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
+        min_coverage: Coverage threshold for outlier removal (0.01–1.0).
-            (0.5–1.0). Default 0.8.
+            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.*,
+        main_only: When True, exclude test files. Default False.
-            etc.). 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
+        method: Scoring method — 'langsize' (default) or 'mdl'.
-            'mdl' (fallback).
+        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()
 

bex/tag_preprocessor/analyze.py

@@ -810,10 +810,26 @@ 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)"
 
+            # 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),
                 }