diff --git a/arborist/qa/corpus.py b/arborist/qa/corpus.py index 9b2e700..28caa57 100644 --- a/arborist/qa/corpus.py +++ b/arborist/qa/corpus.py @@ -134,10 +134,20 @@ def apply_title_boost( return rescored -@dataclass(frozen=True) +@dataclass class Hit: """One retrieval result, backend-agnostic. + Mutable on purpose: the rerank pipeline in + ``arborist.qa.retrieval_routes`` mutates ``.score`` (and tags + ``.extras["source_role"]``) in place as a sequence of multiplier + stages. Frozen would block that and force every rerank to rebuild + the list via ``dataclasses.replace``, doubling allocation and + losing the LCS-on-sort optimization legacy ``query()`` relies on. + Hit objects are short-lived per-request artifacts — they never + enter the persistence layer (that's ChunkRow + the verifier's + evidence map). + ``score`` semantics differ per backend (FTS5 BM25 ranges negative [-20, 0]; sidecar BM25 ranges positive [+25, +60]). Callers that merge across backends must use rank or a normalization layer — diff --git a/arborist/qa/retrieval_routes.py b/arborist/qa/retrieval_routes.py index f85bd97..21c54d3 100644 --- a/arborist/qa/retrieval_routes.py +++ b/arborist/qa/retrieval_routes.py @@ -1,15 +1,16 @@ """Retrieval orchestration helpers shared between legacy ``query()`` -and the unified ``run_query()`` (Phase 1 step 3 of #53). +and the unified ``run_query()`` (#000072 Phase 1). -Functions here take pre-retrieved Hits + a question and return filtered -or reranked Hits. They never open a sqlite connection directly — the -caller does the data fetch via the Corpus protocol and hands the -results in. Backwards-compatible with both the legacy ``_Hit`` -dataclass and the protocol ``arborist.qa.corpus.Hit`` since both -expose ``.title`` and ``.document_root``. +Functions here take pre-retrieved Hits + a question and return +filtered or reranked Hits. They never open a sqlite connection +directly — the caller does the data fetch via the Corpus protocol +and hands the results in. Backwards-compatible with both the legacy +``_Hit`` dataclass and the protocol ``arborist.qa.corpus.Hit`` since +both expose ``.title``, ``.document_root``, and ``.score``. """ from __future__ import annotations +import math as _math import re as _re from typing import Callable, Iterable @@ -170,3 +171,267 @@ def filter_by_title_relevance( if not kept: return hits[: max(1, fallback_top_n)] if hits else [] return kept + + +# =========================================================================== +# Body-density gate + 5 reranks (#000072 Phase 2 prep: parity with legacy +# query()'s downstream suppression of noisy phrase-route hits). +# =========================================================================== + + +def _body_count_with_stem(body: str, t: str) -> int: + """Count mentions of ``t`` in ``body``, falling back to the + lite-stemmed form if the literal didn't match. Returns the LARGER + of the two counts so a query token that appears under both forms + (rare) still scores. Lifted verbatim from legacy + arborist.qa.query._body_count_with_stem.""" + n_literal = body.count(t) + if n_literal: + return n_literal + stem = stem_for_match(t) + if stem != t: + return body.count(stem) + return 0 + + +def body_density_passes( + corpus, document_root: str, qtokens: set[str], + *, min_mentions: int = 3, +) -> bool: + """Body-token CO-OCCURRENCE relevance gate. + + Breadth scales with query length: + ≤ 2 tokens require ALL present in body + 3+ tokens require N - 1 (allow one weak signal token to miss) + Depth: ``total_mentions >= min_mentions`` also enforced. + + Catches the "Girlfriends (TV show) passing a supermans girlfriend + query because body had 'girlfriend' but no 'superman' whatsoever" + failure mode. Stem-tolerant via ``_body_count_with_stem``. + + Uses ``corpus.doc_body(document_root)`` (shipped in Corpus protocol + step 4 / commit e322bbd) so works against any backend. + """ + if not qtokens: + return False + body = corpus.doc_body(document_root) + if not body: + return False + body = body.lower() + counts = {t: _body_count_with_stem(body, t.lower()) for t in qtokens} + distinct_present = sum(1 for n in counts.values() if n > 0) + total_mentions = sum(counts.values()) + if len(qtokens) <= 2: + breadth_threshold = len(qtokens) + else: + breadth_threshold = len(qtokens) - 1 + return distinct_present >= breadth_threshold and total_mentions >= min_mentions + + +def filter_by_body_density( + hits: list, question: str, corpus, *, min_mentions: int = 3, +) -> list: + """Run body_density_passes on each hit; keep only those that + pass. Fall-open if all hits fail (don't strand the LLM with no + context — legacy convention). + """ + from arborist.qa.query import _title_query_tokens + qtokens = _title_query_tokens(question) + if not qtokens: + return hits + kept = [ + h for h in hits + if body_density_passes( + corpus, h.document_root, qtokens, min_mentions=min_mentions, + ) + ] + return kept if kept else hits + + +def rerank_by_source_role(hits: list, question: str) -> list: + """Classify each hit's source_role + rescale .score by + SOURCE_ROLE_RANK_WEIGHTS. + + Mutates ``h.extras["source_role"]`` (and ``h.source_role`` on + legacy _Hit objects) so downstream context-build can reuse the + value. Sort by score (lower = better for FTS5 bm25 negative; + DESC if scores are positive). Default-DESC matches legacy + behavior; FTS5-negative inputs end up sorted as more-negative- + first which is still "best first" if signs stay homogeneous. + """ + from arborist.qa.query import _title_query_tokens + from arborist.qa.source_roles import ( + SOURCE_ROLE_RANK_WEIGHTS, + classify_source_role, + ) + qtokens_stem = { + stem_for_match(t.lower()) + for t in _title_query_tokens(question) + } + for h in hits: + role = classify_source_role( + h.title, qtokens_stem, + document_uri=getattr(h, "document_uri", None), + ) + # Mutate both .source_role (legacy _Hit) and .extras (Hit + # dataclass with frozen-ish fields) without breaking either. + try: + h.source_role = role + except (AttributeError, TypeError): + pass + if hasattr(h, "extras") and isinstance(h.extras, dict): + h.extras["source_role"] = role + weight = SOURCE_ROLE_RANK_WEIGHTS.get(role, 1.0) + # In-place score multiplication. Hit dataclass uses _replace + # for immutable fields but our Hit is plain @dataclass and + # score is mutable. + h.score = h.score * weight + hits.sort(key=lambda h: -h.score) + return hits + + +def rerank_by_title_purity( + hits: list, question: str, *, shards_dir=None, +) -> list: + """Boost titles by purity AND multi-token-match breadth. + + Multiplier ``(1 + overlap_count) * (1 + purity)``: + - purity = |title∩query| / |title_tokens| (rewards clean titles) + - overlap_count = |title∩query| (rewards multi-token coverage) + + Examples: + ``Jurassic Park (film)`` (overlap 2, purity 1.0) → 6.0× + ``Jurassic Park (NES game)`` (overlap 2, purity 0.5) → 4.5× + + Optionally expand qtokens via the STRICT synonym view (manual + + acronym_parens evidence) when ``shards_dir`` provided. Strict + avoids the broad-view link_reciprocity noise tail. + + Stem-aware via stem_for_match (possessive/plural collapse). + """ + from arborist.qa.query import _title_query_tokens + qtokens = _title_query_tokens(question) + if not qtokens: + return hits + if shards_dir is not None: + try: + from arborist.concepts import synonym_expand_strict + qtokens = synonym_expand_strict( + qtokens, shards_dir=shards_dir, + ) or qtokens + except (ImportError, Exception): + pass + qstems = {stem_for_match(t) for t in qtokens} + for h in hits: + title = getattr(h, "title", "") or "" + if not title: + continue + ttokens = _title_query_tokens(title.replace("_", " ")) + if not ttokens: + continue + tstems = {stem_for_match(t) for t in ttokens} + overlap = tstems & qstems + if not overlap: + continue + purity = len(overlap) / len(tstems) + overlap_count = len(overlap) + h.score = h.score * (1.0 + overlap_count) * (1.0 + purity) + hits.sort(key=lambda h: -h.score) + return hits + + +def _ordered_match_length( + query_tokens: list[str], title_tokens: list[str], +) -> int: + """Longest common subsequence length over two token lists. O(N*M). + Both lists are typically <10 in practice so cost is negligible. + """ + if not query_tokens or not title_tokens: + return 0 + n = len(query_tokens) + m = len(title_tokens) + prev = [0] * (m + 1) + for i in range(1, n + 1): + cur = [0] * (m + 1) + for j in range(1, m + 1): + if query_tokens[i - 1] == title_tokens[j - 1]: + cur[j] = prev[j - 1] + 1 + else: + cur[j] = max(cur[j - 1], prev[j]) + prev = cur + return prev[m] + + +def rerank_by_ordered_token_match(hits: list, question: str) -> list: + """Boost titles whose tokens appear in the same order as the + query. Multiplier ``1 + 0.5 * (match_length - 1)`` for + match_length ≥ 2. Stem-aware. No-op on single-token queries + (no order to match). + + Caught the 2026-05-01 "plot of red fish blue fish?" defect: + the Dr. Seuss book's title (ordered-match 4) sits cleanly + above Red Dwarf / Toronto Blue Jays (ordered-match 1). + """ + from arborist.qa.query import ( + _title_query_tokens, _TITLE_TOKEN_RE, _TITLE_STOPWORDS, + ) + qtokens = _title_query_tokens(question) + if len(qtokens) < 2: + return hits + qtokens_ordered: list[str] = [] + seen: set[str] = set() + for tok in _TITLE_TOKEN_RE.findall(question): + t = tok.lower() + if t in _TITLE_STOPWORDS or len(t) <= 1 or t in seen: + continue + seen.add(t) + qtokens_ordered.append(stem_for_match(t)) + if len(qtokens_ordered) < 2: + return hits + for h in hits: + title = getattr(h, "title", "") or "" + if not title: + continue + ttokens_ordered: list[str] = [] + title_seen: set[str] = set() + for tok in _TITLE_TOKEN_RE.findall(title.replace("_", " ")): + t = tok.lower() + if t in _TITLE_STOPWORDS or len(t) <= 1 or t in title_seen: + continue + title_seen.add(t) + ttokens_ordered.append(stem_for_match(t)) + if not ttokens_ordered: + continue + match_len = _ordered_match_length(qtokens_ordered, ttokens_ordered) + if match_len >= 2: + h.score = h.score * (1.0 + 0.5 * (match_len - 1)) + hits.sort(key=lambda h: -h.score) + return hits + + +def rerank_by_body_coverage( + hits: list, question: str, corpus, *, weight: float = 0.6, +) -> list: + """Boost score by per-token body coverage. sqrt-scaling lets long + topical articles meaningfully out-score short tangential ones + without runaway domination by enumerative list pages. + + Counters BM25's short-doc bias. Cost: one body fetch per + surviving candidate via ``corpus.doc_body`` (cheap on local + SQLite; the cloud path pays an HTTP fetch per doc). + """ + from arborist.qa.query import _title_query_tokens + qtokens_lower = {t.lower() for t in _title_query_tokens(question)} + if not qtokens_lower: + return hits + for h in hits: + body = corpus.doc_body(h.document_root) + if not body: + continue + body_lower = body.lower() + coverage = sum( + _math.sqrt(body_lower.count(t)) for t in qtokens_lower + ) + h.score += coverage * weight + hits.sort(key=lambda h: -h.score) + return hits diff --git a/tests/test_qa_corpus.py b/tests/test_qa_corpus.py index 1490542..9c5072d 100644 --- a/tests/test_qa_corpus.py +++ b/tests/test_qa_corpus.py @@ -79,12 +79,17 @@ def shard_conn(tmp_path): # --------------------------------------------------------------------------- -def test_hit_is_frozen_and_equals_by_value(): +def test_hit_is_mutable_and_equals_by_value(): + """Hit is mutable on purpose so the rerank pipeline in + arborist.qa.retrieval_routes can mutate .score in place (matches + legacy _Hit convention). Equality stays value-based.""" h1 = Hit("root1", "uri1", "title1", 1.5) h2 = Hit("root1", "uri1", "title1", 1.5) - with pytest.raises(Exception): - h1.score = 2.0 # type: ignore[misc] assert h1 == h2 # value equality + # Mutation works — score is the lever rerank stages use. + h1.score = 2.0 + assert h1.score == 2.0 + assert h1 != h2 # mutation breaks equality (as expected) # extras carries a dict so Hit isn't hashable by design — callers # dedup by (document_root, ...) tuples, not by put-in-a-set.