# app/services/retriever.py """ Hybrid BM25 + semantic retriever. BSL 1.1 — semantic path requires PAGEPIPER_OLLAMA_URL (BYOK gate). BM25-only path is MIT and has no gate. The semantic half uses Agent-ModernColBERT (pagepiper#8) — a late-interaction retriever scored via MaxSim, replacing the earlier nomic-embed-text bi-encoder + cosine similarity approach. The model runs locally via `pylate`; the BYOK gate below matches product tiering, not a technical dependency on Ollama. """ from __future__ import annotations import logging import sqlite3 from dataclasses import dataclass from app.services.bm25_index import BM25Index from app.services.colbert_index import ColBERTIndex logger = logging.getLogger(__name__) def _fetch_adjacent( hits: list["RetrievedChunk"], db_path: str, window: int = 1, ) -> list["RetrievedChunk"]: """Return chunks immediately before/after each hit that aren't already in the hit set. Definitional passages often start mid-sentence because the EPUB/PDF chunk boundary fell mid-paragraph. Fetching the preceding chunk restores the subject so the LLM can understand 'them' / 'they' references correctly. """ if not hits: return [] existing_keys = {(c.doc_id, c.page_number) for c in hits} needed: dict[str, set[int]] = {} for c in hits: for delta in range(-window, window + 1): if delta == 0: continue adj_page = c.page_number + delta if adj_page > 0 and (c.doc_id, adj_page) not in existing_keys: needed.setdefault(c.doc_id, set()).add(adj_page) if not needed: return [] extra: list[RetrievedChunk] = [] try: conn = sqlite3.connect(db_path) conn.row_factory = sqlite3.Row for doc_id, pages in needed.items(): placeholders = ",".join("?" * len(pages)) rows = conn.execute( f"SELECT id, doc_id, page_number, text FROM page_chunks " f"WHERE doc_id=? AND page_number IN ({placeholders})", [doc_id] + sorted(pages), ).fetchall() for row in rows: extra.append( RetrievedChunk( chunk_id=row["id"], doc_id=row["doc_id"], page_number=row["page_number"], text=row["text"], bm25_score=0.0, vector_score=None, ) ) conn.close() except Exception as exc: logger.warning("Context expansion query failed (non-fatal): %s", exc) return extra @dataclass(frozen=True) class RetrievedChunk: """A chunk returned by the retriever, with source scores.""" chunk_id: str doc_id: str page_number: int text: str bm25_score: float vector_score: float | None class Retriever: def __init__(self, bm25: BM25Index, colbert: ColBERTIndex | None = None) -> None: self._bm25 = bm25 self._colbert = colbert def hybrid_search( self, query: str, top_k: int, doc_ids: list[str] | None, db_path: str, vec_db_path: str, llm, # LLMRouter | None — caller must pass ) -> list[RetrievedChunk]: """ Merge BM25 and semantic (ColBERT) results. Falls back to BM25-only if llm is None or no ColBERT index is configured. """ if llm is None or self._colbert is None: return self._bm25_only(query, top_k, doc_ids, db_path) self._bm25.ensure_fresh(db_path) bm25_hits = { r.chunk_id: r for r in self._bm25.query(query, top_k=top_k * 2, doc_ids=doc_ids) } try: self._colbert.ensure_fresh(db_path) # ColBERTIndex.query already oversamples internally when doc_ids is set — # see app/services/colbert_index.py. colbert_hits = self._colbert.query(query, top_k=top_k, doc_ids=doc_ids) except Exception as exc: logger.warning("ColBERT retrieval failed, falling back to BM25-only: %s", exc) return self._bm25_only(query, top_k, doc_ids, db_path) # Merge: BM25 hits take priority; ColBERT hits fill in additional results merged: dict[str, RetrievedChunk] = {} for cid, r in bm25_hits.items(): merged[cid] = RetrievedChunk( chunk_id=cid, doc_id=r.doc_id, page_number=r.page_number, text=r.text, bm25_score=r.score, vector_score=None, ) # ColBERT MaxSim scores are unbounded (roughly num_query_tokens * # max_per_token_similarity), unlike BM25's already-comparable range — # min-max normalize within this result batch before combining. max_colbert_score = max((h["score"] for h in colbert_hits), default=0.0) for h in colbert_hits: cid = h["chunk_id"] norm_score = (h["score"] / max_colbert_score) if max_colbert_score > 0 else 0.0 if cid in merged: existing = merged[cid] merged[cid] = RetrievedChunk( chunk_id=existing.chunk_id, doc_id=existing.doc_id, page_number=existing.page_number, text=existing.text, bm25_score=existing.bm25_score, vector_score=norm_score, ) else: merged[cid] = RetrievedChunk( chunk_id=cid, doc_id=h["doc_id"], page_number=h["page_number"], text=h["text"], bm25_score=0.0, vector_score=norm_score, ) def _combined(r: RetrievedChunk) -> float: bm25 = r.bm25_score vec = r.vector_score if r.vector_score is not None else 0.0 return bm25 * 0.5 + vec * 0.5 all_ranked = sorted(merged.values(), key=_combined, reverse=True) # Discard results where the best match is pure noise (neither BM25 term # overlap nor vector similarity exceeded the minimum signal threshold). # This lets the caller's empty-result guard fire instead of sending # low-confidence chunks to the LLM where it fills gaps with training data. MIN_SIGNAL = 0.01 if all_ranked and _combined(all_ranked[0]) < MIN_SIGNAL: return [] # Cap per-document contribution to max_per_doc of top_k so that one book # does not crowd out all slots when the query matches it heavily by name # alone (e.g. a character name that appears in every chapter). max_per_doc = max(2, top_k // 3) ranked: list[RetrievedChunk] = [] doc_counts: dict[str, int] = {} for r in all_ranked: if len(ranked) >= top_k: break count = doc_counts.get(r.doc_id, 0) if count < max_per_doc: ranked.append(r) doc_counts[r.doc_id] = count + 1 adjacent = _fetch_adjacent(ranked, db_path) return ranked + adjacent def _bm25_only( self, query: str, top_k: int, doc_ids: list[str] | None, db_path: str ) -> list[RetrievedChunk]: self._bm25.ensure_fresh(db_path) hits = [ RetrievedChunk( chunk_id=r.chunk_id, doc_id=r.doc_id, page_number=r.page_number, text=r.text, bm25_score=r.score, vector_score=None, ) for r in self._bm25.query(query, top_k=top_k, doc_ids=doc_ids) ] MIN_SIGNAL = 0.01 if hits and hits[0].bm25_score < MIN_SIGNAL: return [] adjacent = _fetch_adjacent(hits, db_path) return hits + adjacent