steve
6dfc13be1d
Implements five of the six enhancements from docs/kb-enhancements-proposal.htm, closing the retrieval-quality gap identified in the qmd review. - Cross-encoder reranking: new kb/reranker.py loads an optional reranking model at startup (KB_RERANK_ENABLED, KB_RERANKER_MODEL, KB_RERANK_CANDIDATES). Search degrades gracefully to plain hybrid retrieval when the model is absent. Exposed via a "rerank" block in /status, a rerank flag on search, and --no-rerank in the CLI. - RRF rank fusion: FTS and vector lists now merge by reciprocal rank fusion with a top-rank bonus, replacing the old score blend. Scores are comparable across queries. - Bench harness and explain traces: kb bench runs a query fixture against each backend and reports precision@k, recall and MRR. --explain returns a per-result score breakdown. - Tag context descriptions: tags carry an optional one-line description (kb tag-describe), returned as tag_contexts with search results. Adds a tags.description column migration. - Structured data ingestion: .json/.yaml/.toml files ingest as text via the new "data" doc type, pretty-printing minified JSON before chunking. Query expansion (proposal item 5) is deliberately left out pending bench results. Requires engine v3.3.0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
274 lines
9.6 KiB
Python
274 lines
9.6 KiB
Python
"""Tests for hybrid search: explain traces, document_id, and RRF merging."""
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from kb.database import (
|
|
get_connection,
|
|
init_schema,
|
|
insert_chunk,
|
|
insert_document,
|
|
insert_embedding,
|
|
tag_document,
|
|
)
|
|
from kb.search import _blend_rerank, _rank_by_score, _rrf_merge, hybrid_search
|
|
|
|
DIM = 4
|
|
|
|
# Chunk vectors are axis-aligned so we can steer vector ranking exactly:
|
|
# a query of [1,0,0,0] has distance 0 to chunk A, sqrt(2) to chunk B.
|
|
VEC_A = [1.0, 0.0, 0.0, 0.0]
|
|
VEC_B = [0.0, 1.0, 0.0, 0.0]
|
|
QUERY_VEC = [1.0, 0.0, 0.0, 0.0]
|
|
|
|
|
|
class _Cfg:
|
|
search_threshold = 0.0
|
|
rerank_enabled = False
|
|
reranker_model = "test-reranker"
|
|
rerank_candidates = 40
|
|
|
|
|
|
class _Db:
|
|
def __init__(self, conn, ids):
|
|
self.conn = conn
|
|
self.ids = ids
|
|
|
|
|
|
@pytest.fixture
|
|
def db(tmp_path, monkeypatch):
|
|
"""Real schema (FTS5 + sqlite-vec) with two docs and stubbed embeddings."""
|
|
fake = types.ModuleType("kb.embeddings")
|
|
fake.embed_texts = lambda texts: [QUERY_VEC for _ in texts]
|
|
monkeypatch.setitem(sys.modules, "kb.embeddings", fake)
|
|
|
|
conn = get_connection(str(tmp_path / "kb.db"))
|
|
init_schema(conn, embedding_dim=DIM)
|
|
|
|
doc_a = insert_document(conn, "Alpha doc", "/src/a.md", "hash-a", "markdown")
|
|
doc_b = insert_document(conn, "Bravo doc", "/src/b.md", "hash-b", "markdown")
|
|
chunk_a = insert_chunk(conn, doc_a, 0, "alpha network switch configuration")
|
|
chunk_b = insert_chunk(conn, doc_b, 0, "bravo unrelated cooking recipe")
|
|
insert_embedding(conn, chunk_a, VEC_A)
|
|
insert_embedding(conn, chunk_b, VEC_B)
|
|
tag_document(conn, doc_a, ["ops"])
|
|
|
|
ids = {"doc_a": doc_a, "doc_b": doc_b, "chunk_a": chunk_a, "chunk_b": chunk_b}
|
|
yield _Db(conn, ids)
|
|
conn.close()
|
|
|
|
|
|
def test_results_include_document_id(db):
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg())
|
|
assert result["results"], "expected at least one hit"
|
|
top_hit = result["results"][0]
|
|
assert top_hit["document_id"] == db.ids["doc_a"]
|
|
assert top_hit["chunk_id"] == db.ids["chunk_a"]
|
|
|
|
|
|
def test_explain_absent_by_default(db):
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg())
|
|
assert all("explain" not in r for r in result["results"])
|
|
|
|
|
|
def test_explain_hybrid_breakdown(db):
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg(), explain=True)
|
|
top_hit = result["results"][0]
|
|
exp = top_hit["explain"]
|
|
|
|
# Chunk A is rank 1 in both arms: FTS matches "alpha"/"switch", vector
|
|
# distance is 0 (similarity 1.0).
|
|
assert exp["fts_rank"] == 1
|
|
assert exp["vec_rank"] == 1
|
|
assert exp["fts_score"] > 0
|
|
assert exp["vec_score"] == pytest.approx(1.0)
|
|
assert exp["rrf_fts"] == pytest.approx(1.0 / 61, abs=1e-6)
|
|
assert exp["rrf_vec"] == pytest.approx(1.0 / 61, abs=1e-6)
|
|
assert exp["final_score"] == pytest.approx(
|
|
exp["rrf_fts"] + exp["rrf_vec"] + exp["bonus"], abs=1e-5
|
|
)
|
|
assert exp["final_score"] == pytest.approx(top_hit["score"], abs=1e-5)
|
|
|
|
|
|
def test_explain_single_arm_when_vec_misses(db):
|
|
"""A chunk found only by vector search has null FTS fields."""
|
|
result = hybrid_search(db.conn, "zzz-no-fts-match", _Cfg(), explain=True)
|
|
for r in result["results"]:
|
|
exp = r["explain"]
|
|
assert exp["fts_score"] is None
|
|
assert exp["fts_rank"] is None
|
|
assert exp["rrf_fts"] is None
|
|
assert exp["vec_rank"] is not None
|
|
|
|
|
|
def test_explain_fts_only_shape(db):
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg(), fts_only=True, explain=True)
|
|
top_hit = result["results"][0]
|
|
exp = top_hit["explain"]
|
|
assert exp["fts_rank"] == 1
|
|
assert exp["final_score"] == pytest.approx(exp["fts_score"])
|
|
assert "vec_score" not in exp
|
|
|
|
|
|
def test_explain_vec_only_shape(db):
|
|
result = hybrid_search(db.conn, "anything", _Cfg(), vec_only=True, explain=True)
|
|
top_hit = result["results"][0]
|
|
exp = top_hit["explain"]
|
|
assert exp["vec_rank"] == 1
|
|
assert exp["final_score"] == pytest.approx(exp["vec_score"])
|
|
assert "fts_score" not in exp
|
|
|
|
|
|
def test_rrf_merge_arithmetic():
|
|
fts = {1: 10.0, 2: 5.0}
|
|
vec = {2: 0.9, 3: 0.8}
|
|
scores, details = _rrf_merge(fts, vec)
|
|
by_id = dict(scores)
|
|
|
|
# fts rank 1 → 1/61 + 0.05 bonus
|
|
assert by_id[1] == pytest.approx(1 / 61 + 0.05)
|
|
# fts rank 2 (+0.02) and vec rank 1 (+0.05) — bonuses stack across arms
|
|
assert by_id[2] == pytest.approx(1 / 62 + 1 / 61 + 0.07)
|
|
# vec rank 2 → 1/62 + 0.02
|
|
assert by_id[3] == pytest.approx(1 / 62 + 0.02)
|
|
# Chunk 2 appears in both arms, so it must win.
|
|
assert scores[0][0] == 2
|
|
|
|
assert details[2]["fts_rank"] == 2
|
|
assert details[2]["vec_rank"] == 1
|
|
assert details[2]["bonus"] == pytest.approx(0.07)
|
|
assert details[1]["vec_rank"] is None
|
|
assert details[3]["rrf_fts"] is None
|
|
|
|
|
|
def test_top_rank_bonus_tiers():
|
|
from kb.search import _top_rank_bonus
|
|
|
|
assert _top_rank_bonus(1) == 0.05
|
|
assert _top_rank_bonus(2) == 0.02
|
|
assert _top_rank_bonus(3) == 0.02
|
|
assert _top_rank_bonus(4) == 0.0
|
|
assert _top_rank_bonus(None) == 0.0
|
|
|
|
|
|
def test_bonus_preserves_top_of_arm():
|
|
"""A chunk at rank 1 of one arm beats a chunk at mid-rank in both arms."""
|
|
fts = {10: 100.0, 11: 90.0, 12: 80.0, 13: 70.0, 14: 60.0}
|
|
vec = {20: 0.9, 11: 0.8, 12: 0.7, 13: 0.6, 14: 0.5}
|
|
scores, _ = _rrf_merge(fts, vec)
|
|
order = [cid for cid, _ in scores]
|
|
# Chunk 10 (fts #1, absent from vec): 1/61 + 0.05 ≈ 0.0664.
|
|
# Chunk 13 (rank 4 in fts, rank 4 in vec): 2/64 ≈ 0.031, no bonus.
|
|
assert order.index(10) < order.index(13)
|
|
|
|
|
|
def test_rank_by_score():
|
|
assert _rank_by_score({7: 0.5, 8: 0.9, 9: 0.1}) == {8: 1, 7: 2, 9: 3}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reranking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _enable_fake_reranker(monkeypatch, score_fn):
|
|
from kb import reranker
|
|
|
|
monkeypatch.setattr(reranker, "is_available", lambda: True)
|
|
monkeypatch.setattr(reranker, "rerank_scores", score_fn)
|
|
|
|
|
|
def test_rerank_flags_response_and_explain(db, monkeypatch):
|
|
_enable_fake_reranker(monkeypatch, lambda q, texts: [0.9] * len(texts))
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg(), explain=True, rerank=True)
|
|
|
|
assert result["reranked"] is True
|
|
top_hit = result["results"][0]
|
|
exp = top_hit["explain"]
|
|
assert exp["rerank_score"] == pytest.approx(0.9)
|
|
assert exp["pre_rerank_rank"] == 1
|
|
assert exp["blend_weight"] == 0.75
|
|
assert exp["final_score"] == pytest.approx(top_hit["score"], abs=1e-5)
|
|
# rank 1: 75% retrieval (norm=1.0 for the top candidate) + 25% rerank
|
|
assert top_hit["score"] == pytest.approx(0.75 * 1.0 + 0.25 * 0.9, abs=1e-5)
|
|
|
|
|
|
def test_rerank_can_reorder(db, monkeypatch):
|
|
"""A strong rerank score rescues a lower-retrieval-ranked chunk."""
|
|
|
|
def favour_chunk_b(query, texts):
|
|
return [1.0 if "cooking" in t else 0.0 for t in texts]
|
|
|
|
_enable_fake_reranker(monkeypatch, favour_chunk_b)
|
|
# Neutral-ish query: both chunks retrieved, chunk A ranked first.
|
|
result = hybrid_search(db.conn, "alpha cooking", _Cfg(), rerank=True)
|
|
assert result["reranked"] is True
|
|
assert len(result["results"]) == 2
|
|
|
|
|
|
def test_rerank_false_bypasses(db, monkeypatch):
|
|
called = []
|
|
_enable_fake_reranker(monkeypatch, lambda q, t: called.append(1) or [0.5] * len(t))
|
|
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg(), rerank=False)
|
|
assert result["reranked"] is False
|
|
assert not called
|
|
|
|
|
|
def test_rerank_unavailable_degrades_gracefully(db):
|
|
# No reranker loaded: rerank=True must not error.
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg(), rerank=True, explain=True)
|
|
assert result["reranked"] is False
|
|
assert "rerank_score" not in result["results"][0]["explain"]
|
|
|
|
|
|
def test_rerank_skipped_for_single_arm(db, monkeypatch):
|
|
called = []
|
|
_enable_fake_reranker(monkeypatch, lambda q, t: called.append(1) or [0.5] * len(t))
|
|
|
|
result = hybrid_search(db.conn, "alpha switch", _Cfg(), fts_only=True, rerank=True)
|
|
assert result["reranked"] is False
|
|
assert not called
|
|
|
|
|
|
def test_rerank_default_follows_cfg(db, monkeypatch):
|
|
_enable_fake_reranker(monkeypatch, lambda q, texts: [0.5] * len(texts))
|
|
|
|
class _RerankCfg(_Cfg):
|
|
rerank_enabled = True
|
|
|
|
result = hybrid_search(db.conn, "alpha switch", _RerankCfg())
|
|
assert result["reranked"] is True
|
|
|
|
|
|
def test_blend_rerank_weights_by_position():
|
|
# 12 candidates, retrieval scores 12 down to 1 → norms 1.0 down to 0.0.
|
|
candidates = [(cid, float(12 - i)) for i, cid in enumerate(range(100, 112))]
|
|
details = {cid: {} for cid, _ in candidates}
|
|
rr = [1.0] * len(candidates)
|
|
|
|
blended = dict(_blend_rerank(candidates, rr, details))
|
|
|
|
assert details[100]["blend_weight"] == 0.75 # rank 1
|
|
assert details[102]["blend_weight"] == 0.75 # rank 3
|
|
assert details[103]["blend_weight"] == 0.60 # rank 4
|
|
assert details[109]["blend_weight"] == 0.60 # rank 10
|
|
assert details[110]["blend_weight"] == 0.40 # rank 11
|
|
|
|
# rank 1: norm 1.0 → 0.75*1.0 + 0.25*1.0 = 1.0
|
|
assert blended[100] == pytest.approx(1.0)
|
|
# rank 4: norm 8/11 → 0.6*(8/11) + 0.4*1.0
|
|
assert blended[103] == pytest.approx(0.6 * (8 / 11) + 0.4)
|
|
# rank 11: norm 1/11 → 0.4*(1/11) + 0.6*1.0
|
|
assert blended[110] == pytest.approx(0.4 * (1 / 11) + 0.6)
|
|
|
|
|
|
def test_blend_rerank_constant_retrieval_scores():
|
|
"""Zero span (all candidates same retrieval score) must not divide by zero."""
|
|
candidates = [(1, 0.5), (2, 0.5)]
|
|
details = {1: {}, 2: {}}
|
|
blended = dict(_blend_rerank(candidates, [0.2, 0.8], details))
|
|
assert blended[1] == pytest.approx(0.75 * 1.0 + 0.25 * 0.2)
|
|
assert blended[2] == pytest.approx(0.75 * 1.0 + 0.25 * 0.8)
|