diff --git a/README.md b/README.md index 7bd77f9..d0eda93 100644 --- a/README.md +++ b/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 | diff --git a/app/api/chat.py b/app/api/chat.py index 9226578..34588cf 100644 --- a/app/api/chat.py +++ b/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, diff --git a/app/api/library.py b/app/api/library.py index 5a6afb9..1eec5a8 100644 --- a/app/api/library.py +++ b/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=?", diff --git a/app/deps.py b/app/deps.py index 4826f59..f21ae5f 100644 --- a/app/deps.py +++ b/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), ) diff --git a/app/main.py b/app/main.py index 7a82f95..61f8b20 100644 --- a/app/main.py +++ b/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 +# /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 diff --git a/app/services/colbert_index.py b/app/services/colbert_index.py new file mode 100644 index 0000000..b8991b2 --- /dev/null +++ b/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 diff --git a/app/services/retriever.py b/app/services/retriever.py index 9616e3c..04d58b5 100644 --- a/app/services/retriever.py +++ b/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. - Falls back to BM25-only if llm is None. + 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: + 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 - # 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 + # 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( @@ -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 - text = next((c["text"] for c in self._bm25._chunks if c["id"] == vh.entry_id), "") - if vh.entry_id in merged: - existing = merged[vh.entry_id] - merged[vh.entry_id] = RetrievedChunk( + + # 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=vh.score, + vector_score=norm_score, ) else: - merged[vh.entry_id] = RetrievedChunk( - chunk_id=vh.entry_id, - doc_id=vh.metadata.get("doc_id", ""), - page_number=int(vh.metadata.get("page_number", 0)), - text=text, + 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=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 = (1.0 / (1.0 + r.vector_score)) if r.vector_score is not None else 0.0 + 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) diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md index 45e4bac..cd2697c 100644 --- a/docs/reference/environment-variables.md +++ b/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_EMBED_MODEL` | `nomic-embed-text` | Ollama embedding model | -| `PAGEPIPER_EMBED_DIMS` | `1024` | Embedding dimensions (must match the model) | +| `PAGEPIPER_OLLAMA_URL` | _(unset)_ | Ollama base URL, e.g. `http://localhost:11434`. Enables hybrid search and chat (BYOK gate — see below). | | `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) diff --git a/docs/user-guide/search.md b/docs/user-guide/search.md index e8f945f..50f1c05 100644 --- a/docs/user-guide/search.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 56362a5..d841540 100644 --- a/pyproject.toml +++ b/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] diff --git a/tests/conftest.py b/tests/conftest.py index 65fe8de..8b551c7 100644 --- a/tests/conftest.py +++ b/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(): diff --git a/tests/test_colbert_index.py b/tests/test_colbert_index.py new file mode 100644 index 0000000..6e207ab --- /dev/null +++ b/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" diff --git a/tests/test_retriever.py b/tests/test_retriever.py new file mode 100644 index 0000000..16b6d3f --- /dev/null +++ b/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 == []