diff --git a/graphify/serve.py b/graphify/serve.py index 990a32de0..06e8cf715 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -4,6 +4,7 @@ import math import re import sys +from array import array from pathlib import Path import networkx as nx from networkx.readwrite import json_graph @@ -135,6 +136,109 @@ def _compute_idf(G: nx.Graph, terms: list[str]) -> dict[str, float]: return {t: cache.get(t, math.log(1 + N)) for t in terms} +def _trigrams(text: str) -> set[str]: + """Character trigrams of `text`; for <3-char text the whole string is the key.""" + if len(text) < 3: + return {text} if text else set() + return {text[i:i + 3] for i in range(len(text) - 2)} + + +def _node_search_text(data: dict, nid: str) -> str: + """Concatenate every field _score_nodes / _find_node match a query against, so + one trigram index over this text is a complete candidate generator for both. + + - `norm_label` and `source_file` feed _score_nodes' per-term substring tiers. + - `label_tokens` (the space-joined token form) feeds _find_node's + `term in label_tokens` branch, where a multi-word `term` can span a token + boundary that punctuation hides in `norm_label` (e.g. query "foo bar" matches + label "foo.bar" only via its tokenized form). + - `nid` feeds the whole-query `joined == nid_lower` tier. + + NUL separators stop a trigram from spanning two fields (a query never contains + NUL, so a cross-field trigram can never be a real match). + """ + norm_label = data.get("norm_label") or _strip_diacritics(data.get("label") or "").lower() + label_tokens = " ".join(_search_tokens(data.get("label") or "")) + source = (data.get("source_file") or "").lower() + return "\x00".join((norm_label, label_tokens, str(nid).lower(), source)) + + +def _get_trigram_index(G: nx.Graph) -> dict: + """Lazily build and cache a trigram -> node-position postings map on the graph. + + Cached on `G.graph` so it auto-invalidates when _maybe_reload() swaps in a + fresh graph object, exactly like `_idf_cache`. `set_cache` memoizes per-trigram + id-sets across queries within one graph generation. + """ + idx = G.graph.get("_trigram_index") + if idx is not None: + return idx + ids = list(G.nodes()) + postings: dict[str, array] = {} + for i, nid in enumerate(ids): + for g in _trigrams(_node_search_text(G.nodes[nid], nid)): + bucket = postings.get(g) + if bucket is None: + bucket = array("i") + postings[g] = bucket + bucket.append(i) + idx = {"ids": ids, "postings": postings, "set_cache": {}} + G.graph["_trigram_index"] = idx + return idx + + +def _trigram_candidates(G: nx.Graph, needles: list[str], *, guard_frac: float = 0.10) -> list[str] | None: + """Node IDs whose text could contain any `needle` as a substring, via the + trigram index — a *superset* the caller then re-scores with the exact predicates. + + Returns candidates in graph-iteration order (so order-sensitive callers like + _find_node stay byte-identical to a full scan), or **None** when the index isn't + worth it — a needle is too short to trigram, or its rarest trigram is still + common enough that the candidate set would approach the whole graph. The caller + falls back to the full scan, preserving the never-worse contract. The guard is + cheap: postings-length lookups only, no set intersection. + """ + idx = _get_trigram_index(G) + ids, postings, set_cache = idx["ids"], idx["postings"], idx["set_cache"] + n = len(ids) + if n == 0: + return [] + needles = [s for s in needles if s] + thresh = int(n * guard_frac) + for s in needles: + tgs = _trigrams(s) + if not tgs or any(len(g) < 3 for g in tgs): + return None # too short to trigram-filter + present = [len(postings[g]) for g in tgs if g in postings] + if not present: + continue # this needle matches nothing — contributes no candidates + if min(present) > thresh: + return None # rarest trigram still too common -> not worth the index + cand: set[int] = set() + for s in needles: + sets: list[set] | None = [] + for g in _trigrams(s): + bucket = postings.get(g) + if bucket is None: + sets = None # a trigram absent everywhere -> needle matches nothing + break + cached = set_cache.get(g) + if cached is None: + cached = set(bucket) + set_cache[g] = cached + sets.append(cached) + if not sets: + continue + sets.sort(key=len) # intersect smallest-first + hit = set(sets[0]) + for other in sets[1:]: + hit &= other + if not hit: + break + cand |= hit + return [ids[i] for i in sorted(cand)] + + def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]: scored = [] norm_terms = [tok for t in terms for tok in _search_tokens(t)] @@ -144,7 +248,16 @@ def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]: # Weight the full-query bonus by the rarest constituent term so a specific # multi-word label still outweighs common-token noise; floor at 1.0. joined_w = max((idf.get(t, 1.0) for t in norm_terms), default=1.0) - for nid, data in G.nodes(data=True): + # Trigram prefilter: score only nodes whose text could match a term, falling + # back to the whole graph when the index isn't selective. The result is + # identical either way — the per-node scoring below is unchanged and a + # non-candidate node always scores 0. (IDF above stays a whole-graph statistic.) + candidate_ids = _trigram_candidates(G, norm_terms + ([joined] if joined else [])) + node_iter = ( + G.nodes(data=True) if candidate_ids is None + else ((nid, G.nodes[nid]) for nid in candidate_ids) + ) + for nid, data in node_iter: norm_label = data.get("norm_label") or _strip_diacritics(data.get("label") or "").lower() bare_label = norm_label.rstrip("()") # Tokenized form of the label (punctuation stripped, same transform as the @@ -462,7 +575,14 @@ def _find_node(G: nx.Graph, label: str) -> list[str]: exact: list[str] = [] prefix: list[str] = [] substring: list[str] = [] - for nid, d in G.nodes(data=True): + # Trigram prefilter (graph-iteration order preserved so exact/prefix/substring + # ordering — and thus matches[0] — is byte-identical to the full scan). + candidate_ids = _trigram_candidates(G, [term]) + node_iter = ( + G.nodes(data=True) if candidate_ids is None + else ((nid, G.nodes[nid]) for nid in candidate_ids) + ) + for nid, d in node_iter: norm_label = d.get("norm_label") or _strip_diacritics(d.get("label") or "").lower() bare_label = norm_label.rstrip("()") label_tokens = " ".join(_search_tokens(d.get("label") or "")) @@ -530,6 +650,11 @@ def _build_server(graph_path: str): G = _load_graph(graph_path) communities = _communities_from_graph(G) + # Build the trigram query index eagerly so the first query doesn't pay the + # one-time build (a few seconds on a large graph). serve exists to answer + # queries, so absorbing the cost at startup beats spiking a random first + # request; it's cached on G.graph, so queries just reuse it. + _get_trigram_index(G) # Hot-reload state: mtime+size key lets us detect graph.json changes without # polling. Initialised from the file stat at startup so the first tool call @@ -562,6 +687,8 @@ def _maybe_reload() -> None: new_G = _load_graph(graph_path) except SystemExit: return # keep serving stale graph on transient read error + _get_trigram_index(new_G) # warm before exposing, so the first + # post-reload query is fast too (same rationale as startup) G = new_G communities = _communities_from_graph(new_G) _reload_state["mtime_ns"], _reload_state["size"] = key diff --git a/tests/test_serve.py b/tests/test_serve.py index 5b5926b96..a0432a4eb 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -12,6 +12,10 @@ _bfs, _dfs, _find_node, + _trigrams, + _node_search_text, + _get_trigram_index, + _trigram_candidates, _filter_graph_by_context, _infer_context_filters, _query_terms, @@ -130,6 +134,107 @@ def test_find_node_matches_full_punctuated_unicode_label(): assert _find_node(G, "Skill /auditar — Auditoría inquisitiva de enlaces") == ["n1"] +# --- trigram candidate prefilter (the trigram index that shrinks the O(N) scan) --- + + +def _force_full_scan(monkeypatch): + """Disable the prefilter so a call exercises the original full-node scan.""" + monkeypatch.setattr("graphify.serve._trigram_candidates", lambda *a, **k: None) + + +def _make_big_graph(n: int = 150) -> nx.Graph: + """A graph large enough that the selectivity guard lets the fast-path fire for + rare terms and fall back for common ones. Most labels share the 'item'/'node' + stem (common), plus a few distinctive rare labels and one punctuated label.""" + G = nx.Graph() + for i in range(n): + G.add_node(f"id{i}", label=f"item node {i}", source_file=f"pkg/item_{i}.py") + G.add_node("rareA", label="ZebraQuokkaWidget", source_file="zoo/zqw.py") + G.add_node("rareB", label="MarmosetGadget handler", source_file="zoo/marmoset.py") + G.add_node("punct", label="Foo.Bar:Baz", source_file="pkg/foobar.py") + return G + + +def test_trigrams_basic(): + assert _trigrams("foobar") == {"foo", "oob", "oba", "bar"} + assert _trigrams("ab") == {"ab"} # <3 chars -> whole string is the key + assert _trigrams("") == set() + + +def test_node_search_text_includes_all_matched_fields(): + G = _make_big_graph() + text = _node_search_text(G.nodes["punct"], "punct") + # norm_label, the tokenized label (label_tokens), nid, and source are all present, + # NUL-separated so trigrams can't span fields. + parts = text.split("\x00") + assert parts[0] == "foo.bar:baz" # norm_label (punctuation kept) + assert parts[1] == "foo bar baz" # label_tokens (tokenized) + assert parts[2] == "punct" # nid + assert parts[3] == "pkg/foobar.py" # source_file + + +def test_trigram_candidates_fast_path_fires_for_rare_term(): + G = _make_big_graph() + cand = _trigram_candidates(G, ["zebraquokkawidget"]) + assert cand is not None # selective -> fast-path used + assert "rareA" in cand + assert len(cand) < G.number_of_nodes() # a real shrink, not the whole graph + + +def test_trigram_candidates_falls_back_on_common_term(): + G = _make_big_graph() + # 'item' is in the label of every one of the 150 'item node N' nodes -> the + # rarest trigram is still common -> guard returns None (full-scan fallback). + assert _trigram_candidates(G, ["item"]) is None + + +def test_trigram_candidates_falls_back_on_short_token(): + G = _make_big_graph() + assert _trigram_candidates(G, ["ab"]) is None # <3 chars -> can't trigram-filter + + +def test_score_nodes_prefilter_is_identical_to_full_scan(monkeypatch): + G = _make_big_graph() + queries = ["zebraquokkawidget", "marmosetgadget handler", "foo bar baz", + "item", "node 42", "nonexistentxyz"] + for q in queries: + terms = _query_terms(q) + fast = _score_nodes(G, terms) + _force_full_scan(monkeypatch) + full = _score_nodes(G, terms) + monkeypatch.undo() + assert fast == full, f"prefilter diverged from full scan for {q!r}" + + +def test_find_node_prefilter_is_identical_to_full_scan(monkeypatch): + G = _make_big_graph() + # includes the punctuated label, exercised via its tokenized (label_tokens) form + for label in ["ZebraQuokkaWidget", "MarmosetGadget handler", "Foo Bar Baz", + "item node 7", "missing"]: + fast = _find_node(G, label) + _force_full_scan(monkeypatch) + full = _find_node(G, label) + monkeypatch.undo() + assert fast == full, f"_find_node prefilter diverged (order!) for {label!r}" + + +def test_find_node_label_tokens_branch_covered_by_index(): + # "foo bar baz" matches label "Foo.Bar:Baz" only via the tokenized label_tokens + # form (the dotted/colon norm_label never contains the spaced query). The index + # must surface this node as a candidate, or the prefilter would silently drop it. + G = _make_big_graph() + assert _find_node(G, "Foo Bar Baz") == ["punct"] + + +def test_trigram_index_cached_and_rebuilt_per_graph(): + G = _make_big_graph() + idx1 = _get_trigram_index(G) + assert idx1 is _get_trigram_index(G) # cached on the same graph object + assert G.graph["_trigram_index"] is idx1 + G2 = _make_big_graph() + assert _get_trigram_index(G2) is not idx1 # a fresh graph rebuilds (reload safety) + + def test_query_terms_strips_search_punctuation(): assert _query_terms("what calls extract?") == ["what", "calls", "extract"]