Merge branch 'feat/colbert-retriever' into freeze/pr-12-14-15-integration

# Conflicts:
#	README.md
#	app/api/library.py
#	pyproject.toml

README.md

@@ -31,7 +31,7 @@ No cloud required. Your files stay on your machine.
 - **Your library, not ours.** Documents are indexed and stored locally. Nothing is sent to a third-party service unless you explicitly configure a cloud LLM.
 - **Works without an LLM.** BM25 full-text search runs entirely inside the Docker container. No Ollama, no API key, no GPU required for keyword search.
 - **Answers cite their sources.** Every LLM response includes the document name and page number it drew from. You can verify or dispute every answer.
-- **Hybrid search when you want it.** Connect a local Ollama instance to unlock semantic (vector) search that finds relevant passages even when your question doesn't use the exact words in the text.
+- **Hybrid search when you want it.** Connect a local Ollama instance to unlock hybrid search — BM25 merged with Agent-ModernColBERT, a late-interaction retriever that scores passages by token-level relevance instead of collapsing your whole question into one vector, so multi-part questions find the right passage even when it doesn't use your exact words.
 - **Open shelve pipeline.** The indexing and search layer is MIT-licensed. Add support for new formats, improve the PDF parser, contribute — the community benefits directly.
 
 ---
@@ -98,7 +98,7 @@ PAGEPIPER_EMBED_MODEL=nomic-embed-text
 |-------|-----------|
 | Backend API | FastAPI + SQLite |
 | Full-text search | BM25 (custom index, no external service) |
-| Vector search | sqlite-vec + Ollama embeddings (optional) |
+| Semantic search | Agent-ModernColBERT late-interaction retrieval, via `pylate` (optional, BYOK-gated) |
 | LLM synthesis | Ollama (local, any model) |
 | Frontend | Vue 3 SPA served by nginx |
 | Deployment | Docker Compose |

app/api/chat.py

@@ -80,7 +80,7 @@ def _build_llm_for_alloc(alloc) -> "LLMRouter":
 
 
 def _run_chat(req: "ChatRequest", ctx: "UserCtx", llm) -> "ChatResponse":
