fix(benchmark): count the tokens of the answer query_graph returns

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01B8gwEpjajiWHVW3c9smuKf

graphify/benchmark.py

@@ -3,8 +3,7 @@ from __future__ import annotations
 import sys
 import networkx as nx
 
-from graphify.build import edge_data
-from graphify.serve import _query_terms
+from graphify.serve import _query_graph_text, _read_graph
 from graphify.paths import default_graph_json as _default_graph_json
 
 
@@ -34,43 +33,14 @@ def _estimate_tokens(text: str) -> int:
     return max(1, len(text) // _CHARS_PER_TOKEN)
 
 
-def _query_subgraph_tokens(G: nx.Graph, question: str, depth: int = 3) -> int:
-    """Run BFS from best-matching nodes and return estimated tokens in the subgraph context."""
-    terms = _query_terms(question)
-    scored = []
-    for nid, data in G.nodes(data=True):
-        label = (data.get("label") or "").lower()
-        score = sum(1 for t in terms if t in label)
-        if score > 0:
-            scored.append((score, nid))
-    scored.sort(reverse=True)
-    start_nodes = [nid for _, nid in scored[:3]]
-    if not start_nodes:
+def _query_subgraph_tokens(G: nx.Graph, question: str, depth: int = 3, token_budget: int = 2000) -> int:
+    """Estimated tokens of the answer query_graph returns for ``question`` (BFS,
+    the MCP tool's defaults), 0 when nothing in the graph matches it. The
+    answer is the real one an agent receives: seeds, budget cut and notices."""
+    text = _query_graph_text(G, question, mode="bfs", depth=depth, token_budget=token_budget)
+    if text == "No matching nodes found.":
         return 0
-
-    visited: set[str] = set(start_nodes)
-    frontier = set(start_nodes)
-    edges_seen: list[tuple] = []
-    for _ in range(depth):
-        next_frontier: set[str] = set()
-        for n in frontier:
-            for neighbor in G.neighbors(n):
-                if neighbor not in visited:
-                    next_frontier.add(neighbor)
-                    edges_seen.append((n, neighbor))
-        visited.update(next_frontier)
-        frontier = next_frontier
-
-    lines = []
-    for nid in visited:
-        d = G.nodes[nid]
-        lines.append(f"NODE {d.get('label', nid)} src={d.get('source_file', '')} loc={d.get('source_location', '')}")
-    for u, v in edges_seen:
-        if u in visited and v in visited:
-            d = edge_data(G, u, v)
-            lines.append(f"EDGE {G.nodes[u].get('label', u)} --{d.get('relation', '')}--> {G.nodes[v].get('label', v)}")
-
-    return _estimate_tokens("\n".join(lines))
+    return _estimate_tokens(text)
 
 
 _SAMPLE_QUESTIONS = [
@@ -97,11 +67,10 @@ def run_benchmark(
     Returns dict with: corpus_tokens, avg_query_tokens, reduction_ratio, per_question
     """
     graph_path = graph_path or _default_graph_json()
-    # Size-cap check + links/edges normalization + node-link parse. A raw
-    # --no-cluster graph stores edges under "edges" and used to KeyError
-    # here (#2212).
-    from graphify.paths import load_node_link_graph
-    G = load_node_link_graph(graph_path)
+    # Loaded exactly as the MCP server loads it (size cap, links/edges
+    # normalization for a raw --no-cluster graph, #2212), so every answer
+    # measured is the one query_graph returns.
+    G = _read_graph(str(graph_path))
 
     if corpus_words is None:
         # Rough estimate: each node label is ~3 words, plus source context

graphify/serve.py

@@ -41,53 +41,60 @@ class ToolError(Exception):
     """
 
 
+def _read_graph(graph_path: str) -> nx.Graph:
+    """Load graph.json exactly as the MCP server does, raising on a bad path
+    or file (ValueError, FileNotFoundError, json.JSONDecodeError) so a library
+    caller such as the benchmark can handle it."""
+    resolved = Path(graph_path).resolve()
+    if resolved.suffix != ".json":
+        raise ValueError(f"Graph path must be a .json file, got: {graph_path!r}")
+    if not resolved.exists():
+        raise FileNotFoundError(f"Graph file not found: {resolved}")
+    check_graph_file_size_cap(resolved)
+    safe = resolved
+    data = json.loads(safe.read_text(encoding="utf-8"))
+    if "links" not in data and "edges" in data:
+        data = dict(data, links=data["edges"])
+    # Stash the on-disk logical flag before the load-time override below:
+    # `directed: True` exists only so renderers can recover stored arc
+    # order (#2309); tools that care about logical direction (#2487) must
+    # not mistake the override for graph truth.
+    _logical_directed = bool(data.get("directed", False))
+    data = {**data, "directed": True}
+    try:
+        from graphify.build import graph_has_legacy_ids as _legacy
+        if _legacy(data.get("nodes", [])):
+            print(
+                "[graphify] note: this graph uses the pre-#1504 node-ID scheme; "
+                "rebuild with `graphify extract --force` for path-qualified IDs.",
+                file=sys.stderr,
+            )
+    except Exception:
+        pass
+    try:
+        G = json_graph.node_link_graph(data, edges="links")
+    except TypeError:
+        G = json_graph.node_link_graph(data)
+    G.graph["_logical_directed"] = _logical_directed
+    # node_link_graph drops top-level keys other than data["graph"], so the
+    # build commit export.to_json writes there is stashed by hand (#3354).
+    _commit = data.get("built_at_commit")
+    if isinstance(_commit, str) and _commit.strip():
+        G.graph["_built_at_commit"] = _commit.strip()
+    # Attach the work-memory overlay (derived sidecar next to graph.json) so
+    # the query/MCP read surface can annotate NODE lines display-only. Empty
+    # when no sidecar exists, leaving un-annotated output byte-identical.
+    try:
+        from graphify.reflect import load_learning_overlay as _llo
+        G.graph["_learning_overlay"] = _llo(resolved)
+    except Exception:
+        G.graph["_learning_overlay"] = {}
+    return G
+
+
 def _load_graph(graph_path: str) -> nx.Graph:
     try:
-        resolved = Path(graph_path).resolve()
-        if resolved.suffix != ".json":
-            raise ValueError(f"Graph path must be a .json file, got: {graph_path!r}")
-        if not resolved.exists():
-            raise FileNotFoundError(f"Graph file not found: {resolved}")
-        check_graph_file_size_cap(resolved)
-        safe = resolved
-        data = json.loads(safe.read_text(encoding="utf-8"))
-        if "links" not in data and "edges" in data:
-            data = dict(data, links=data["edges"])
-        # Stash the on-disk logical flag before the load-time override below:
-        # `directed: True` exists only so renderers can recover stored arc
-        # order (#2309); tools that care about logical direction (#2487) must
-        # not mistake the override for graph truth.
-        _logical_directed = bool(data.get("directed", False))
-        data = {**data, "directed": True}
-        try:
-            from graphify.build import graph_has_legacy_ids as _legacy
-            if _legacy(data.get("nodes", [])):
-                print(
-                    "[graphify] note: this graph uses the pre-#1504 node-ID scheme; "
-                    "rebuild with `graphify extract --force` for path-qualified IDs.",
-                    file=sys.stderr,
-                )
-        except Exception:
-            pass
-        try:
-            G = json_graph.node_link_graph(data, edges="links")
-        except TypeError:
-            G = json_graph.node_link_graph(data)
-        G.graph["_logical_directed"] = _logical_directed
-        # node_link_graph drops top-level keys other than data["graph"], so the
-        # build commit export.to_json writes there is stashed by hand (#3354).
-        _commit = data.get("built_at_commit")
-        if isinstance(_commit, str) and _commit.strip():
-            G.graph["_built_at_commit"] = _commit.strip()
-        # Attach the work-memory overlay (derived sidecar next to graph.json) so
-        # the query/MCP read surface can annotate NODE lines display-only. Empty
-        # when no sidecar exists, leaving un-annotated output byte-identical.
-        try:
-            from graphify.reflect import load_learning_overlay as _llo
-            G.graph["_learning_overlay"] = _llo(resolved)
-        except Exception:
-            G.graph["_learning_overlay"] = {}
-        return G
+        return _read_graph(graph_path)
     except json.JSONDecodeError as exc:
         print(f"error: graph.json is corrupted ({exc}). Re-run /graphify to rebuild.", file=sys.stderr)
         sys.exit(1)

tests/test_benchmark_real_query.py

@@ -0,0 +1,74 @@
+"""`graphify benchmark` counts the answer query_graph really returns."""
+from __future__ import annotations
+
+import json
+
+import networkx as nx
+import pytest
+from networkx.readwrite import json_graph
+
+from graphify.benchmark import _estimate_tokens, _query_subgraph_tokens, run_benchmark
+from graphify.serve import _load_graph, _query_graph_text, _read_graph
+
+
+def _write(G: nx.Graph, path) -> str:
+    path.write_text(json.dumps(json_graph.node_link_data(G, edges="links")), encoding="utf-8")
+    return str(path)
+
+
+def _small_graph() -> nx.DiGraph:
+    G = nx.DiGraph()
+    G.add_node("auth", label="authentication()", source_file="auth.py", source_location="L1")
+    G.add_node("login", label="login()", source_file="views.py", source_location="L9")
+    G.add_node("user", label="get_user()", source_file="db.py", source_location="L3")
+    G.add_edge("login", "auth", relation="calls", confidence="EXTRACTED")
+    G.add_edge("auth", "user", relation="calls", confidence="EXTRACTED")
+    return G
+
+
+def _hub_graph(spokes: int = 400) -> nx.DiGraph:
+    G = nx.DiGraph()
+    G.add_node("auth", label="authentication()", source_file="auth.py", source_location="L1")
+    for i in range(spokes):
+        G.add_node(f"n{i}", label=f"helper_function_number_{i}_with_a_long_name()",
+                   source_file=f"pkg/module_{i}.py", source_location=f"L{i}")
+        G.add_edge("auth", f"n{i}", relation="calls", confidence="EXTRACTED")
+    return G
+
+
+def test_tokens_are_those_of_the_mcp_query_graph_answer(tmp_path):
+    path = _write(_small_graph(), tmp_path / "graph.json")
+    question = "how does authentication work"
+    answer = _query_graph_text(_load_graph(path), question, mode="bfs", depth=3, token_budget=2000)
+    result = run_benchmark(path, corpus_words=10_000, questions=[question])
+    assert result["per_question"][0]["query_tokens"] == _estimate_tokens(answer)
+
+
+def test_tokens_stop_at_the_query_budget(tmp_path):
+    # 400 callees: query_graph cuts the answer at its ~2000-token budget, the
+    # benchmark must count that cut answer, not every node BFS reaches.
+    path = _write(_hub_graph(), tmp_path / "graph.json")
+    G = _read_graph(path)
+    full = _query_graph_text(G, "authentication", token_budget=10**6)
+    tokens = _query_subgraph_tokens(G, "authentication")
+    assert "[!] TRUNCATED" in _query_graph_text(G, "authentication")
+    assert tokens <= 2000 < _estimate_tokens(full)
+
+
+def test_a_larger_budget_buys_a_larger_answer(tmp_path):
+    G = _read_graph(_write(_hub_graph(), tmp_path / "graph.json"))
+    assert _query_subgraph_tokens(G, "authentication", token_budget=500) < _query_subgraph_tokens(G, "authentication")
+
+
+def test_no_answer_counts_zero_and_is_left_out(tmp_path):
+    path = _write(_small_graph(), tmp_path / "graph.json")
+    assert _query_subgraph_tokens(_read_graph(path), "xyzzy plugh") == 0
+    result = run_benchmark(path, corpus_words=10_000, questions=["xyzzy plugh", "authentication"])
+    assert [p["question"] for p in result["per_question"]] == ["authentication"]
+
+
+def test_bad_graph_path_raises_instead_of_exiting(tmp_path):
+    with pytest.raises(FileNotFoundError):
+        run_benchmark(str(tmp_path / "missing.json"))
+    with pytest.raises(ValueError):
+        run_benchmark(str(tmp_path / "graph.txt"))