diff --git a/aborist/qa/query.py b/aborist/qa/query.py index c97d748..7e6858f 100644 --- a/aborist/qa/query.py +++ b/aborist/qa/query.py @@ -6,13 +6,17 @@ asks Hermes, caches the answer. The flow: 1. FTS5 search across all shards (chunks_fts can't be UNION'd in views, so we query each shard's index independently and merge by score). - 2. Pick top-K distinct documents within a character budget. - 3. Compute `context_root` = Merkle root over the sorted source + 2. Title-boost rerank: hits whose title contains query tokens get a + score bump. Title is a strong topical signal that BM25 alone misses + (BM25 favors short docs with rare body tokens — Tell_(poker) outranks + Back_to_the_Future without a title boost). + 3. Pick top-K distinct documents within a character budget. + 4. Compute `context_root` = Merkle root over the sorted source document_roots — that's the "source" dimension of v9.8's 8-dim cache_key for this multi-source answer. - 4. Cache lookup; hit returns STRICT immediately. - 5. Miss calls Hermes via the OpenAI-compatible client. - 6. Persist record with merkle_proof = {context_root, sources: [...]} + 5. Cache lookup; hit returns STRICT immediately. + 6. Miss calls Hermes via the OpenAI-compatible client. + 7. Persist record with merkle_proof = {context_root, sources: [...]} so a verifier can recompute context_root and check each source document_root against its shard. @@ -24,6 +28,7 @@ chunk's proof on demand. from __future__ import annotations import json +import re import time from dataclasses import dataclass from pathlib import Path @@ -51,6 +56,42 @@ from aborist.store import ( ) +_TITLE_TOKEN_RE = re.compile(r"[A-Za-z][A-Za-z0-9]*") +_TITLE_STOPWORDS = frozenset( + """ + the a an is are was were be been being of to in on at for with by from + as about into through during and or but not no nor so yet what who where + when why how which this that these those such i you he she it we they + me him her us them do does did have has had can could should would will + """.split() +) + + +def _title_query_tokens(s: str) -> set[str]: + return { + t.lower() + for t in _TITLE_TOKEN_RE.findall(s) + if t.lower() not in _TITLE_STOPWORDS and len(t) > 1 + } + + +def _rerank_by_title(hits: list, question: str, boost: float = 8.0) -> list: + """Boost hits whose title contains query tokens — title is high signal.""" + qtokens = _title_query_tokens(question) + if not qtokens: + return hits + for h in hits: + if not h.title: + continue + # Replace _ with space: "Back_to_the_Future" -> ["back", "future"] + ttokens = _title_query_tokens(h.title.replace("_", " ")) + overlap = len(qtokens & ttokens) + if overlap: + h.score += overlap * boost + hits.sort(key=lambda h: -h.score) + return hits + + DEFAULT_QUERY_POLICY = { "system_prompt": ( "You are answering a question using ONLY the sources provided below. " @@ -130,6 +171,11 @@ def _search_corpus( return out +def _rerank(hits: list[_Hit], question: str) -> list[_Hit]: + """Public-private rerank: title-boost. Easy to extend with other signals.""" + return _rerank_by_title(hits, question) + + def _load_doc_text(shard_path: str, document_root: str) -> str | None: """Concatenate all hot chunks of a document. Returns None if cold or missing.""" conn = connect(shard_path) @@ -183,6 +229,9 @@ def query( hits = _search_corpus(shards_dir, single_db, question, over_fetch) if not hits: return {"status": "no_sources", "msg": "FTS5 search returned no hits"} + # 1b. Rerank — title-boost gives topically-named docs priority over + # body-only matches like Tell_(poker) for "tell me about ...". + hits = _rerank(hits, question) # 2. Pull doc texts within budget. chosen: list[_Hit] = [] diff --git a/aborist/sources/wikipedia.py b/aborist/sources/wikipedia.py index 876b025..b202fd0 100644 --- a/aborist/sources/wikipedia.py +++ b/aborist/sources/wikipedia.py @@ -244,6 +244,7 @@ class WikipediaSqlDump(Source): base_uri: str = "https://en.wikipedia.org/wiki/", shard: tuple[int, int] | None = None, start_id: int = 0, + encoding: str = "latin-1", ): if table not in ("cur", "old"): raise ValueError("table must be 'cur' or 'old'") @@ -265,12 +266,17 @@ class WikipediaSqlDump(Source): # High-water mark observed during iteration. Caller reads this back # after each batch and persists it via store.set_meta(). self.last_id: int = start_id + # 2003-era MediaWiki dumps mix Latin-1 raw bytes (e.g., 0xE9 for é) + # with HTML entities. Latin-1 decode is lossless on every byte and + # gives the right code point for the raw-byte cases. Modern dumps + # are UTF-8 — pass encoding="utf-8" for those. + self.encoding = encoding def iter_documents(self) -> Iterator[Document]: opener = bz2.open if str(self.path).endswith(".bz2") else open rank, total = (0, 1) if self.shard is None else self.shard idx = 0 - with opener(self.path, "rt", encoding="utf-8", errors="replace") as f: + with opener(self.path, "rt", encoding=self.encoding, errors="replace") as f: for stmt in _iter_insert_statements(f): head = stmt.lstrip() if not head.upper().startswith("INSERT INTO"):