From dafa26cf558cce74ee34a0051be98800d553a4c9 Mon Sep 17 00:00:00 2001 From: Dmitry Teryaev Date: Sat, 11 Jul 2026 14:57:56 +0300 Subject: [PATCH] feat(search): BM25 lexical fallback via LadybugDB FTS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the macOS-Intel heuristic keyword scan with Okapi BM25 over a LadybugDB FTS index on Symbol.search_text, fetched DB-side — killing the bounded 5000-row Python scan that silently missed matches past the cap on large repos. The hand-rolled name/type/fqn/role heuristic now re-ranks the BM25 candidates, and remains the fallback when the FTS index or extension is unavailable (older graph, or an offline first run where INSTALL FTS can't fetch from extension.ladybugdb.com). LadybugDB (the embedded graph DB that is kuzu's ancestor, not a kuzu wrapper) ships this natively in the same code_graph.lbug the graph lives in — no new dependency, no sidecar store. The PR #403 note "upgradable to SQLite FTS5" pointed at the wrong store. Build (build_ast_graph): - Symbol.search_text column: camelCase-split token soup of name + fqn + signature + annotations + capabilities (same _split_identifier the query path uses, so index/query tokenization agree). - sym_fts FTS index (Okapi BM25, porter stemmer): ensured at the end of every full build and as a self-heal on incremental. Auto-maintains on COPY/MERGE/ DELETE, so increment stays fresh without drop+recreate. - _drop_all drops the FTS index first (LadybugDB refuses DROP TABLE on a table referenced by an FTS index); incremental_rebuild and _delete_file_scope LOAD EXTENSION FTS so Symbol DML can maintain the index. - ONTOLOGY_VERSION 18 -> 19 (semantic change; one-time reindex). Query (search_lexical): BM25-first — QUERY_FTS_INDEX fetches top-K candidates, re-MATCH re-applies the NodeFilter (role/module/path/kind), then the heuristic re-ranks and dedup_by_fqn runs as before. Per-connection FTS-load cache keyed by object (WeakSet), not id() — id() is reused after GC and broke under test batching. Same SearchHit/SearchOutput/NodeFilter contract; _score_components gains a bm25 entry. Bench (~/jrag-bench/shopizer, master vs branch): build phases flat (init 34.65->30.19s, reprocess --graph-only 7.63->7.21s, increment 18.98->17.30s — within run variance; vectors phase untouched); query mean 82.8ms -> 23.6ms (3.5x), max 111.7ms -> 25.0ms. Full suite green. Fork B (BM25 + vector hybrid via RRF on the primary path) tracked in #431. Co-Authored-By: Claude --- README.md | 2 +- docs/AGENT-GUIDE.md | 2 +- docs/ARCHITECTURE.md | 10 +- docs/CONFIGURATION.md | 6 +- src/java_codebase_rag/ast/ast_java.py | 3 +- .../graph/build_ast_graph.py | 127 +++++++++++++- .../search/search_lexical.py | 158 ++++++++++++++---- .../search/search_scoring.py | 5 + tests/fixtures/graph_baseline_bank_chat.json | 2 +- tests/integration/test_incremental_graph.py | 6 +- tests/package/test_jrag_status.py | 2 +- tests/search/test_search_lexical.py | 57 ++++++- 12 files changed, 329 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 075c0e2..a34e72c 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ The rest of this README is the install, the tool/command orientation, and the re pip install java-codebase-rag ``` -Python **3.11+** required, on **Linux, macOS, and Windows**. On Linux, Windows, and **Apple Silicon** Macs every native dependency (LanceDB, LadybugDB/kuzu, CocoIndex) ships a wheel and you get the full semantic + graph search. **Intel Macs (x86_64) install graph-only**: PyTorch ≥2.3 and LanceDB ≥0.26 dropped macOS Intel wheels, so the vector stack is auto-excluded via PEP 508 markers — `pip install java-codebase-rag` works out of the box, the graph layer (`find` / `describe` / `neighbors` / `resolve`) is fully usable, and the `search` tool falls back to **lexical (keyword) search** over the symbol graph (same tool contract, keyword-ranked instead of semantic; an advisory notes the mode). Semantic/vector search needs Apple Silicon, Linux, or Windows. After install, `java-codebase-rag --help` should print the CLI groups. +Python **3.11+** required, on **Linux, macOS, and Windows**. On Linux, Windows, and **Apple Silicon** Macs every native dependency (LanceDB, LadybugDB, CocoIndex) ships a wheel and you get the full semantic + graph search. **Intel Macs (x86_64) install graph-only**: PyTorch ≥2.3 and LanceDB ≥0.26 dropped macOS Intel wheels, so the vector stack is auto-excluded via PEP 508 markers — `pip install java-codebase-rag` works out of the box, the graph layer (`find` / `describe` / `neighbors` / `resolve`) is fully usable, and the `search` tool falls back to **lexical search** over the symbol graph — BM25-ranked over a LadybugDB full-text index (same tool contract, keyword-ranked instead of semantic; an advisory notes the mode). Semantic/vector search needs Apple Silicon, Linux, or Windows. After install, `java-codebase-rag --help` should print the CLI groups. The package includes the CocoIndex lifecycle dependency used by `init`, `increment`, `reprocess`, and `erase` on platforms that have it (it is absent on Intel Mac). ### Interactive setup (recommended) diff --git a/docs/AGENT-GUIDE.md b/docs/AGENT-GUIDE.md index 49208df..860b1be 100644 --- a/docs/AGENT-GUIDE.md +++ b/docs/AGENT-GUIDE.md @@ -207,7 +207,7 @@ Prefer **`resolve` → `describe(id=…)`** over **`describe(fqn=…)`** when an Ranked chunk retrieval. Args: `query`, `table` (`java`|`sql`|`yaml`|`all`, default `java`), `hybrid` (bool), `limit` (default 5), `offset`, `path_contains`, optional `filter` (symbol-applicable `NodeFilter` only), optional `chunks` (bool, default `false`). Returns one row per `primary_type_fqn` (symbol/type) by default; set `chunks=true` to restore chunk-level output. When deduped, each hit includes a `chunks` field (≥1) indicating how many chunks were collapsed into that hit. -> **Intel Mac (graph-only) installs:** `search` runs the **lexical backend** — keyword relevance over the symbol graph instead of embeddings, behind this same contract. Same `query`/`table`/`filter`/`limit`/`chunks` behavior; results are keyword-ranked (not semantic), `hybrid` is ignored, `sql`/`yaml` tables aren't indexed (only Java symbols), and an `advisories` entry + `lexical_mode=true` flag note the mode. Structural discovery (`find`/`describe`/`neighbors`/`resolve`) is unaffected. +> **Intel Mac (graph-only) installs:** `search` runs the **lexical backend** — BM25 keyword ranking over the symbol graph's LadybugDB full-text index instead of embeddings, behind this same contract. Same `query`/`table`/`filter`/`limit`/`chunks` behavior; results are keyword-ranked (not semantic), `hybrid` is ignored, `sql`/`yaml` tables aren't indexed (only Java symbols), and an `advisories` entry + `lexical_mode=true` flag note the mode. Structural discovery (`find`/`describe`/`neighbors`/`resolve`) is unaffected. #### `find` diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index eae7d07..2ee34c6 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -68,7 +68,7 @@ java_codebase_rag/pipeline.py ``` MCP tool call (server.py) ──asyncio.to_thread──▶ mcp_v2.* ├─ search ─▶ search_lancedb.run_search (vector / hybrid; optional graph-expand + RRF rank fusion) - │ └─ lancedb import absent (Intel Mac) → search_lexical (keyword over graph) + │ └─ lancedb import absent (Intel Mac) → search_lexical (BM25 over Symbol FTS index; heuristic scan fallback) ├─ find / describe / neighbors ─▶ ladybug_queries.LadybugGraph (Cypher) └─ resolve ─▶ resolve_service.resolve_v2 (cascade → status one | many | none) on empty ─▶ absence_diagnosis.diagnose → verdict + (optional) proof @@ -77,13 +77,13 @@ MCP tool call (server.py) ──asyncio.to_thread──▶ mcp_v2.* | Tool | Backing | Notes | | --- | --- | --- | -| `search` | Lance vector/hybrid, or lexical fallback | dedup by FQN; role weights via `search_scoring` | +| `search` | Lance vector/hybrid, or BM25 lexical fallback | dedup by FQN; role weights via `search_scoring` | | `find` | Ladybug Cypher | required `NodeFilter`; strict per-kind frame | | `describe` | Ladybug Cypher | node record + `edge_summary` (composed/override rollups) | | `neighbors` | Ladybug Cypher | one hop; `direction` + `edge_types` required; dot-key composed edges | | `resolve` | Ladybug Cypher | per-kind generators exact→fuzzy; cap 10 candidates | -**Lexical fallback** is selected by import availability (`mcp_v2` guards `from search_lancedb import …`): same row contract, flagged via `lexical_mode` + advisory. **`jrag` CLI** calls the same `mcp_v2.*` functions — identical backends, only rendering differs. +**Lexical fallback** is selected by import availability (`mcp_v2` guards `from search_lancedb import …`): same row contract, flagged via `lexical_mode` + advisory. It is **BM25-first**: `build_ast_graph` indexes `Symbol.search_text` (camelCase-split token soup) under a LadybugDB FTS index (`sym_fts`, Okapi BM25), and `search_lexical` fetches top-K candidates via `QUERY_FTS_INDEX` then re-ranks them with the name/type/fqn/role heuristic in `search_scoring`. The FTS index auto-maintains on `increment`; the heuristic scan is the fallback when the index/extension is absent (older graph, offline first run). **`jrag` CLI** calls the same `mcp_v2.*` functions — identical backends, only rendering differs. ## Stores @@ -108,7 +108,7 @@ Dev workflow (editable install, test-reset ritual, full-suite discipline) — se | Constant | Value / location | | --- | --- | -| `ONTOLOGY_VERSION` | `18` — `ast_java.py:87` | +| `ONTOLOGY_VERSION` | `19` — `ast_java.py:87` | | `LANCE_TABLE_NAMES` | 3 tables — `java_codebase_rag/lance_optimize.py:35` | | Graph passes | 6 (labels `build_ast_graph.py:83`) | | Incremental cap | `expansion_cap=50` — `build_ast_graph.py:3800` | @@ -117,4 +117,4 @@ Dev workflow (editable install, test-reset ritual, full-suite discipline) — se ## TL;DR -Two stores built in lockstep — LanceDB vectors via CocoIndex, LadybugDB graph via a 6-pass tree-sitter build — queried by 5 MCP tools that split cleanly: `search` → vector/lexical, `find`/`describe`/`neighbors`/`resolve` → Cypher. Hints and absence wrap every response; `ONTOLOGY_VERSION=18` is the rebuild/staleness contract. Contributors extend via `EDGE_SCHEMA` + builder passes, and bump the version on any semantic change. +Two stores built in lockstep — LanceDB vectors via CocoIndex, LadybugDB graph via a 6-pass tree-sitter build — queried by 5 MCP tools that split cleanly: `search` → vector/lexical, `find`/`describe`/`neighbors`/`resolve` → Cypher. Hints and absence wrap every response; `ONTOLOGY_VERSION=19` is the rebuild/staleness contract. Contributors extend via `EDGE_SCHEMA` + builder passes, and bump the version on any semantic change. diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index ceca689..215334b 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -437,7 +437,9 @@ Combined, these pull `processClientMessage` / `pickEligibleOperator` / `onOperat #### Graph-only (macOS Intel) lexical ranking -On Intel Mac installs the vector stack is absent (see `README.md`), so `search` runs the **lexical backend** — keyword relevance over the symbol graph instead of embeddings, behind the same tool contract. Ranking components, normalized into a `[0,1]` score: +On Intel Mac installs the vector stack is absent (see `README.md`), so `search` runs the **lexical backend** — keyword ranking over the symbol graph instead of embeddings, behind the same tool contract. It is **BM25-first**: at index time every `Symbol` gets a `search_text` column (camelCase-split tokens of name + fqn + signature + annotations + capabilities, tokenized with the same splitter the query path uses) and a LadybugDB FTS index (`sym_fts`, Okapi BM25, porter stemmer) over it. At query time `QUERY_FTS_INDEX` fetches the top-K candidates DB-side, which are then re-ranked in Python by the heuristic below and deduped by FQN. + +The heuristic re-rank decides final order (what `--explain` reports as `relevance=` / `name=` / `type=` / `fqn=`, plus a `bm25=` component for the index score): - **Name token overlap** — strongest signal (weight `0.45`); full query-token coverage of the declaration name scores `1.0`. - **Type-name overlap** — `+0.05`/hit, capped `+0.10` (same convention as the vector symbol bonus above). @@ -445,7 +447,7 @@ On Intel Mac installs the vector stack is absent (see `README.md`), so `search` - **Signature / annotation / capability text overlap** — weight `0.15`. - **Role weights** — the same table above applies as a tie-breaker/booster. -Rows with **no keyword overlap** are dropped — role alone never qualifies a hit (it only reorders matches). Locking `role=` / `exclude_roles` still skips the role weight. `--explain` surfaces `name=` / `type=` / `fqn=` / `relevance=` components. `sql` / `yaml` tables aren't indexed in graph-only mode (only Java symbols are), and `hybrid` is ignored (lexical-only). +BM25 only selects the candidate pool (so large repos no longer hit the old bounded Python scan); the heuristic decides order. If the FTS index or extension is unavailable (older graph, or an offline first run where `INSTALL FTS` can't fetch it), the backend falls back to that bounded scan over `Symbol` rows using the same heuristic. Locking `role=` / `exclude_roles` skips the role weight. `sql` / `yaml` tables aren't indexed in graph-only mode (only Java symbols are), and `hybrid` is ignored (lexical-only). ### Debugging empty `context_before` / `context_after` diff --git a/src/java_codebase_rag/ast/ast_java.py b/src/java_codebase_rag/ast/ast_java.py index c10e8e7..386649d 100644 --- a/src/java_codebase_rag/ast/ast_java.py +++ b/src/java_codebase_rag/ast/ast_java.py @@ -84,7 +84,8 @@ # Phase 11: `EDGE_SCHEMA` in `java_ontology.py` (canonical edge navigation schema; v14 re-index). # Phase 12: CALLS `callee_declaring_role`, supertype-walk dedup, pass3 unresolved counters (v15 re-index). # Bumps whenever extraction / enrichment semantics change. -ONTOLOGY_VERSION = 18 +# Phase 13: Symbol.search_text + LadybugDB FTS (Okapi BM25) index for lexical search (v19 re-index). +ONTOLOGY_VERSION = 19 ROLE_ANNOTATIONS: dict[str, str] = { # Spring Web diff --git a/src/java_codebase_rag/graph/build_ast_graph.py b/src/java_codebase_rag/graph/build_ast_graph.py index d4a041b..54c7a1d 100644 --- a/src/java_codebase_rag/graph/build_ast_graph.py +++ b/src/java_codebase_rag/graph/build_ast_graph.py @@ -68,6 +68,7 @@ symbol_id, ) from java_codebase_rag.graph.path_filtering import LayeredIgnore, iter_java_source_files +from java_codebase_rag.search.search_scoring import SYMBOL_FTS_INDEX as _SYMBOL_FTS_INDEX, _split_identifier from java_codebase_rag.graph.java_ontology import ( CLIENT_KIND_FEIGN_METHOD, CLIENT_KIND_REST_TEMPLATE, @@ -801,6 +802,12 @@ def _delete_file_scope( Producer nodes use DETACH DELETE as a safety net for any edges missed in Phase 1. """ + # Symbol DELETEs below maintain the FTS index, which needs the extension loaded on + # this connection. Idempotent + best-effort (no-op when no index exists / FTS absent). + try: + conn.execute("LOAD EXTENSION FTS") + except Exception: + pass scope_files = changed_files | dependent_files scope_list = list(scope_files) changed_list = list(changed_files) @@ -2927,7 +2934,8 @@ def _micro_factor(member: MemberEntry | None) -> float: "start_byte INT64, end_byte INT64, " "modifiers STRING[], annotations STRING[], capabilities STRING[], " "role STRING, signature STRING, parent_id STRING, resolved BOOLEAN, " - "generated BOOLEAN, generated_by STRING" + "generated BOOLEAN, generated_by STRING, " + "search_text STRING" ")" ) @@ -3050,7 +3058,28 @@ def _micro_factor(member: MemberEntry | None) -> float: ) +def _drop_symbol_fts_index_if_present(conn: ladybug.Connection) -> None: + """Drop the Symbol FTS index so the table drop below succeeds. + + LadybugDB refuses ``DROP TABLE Symbol`` while an FTS index references it ("Cannot + delete node table ... referenced by index"), so the index must go first. No-op when + FTS is unavailable (offline first-run) or the index was never created — in those + cases nothing blocks the table drop. Best-effort: any failure is swallowed. + """ + try: + conn.execute("LOAD EXTENSION FTS") + existing = conn.execute("CALL SHOW_INDEXES() RETURN index_name") + names: set[str] = set() + while existing.has_next(): + names.add(existing.get_next()[0]) + if _SYMBOL_FTS_INDEX in names: + conn.execute(f"CALL DROP_FTS_INDEX('Symbol', '{_SYMBOL_FTS_INDEX}')") + except Exception: + pass + + def _drop_all(conn: ladybug.Connection) -> None: + _drop_symbol_fts_index_if_present(conn) for stmt in ( "DROP TABLE IF EXISTS DECLARES_CLIENT", "DROP TABLE IF EXISTS DECLARES_PRODUCER", @@ -3101,6 +3130,80 @@ def _create_schema(conn: ladybug.Connection) -> None: conn.execute(stmt) +_IDENT_RE = re.compile(r"[A-Za-z][A-Za-z0-9]*") + + +def _compute_symbol_search_text( + name: str, + fqn: str, + signature: str, + annotations: list[str], + capabilities: list[str], + package: str = "", +) -> str: + """Build the BM25-indexable token soup for a Symbol (fork A). + + camelCase / snake_case identifiers are split into lowercase tokens via the SAME + ``_split_identifier`` the query path uses (index- and query-time tokenization + must agree), then space-joined into one STRING that LadybugDB FTS porter-stems. + Name + fqn carry the strongest discovery signal; signature / annotations / + capabilities / package segments are weaker corroborators. Members reach this via + the same ``_node_row`` constructor as types. + """ + toks: list[str] = [] + for field in (name, fqn, signature, package): + for m in _IDENT_RE.findall(str(field or "")): + toks.extend(_split_identifier(m)) + for lst in (annotations, capabilities): + for item in (lst or []): + for m in _IDENT_RE.findall(str(item)): + toks.extend(_split_identifier(m)) + seen: set[str] = set() + out: list[str] = [] + for t in toks: + if len(t) >= 2 and t not in seen: + seen.add(t) + out.append(t) + return " ".join(out) + + +def _ensure_symbol_fts_index(conn: ladybug.Connection, *, verbose: bool) -> None: + """Best-effort create the Symbol FTS (Okapi BM25) index if absent (fork A). + + Idempotent: skips when the index already exists. The FTS extension is fetched + from extension.ladybugdb.com on first ``INSTALL FTS`` (cached locally after); an + offline first run fails softly — ``run_lexical_search`` then falls back to the + heuristic over the bare Symbol scan. The index auto-maintains on later + COPY / MERGE / DELETE (verified), so this only runs on full builds and as a + cheap self-heal at the end of an incremental rebuild. + """ + try: + try: + conn.execute("INSTALL FTS") # already-installed raises; swallow + except Exception: + pass + conn.execute("LOAD EXTENSION FTS") + existing = conn.execute("CALL SHOW_INDEXES() RETURN index_name") + names: set[str] = set() + while existing.has_next(): + names.add(existing.get_next()[0]) + if _SYMBOL_FTS_INDEX not in names: + conn.execute( + f"CALL CREATE_FTS_INDEX('Symbol', '{_SYMBOL_FTS_INDEX}', " + "['search_text'], stemmer := 'porter')" + ) + if verbose: + _verbose_stderr_line( + f"[graph] fts · created Symbol {_SYMBOL_FTS_INDEX} (BM25 over search_text)" + ) + except Exception as e: + if verbose: + _verbose_stderr_line( + f"[graph] fts · unavailable ({type(e).__name__}: {e}); " + "lexical search will use the heuristic scan" + ) + + def _node_row(**kwargs) -> dict: base = { "kind": "", "name": "", "fqn": "", "package": "", @@ -3109,9 +3212,16 @@ def _node_row(**kwargs) -> dict: "start_byte": 0, "end_byte": 0, "modifiers": [], "annotations": [], "capabilities": [], "role": "OTHER", "signature": "", "parent_id": "", "resolved": True, - "generated": False, "generated_by": None, + "generated": False, "generated_by": None, "search_text": "", } base.update(kwargs) + # Derive the BM25 token soup unless the caller set it explicitly. + if not base.get("search_text"): + base["search_text"] = _compute_symbol_search_text( + base.get("name", ""), base.get("fqn", ""), base.get("signature", ""), + base.get("annotations", []), base.get("capabilities", []), + base.get("package", ""), + ) return base @@ -3161,7 +3271,7 @@ def _existing_node_ids(conn: ladybug.Connection) -> set[str]: "id", "kind", "name", "fqn", "package", "module", "microservice", "filename", "start_line", "end_line", "start_byte", "end_byte", "modifiers", "annotations", "capabilities", "role", "signature", "parent_id", "resolved", - "generated", "generated_by" + "generated", "generated_by", "search_text" ] # Type declaration kinds. Tuple (not set) so the rendered SQL `IN` clause is @@ -3185,7 +3295,7 @@ def _existing_node_ids(conn: ladybug.Connection) -> set[str]: "n.modifiers = $modifiers, n.annotations = $annotations, " "n.capabilities = $capabilities, n.role = $role, " "n.signature = $signature, n.parent_id = $parent_id, n.resolved = $resolved, " - "n.generated = $generated, n.generated_by = $generated_by" + "n.generated = $generated, n.generated_by = $generated_by, n.search_text = $search_text" ) # Refresh every mutable Route field on an existing Route node by id. Mirrors the @@ -3835,6 +3945,13 @@ def incremental_rebuild( db = ladybug.Database(str(ladybug_path)) conn = ladybug.Connection(db) + # If a Symbol FTS index exists, the scoped DELETE/MERGE/COPY below must maintain it, + # which needs the FTS extension loaded on THIS connection. Best-effort: when FTS is + # unavailable no index exists, so there's nothing to maintain and DML proceeds fine. + try: + conn.execute("LOAD EXTENSION FTS") + except Exception: + pass # Check ontology version try: @@ -4007,6 +4124,7 @@ def incremental_rebuild( # Update GraphMeta _write_meta(conn, tables_for_global, source_root) + _ensure_symbol_fts_index(conn, verbose=verbose) # Remove crash marker crash_marker_path.unlink(missing_ok=True) @@ -4236,6 +4354,7 @@ def write_ladybug( if verbose: _verbose_stderr_line(f"[graph] writing · routes/exposes written in {time.time() - t2:.2f}s") _write_meta(conn, tables, source_root) + _ensure_symbol_fts_index(conn, verbose=verbose) conn.close() db.close() diff --git a/src/java_codebase_rag/search/search_lexical.py b/src/java_codebase_rag/search/search_lexical.py index 0d1ca17..1bf795b 100644 --- a/src/java_codebase_rag/search/search_lexical.py +++ b/src/java_codebase_rag/search/search_lexical.py @@ -18,11 +18,13 @@ from __future__ import annotations import os +import weakref from pathlib import Path from typing import TYPE_CHECKING, Any from java_codebase_rag.graph.ladybug_queries import LadybugGraph from java_codebase_rag.search.search_scoring import ( + SYMBOL_FTS_INDEX, _ROLE_SCORE_WEIGHTS, _TYPE_MATCH_BONUS_CAP, _TYPE_MATCH_BONUS_PER_HIT, @@ -170,6 +172,88 @@ def _token_overlap(haystack_toks: set[str], needle_toks: set[str]) -> float: return len(needle_toks & haystack_toks) / len(needle_toks) +# BM25 candidate fetch via the LadybugDB FTS index (fork A). DB-side indexed ranking +# replaces the heuristic's bounded Python scan; the heuristic below still scores the +# fetched candidates (name/type/fqn/role) and is the fallback when the FTS index or +# extension is unavailable (older graph, offline first run). +_FTS_CANDIDATE_K = 200 # top-K BM25 candidates; re-filtered by NodeFilter before ranking +# Connections that have run LOAD EXTENSION FTS. Keyed by the connection OBJECT (WeakSet), +# NOT id() — id() is reused after GC, which would let a fresh connection skip LOAD and then +# fail at QUERY_FTS_INDEX under test batching. Entries die with the connection. +_FTS_LOADED_CONNS: "weakref.WeakSet[object]" = weakref.WeakSet() + + +def _ensure_fts_loaded(g: LadybugGraph) -> bool: + """LOAD EXTENSION FTS on the graph's (read-only) connection, once per connection. + + Returns False if the extension can't be loaded (absent / offline) so the caller + falls back to the heuristic scan. + """ + conn = g._conn # noqa: SLF001 + try: + if conn in _FTS_LOADED_CONNS: + return True + except Exception: # connection not weakref-able → LOAD every call (correct, slow) + pass + try: + g._rows("LOAD EXTENSION FTS") # noqa: SLF001 + try: + _FTS_LOADED_CONNS.add(conn) + except Exception: + pass + return True + except Exception: + return False + + +def _try_fts_candidates( + g: LadybugGraph, + query: str, + filter: NodeFilter | None, + path_contains: str | None, +) -> dict | None: + """Fetch BM25-ranked Symbol candidates via the FTS index; re-apply NodeFilter. + + Returns ``{"rows": [...], "scores": {id: bm25}}`` (rows are the same shape the + heuristic scan yields), or ``None`` when FTS is unavailable (extension won't load, + or the index isn't present on this graph) so the caller falls back. + + Two-step: (1) ``QUERY_FTS_INDEX`` returns the top-K node ids by Okapi BM25 over + ``Symbol.search_text``; (2) re-MATCH those ids with the full ``_lexical_where`` + predicates (role / module / path / kind≠file,package) so the filter logic stays + defined in one place. ``search_text`` is built at index time by ``build_ast_graph`` + from the same ``_split_identifier`` the re-rank below uses, so index- and query-time + tokenization agree. + """ + if not _ensure_fts_loaded(g): + return None + idx_rows = g._rows("CALL SHOW_INDEXES() RETURN index_name") # noqa: SLF001 + names = {row.get("index_name") for row in idx_rows} + if SYMBOL_FTS_INDEX not in names: + return None + fts = g._rows( # noqa: SLF001 + f"CALL QUERY_FTS_INDEX('Symbol', '{SYMBOL_FTS_INDEX}', $q, top := $k) " + "RETURN node.id AS id, score", + {"q": query, "k": _FTS_CANDIDATE_K}, + ) + if not fts: + return {"rows": [], "scores": {}} + scores = {row["id"]: float(row.get("score") or 0.0) for row in fts} + ids = list(scores.keys()) + + # Re-MATCH the K ids with the SAME predicates the heuristic pushes down, so + # NodeFilter / path / structural-kind filtering is defined exactly once. + where, params = _lexical_where(filter, path_contains=path_contains) + struct_pred = "(s.kind <> 'file' AND s.kind <> 'package')" + if not where: + where = f"WHERE s.id IN $ids AND {struct_pred}" + else: + where = where.replace("WHERE ", f"WHERE s.id IN $ids AND {struct_pred} AND ", 1) + params["ids"] = ids + rows = g._rows(f"MATCH (s:Symbol) {where} RETURN {_SYMBOL_RETURN}", params) # noqa: SLF001 + return {"rows": rows, "scores": scores} + + def run_lexical_search( query: str, *, @@ -185,6 +269,12 @@ def run_lexical_search( ) -> list[dict]: """Keyword search over Symbol nodes; returns ``run_search``-shaped row-dicts. + BM25-first (fork A): when the LadybugDB ``sym_fts`` index exists, candidates are + fetched DB-side via Okapi BM25 over ``Symbol.search_text`` (killing the bounded + Python scan that silently missed matches past the cap on large repos) and then + re-ranked here by the name/type/fqn/role heuristic. Falls back to that heuristic + scan when the FTS index or extension is unavailable (older graph, offline first run). + Raises ``RuntimeError`` (message contains "lexical search unavailable") if no symbol graph exists — the caller maps that to a clean failure envelope. Returns ``[]`` for ``table in ("sql", "yaml")`` (those LanceDB tables aren't built in @@ -201,33 +291,38 @@ def run_lexical_search( ) g = graph or LadybugGraph.get() - where, params = _lexical_where(filter, path_contains=path_contains) - # Always exclude structural Symbol nodes. Files and packages are :Symbol-labeled - # (kind='file'/'package') but aren't searchable code declarations — without this - # a token that appears in a filename (e.g. 'distribution' in - # 'DistributionChunkService.java') would surface the file node as a hit. - struct_pred = "(s.kind <> 'file' AND s.kind <> 'package')" - where = f"WHERE {struct_pred}" if not where else where.replace("WHERE ", f"WHERE {struct_pred} AND ", 1) - # Lexical ranking is done in Python (LadybugDB/kuzu has no keyword ranking without - # FTS5, which is deferred), and the MATCH scan returns rows in storage order — there - # is NO DB-side relevance ORDER BY. So fetch the FULL candidate pool up to the safety - # cap and rank here. A pagination-derived LIMIT (4x the page) is correct on the vector - # path where LanceDB returns rows pre-ranked by similarity, but on this unordered scan - # it would return only the first ~N symbols in arbitrary storage order and silently - # miss the best match on any non-trivial repo. - params["lim"] = _CANDIDATE_LIMIT_CAP - cypher = f"MATCH (s:Symbol) {where} RETURN {_SYMBOL_RETURN} LIMIT $lim" - rows = g._rows(cypher, params) # noqa: SLF001 — de facto public read API (see find_v2) - # If the fetch hit the safety cap, deeper matches were never ranked (the scan has no - # ORDER BY — kuzu returns an arbitrary, storage-order-dependent subset). Surface it so - # a user on a large repo isn't silently shown an incomplete result set; refining the - # query or adding a filter narrows the pool below the cap. Raising the cap / FTS5 is - # the deferred long-term fix (see the plan's "Out of scope" note). - if advisories is not None and len(rows) >= _CANDIDATE_LIMIT_CAP: - advisories.append( - f"lexical search scanned the first {_CANDIDATE_LIMIT_CAP} matching symbols " - "(repo cap); deeper matches were not ranked — refine the query or add a filter" - ) + # --- candidate fetch: BM25 (FTS) preferred, heuristic scan fallback --- + bm25_scores: dict[str, float] = {} + use_fts = False + fts = _try_fts_candidates(g, query, filter, path_contains) + if fts is not None: + rows = fts["rows"] + bm25_scores = fts["scores"] + use_fts = True + else: + where, params = _lexical_where(filter, path_contains=path_contains) + # Always exclude structural Symbol nodes. Files and packages are :Symbol-labeled + # (kind='file'/'package') but aren't searchable code declarations — without this + # a token that appears in a filename (e.g. 'distribution' in + # 'DistributionChunkService.java') would surface the file node as a hit. + struct_pred = "(s.kind <> 'file' AND s.kind <> 'package')" + where = f"WHERE {struct_pred}" if not where else where.replace("WHERE ", f"WHERE {struct_pred} AND ", 1) + # The heuristic scan returns rows in storage order — there is NO DB-side relevance + # ORDER BY without the FTS index — so fetch the FULL candidate pool up to the safety + # cap and rank here. A pagination-derived LIMIT (4x the page) is correct on the + # vector path where LanceDB returns rows pre-ranked, but on this unordered scan it + # would return only the first ~N symbols in arbitrary storage order and silently + # miss the best match on any non-trivial repo. The BM25 (FTS) path above has no cap. + params["lim"] = _CANDIDATE_LIMIT_CAP + cypher = f"MATCH (s:Symbol) {where} RETURN {_SYMBOL_RETURN} LIMIT $lim" + rows = g._rows(cypher, params) # noqa: SLF001 — de facto public read API (see find_v2) + # If the fetch hit the safety cap, deeper matches were never ranked. Surface it so + # a user on a large repo isn't silently shown an incomplete result set. + if advisories is not None and len(rows) >= _CANDIDATE_LIMIT_CAP: + advisories.append( + f"lexical search scanned the first {_CANDIDATE_LIMIT_CAP} matching symbols " + "(repo cap); deeper matches were not ranked — refine the query or add a filter" + ) query_toks = _query_tokens(query) source_root = _resolve_source_root(g) @@ -265,9 +360,10 @@ def run_lexical_search( text_match = text_overlap * _TEXT_MATCH_WEIGHT # A keyword search must require at least one lexical hit — role alone never - # qualifies a row (it only boosts/reorders matches). Degenerate queries with - # no usable tokens fall through to role-ranked listing. - if query_toks and not (name_overlap or type_hits or fqn_match or text_overlap): + # qualifies a row (it only boosts/reorders matches). On the BM25 path the FTS + # index already established textual relevance, so the qualifier is heuristic-only. + # Degenerate queries with no usable tokens fall through to role-ranked listing. + if query_toks and not use_fts and not (name_overlap or type_hits or fqn_match or text_overlap): continue role_w = 0.0 if role_locked else _ROLE_SCORE_WEIGHTS.get(role_raw.upper(), 0.0) @@ -291,6 +387,8 @@ def run_lexical_search( "lexical_relevance": round(raw, 4), "role_weight": role_w, } + if use_fts: + comps["bm25"] = round(float(bm25_scores.get(r.get("id"), 0.0)), 4) sl, el, sb, eb = r.get("start_line"), r.get("end_line"), r.get("start_byte"), r.get("end_byte") out.append( diff --git a/src/java_codebase_rag/search/search_scoring.py b/src/java_codebase_rag/search/search_scoring.py index 3a30546..8d539ec 100644 --- a/src/java_codebase_rag/search/search_scoring.py +++ b/src/java_codebase_rag/search/search_scoring.py @@ -14,6 +14,11 @@ import json import re +# Name of the LadybugDB FTS (Okapi BM25) index over Symbol.search_text (fork A). +# Shared by the build path (build_ast_graph._ensure_symbol_fts_index) and the +# query path (search_lexical.run_lexical_search) so the two never drift. +SYMBOL_FTS_INDEX = "sym_fts" + # Over-fetch multiplier for dedup: fetch 4x to absorb per-FQN chunk multiplicity # so that after collapsing by primary_type_fqn, a page stays full and the +1 # truncation sentinel survives. The formula: need = max((limit + offset) * 4, limit + offset + 1) diff --git a/tests/fixtures/graph_baseline_bank_chat.json b/tests/fixtures/graph_baseline_bank_chat.json index b9357e2..7d52e51 100644 --- a/tests/fixtures/graph_baseline_bank_chat.json +++ b/tests/fixtures/graph_baseline_bank_chat.json @@ -10,7 +10,7 @@ "UNRESOLVED_AT": 227 }, "graph_meta": { - "ontology_version": 18, + "ontology_version": 19, "built_at": 1782110216, "source_root": "/Users/dmitry/Desktop/CursorProjects/java-enterprise-codebase-rag/tests/bank-chat-system", "counts_json": "{packages: 29, files: 130, types: 140, members: 606, phantoms: 54, extends: 18, implements: 21, injects: 94, declares: 606, overrides: 38, calls: 684, routes: 29, exposes: 15, clients: 8, declares_client: 8, producers: 9, declares_producer: 9, http_calls: 8, async_calls: 9}" diff --git a/tests/integration/test_incremental_graph.py b/tests/integration/test_incremental_graph.py index 67e11f9..25b0091 100644 --- a/tests/integration/test_incremental_graph.py +++ b/tests/integration/test_incremental_graph.py @@ -222,9 +222,9 @@ def test_source_file_value_matches_symbol_filename(self, tmp_path: Path) -> None sub_filename, edge_source_file = result.get_next() assert sub_filename == edge_source_file - def test_ontology_version_bumped_to_18(self) -> None: - """ONTOLOGY_VERSION == 18.""" - assert ONTOLOGY_VERSION == 18 + def test_ontology_version_bumped_to_19(self) -> None: + """ONTOLOGY_VERSION == 19 (v19: Symbol.search_text + LadybugDB FTS index, fork A).""" + assert ONTOLOGY_VERSION == 19 class TestIncrementalOrchestrator: diff --git a/tests/package/test_jrag_status.py b/tests/package/test_jrag_status.py index dfbdbdf..4655905 100644 --- a/tests/package/test_jrag_status.py +++ b/tests/package/test_jrag_status.py @@ -70,7 +70,7 @@ def test_status_reports_ontology_version_and_counts( payload = json.loads(proc.stdout) assert payload["status"] == "ok" index = payload["nodes"]["index"] - assert index["ontology_version"] == 18 + assert index["ontology_version"] == 19 # Counts is a top-level nested dict on the index node (the generic # nested-sections dispatch signal - any dict-typed value renders as an # indented alphabetical section; edge_summary is NOT used as the dispatch diff --git a/tests/search/test_search_lexical.py b/tests/search/test_search_lexical.py index 924f954..8bb5bab 100644 --- a/tests/search/test_search_lexical.py +++ b/tests/search/test_search_lexical.py @@ -226,13 +226,16 @@ def _spy(cypher, params): # capture the LIMIT param the backend sends to the gr def test_cap_truncation_advisory_fires_at_cap(monkeypatch, tmp_path: Path) -> None: - """Correctness review: when the fetch hits the safety cap, deeper matches are never + """Heuristic-fallback path: when the fetch hits the safety cap, deeper matches are never ranked (the scan has no ORDER BY). Surface an advisory instead of silently returning a - storage-order-dependent subset. Lower the cap so a small fixture triggers it.""" + storage-order-dependent subset. Lower the cap so a small fixture triggers it. The BM25 + (FTS) path has no cap, so this forces the heuristic scan via ``_try_fts_candidates → + None`` to keep the cap + advisory exercised (fork A).""" from java_codebase_rag.search import search_lexical db = _build_corpus(tmp_path) g = _graph(db) + monkeypatch.setattr(search_lexical, "_try_fts_candidates", lambda *a, **k: None) monkeypatch.setattr(search_lexical, "_CANDIDATE_LIMIT_CAP", 3) adv: list[str] = [] run_lexical_search("distribution", limit=5, graph=g, advisories=adv) @@ -337,3 +340,53 @@ def test_search_v2_lexical_dispatch_end_to_end(monkeypatch, tmp_path: Path) -> N out2 = mcp_v2.search_v2(query="distribution chunk", graph=g, explain=True) assert out2.success and out2.results assert out2.results[0].score_components, "explain should populate score_components" + + +def test_bm25_path_active_when_fts_index_present(tmp_path: Path) -> None: + """Fork A: with a built FTS index (the default after init/reprocess), run_lexical_search + fetches candidates DB-side via Okapi BM25 and records a bm25 score component on each hit. + Ranking is still driven by the heuristic re-rank (lexical_relevance), not by bm25.""" + db = _build_corpus(tmp_path) + rows = run_lexical_search("distribution chunk service", limit=5, graph=_graph(db)) + assert rows, "expected hits" + assert all("bm25" in r["_score_components"] for r in rows), ( + f"BM25 path should annotate each hit with a bm25 component; got {rows[0]['_score_components']}" + ) + assert rows[0]["fqn"] == "svc.DistributionChunkService" + + +def test_bm25_falls_back_to_heuristic_when_fts_index_absent(tmp_path: Path) -> None: + """Fork A: if the FTS index is missing (older graph, or an offline first-run where + INSTALL FTS failed), run_lexical_search falls back to the heuristic scan — it still + returns hits, just without a bm25 component.""" + import ladybug + from java_codebase_rag.search.search_scoring import SYMBOL_FTS_INDEX + + db = _build_corpus(tmp_path) + # Drop the FTS index to simulate its absence. + rw_db = ladybug.Database(str(db), read_only=False) + rw_conn = ladybug.Connection(rw_db) + try: + rw_conn.execute("LOAD EXTENSION FTS") + rw_conn.execute(f"CALL DROP_FTS_INDEX('Symbol', '{SYMBOL_FTS_INDEX}')") + finally: + rw_conn.close() + rw_db.close() + + rows = run_lexical_search("distribution chunk service", limit=5, graph=_graph(db)) + assert rows, "heuristic fallback must still return hits" + assert all("bm25" not in r["_score_components"] for r in rows), ( + "heuristic fallback must not annotate a bm25 component" + ) + + +def test_bm25_path_respects_role_filter(tmp_path: Path) -> None: + """Fork A: the BM25 fetch re-applies the NodeFilter on the re-MATCH (role pushdown), so a + role=CONTROLLER filter surfaces the @RestController hit and excludes the @Service one.""" + db = _build_corpus(tmp_path) + rows = run_lexical_search( + "operator session", limit=5, graph=_graph(db), filter=NodeFilter(role="CONTROLLER") + ) + assert rows, "expected the controller hit through the role filter" + assert all(r["role"] == "CONTROLLER" for r in rows), [r["role"] for r in rows] + assert rows[0]["fqn"] == "ctrl.OperatorSessionController"