diff --git a/graphify/__main__.py b/graphify/__main__.py index 74ffb8b69..048206a4a 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1223,7 +1223,7 @@ def main() -> None: sys.exit(0) start = [nid for _, nid in scored[:5]] nodes, edges = (_dfs if use_dfs else _bfs)(G, start, depth=2) - print(_subgraph_to_text(G, nodes, edges, token_budget=budget)) + print(_subgraph_to_text(G, nodes, edges, token_budget=budget, seeds=start)) elif cmd == "save-result": # graphify save-result --question Q --answer A --type T [--nodes N1 N2 ...] import argparse as _ap diff --git a/graphify/serve.py b/graphify/serve.py index 361dec3c0..02bf88cef 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -45,13 +45,25 @@ def _strip_diacritics(text: str) -> str: return "".join(c for c in nfkd if not unicodedata.combining(c)) +# Identifier-style queries (e.g. a function name) often tie at score 1 against +# every hub module that mentions the symbol, then lose tie-breaks to high-degree +# nodes during _subgraph_to_text rendering. The bonus lifts an exact-token match +# above any combination of substring + source-file scoring so the actual symbol +# always seeds the BFS and renders first. +EXACT_MATCH_BONUS = 100.0 + + def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]: scored = [] norm_terms = [_strip_diacritics(t).lower() for t in terms] for nid, data in G.nodes(data=True): norm_label = data.get("norm_label") or _strip_diacritics(data.get("label") or "").lower() + # AST extractors emit function labels with trailing parens (e.g. "foo()"). + # Treat the bare identifier as an exact match too. + norm_label_bare = norm_label.rstrip().rstrip("()") source = (data.get("source_file") or "").lower() score = sum(1 for t in norm_terms if t in norm_label) + sum(0.5 for t in norm_terms if t in source) + score += sum(EXACT_MATCH_BONUS for t in norm_terms if t == norm_label or t == norm_label_bare) if score > 0: scored.append((score, nid)) return sorted(scored, reverse=True) @@ -89,11 +101,31 @@ def _dfs(G: nx.Graph, start_nodes: list[str], depth: int) -> tuple[set[str], lis return visited, edges_seen -def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_budget: int = 2000) -> str: - """Render subgraph as text, cutting at token_budget (approx 3 chars/token).""" +def _subgraph_to_text( + G: nx.Graph, + nodes: set[str], + edges: list[tuple], + token_budget: int = 2000, + seeds: list[str] | None = None, +) -> str: + """Render subgraph as text, cutting at token_budget (approx 3 chars/token). + + When ``seeds`` is provided, those node IDs render first in the order given + (preserving the caller's ranking from _score_nodes). Remaining nodes follow + sorted by degree desc — without seeds, this matches the legacy ordering. + """ char_budget = token_budget * 3 lines = [] - for nid in sorted(nodes, key=lambda n: G.degree(n), reverse=True): + seeds = seeds or [] + seen: set[str] = set() + ordered: list[str] = [] + for nid in seeds: + if nid in nodes and nid not in seen: + ordered.append(nid) + seen.add(nid) + for nid in sorted((n for n in nodes if n not in seen), key=lambda n: G.degree(n), reverse=True): + ordered.append(nid) + for nid in ordered: d = G.nodes[nid] line = f"NODE {sanitize_label(d.get('label', nid))} [src={d.get('source_file', '')} loc={d.get('source_location', '')} community={d.get('community', '')}]" lines.append(line) @@ -246,7 +278,7 @@ def _tool_query_graph(arguments: dict) -> str: return "No matching nodes found." nodes, edges = _dfs(G, start_nodes, depth) if mode == "dfs" else _bfs(G, start_nodes, depth) header = f"Traversal: {mode.upper()} depth={depth} | Start: {[G.nodes[n].get('label', n) for n in start_nodes]} | {len(nodes)} nodes found\n\n" - return header + _subgraph_to_text(G, nodes, edges, budget) + return header + _subgraph_to_text(G, nodes, edges, budget, seeds=start_nodes) def _tool_get_node(arguments: dict) -> str: label = arguments["label"].lower() diff --git a/tests/test_serve.py b/tests/test_serve.py index 6457ac501..3eebe639f 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -11,6 +11,7 @@ _dfs, _subgraph_to_text, _load_graph, + EXACT_MATCH_BONUS, ) @@ -151,3 +152,101 @@ def test_load_graph_missing_file(tmp_path): graphify_dir.mkdir() with pytest.raises(SystemExit): _load_graph(str(graphify_dir / "nonexistent.json")) + + +# --- exact-match bonus + seed-first rendering --- +# +# Single-token identifier queries (e.g. a function name) used to tie at score 1 +# against every node containing the substring, then lose tie-breaks during +# _subgraph_to_text rendering, where high-degree hubs (app.js, controller.js) +# always rendered first. The two changes below — EXACT_MATCH_BONUS in +# _score_nodes and the optional `seeds` arg in _subgraph_to_text — fix that. + +def _make_hub_graph() -> nx.Graph: + """Graph with a low-degree exact-match seed and a high-degree hub.""" + G = nx.Graph() + # Seed: the function we're querying for. Degree 1. + G.add_node("seed", label="pasteFromClipboard()", source_file="frontend/clipboard.js", + source_location="L42", community=1) + # High-degree hub. The substring "pasteFromClipboard" also appears in app.js + # via call sites, so it scores 1 on a single-token query without the bonus. + G.add_node("hub", label="app.js", source_file="frontend/app.js", + source_location="L1", community=1) + G.add_edge("seed", "hub", relation="defined_in", confidence="EXTRACTED") + # Distractors that mention the substring but aren't exact matches. + for i, name in enumerate(["pasteFromClipboard_handler", "wrap_pasteFromClipboard", "_pasteFromClipboard_inner"]): + nid = f"sub{i}" + G.add_node(nid, label=name, source_file="frontend/app.js", + source_location=f"L{100+i}", community=1) + G.add_edge("hub", nid, relation="defines", confidence="EXTRACTED") + # Pad the hub up to degree ~10 so degree-sort would otherwise float it to top. + for i in range(7): + nid = f"pad{i}" + G.add_node(nid, label=f"helper{i}", source_file="frontend/app.js", + source_location=f"L{200+i}", community=1) + G.add_edge("hub", nid, relation="defines", confidence="EXTRACTED") + return G + + +def test_score_nodes_exact_match_beats_substring(): + G = _make_hub_graph() + scored = _score_nodes(G, ["pastefromclipboard"]) + assert scored, "expected at least one scoring node" + top_score, top_nid = scored[0] + assert top_nid == "seed", f"exact match should win; got {top_nid}" + # Bonus should dominate any substring sum. + assert top_score >= EXACT_MATCH_BONUS + # The substring-only matches must rank well below the seed. + sub_scores = [s for s, nid in scored if nid != "seed"] + assert all(s < EXACT_MATCH_BONUS for s in sub_scores) + + +def test_score_nodes_exact_match_strips_function_parens(): + """Labels emitted by the AST extractor often carry trailing parens (foo()).""" + G = nx.Graph() + G.add_node("a", label="saveDiagram()", source_file="x.js", source_location="L1", community=0) + G.add_node("b", label="saveDiagram_helper", source_file="x.js", source_location="L2", community=0) + scored = _score_nodes(G, ["savediagram"]) + assert scored[0][1] == "a" + assert scored[0][0] >= EXACT_MATCH_BONUS + + +def test_score_nodes_exact_match_no_false_positive(): + """Unrelated query must not trigger the bonus.""" + G = _make_hub_graph() + scored = _score_nodes(G, ["xyzzy"]) + assert scored == [] + + +def test_subgraph_to_text_seeds_render_first(): + G = _make_hub_graph() + nodes = {"seed", "hub", "sub0", "sub1", "sub2"} | {f"pad{i}" for i in range(7)} + edges = list(G.edges()) + text = _subgraph_to_text(G, nodes, edges, token_budget=4000, seeds=["seed"]) + seed_pos = text.index("pasteFromClipboard") + hub_pos = text.index("app.js") + assert seed_pos < hub_pos, "seed must render before the high-degree hub" + + +def test_subgraph_to_text_no_seeds_preserves_legacy_order(): + """Without seeds, ordering still falls back to degree desc (back-compat).""" + G = _make_hub_graph() + nodes = {"seed", "hub", "sub0", "sub1", "sub2"} | {f"pad{i}" for i in range(7)} + edges = list(G.edges()) + legacy = _subgraph_to_text(G, nodes, edges, token_budget=4000) + explicit_none = _subgraph_to_text(G, nodes, edges, token_budget=4000, seeds=None) + assert legacy == explicit_none + # Hub has degree ~11 vs seed degree 1, so without seeds the hub renders first. + assert legacy.index("app.js") < legacy.index("pasteFromClipboard") + + +def test_query_pipeline_exact_match_ranks_above_hub(): + """End-to-end: _score_nodes -> _bfs -> _subgraph_to_text(seeds=...).""" + G = _make_hub_graph() + terms = ["pastefromclipboard"] + scored = _score_nodes(G, terms) + start = [nid for _, nid in scored[:5]] + assert start[0] == "seed" + nodes, edges = _bfs(G, start, depth=2) + text = _subgraph_to_text(G, nodes, edges, token_budget=4000, seeds=start) + assert text.index("pasteFromClipboard") < text.index("app.js")