pyr0ball
89a58ec9b0
Bi-encoder embeddings collapse a whole query into one vector, losing multi-part reasoning structure — queries like "the procedure for setting an IP on an AVC-X" or "what is the action economy for a fighter casting a spell while prone" lose nuance. Agent-ModernColBERT is a late-interaction retriever: per-token embeddings, scored via MaxSim at query time, built specifically for agentic/multi-hop queries. Implements Option A from the issue (in-process, via `pylate`) rather than Option B (managed cf-orch service) — cf-orch already has `agent-moderncolbert` registered in model_registry.yaml with a `pagepiper/retrieve` assignment in assignments.yaml pointing at it and referencing this issue directly, someone had already pre-wired that side. - app/services/colbert_index.py: new ColBERTIndex class, mirrors BM25Index's dirty-flag/rebuild-from-SQLite pattern exactly — no separate per-shelve indexing step needed, just mark_dirty() on the same callback that already marks BM25 dirty. - app/services/retriever.py: hybrid_search's semantic half now merges BM25 with ColBERT MaxSim scores (min-max normalized per-batch, since MaxSim is unbounded unlike the old sqlite-vec L2-distance path) instead of Ollama-embed + sqlite-vec cosine. BM25 merge/rank/per-doc-cap/ adjacent-chunk-window logic is unchanged. - app/main.py / app/deps.py: per-user ColBERTIndex registry, same pattern as the existing per-user BM25Index registry. - Existing BYOK tier gate preserved exactly (llm is None check) — this is a retrieval-technology swap, not a tier/licensing change. The ColBERT model runs locally via pylate with no Ollama dependency, but gating still follows product tiering. - 12 new tests. pylate is intentionally NOT installed in the dev/test env — see the cf-sysadmin skill's "Known Gotchas" for why (installing it directly into the shared `cf` conda env broke several other services' torch/transformers pins on 2026-07-10). Tests inject fake pylate modules via sys.modules instead. Known follow-up (not addressed here): shelve scripts still compute and store Ollama embeddings into `page_vecs` at shelve time — that table is no longer read by search/chat now that retrieval uses the ColBERT index. Removing the now-redundant embedding step is separate cleanup. Closes: #8
113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
# 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 == []
|