diff --git a/bex/mcp_server.py b/bex/mcp_server.py index 86b04bf..0626fef 100644 --- a/bex/mcp_server.py +++ b/bex/mcp_server.py @@ -1,12 +1,17 @@ """Dervish — MCP server. -Provides tools to infer regular expression grammars from example sequences. +Provides tools to infer regular expression grammars from example sequences, +and to look up the right grammar at code generation time. + Run as: python -m bex.mcp_server """ +import os + from mcp.server.fastmcp import FastMCP from .ensemble import infer_ensemble, _matches +from .grammar_index import load_grammar_index from .tag_preprocessor.analyze import ( analyze_directory as _analyze_directory, _build_yaml_output, @@ -16,6 +21,10 @@ from .tag_preprocessor.analyze import ( mcp = FastMCP("grammar-inference", log_level="ERROR") +# --------------------------------------------------------------------------- +# Inference tools +# --------------------------------------------------------------------------- + @mcp.tool() def infer_best_grammar( sequences: list[list[str]], @@ -87,7 +96,7 @@ def infer_best_grammar( def analyze_directory( directory: str, slice: str = "package", - min_coverage: float = 0.8, + min_coverage: float = 0.05, prefer: str = "", kmax: int = 2, include: str = "", @@ -96,10 +105,13 @@ def analyze_directory( max_mdl: float = 200, persist: bool = True, method: str = "langsize", + min_methods: int = 3, + min_structure: float = 0.2, + split_mixed: bool = True, ) -> str: """Scan a source code directory and infer behavioral conventions (regular expression grammars) per package. Returns compact patterns - grouped by module, sorted by quality (MDL score). + grouped by module, sorted by quality. Use this when you need to understand the calling conventions in a codebase — what patterns new code should follow. The grammar @@ -107,27 +119,30 @@ def analyze_directory( regular expression. Auto-persists results to {directory}/.dervish/grammars.yml unless - persist=False. + persist=False. An agent can later look up the right grammar via + get_grammar() during code generation. Args: directory: Path to the source code directory to analyze. slice: Grouping strategy — 'package' (per directory, default) or 'flat' (one per language). - min_coverage: BEX core coverage threshold for outlier removal - (0.5–1.0). Default 0.8. + min_coverage: Coverage threshold for outlier removal (0.01–1.0). + Lower = see more symbols. Default 0.05. prefer: Optional — 'crx' for full vocabulary, 'idregex' for minimal core. Omit to auto-pick by scoring. kmax: Context depth for k-ORE inference. Default 2. include: Glob pattern to include only matching files. exclude: Glob pattern to skip matching files. - main_only: When True, exclude test files (src/test/**, *Test.*, - etc.). Default False. + main_only: When True, exclude test files. Default False. max_mdl: Drop groups with score above this threshold. Default 200. - Lower = tighter patterns only. Set higher to see noisier groups. persist: When True (default), write results to {directory}/.dervish/grammars.yml. - method: Scoring method — 'langsize' (default, Bex et al.) or - 'mdl' (fallback). + method: Scoring method — 'langsize' (default) or 'mdl'. + min_methods: Minimum methods per group to attempt inference. Default 3. + min_structure: Minimum grammar structure score (0.0–1.0). Groups + below this produce flat bags. Default 0.2. + split_mixed: When True (default), recursively split groups with + diverse first symbols into uniform sub-groups before inference. Returns: YAML string with grammars grouped by top-level module, sorted @@ -143,6 +158,9 @@ def analyze_directory( exclude=exclude or None, main_only=main_only, method=method, + min_methods=min_methods, + split_mixed=split_mixed, + min_structure=min_structure, ) yaml_content = _build_yaml_output(results, directory, max_mdl=max_mdl) if persist: @@ -150,6 +168,90 @@ def analyze_directory( return yaml_content +# --------------------------------------------------------------------------- +# Runtime lookup tools +# --------------------------------------------------------------------------- + +@mcp.tool() +def get_grammar( + directory: str, + file_path: str, + context_symbol: str = "", +) -> str: + """Get the grammar constraint for generating code in a specific file. + + Call this before generating code to get the GBNF grammar that matches + the file's package and the calling context. The agent should pass the + returned grammar to the LLM's constrained generation backend. + + Workflow: + 1. Agent is about to write code in `file_path` + 2. Call get_grammar(directory, file_path, context_symbol="return") + 3. Pass the returned GBNF to the LLM's grammar-constrained sampler + 4. LLM generates code that follows the package's convention + + Args: + directory: Project root (must match the directory used in + analyze_directory). + file_path: Path to the file being written (absolute or relative + to directory). + context_symbol: Optional — the first symbol of the code being + generated (e.g. "return", "if", "try"). When provided, returns + the leaf grammar for that specific context. When empty, returns + the best grammar for the file's package. + + Returns: + GBNF grammar string, or a message explaining why no grammar was found. + """ + idx = load_grammar_index(directory) + if not idx.get_all(): + return f"No grammars indexed for {directory}. Run analyze_directory first." + + ctx = context_symbol if context_symbol else None + grammar = idx.get(file_path, context_symbol=ctx) + if grammar is None: + return f"No grammar found for {file_path} (context={context_symbol or 'best'})." + + return grammar + + +@mcp.tool() +def get_package_grammars( + directory: str, + file_path: str, +) -> str: + """Get all available grammars for a file's package, ranked by quality. + + Use this to explore what conventions exist for a package before + choosing which context to generate in. Returns all leaf grammars + (one per calling context like "return", "if", "try", etc.). + + Args: + directory: Project root. + file_path: Path to the file (absolute or relative to directory). + + Returns: + Formatted list of (context_symbol, grammar, score, methods) tuples. + """ + idx = load_grammar_index(directory) + if not idx.get_all(): + return f"No grammars indexed for {directory}. Run analyze_directory first." + + entries = idx.get_package(file_path) + if not entries: + return f"No grammars found for package of {file_path}." + + lines = [f"Grammars for {file_path}:"] + for sym, grammar, score, methods in entries: + label = f"[{sym}]" if sym else "(best)" + lines.append(f" {label:25s} score={score:.3f} {methods}m {grammar[:70]}") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + def main(): mcp.run() diff --git a/bex/tag_preprocessor/analyze.py b/bex/tag_preprocessor/analyze.py index f25d462..0c12744 100644 --- a/bex/tag_preprocessor/analyze.py +++ b/bex/tag_preprocessor/analyze.py @@ -810,16 +810,32 @@ def _build_yaml_output(results, dir_path, max_mdl=200.0): parts = label.replace(os.sep, "/").split("/") module = parts[0] if len(parts) > 1 else "(root)" - entry = { - "package": label, - "methods": count, - "grammar": best["grammar"], - "algorithm": best["algorithm"], - "mdl": round(best["mdl_score"], 1), - } + # When recursive split produced leaf grammars, output each one + all_grammars = meta.get("all_grammars", []) + if all_grammars: + for leaf_label, leaf_grammar, leaf_score, leaf_count in all_grammars: + entry = { + "package": leaf_label, + "methods": leaf_count, + "grammar": leaf_grammar, + "score": round(leaf_score, 3), + "algorithm": best["algorithm"], + "mdl": round(best["mdl_score"], 1), + } + modules.setdefault(module, []).append(entry) + total_patterns += 1 + else: + entry = { + "package": label, + "methods": count, + "grammar": best["grammar"], + "score": round(best.get("mdl_score", 0), 3), + "algorithm": best["algorithm"], + "mdl": round(best["mdl_score"], 1), + } - modules.setdefault(module, []).append(entry) - total_patterns += 1 + modules.setdefault(module, []).append(entry) + total_patterns += 1 # Sort entries within each module by MDL for module in modules: