From ee60b628548531b0d8b4dc6a9737ced32d7da7ed Mon Sep 17 00:00:00 2001
From: tobjend <tobend85@gmail.com>
Date: Sun, 12 Jul 2026 19:53:24 +0200
Subject: [PATCH] fix: YAML expansion bug and add GBNF to output
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Removed all_grammars expansion in _build_yaml_output (226→54 RAGSAK, 384→87 FastAPI)
- Added gbnf field to each YAML entry via to_gbnf()
- Updated test_roundtrip_persisted_ragsak to use correct file path

Closes #66
---
 bex/tag_preprocessor/analyze.py | 41 +++++++++++----------------------
 tests/test_grammar_index.py     |  8 +++----
 2 files changed, 17 insertions(+), 32 deletions(-)

diff --git a/bex/tag_preprocessor/analyze.py b/bex/tag_preprocessor/analyze.py
index fb5bcdd..fa9432e 100644
--- a/bex/tag_preprocessor/analyze.py
+++ b/bex/tag_preprocessor/analyze.py
@@ -20,7 +20,7 @@ import pathspec
 
 from .code import preprocess_by_method, extract_arg_info, _summarize_arg_info
 from bex.ensemble import infer_ensemble
-from bex.gbnf import validate_sore, grammar_structure_score
+from bex.gbnf import validate_sore, grammar_structure_score, to_gbnf
 from bex.distributional import distributional_split
 from bex.decompose import decompose_with_coverage, get_decomposition_stats
 
@@ -916,34 +916,19 @@ def _build_yaml_output(results, dir_path, max_mdl=200.0, min_structure=0.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:
-                    if grammar_structure_score(leaf_grammar) < min_structure:
-                        continue
-                    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),
-                }
+            # Output the best grammar per parent group (don't expand leaves)
+            entry = {
+                "package": label,
+                "methods": count,
+                "grammar": best["grammar"],
+                "gbnf": to_gbnf(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:
diff --git a/tests/test_grammar_index.py b/tests/test_grammar_index.py
index 8154e23..904d79f 100644
--- a/tests/test_grammar_index.py
+++ b/tests/test_grammar_index.py
@@ -153,10 +153,10 @@ class TestLoadGrammarIndex:
             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
+        # Should be able to resolve a file in a package with a grammar
         g = idx.get(
-            "/home/tobi/Desktop/kesai/RAGSAK/infrastructure/adapters/search/src/main/kotlin/eu/corentic/springrag/service/storage/GcsStorage.kt",
-            context_symbol="return",
+            "/home/tobi/Desktop/kesai/RAGSAK/agents/capability/support/src/main/kotlin/eu/corentic/springrag/agent/capability/support/AgentCapabilityConfiguration.kt",
+            context_symbol="id",
         )
         assert g is not None
-        assert "return" in g
+        assert "id" in g