diff --git a/arborist/qa/corpus.py b/arborist/qa/corpus.py index a6b8bb6..4114436 100644 --- a/arborist/qa/corpus.py +++ b/arborist/qa/corpus.py @@ -342,13 +342,38 @@ class SqliteShardCorpus: ] def fts_title(self, query: str, *, limit: int = 8) -> list[Hit]: - """Title-only bm25() via the documents_fts virtual table that - lives alongside chunks_fts in every shard. Uses the SAME - sanitizer as fts_body so a user query with punctuation / - stopwords gets the same shape on both routes.""" - from arborist.wallet.bucket import _to_fts5 - match_expr = _to_fts5(query) - if not match_expr.strip(): + """Title-only bm25() via the documents_fts virtual table. + + Per-token fan-out (one MATCH per non-stopword token, dedupe + by doc_root keeping best bm25). Single-MATCH ``tok1 OR tok2`` + is biased toward the rarer token — for "what is the capitol + of paris", "capitol" (DF ~1k titles) outranks "paris" (DF ~5k + titles) so a unioned MATCH returns 12 Capitol-* titles and + zero Paris-* hits, squeezing the actual "Paris" article out + of the top-K entirely. Per-token MATCH gives each token its + own candidate space; the merge then takes the best bm25 each + doc earned in any token's run. + """ + from arborist.wallet.bucket import ( + _to_fts5, + _FTS5_STOPWORDS, + _FTS5_WORD_RE, + ) + import unicodedata + # Single-token MATCH list — preserve FTS5-safe tokenization + # via the same sanitizer fts_body uses (NFKD-fold, drop + # stopwords, ≥2 char filter). + if not _to_fts5(query).strip(): + return [] + folded = "".join( + c for c in unicodedata.normalize("NFKD", query) + if not unicodedata.combining(c) + ) + tokens = [ + t.lower() for t in _FTS5_WORD_RE.findall(folded) + if t and t.lower() not in _FTS5_STOPWORDS and len(t) > 1 + ] + if not tokens: return [] sql = ( "SELECT d.document_root, d.document_uri, d.title, " @@ -358,16 +383,26 @@ class SqliteShardCorpus: "WHERE documents_fts MATCH ? " "ORDER BY score LIMIT ?" ) - rows = list(self._conn.execute(sql, (match_expr, limit))) + best: dict[str, tuple] = {} + for tok in tokens: + for r in self._conn.execute(sql, (tok, limit)): + droot = r["document_root"] + cur = best.get(droot) + if cur is None or r["score"] < cur[0]: + best[droot] = ( + r["score"], r["document_uri"] or "", + r["title"] or "", + ) + ranked = sorted(best.items(), key=lambda kv: kv[1][0])[:limit] return [ Hit( - document_root=r["document_root"], - document_uri=r["document_uri"] or "", - title=r["title"] or "", - score=r["score"], + document_root=droot, + document_uri=uri, + title=title, + score=score, shard_id=None, ) - for r in rows + for droot, (score, uri, title) in ranked ] def fts_phrase(self, ngrams, *, limit: int = 8) -> list[Hit]: @@ -618,8 +653,107 @@ class MultiShardSqliteCorpus: return merged[:limit] def fts_title(self, query: str, *, limit: int = 8) -> list[Hit]: - """Per-shard fts_title merged by raw bm25 score (ASC).""" - return self._fanout("fts_title", limit, query=query) + """Per-token title fan-out across all shards. + + Single-MATCH ``tok1 OR tok2`` over documents_fts is biased + toward whichever token has lower corpus DF — rarer tokens get + higher bm25 per match, so a single-MATCH return is dominated + by their hits. For "what is the capitol of paris" that meant + Capitol-* titles filled top-16 and the "Paris" article never + made it into the merge pool. Per-shard per-token MATCH gets + balanced per-token at the shard level, BUT MultiShard's + generic _fanout then re-sorts globally by bm25 — undoing the + balance. Fix: run per-token here too, giving each token its + own slot in the final return. + """ + from arborist.wallet.bucket import ( + _to_fts5, _FTS5_STOPWORDS, _FTS5_WORD_RE, + ) + from concurrent.futures import ThreadPoolExecutor + import unicodedata + + if not _to_fts5(query).strip(): + return [] + folded = "".join( + c for c in unicodedata.normalize("NFKD", query) + if not unicodedata.combining(c) + ) + tokens = [ + t.lower() for t in _FTS5_WORD_RE.findall(folded) + if t and t.lower() not in _FTS5_STOPWORDS and len(t) > 1 + ] + if not tokens: + return [] + + # Per-token: fan out across shards in parallel; each token + # gets its OWN top-(limit*2) bucket. The final return is the + # union of all tokens' buckets — guarantees every token has + # surface area even when one token (e.g. "capitol", low-DF) + # outranks another ("paris", higher-DF) on raw bm25. + def _shard_token(item, token): + sh_path, sc = item + try: + hits = sc.fts_title(token, limit=limit * 2) + return sh_path, hits + except Exception: + return sh_path, [] + + per_token_buckets: list[list[Hit]] = [] + with ThreadPoolExecutor( + max_workers=max(1, len(self._shards)) + ) as ex: + for tok in tokens: + futures = [ + ex.submit(_shard_token, item, tok) + for item in self._shards + ] + bucket: list[Hit] = [] + for f in futures: + sh_path, hits = f.result() + for h in hits: + bucket.append(Hit( + document_root=h.document_root, + document_uri=h.document_uri, + title=h.title, + score=h.score, + shard_id=sh_path, + extras=dict(h.extras), + )) + bucket.sort(key=lambda h: h.score) + per_token_buckets.append(bucket[: limit * 2]) + # Union: each token gets equal allocation. Dedupe by + # document_root keeping best bm25 across tokens. The final + # sort is per-bucket-aware: walk all buckets, take the + # top-N from each round-robin, until we hit `limit`. + best: dict[str, Hit] = {} + for bucket in per_token_buckets: + for h in bucket: + cur = best.get(h.document_root) + if cur is None or h.score < cur.score: + best[h.document_root] = h + # Round-robin across token buckets so each token contributes + # representatives even when one token's bm25 dominates the + # global sort. Within a bucket, walk by bm25-ascending order. + bucket_iters = [iter(b) for b in per_token_buckets] + seen_roots: set[str] = set() + ranked: list[Hit] = [] + active = True + while active and len(ranked) < limit: + active = False + for it in bucket_iters: + if len(ranked) >= limit: + break + for h in it: + active = True + if h.document_root in seen_roots: + continue + seen_roots.add(h.document_root) + # Use the merged best-score (in case the same + # doc surfaced from multiple tokens with better + # bm25 elsewhere). + ranked.append(best[h.document_root]) + break + return ranked def fts_phrase(self, ngrams, *, limit: int = 8) -> list[Hit]: """Per-shard fts_phrase merged by raw bm25 score (ASC). diff --git a/arborist/qa/corpus_query.py b/arborist/qa/corpus_query.py index 437d296..3df2a04 100644 --- a/arborist/qa/corpus_query.py +++ b/arborist/qa/corpus_query.py @@ -134,6 +134,7 @@ def run_query( max_tokens: int = 512, warrant_check_enabled: bool = False, policy: dict | None = None, + precomputed_hits: list | None = None, ) -> dict: """End-to-end claim-lattice query over any Corpus. @@ -191,7 +192,12 @@ def run_query( # cut before reranking sees it). ts = _time.time() multi_route = bool(policy and policy.get("multi_route", False)) - if multi_route: + if precomputed_hits is not None: + # Caller (providence_query) already ran retrieval; skip the + # in-function body-only fetch and the multi-route fan-out so + # the title-aware union the caller computed survives. + hits = list(precomputed_hits) + elif multi_route: from arborist.qa.retrieval_routes import ( filter_by_title_relevance, question_phrases, @@ -235,10 +241,14 @@ def run_query( # Shared title-boost rerank — lifts "Homer Simpson" main article # above "You Only Move Twice" sibling on the homer query, etc. # Same logic for SqliteShardCorpus + SidecarBucketCorpus. - hits = apply_title_boost( - hits, question, - higher_is_better=getattr(corpus, "higher_is_better", False), - )[:top_k] + # Skip when caller already ran boost on a precomputed set. + if precomputed_hits is None: + hits = apply_title_boost( + hits, question, + higher_is_better=getattr(corpus, "higher_is_better", False), + )[:top_k] + else: + hits = hits[:top_k] timings["search"] = _time.time() - ts if not hits: diff --git a/arborist/qa/providence_query.py b/arborist/qa/providence_query.py index 7109a04..9916148 100644 --- a/arborist/qa/providence_query.py +++ b/arborist/qa/providence_query.py @@ -317,7 +317,6 @@ def providence_query( # The merge picks whichever route ranked a doc best. Title-boost # then runs on the union. from arborist.qa.corpus import NotSupportedError, apply_title_boost - from concurrent.futures import ThreadPoolExecutor progress.emit("search.start", top_k=top_k) t_ret = _time.time() @@ -329,14 +328,14 @@ def providence_query( return fn(question, limit=top_k * 4) except NotSupportedError: return [] - except Exception: - return [] - with ThreadPoolExecutor(max_workers=2) as ex: - f_body = ex.submit(_safe, "fts_body") - f_title = ex.submit(_safe, "fts_title") - body_hits = f_body.result() - title_hits = f_title.result() + # Sequential: MultiShardSqliteCorpus.fts_body / fts_title already + # fan out across shards internally with a ThreadPoolExecutor. An + # outer pool here would double-parallelize, and single-shard + # SqliteShardCorpus connections are not thread-safe by default + # (check_same_thread=True). Run body then title. + body_hits = _safe("fts_body") + title_hits = _safe("fts_title") # Merge by document_root, keeping the best bm25 score across # routes. FTS5 bm25 is "more negative = better" — title typically @@ -476,11 +475,12 @@ def providence_query( "burned_existing": 0, "elapsed_s": round(_time.time() - t_total, 3), } - # 5. Miss — call run_query (which re-does retrieval + does - # LLM + verify). The retrieval re-run wastes ~sub-second but - # keeps the orchestrator's prompt/verify path single-sourced; - # a future refactor can pass `precomputed_hits=hits` in to - # eliminate the redundant retrieval. + # 5. Miss — call run_query for LLM + verify. Pass our hits + # through so run_query doesn't redo retrieval and discard our + # title-aware union (previously dropped Paris off the "capitol + # of paris" query because run_query re-did body-only retrieval + # internally and high-DF body noise outranked the canonical + # title-hit primary). progress.emit("cache.miss", lookup_path="primary") progress.emit("llm.start", model=model_id, top_k=top_k) t_llm = _time.time() @@ -490,6 +490,7 @@ def providence_query( max_context_chars=max_context_chars, temperature=temperature, max_tokens=max_tokens, policy=policy, + precomputed_hits=hits, ) progress.emit( "llm.done",