7b54a5ef77
Add GrammarIndex class that loads grammars.yml and provides fast lookup by (package, context_symbol). Enables agents to get the right GBNF grammar at code generation time. API: load_grammar_index(project_root) → GrammarIndex idx.get(file_path, context_symbol=None) → gbnf_string idx.get_package(file_path) → [(symbol, grammar, score, methods)] 14 tests, all passing.
139 lines
4.4 KiB
Python
139 lines
4.4 KiB
Python
"""Grammar index for runtime lookup.
|
|
|
|
Loads grammars.yml and builds a lookup index mapping
|
|
(package, context_symbol) → GBNF grammar string.
|
|
|
|
Usage:
|
|
from bex.grammar_index import load_grammar_index
|
|
|
|
idx = load_grammar_index("/path/to/project")
|
|
grammar = idx.get(file_path="src/services/Chat.kt", context_symbol="return")
|
|
all_grammars = idx.get_package("src/services")
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import yaml
|
|
|
|
|
|
_LABEL_RE = re.compile(r"^(.+)\s+\[([^\]]+)\]$")
|
|
|
|
|
|
class GrammarIndex:
|
|
"""In-memory index of grammars keyed by (package, context_symbol)."""
|
|
|
|
def __init__(self, project_root, entries):
|
|
"""
|
|
Args:
|
|
project_root: absolute path to the project root.
|
|
entries: list of dicts with keys: package, grammar, score, methods, algorithm.
|
|
"""
|
|
self.project_root = project_root
|
|
self._by_package = {} # package → [(symbol, grammar, score)]
|
|
self._all = entries
|
|
|
|
for e in entries:
|
|
label = e["package"]
|
|
m = _LABEL_RE.match(label)
|
|
if m:
|
|
pkg = m.group(1).rstrip("/")
|
|
symbol = m.group(2)
|
|
else:
|
|
pkg = label.rstrip("/")
|
|
symbol = ""
|
|
|
|
self._by_package.setdefault(pkg, []).append((
|
|
symbol,
|
|
e["grammar"],
|
|
e.get("score", 0),
|
|
e.get("methods", 0),
|
|
))
|
|
|
|
# Sort each package's entries by score descending (best first)
|
|
for pkg in self._by_package:
|
|
self._by_package[pkg].sort(key=lambda x: -x[2])
|
|
|
|
def get(self, file_path, context_symbol=None):
|
|
"""Get the best grammar for a file, optionally filtered by context.
|
|
|
|
Args:
|
|
file_path: path to the source file (absolute or relative to project_root).
|
|
context_symbol: if provided, match the leaf grammar whose first symbol
|
|
is this. If None, return the best grammar for the package.
|
|
|
|
Returns:
|
|
GBNF grammar string, or None if no match.
|
|
"""
|
|
pkg = self._resolve_package(file_path)
|
|
entries = self._by_package.get(pkg, [])
|
|
if not entries:
|
|
return None
|
|
|
|
if context_symbol:
|
|
for sym, grammar, score, methods in entries:
|
|
if sym == context_symbol:
|
|
return grammar
|
|
|
|
# Fall back to best grammar for the package
|
|
return entries[0][1] if entries else None
|
|
|
|
def get_package(self, file_path):
|
|
"""Get all grammars for a file's package.
|
|
|
|
Returns:
|
|
list of (context_symbol, grammar, score, methods) tuples,
|
|
sorted by score descending. Empty list if no match.
|
|
"""
|
|
pkg = self._resolve_package(file_path)
|
|
return list(self._by_package.get(pkg, []))
|
|
|
|
def get_all(self):
|
|
"""Return all entries as a flat list."""
|
|
return list(self._all)
|
|
|
|
def packages(self):
|
|
"""Return sorted list of all indexed packages."""
|
|
return sorted(self._by_package.keys())
|
|
|
|
def _resolve_package(self, file_path):
|
|
"""Map a file path to its package."""
|
|
# Make absolute if relative
|
|
if not os.path.isabs(file_path):
|
|
file_path = os.path.join(self.project_root, file_path)
|
|
|
|
# Get directory of file relative to project root
|
|
try:
|
|
rel = os.path.relpath(os.path.dirname(file_path), self.project_root)
|
|
except ValueError:
|
|
# Different drives on Windows
|
|
return ""
|
|
if rel == ".":
|
|
return ""
|
|
return rel
|
|
|
|
def __repr__(self):
|
|
return f"GrammarIndex({self.project_root}, {len(self._all)} entries, {len(self._by_package)} packages)"
|
|
|
|
|
|
def load_grammar_index(project_root):
|
|
"""Load grammars.yml from {project_root}/.dervish/grammars.yml.
|
|
|
|
Returns GrammarIndex, or empty index if no file found.
|
|
"""
|
|
yml_path = os.path.join(project_root, ".dervish", "grammars.yml")
|
|
if not os.path.exists(yml_path):
|
|
return GrammarIndex(project_root, [])
|
|
|
|
with open(yml_path) as f:
|
|
data = yaml.safe_load(f)
|
|
|
|
entries = []
|
|
if isinstance(data, dict):
|
|
for module, items in data.items():
|
|
if not isinstance(items, list):
|
|
continue
|
|
for item in items:
|
|
if isinstance(item, dict) and "package" in item and "grammar" in item:
|
|
entries.append(item)
|
|
|
|
return GrammarIndex(project_root, entries)
|