-    retriever = Retriever(ctx.bm25)
+    retriever = Retriever(ctx.bm25, ctx.colbert)
     chunks = retriever.hybrid_search(
         query=req.message,
         top_k=req.top_k,

app/api/library.py

@@ -46,6 +46,12 @@ _SHELVE_RUNNERS = {
 }
 
 
+def _mark_indexes_dirty(ctx: UserCtx) -> None:
+    """Mark both BM25 and ColBERT indexes dirty — call after any document is shelved."""
+    ctx.bm25.mark_dirty()
+    ctx.colbert.mark_dirty()
+
+
 def _dispatch_shelve(
     doc_id: str,
     file_path: str,
@@ -175,7 +181,7 @@ def scan_library(
             db.commit()
 
         task_id = _dispatch_shelve(
-            doc_id, path_str, background_tasks, ctx.data_dir, ctx.bm25.mark_dirty
+            doc_id, path_str, background_tasks, ctx.data_dir, lambda: _mark_indexes_dirty(ctx)
         )
         db.execute(
             "UPDATE documents SET status='processing', task_id=? WHERE id=?",
@@ -199,7 +205,7 @@ def reshelve_document(
         raise HTTPException(status_code=404, detail="Document not found")
 
     task_id = _dispatch_shelve(
-        doc_id, row["file_path"], background_tasks, ctx.data_dir, ctx.bm25.mark_dirty
+        doc_id, row["file_path"], background_tasks, ctx.data_dir, lambda: _mark_indexes_dirty(ctx)
     )
     db.execute(
         "UPDATE documents SET status='processing', task_id=?, error_msg=NULL WHERE id=?",
@@ -231,7 +237,7 @@ def delete_document(
     except Exception as exc:
         logger.warning("Could not remove vectors for doc %s: %s", doc_id, exc)
 
-    ctx.bm25.mark_dirty()
+    _mark_indexes_dirty(ctx)
 
 
 def _get_vec_count(doc_id: str, vec_db_path: str) -> int:
@@ -308,7 +314,7 @@ def upload_document(
         db.commit()
 
     task_id = _dispatch_shelve(
-        doc_id, path_str, background_tasks, ctx.data_dir, ctx.bm25.mark_dirty
+        doc_id, path_str, background_tasks, ctx.data_dir, lambda: _mark_indexes_dirty(ctx)
     )
     db.execute(
         "UPDATE documents SET status='processing', task_id=? WHERE id=?",

app/deps.py

@@ -11,11 +11,12 @@ from fastapi import Depends, Request
 
 from app.config import DATA_DIR, LOCAL_USER_ID
 from app.services.bm25_index import BM25Index
+from app.services.colbert_index import ColBERTIndex
 
 
 @dataclass
 class UserCtx:
-    """Per-request context routing DB paths and BM25 to the right user."""
+    """Per-request context routing DB paths, BM25, and ColBERT to the right user."""
 
     user_id: str
     db_path: str
@@ -23,6 +24,7 @@ class UserCtx:
     data_dir: Path
     watch_dir: Path
     bm25: BM25Index
+    colbert: ColBERTIndex
 
 
 _user_startup_done: set[str] = set()
@@ -67,6 +69,7 @@ def get_user_ctx(request: Request) -> UserCtx:
         data_dir=user_dir,
         watch_dir=watch_dir,
         bm25=_main._get_bm25_for(user_id),
+        colbert=_main._get_colbert_for(user_id, user_dir),
     )
 
 

app/main.py

@@ -10,12 +10,17 @@ from fastapi import FastAPI
 
 from app.config import DB_PATH, VEC_DB_PATH, VEC_DIMENSIONS
 from app.services.bm25_index import BM25Index
+from app.services.colbert_index import ColBERTIndex
 
 logger = logging.getLogger("pagepiper")
 
 # Per-user BM25 registry — keyed by user_id; "__local__" for single-user mode
 _bm25_map: dict[str, BM25Index] = {}
 
+# Per-user ColBERT registry — keyed by user_id; index files live under
+# <user_dir>/colbert_index/ (see _get_colbert_for)
+_colbert_map: dict[str, ColBERTIndex] = {}
+
 
 def _get_bm25_for(user_id: str) -> BM25Index:
     if user_id not in _bm25_map:
@@ -23,6 +28,13 @@ def _get_bm25_for(user_id: str) -> BM25Index:
     return _bm25_map[user_id]
 
 
+def _get_colbert_for(user_id: str, user_dir) -> ColBERTIndex:
+    if user_id not in _colbert_map:
+        index_dir = str(user_dir / "colbert_index")
+        _colbert_map[user_id] = ColBERTIndex(index_dir=index_dir)
+    return _colbert_map[user_id]
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     from app.cloud_session import CLOUD_MODE
@@ -38,9 +50,12 @@ async def lifespan(app: FastAPI):
         warn_if_unencrypted(str(DATA_DIR))
     else:
         # In cloud mode, per-user migration and vec schema check run on first request (deps.py).
+        from app.config import DATA_DIR
+
         apply_migrations(DB_PATH)
         check_and_rebuild_vec_schema(VEC_DB_PATH, VEC_DIMENSIONS, DB_PATH)
         _get_bm25_for(LOCAL_USER_ID).mark_dirty()
+        _get_colbert_for(LOCAL_USER_ID, DATA_DIR).mark_dirty()
 
     yield
 

app/services/colbert_index.py

@@ -0,0 +1,137 @@
+# app/services/colbert_index.py
+"""
+ColBERT late-interaction index (pagepiper#8).
+
+Replaces nomic-embed-text bi-encoder + cosine similarity for the semantic
+half of hybrid search with Agent-ModernColBERT (lightonai/Agent-ModernColBERT),
+a late-interaction retriever that keeps per-token embeddings and scores via
+MaxSim at query time — better suited to multi-part rulebook questions than a
+single collapsed query vector.
+
+BSL 1.1 — same BYOK gate as the rest of hybrid search (Retriever only reaches
+this index when an LLM is configured). The model itself runs locally via
+`pylate`, no Ollama call required — the gate matches product tiering, not a
+technical dependency on Ollama.
+
+Mirrors BM25Index's dirty-flag, rebuild-from-SQLite pattern: no separate
+per-shelve indexing step is needed. `mark_dirty()` is called by the same
+callback that already marks BM25 dirty on shelve completion; the next query
+triggers a full rebuild from `page_chunks`.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sqlite3
+import threading
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MODEL = "lightonai/Agent-ModernColBERT"
+
+
+class ColBERTIndex:
+    def __init__(self, index_dir: str, model_name: str | None = None) -> None:
+        self._index_dir = index_dir
+        self._model_name = model_name or os.environ.get("PAGEPIPER_COLBERT_MODEL", _DEFAULT_MODEL)
+        self._model = None
+        self._index = None
+        self._chunks: dict[str, dict] = {}
+        self._dirty = True
+        self._lock = threading.Lock()
+
+    def mark_dirty(self) -> None:
+        """Signal that the index needs rebuilding (call after any document is shelved)."""
+        self._dirty = True
+
+    def _get_model(self):
+        if self._model is None:
+            from pylate import models
+
+            logger.info("Loading ColBERT model %s", self._model_name)
+            self._model = models.ColBERT(model_name_or_path=self._model_name)
+        return self._model
+
+    def ensure_fresh(self, db_path: str) -> None:
+        """Rebuild from SQLite if dirty."""
+        if not self._dirty:
+            return
+        with self._lock:
+            if not self._dirty:
+                return
+            try:
+                conn = sqlite3.connect(db_path)
+                conn.row_factory = sqlite3.Row
+                try:
+                    rows = conn.execute(
+                        "SELECT id, doc_id, page_number, text FROM page_chunks ORDER BY doc_id, page_number"
+                    ).fetchall()
+                finally:
+                    conn.close()
+            except sqlite3.Error as exc:
+                logger.error("ColBERT index rebuild failed: %s", exc)
+                return
+
+            self._chunks = {str(r["id"]): dict(r) for r in rows}
+
+            if not rows:
+                self._index = None
+                self._dirty = False
+                return
+
+            from pylate import indexes
+
+            model = self._get_model()
+            ids = [str(r["id"]) for r in rows]
+            texts = [r["text"] for r in rows]
+            embeddings = model.encode(texts, is_query=False, show_progress_bar=False)
+
+            os.makedirs(self._index_dir, exist_ok=True)
+            index = indexes.Voyager(
+                index_folder=self._index_dir, index_name="colbert", override=True
+            )
+            index.add_documents(documents_ids=ids, documents_embeddings=embeddings)
+            self._index = index
+            self._dirty = False
+            logger.info("ColBERT index rebuilt: %d chunks", len(rows))
+
+    def query(
+        self,
+        query_text: str,
+        top_k: int = 10,
+        doc_ids: list[str] | None = None,
+    ) -> list[dict]:
+        """Search the corpus. Returns results sorted by descending MaxSim score."""
+        if self._index is None:
+            return []
+
+        from pylate import retrieve
+
+        model = self._get_model()
+        query_embeddings = model.encode([query_text], is_query=True, show_progress_bar=False)
+        retriever = retrieve.ColBERT(index=self._index)
+
+        # Oversample when filtering to a doc subset — same pattern as the
+        # sqlite-vec path this replaces (see app/services/retriever.py).
+        k = top_k * 20 if doc_ids else top_k * 2
+        results = retriever.retrieve(queries_embeddings=query_embeddings, k=k)[0]
+
+        hits: list[dict] = []
+        for r in results:
+            chunk = self._chunks.get(str(r["id"]))
+            if not chunk:
+                continue
+            if doc_ids is not None and chunk["doc_id"] not in doc_ids:
+                continue
+            hits.append(
+                {
+                    "chunk_id": chunk["id"],
+                    "doc_id": chunk["doc_id"],
+                    "page_number": chunk["page_number"],
+                    "text": chunk["text"],
+                    "score": r["score"],
+                }
+            )
+            if len(hits) >= top_k:
+                break
+        return hits

app/services/retriever.py

@@ -4,6 +4,11 @@ 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
 
@@ -12,6 +17,7 @@ import sqlite3
 from dataclasses import dataclass
 
 from app.services.bm25_index import BM25Index
+from app.services.colbert_index import ColBERTIndex
 
 logger = logging.getLogger(__name__)
 
@@ -85,8 +91,9 @@ class RetrievedChunk:
 
 
 class Retriever:
-    def __init__(self, bm25: BM25Index) -> None:
+    def __init__(self, bm25: BM25Index, colbert: ColBERTIndex | None = None) -> None:
         self._bm25 = bm25
+        self._colbert = colbert
 
     def hybrid_search(
         self,
@@ -98,14 +105,12 @@ class Retriever:
         llm,  # LLMRouter | None — caller must pass
     ) -> list[RetrievedChunk]:
         """
-        Merge BM25 and semantic results.
+        Merge BM25 and semantic (ColBERT) results.
-        Falls back to BM25-only if llm is None.
+        Falls back to BM25-only if llm is None or no ColBERT index is configured.
         """
-        if llm is None:
+        if llm is None or self._colbert is None:
             return self._bm25_only(query, top_k, doc_ids, db_path)
 
-        from circuitforge_core.vector.sqlite_vec import LocalSQLiteVecStore
-
         self._bm25.ensure_fresh(db_path)
         bm25_hits = {
             r.chunk_id: r
@@ -113,24 +118,15 @@ class Retriever:
         }
 
         try:
-            vec = llm.embed([query])[0]
+            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("Embed failed, falling back to BM25-only: %s", exc)
+            logger.warning("ColBERT retrieval failed, falling back to BM25-only: %s", exc)
             return self._bm25_only(query, top_k, doc_ids, db_path)
-        from app.config import VEC_DIMENSIONS
-        store = LocalSQLiteVecStore(db_path=vec_db_path, table="page_vecs", dimensions=VEC_DIMENSIONS)
 
-        # sqlite-vec applies filter_metadata as a Python post-filter after fetching k
+        # Merge: BM25 hits take priority; ColBERT hits fill in additional results
-        # nearest globally. When the corpus spans many documents and only a subset is
-        # selected, most of those k candidates are from non-target docs and get dropped,
-        # leaving too few vector hits. Oversample heavily and filter in Python instead.
-        if doc_ids:
-            vec_candidates = store.query(vec, top_k=top_k * 20)
-            vec_hits = [h for h in vec_candidates if h.metadata.get("doc_id") in doc_ids]
-        else:
-            vec_hits = store.query(vec, top_k=top_k * 2)
-
-        # Merge: BM25 hits take priority; vector hits fill in additional results
         merged: dict[str, RetrievedChunk] = {}
         for cid, r in bm25_hits.items():
             merged[cid] = RetrievedChunk(
@@ -141,33 +137,37 @@ class Retriever:
                 bm25_score=r.score,
                 vector_score=None,
             )
-        for vh in vec_hits:
+
-            # _chunks is the loaded list of dicts from BM25Index; no public accessor exists
+        # ColBERT MaxSim scores are unbounded (roughly num_query_tokens *
-            text = next((c["text"] for c in self._bm25._chunks if c["id"] == vh.entry_id), "")
+        # max_per_token_similarity), unlike BM25's already-comparable range —
-            if vh.entry_id in merged:
+        # min-max normalize within this result batch before combining.
-                existing = merged[vh.entry_id]
+        max_colbert_score = max((h["score"] for h in colbert_hits), default=0.0)
-                merged[vh.entry_id] = RetrievedChunk(
+        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=vh.score,
+                    vector_score=norm_score,
                 )
             else:
-                merged[vh.entry_id] = RetrievedChunk(
+                merged[cid] = RetrievedChunk(
-                    chunk_id=vh.entry_id,
+                    chunk_id=cid,
-                    doc_id=vh.metadata.get("doc_id", ""),
+                    doc_id=h["doc_id"],
-                    page_number=int(vh.metadata.get("page_number", 0)),
+                    page_number=h["page_number"],
-                    text=text,
+                    text=h["text"],
                     bm25_score=0.0,
-                    vector_score=vh.score,
+                    vector_score=norm_score,
                 )
 
         def _combined(r: RetrievedChunk) -> float:
             bm25 = r.bm25_score
-            # sqlite-vec returns L2 distance (lower=better); invert to [0,1] higher-is-better
+            vec = r.vector_score if r.vector_score is not None else 0.0
-            vec = (1.0 / (1.0 + 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)

docs/reference/environment-variables.md

@@ -14,10 +14,18 @@ Copy `.env.example` to `.env` and configure as needed.
 
 | Variable | Default | Description |
 |----------|---------|-------------|
-| `PAGEPIPER_OLLAMA_URL` | _(unset)_ | Ollama base URL, e.g. `http://localhost:11434`. Enables hybrid search and chat. |
+| `PAGEPIPER_OLLAMA_URL` | _(unset)_ | Ollama base URL, e.g. `http://localhost:11434`. Enables hybrid search and chat (BYOK gate — see below). |
-| `PAGEPIPER_EMBED_MODEL` | `nomic-embed-text` | Ollama embedding model |
-| `PAGEPIPER_EMBED_DIMS` | `1024` | Embedding dimensions (must match the model) |
 | `PAGEPIPER_CHAT_MODEL` | `mistral:7b` | Ollama chat/completion model |
+| `PAGEPIPER_EMBED_MODEL` | `nomic-embed-text` | Ollama embedding model — used for shelve-time embeddings only (`page_vecs`), not for search retrieval (see ColBERT below) |
+| `PAGEPIPER_EMBED_DIMS` | `1024` | Embedding dimensions (must match `PAGEPIPER_EMBED_MODEL`) |
+
+## Semantic search (ColBERT)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `PAGEPIPER_COLBERT_MODEL` | `lightonai/Agent-ModernColBERT` | HuggingFace model used for hybrid search's semantic half — a late-interaction retriever, runs locally via `pylate`, no Ollama call required. Gated behind the same BYOK check as the rest of hybrid search (`PAGEPIPER_OLLAMA_URL` or `CF_ORCH_URL`/`GPU_SERVER_URL` must be set). |
+
+**Note:** `page_vecs` (the sqlite-vec table populated at shelve time using `PAGEPIPER_EMBED_MODEL`) is no longer read by search or chat — retrieval was switched to the ColBERT index above (pagepiper#8). Shelving still computes and stores those embeddings for now; removing that redundant work is tracked as a follow-up.
 
 ## GPU server / cf-orch (managed deployments)
 

docs/user-guide/search.md

@@ -21,4 +21,4 @@ BM25 (Best Match 25) ranks pages by term frequency weighted against how rare eac
 
 ## Hybrid search (requires Ollama)
 
-When Ollama is configured, the Chat endpoint uses hybrid search behind the scenes: BM25 results are merged with semantic vector results using a 50/50 score blend. The Search page always uses BM25 only.
+When Ollama is configured, the Chat endpoint uses hybrid search behind the scenes: BM25 results are merged with Agent-ModernColBERT late-interaction results using a 50/50 score blend. The Search page always uses BM25 only.

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "python-docx>=1.0",
     "odfpy>=1.4",
     "openpyxl>=3.1",
+    "pylate[voyager]>=1.6",
 ]
 
 [tool.setuptools.packages.find]

tests/conftest.py

@@ -31,6 +31,7 @@ def client(test_db, tmp_path, monkeypatch):
     from app.deps import UserCtx, get_db, get_user_ctx
     from app.main import app
     from app.services.bm25_index import BM25Index
+    from app.services.colbert_index import ColBERTIndex
     from app.startup import apply_migrations, check_and_rebuild_vec_schema
 
     monkeypatch.setattr(_main_module, "_apply_migrations", lambda: None, raising=False)
@@ -43,6 +44,7 @@ def client(test_db, tmp_path, monkeypatch):
 
     test_bm25 = BM25Index()
     test_bm25.mark_dirty()
+    test_colbert = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
 
     def override_user_ctx():
         return UserCtx(
@@ -52,6 +54,7 @@ def client(test_db, tmp_path, monkeypatch):
             data_dir=Path(tmp_path),
             watch_dir=Path(tmp_path) / "books",
             bm25=test_bm25,
+            colbert=test_colbert,
         )
 
     def override_db():

tests/test_colbert_index.py

@@ -0,0 +1,168 @@
+# tests/test_colbert_index.py
+"""Tests for app.services.colbert_index.
+
+pylate is NOT installed in the dev/test env by design (see cf-sysadmin skill's
+"Known Gotchas" — installing it directly into the shared `cf` conda env broke
+several other services' pinned torch/transformers versions on 2026-07-10).
+These tests inject fake `pylate`/`pylate.models`/`pylate.indexes`/`pylate.retrieve`
+modules via sys.modules so ColBERTIndex's lazy imports resolve to mocks without
+pylate ever needing to be installed here.
+"""
+from __future__ import annotations
+
+import sqlite3
+import sys
+import types
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from app.services.colbert_index import ColBERTIndex
+
+
+@pytest.fixture
+def fake_pylate(monkeypatch):
+    fake_models_mod = types.ModuleType("pylate.models")
+    fake_indexes_mod = types.ModuleType("pylate.indexes")
+    fake_retrieve_mod = types.ModuleType("pylate.retrieve")
+    fake_pylate_mod = types.ModuleType("pylate")
+    fake_pylate_mod.models = fake_models_mod
+    fake_pylate_mod.indexes = fake_indexes_mod
+    fake_pylate_mod.retrieve = fake_retrieve_mod
+
+    mock_model = MagicMock()
+    mock_model.encode.return_value = [[0.1, 0.2], [0.3, 0.4]]
+    fake_models_mod.ColBERT = MagicMock(return_value=mock_model)
+
+    mock_index = MagicMock()
+    fake_indexes_mod.Voyager = MagicMock(return_value=mock_index)
+
+    mock_retriever = MagicMock()
+    fake_retrieve_mod.ColBERT = MagicMock(return_value=mock_retriever)
+
+    monkeypatch.setitem(sys.modules, "pylate", fake_pylate_mod)
+    monkeypatch.setitem(sys.modules, "pylate.models", fake_models_mod)
+    monkeypatch.setitem(sys.modules, "pylate.indexes", fake_indexes_mod)
+    monkeypatch.setitem(sys.modules, "pylate.retrieve", fake_retrieve_mod)
+
+    return types.SimpleNamespace(
+        model_cls=fake_models_mod.ColBERT,
+        model=mock_model,
+        index_cls=fake_indexes_mod.Voyager,
+        index=mock_index,
+        retriever_cls=fake_retrieve_mod.ColBERT,
+        retriever=mock_retriever,
+    )
+
+
+@pytest.fixture
+def seeded_db(tmp_path) -> str:
+    db_path = str(tmp_path / "test.db")
+    schema = Path("migrations/001_initial_schema.sql").read_text()
+    conn = sqlite3.connect(db_path)
+    conn.executescript(schema)
+    conn.execute(
+        "INSERT INTO documents(id, title, file_path, status) VALUES ('d1','Test','test.pdf','ready')"
+    )
+    conn.execute(
+        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
+        "VALUES ('c1','d1',1,'Setting the IP on the AVC-X','text',6)"
+    )
+    conn.execute(
+        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
+        "VALUES ('c2','d1',2,'Filter cartridge replacement steps','text',5)"
+    )
+    conn.commit()
+    conn.close()
+    return db_path
+
+
+def test_ensure_fresh_builds_index_from_sqlite(fake_pylate, seeded_db, tmp_path):
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(seeded_db)
+
+    fake_pylate.model.encode.assert_called_once()
+    call_args = fake_pylate.model.encode.call_args
+    assert set(call_args[0][0]) == {"Setting the IP on the AVC-X", "Filter cartridge replacement steps"}
+    assert call_args[1]["is_query"] is False
+
+    fake_pylate.index.add_documents.assert_called_once()
+    add_kwargs = fake_pylate.index.add_documents.call_args[1]
+    assert set(add_kwargs["documents_ids"]) == {"c1", "c2"}
+
+
+def test_ensure_fresh_skips_rebuild_when_not_dirty(fake_pylate, seeded_db, tmp_path):
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(seeded_db)
+    idx.ensure_fresh(seeded_db)
+
+    fake_pylate.model.encode.assert_called_once()
+
+
+def test_mark_dirty_forces_rebuild(fake_pylate, seeded_db, tmp_path):
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(seeded_db)
+    idx.mark_dirty()
+    idx.ensure_fresh(seeded_db)
+
+    assert fake_pylate.model.encode.call_count == 2
+
+
+def test_ensure_fresh_with_empty_corpus_leaves_index_none(fake_pylate, tmp_path):
+    db_path = str(tmp_path / "empty.db")
+    schema = Path("migrations/001_initial_schema.sql").read_text()
+    conn = sqlite3.connect(db_path)
+    conn.executescript(schema)
+    conn.commit()
+    conn.close()
+
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(db_path)
+
+    fake_pylate.index_cls.assert_not_called()
+    assert idx.query("anything") == []
+
+
+def test_query_maps_results_back_to_chunks(fake_pylate, seeded_db, tmp_path):
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(seeded_db)
+
+    fake_pylate.retriever.retrieve.return_value = [
+        [{"id": "c1", "score": 13.5}, {"id": "c2", "score": 9.2}]
+    ]
+
+    results = idx.query("how do I set the IP on the AVC-X", top_k=10)
+
+    assert len(results) == 2
+    assert results[0]["chunk_id"] == "c1"
+    assert results[0]["doc_id"] == "d1"
+    assert results[0]["score"] == 13.5
+    assert results[1]["chunk_id"] == "c2"
+
+
+def test_query_filters_by_doc_ids(fake_pylate, seeded_db, tmp_path):
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(seeded_db)
+
+    fake_pylate.retriever.retrieve.return_value = [
+        [{"id": "c1", "score": 13.5}, {"id": "c2", "score": 9.2}]
+    ]
+
+    results = idx.query("query", top_k=10, doc_ids=["other-doc"])
+
+    assert results == []
+
+
+def test_query_respects_top_k(fake_pylate, seeded_db, tmp_path):
+    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
+    idx.ensure_fresh(seeded_db)
+
+    fake_pylate.retriever.retrieve.return_value = [
+        [{"id": "c1", "score": 13.5}, {"id": "c2", "score": 9.2}]
+    ]
+
+    results = idx.query("query", top_k=1)
+
+    assert len(results) == 1
+    assert results[0]["chunk_id"] == "c1"

tests/test_retriever.py

@@ -0,0 +1,113 @@
+# tests/test_retriever.py
+"""Tests for app.services.retriever.Retriever.hybrid_search — the BM25 + ColBERT merge."""
+from __future__ import annotations
+
+import sqlite3
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from app.services.bm25_index import BM25Index
+from app.services.retriever import Retriever
+
+
+@pytest.fixture
+def seeded_db(tmp_path) -> str:
+    db_path = str(tmp_path / "test.db")
+    schema = Path("migrations/001_initial_schema.sql").read_text()
+    conn = sqlite3.connect(db_path)
+    conn.executescript(schema)
+    conn.execute(
+        "INSERT INTO documents(id, title, file_path, status) VALUES ('d1','Test','test.pdf','ready')"
+    )
+    conn.execute(
+        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
+        "VALUES ('c1','d1',1,'Setting the IP address on the AVC-X device','text',7)"
+    )
+    conn.execute(
+        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
+        "VALUES ('c2','d1',2,'Filter cartridge replacement procedure','text',4)"
+    )
+    # Third, unrelated chunk — with only 2 chunks total, a term appearing in
+    # exactly one of them gets an Okapi BM25 IDF of exactly log(1.0) == 0
+    # (log((N-n+0.5)/(n+0.5)) with N=2, n=1), silently zeroing every score.
+    # A third chunk dilutes N enough for real term-overlap scores to surface.
+    conn.execute(
+        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
+        "VALUES ('c3','d1',3,'Warranty terms and annual maintenance schedule','text',5)"
+    )
+    conn.commit()
+    conn.close()
+    return db_path
+
+
+def _seeded_bm25() -> BM25Index:
+    idx = BM25Index()
+    idx._dirty = True
+    return idx
+
+
+def test_hybrid_search_falls_back_to_bm25_only_without_llm(seeded_db):
+    retriever = Retriever(_seeded_bm25(), colbert=MagicMock())
+    results = retriever.hybrid_search(
+        query="IP address", top_k=5, doc_ids=None,
+        db_path=seeded_db, vec_db_path="unused", llm=None,
+    )
+    assert any(r.chunk_id == "c1" for r in results)
+
+
+def test_hybrid_search_falls_back_to_bm25_only_without_colbert(seeded_db):
+    retriever = Retriever(_seeded_bm25(), colbert=None)
+    results = retriever.hybrid_search(
+        query="IP address", top_k=5, doc_ids=None,
+        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
+    )
+    assert any(r.chunk_id == "c1" for r in results)
+
+
+def test_hybrid_search_merges_bm25_and_colbert_hits(seeded_db):
+    fake_colbert = MagicMock()
+    fake_colbert.query.return_value = [
+        {"chunk_id": "c1", "doc_id": "d1", "page_number": 1, "text": "Setting the IP address on the AVC-X device", "score": 15.0},
+        {"chunk_id": "c2", "doc_id": "d1", "page_number": 2, "text": "Filter cartridge replacement procedure", "score": 5.0},
+    ]
+
+    retriever = Retriever(_seeded_bm25(), colbert=fake_colbert)
+    results = retriever.hybrid_search(
+        query="IP address AVC-X", top_k=5, doc_ids=None,
+        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
+    )
+
+    fake_colbert.ensure_fresh.assert_called_once_with(seeded_db)
+    result_ids = {r.chunk_id for r in results}
+    assert "c1" in result_ids
+    c1 = next(r for r in results if r.chunk_id == "c1")
+    assert c1.vector_score == 1.0  # highest colbert score, normalized to max
+
+
+def test_hybrid_search_falls_back_when_colbert_raises(seeded_db):
+    fake_colbert = MagicMock()
+    fake_colbert.query.side_effect = RuntimeError("model not loaded")
+
+    retriever = Retriever(_seeded_bm25(), colbert=fake_colbert)
+    results = retriever.hybrid_search(
+        query="IP address", top_k=5, doc_ids=None,
+        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
+    )
+
+    assert any(r.chunk_id == "c1" for r in results)
+
+
+def test_hybrid_search_discards_pure_noise(seeded_db):
+    fake_colbert = MagicMock()
+    fake_colbert.query.return_value = []
+
+    retriever = Retriever(_seeded_bm25(), colbert=fake_colbert)
+    results = retriever.hybrid_search(
+        query="completely unrelated gibberish xyzzy",
+        top_k=5, doc_ids=None,
+        db_path=seeded_db, vec_db_path="unused", llm=MagicMock(),
+    )
+
+    assert results == []