grammar-inference-engine/tests/test_grammar_index.py
tobjend 7b54a5ef77 feat: grammar index for runtime lookup
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.
2026-07-12 14:08:05 +02:00

162 lines
5.5 KiB
Python

"""Tests for bex.grammar_index — runtime grammar lookup."""
import os
import tempfile
import pytest
import yaml
from bex.grammar_index import GrammarIndex, load_grammar_index
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
def _make_yml(tmp_path, entries):
"""Write a grammars.yml with the given entries grouped by module."""
dervish = tmp_path / ".dervish"
dervish.mkdir()
data = {"test_module": entries}
with open(dervish / "grammars.yml", "w") as f:
yaml.dump(data, f, default_flow_style=False)
SAMPLE_ENTRIES = [
{
"package": "src/services [return]",
"methods": 10,
"grammar": "return.ok+?.error?",
"score": 0.8,
"algorithm": "CRX",
"mdl": 12.0,
},
{
"package": "src/services [if]",
"methods": 5,
"grammar": "if.cond+.then+.else?",
"score": 0.6,
"algorithm": "CRX",
"mdl": 15.0,
},
{
"package": "src/utils [return]",
"methods": 8,
"grammar": "return.value+",
"score": 1.0,
"algorithm": "CRX",
"mdl": 5.0,
},
]
# ---------------------------------------------------------------------------
# GrammarIndex construction
# ---------------------------------------------------------------------------
class TestGrammarIndex:
def test_parse_label_with_symbol(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
pkgs = idx.packages()
assert "src/services" in pkgs
assert "src/utils" in pkgs
def test_parse_label_without_symbol(self, tmp_path):
entries = [{"package": "src/root", "grammar": "a.b", "score": 0.5}]
idx = GrammarIndex(str(tmp_path), entries)
assert "src/root" in idx.packages()
def test_sorted_by_score_desc(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
entries = idx.get_package("src/services")
scores = [s for _, _, s, _ in entries]
assert scores == sorted(scores, reverse=True)
def test_repr(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
assert "3 entries" in repr(idx)
assert "2 packages" in repr(idx)
# ---------------------------------------------------------------------------
# get()
# ---------------------------------------------------------------------------
class TestGet:
def test_get_best_no_context(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
g = idx.get("src/services/main.kt")
assert g == "return.ok+?.error?"
def test_get_by_context_symbol(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
g = idx.get("src/services/main.kt", context_symbol="if")
assert g == "if.cond+.then+.else?"
def test_get_missing_context_falls_back_to_best(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
g = idx.get("src/services/main.kt", context_symbol="nonexistent")
assert g == "return.ok+?.error?"
def test_get_unknown_package_returns_none(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
g = idx.get("src/unknown/file.kt")
assert g is None
def test_get_absolute_path(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
abs_path = str(tmp_path / "src" / "services" / "main.kt")
g = idx.get(abs_path, context_symbol="return")
assert g == "return.ok+?.error?"
# ---------------------------------------------------------------------------
# get_package()
# ---------------------------------------------------------------------------
class TestGetPackage:
def test_returns_all_entries(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
entries = idx.get_package("src/services/main.kt")
assert len(entries) == 2
def test_unknown_package_returns_empty(self, tmp_path):
idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
entries = idx.get_package("src/unknown/file.kt")
assert entries == []
# ---------------------------------------------------------------------------
# load_grammar_index()
# ---------------------------------------------------------------------------
class TestLoadGrammarIndex:
def test_load_existing(self, tmp_path):
_make_yml(tmp_path, SAMPLE_ENTRIES)
idx = load_grammar_index(str(tmp_path))
assert len(idx.get_all()) == 3
assert "src/services" in idx.packages()
def test_load_missing_file(self, tmp_path):
idx = load_grammar_index(str(tmp_path))
assert len(idx.get_all()) == 0
assert idx.packages() == []
def test_roundtrip_persisted_ragsak(self):
"""Integration test: grammars.yml generated by analyze_directory."""
ragsak_yml = "/home/tobi/Desktop/kesai/RAGSAK/.dervish/grammars.yml"
if not os.path.exists(ragsak_yml):
pytest.skip("RAGSAK grammars.yml not generated yet")
idx = load_grammar_index("/home/tobi/Desktop/kesai/RAGSAK")
assert len(idx.get_all()) > 0
# Should be able to resolve a file in the storage package
g = idx.get(
"/home/tobi/Desktop/kesai/RAGSAK/infrastructure/adapters/search/src/main/kotlin/eu/corentic/springrag/service/storage/GcsStorage.kt",
context_symbol="return",
)
assert g is not None
assert "return" in g