diff --git a/aborist/qa/query.py b/aborist/qa/query.py index 3c31573..1c5cac0 100644 --- a/aborist/qa/query.py +++ b/aborist/qa/query.py @@ -911,17 +911,15 @@ def _search_corpus( out-ranks FTS5 body hits so the actual topic article rises to the top. """ qtokens = _title_query_tokens(question) - # Title-LIKE backup search uses ORIGINAL qtokens only — NOT synonym- - # expanded. With expansion, this path becomes O(corpus × |accept|) - # full-scan because LOWER(title) LIKE '%tok%' can't use any index; - # 50 LIKE patterns × 870k docs × 4 shards = ~56s of pure waste. - # Synonym expansion stays in `_filter_by_title_relevance` (post- - # retrieval, in-memory, cheap regardless of accept-set size) where - # it actually does useful work surfacing synonym-related titles - # from the FTS5 BM25 hits. Title-LIKE is a thin backup for docs - # whose body is short on query terms but whose title is the literal - # topic — that case is well-served by the original qtokens. + # Title-LIKE backup uses ORIGINAL qtokens only (LIKE %tok% can't + # use any index — adding synonyms makes it O(corpus × |accept|)). + # Synonym expansion stays useful in two places: (1) the FTS5 + # OR-mode fallback (top-5 longest pool merged with synonyms — long + # topical synonyms like "neurotechnology" surface relevant titles + # without paying for full-scan), and (2) `_filter_by_title_relevance` + # post-retrieval filtering (in-memory, cheap). accept_tokens = set(qtokens) + or_synonym_pool = synonym_expand(qtokens, shards_dir=shards_dir) paths: list[Path] if shards_dir is not None: paths = discover_shards(shards_dir) @@ -948,7 +946,9 @@ def _search_corpus( conn = connect(p) try: backend = FTS5Backend(conn) - for h in backend.search(question, limit=over_fetch): + for h in backend.search( + question, limit=over_fetch, extra_or_tokens=or_synonym_pool + ): raw.append( ( h.score, diff --git a/aborist/search/fts5.py b/aborist/search/fts5.py index 2029bc3..0573efc 100644 --- a/aborist/search/fts5.py +++ b/aborist/search/fts5.py @@ -52,7 +52,12 @@ def _quote(t: str) -> str: return '"' + t.replace('"', '""') + '"' -def _escape_fts5(query: str, *, mode: str = "and") -> str: +def _escape_fts5( + query: str, + *, + mode: str = "and", + extra_or_tokens: set[str] | None = None, +) -> str: """Build a MATCH expression from a free-text query. Tokenizes by alpha runs (so `?` and other punctuation can't break the @@ -67,6 +72,14 @@ def _escape_fts5(query: str, *, mode: str = "and") -> str: clause matches millions of docs and forces BM25 to rank them all — 13s/shard observed pre-cap). + The ``extra_or_tokens`` kwarg accepts synonym-expanded tokens + (e.g. "telepathy" / "neurotechnology" expanded from the query token + "thoughts"). These join the top-N-longest OR pool — long synonym + tokens like "neurotechnology" (15 chars) outrank short query tokens + like "thoughts" (8) by length and surface the right titles. Pure + quality win at OR-mode-fallback time without extra retrieval cost + since the top-N cap still applies to the merged pool. + All-stopword queries fall back to OR over the raw tokens so they still find something instead of crashing FTS5 with an empty MATCH. """ @@ -83,8 +96,20 @@ def _escape_fts5(query: str, *, mode: str = "and") -> str: # corpus). Bounds the per-clause cost so OR-mode terminates # quickly instead of scanning the corpus. sep = " OR " - if len(tokens) > _OR_FALLBACK_MAX_TOKENS: - tokens = sorted(tokens, key=len, reverse=True)[:_OR_FALLBACK_MAX_TOKENS] + candidate_pool = list(tokens) + if extra_or_tokens: + # Synonyms join the OR pool; dedupe lowercase. + seen_lower = {t.lower() for t in candidate_pool} + for t in extra_or_tokens: + tl = t.lower() + if tl and tl not in seen_lower and tl not in _FTS5_STOPWORDS: + candidate_pool.append(tl) + seen_lower.add(tl) + if len(candidate_pool) > _OR_FALLBACK_MAX_TOKENS: + candidate_pool = sorted(candidate_pool, key=len, reverse=True)[ + :_OR_FALLBACK_MAX_TOKENS + ] + tokens = candidate_pool return sep.join(_quote(t) for t in tokens) @@ -153,12 +178,31 @@ class FTS5Backend(SearchBackend): name = "fts5" audit_mode = AuditMode.UNGROUNDED - def search(self, query: str, limit: int = 20) -> list[Hit]: + def search( + self, + query: str, + limit: int = 20, + extra_or_tokens: set[str] | None = None, + ) -> list[Hit]: + """Run FTS5 BM25 over chunk content. + + ``extra_or_tokens`` (synonym-expanded set) is passed through to + the OR-mode fallback only. AND mode stays on original query + tokens (adding synonyms there would relax the AND constraint + and pull in noise). The intended caller is the retrieval + pipeline that has already computed ``synonym_expand(qtokens)`` + — passing it here saves the OR-mode pool from missing topical + synonym terms. + """ if not query.strip(): return [] # Try strict AND first; fall back to OR if it returns nothing. for mode in ("and", "or"): - fts_query = _escape_fts5(query, mode=mode) + fts_query = _escape_fts5( + query, + mode=mode, + extra_or_tokens=extra_or_tokens if mode == "or" else None, + ) try: rows = self.conn.execute( """