grammar-inference-engine/bex/mcp_server.py
tobjend 44415c5b42 feat: runtime grammar lookup via MCP + fix YAML output
MCP server:
- analyze_directory: updated signature to match CLI (split_mixed,
  min_structure, min_methods, min_coverage=0.05)
- get_grammar(directory, file_path, context_symbol): returns the right
  GBNF for constrained generation at code generation time
- get_package_grammars(directory, file_path): lists all grammars for
  a file's package, ranked by quality

YAML output:
- _build_yaml_output now includes leaf grammars from recursive split
  (all_grammars in meta), so each calling context gets its own entry

Workflow for agents:
  1. analyze_directory → grammars.yml persisted
  2. get_grammar(file, 'return') → GBNF for constrained generation
  3. LLM generates code following the package's convention
2026-07-12 14:08:12 +02:00

260 lines
9.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Dervish — MCP server.
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,
_persist_grammars,
)
mcp = FastMCP("grammar-inference", log_level="ERROR")
# ---------------------------------------------------------------------------
# Inference tools
# ---------------------------------------------------------------------------
@mcp.tool()
def infer_best_grammar(
sequences: list[list[str]],
prefer: str = "",
kmax: int = 2,
N: int = 3,
min_coverage: float = 1.0,
method: str = "langsize",
) -> str:
"""Infer a compact grammar from example sequences. Use this when you
have examples of sequential data and want to learn the pattern.
The grammar compresses N examples into ~100 chars — far fewer tokens
than passing all examples. Pass the existing sequences, get back a
pattern you can follow to generate new instances.
Runs CRX + iDRegEx, picks best by scoring. kORE is excluded by
default (slow, rarely wins on real data). Set prefer='koreinference'
to force it.
Args:
sequences: List of sequences, each a list of strings (symbols in
the order they appear). Example: [["file","copy","command"],
["file","template","command"]].
prefer: Optional — 'crx' for full vocabulary (accepts all examples),
'idregex' for deterministic minimal core, 'koreinference' for
k-OA with rwr0 repair (slow). Omit to auto-pick by MDL.
kmax: Context depth for k-ORE inference (iDRegEx, kOREInference).
Default 2.
N: Random trials for k-ORE inference (higher = better, slower).
min_coverage: (Expert) When < 1.0, also runs a **core+outlier analysis**:
iteratively removes outlier sequences (those with rarest symbols)
until at least this fraction remain. Returns the core grammar
for the majority, plus a list of which sequences were removed and why.
Default 1.0 = no core analysis. Set to 0.8 to find the tight
pattern shared by ~80% of examples while flagging the other ~20%
as variations.
Returns:
A formatted string with the best grammar, scores, and explanation.
When min_coverage < 1.0, includes the core grammar and outlier info.
Grammar notation: a.b = a then b, (a+b) = a or b, r? = optional,
r+ = one or more, r+? = zero or more.
"""
pref = prefer if prefer else None
result = infer_ensemble(sequences, kmax=kmax, N=N, prefer=pref, min_coverage=min_coverage, method=method)
if result['best'] is None:
return f"No grammar found. {result['why']}"
lines = [f"Best: {result['best']['algorithm']} (Score {result['best']['mdl_score']})",
f"Grammar: {result['best']['grammar']}",
""]
if len(result['all']) > 1:
for r in result['all']:
m = sum(1 for s in sequences if _matches(r['grammar'], s))
lines.append(f" {r['algorithm']:10s} Score={r['mdl_score']:>8.2f} match={m}/{len(sequences)}")
lines.append("")
lines.append(f"Why: {result['why']}")
if 'core' in result and result['core']:
c = result['core']
lines.append(f"\nCore CRX ({c['coverage']:.0%} coverage, {c['outlier_count']} outliers): {c['grammar']}")
if c['outliers']:
lines.append(f" Outlier sequences:")
for i, o in enumerate(c['outliers'], 1):
lines.append(f" {i}. {''.join(str(x) for x in o[:8])}{'...' if len(o) > 8 else ''}")
return "\n".join(lines)
@mcp.tool()
def analyze_directory(
directory: str,
slice: str = "package",
min_coverage: float = 0.05,
prefer: str = "",
kmax: int = 2,
include: str = "",
exclude: str = "",
main_only: bool = False,
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.
Use this when you need to understand the calling conventions in a
codebase — what patterns new code should follow. The grammar
compresses each package's method call patterns into a compact
regular expression.
Auto-persists results to {directory}/.dervish/grammars.yml unless
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: Coverage threshold for outlier removal (0.011.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. Default False.
max_mdl: Drop groups with score above this threshold. Default 200.
persist: When True (default), write results to
{directory}/.dervish/grammars.yml.
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.01.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
by score (tightest/most useful first).
"""
results = _analyze_directory(
directory,
min_coverage=min_coverage,
prefer=prefer or None,
kmax=kmax,
slice=slice,
include=include or None,
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:
_persist_grammars(yaml_content, 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()
if __name__ == "__main__":
main()