From 687cedd4077dc77ccdc7e557f85acd523faeb11c Mon Sep 17 00:00:00 2001 From: rohit-jsfreaky Date: Mon, 12 Oct 2026 01:39:17 +0530 Subject: [PATCH] fix(benchmark): count the tokens of the answer query_graph returns Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01B8gwEpjajiWHVW3c9smuKf --- graphify/benchmark.py | 55 ++++------------- graphify/serve.py | 97 ++++++++++++++++-------------- tests/test_benchmark_real_query.py | 74 +++++++++++++++++++++++ 3 files changed, 138 insertions(+), 88 deletions(-) create mode 100644 tests/test_benchmark_real_query.py diff --git a/graphify/benchmark.py b/graphify/benchmark.py index a7edf5d08..7a117f6fe 100644 --- a/graphify/benchmark.py +++ b/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 diff --git a/graphify/serve.py b/graphify/serve.py index f2fbfe231..165c2466e 100644 --- a/graphify/serve.py +++ b/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) diff --git a/tests/test_benchmark_real_query.py b/tests/test_benchmark_real_query.py new file mode 100644 index 000000000..9be6c6422 --- /dev/null +++ b/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"))