diff --git a/src/vouch/context.py b/src/vouch/context.py index fa7e3f16..64ec2e3e 100644 --- a/src/vouch/context.py +++ b/src/vouch/context.py @@ -44,6 +44,7 @@ ContextItemKind = Literal["claim", "page", "entity", "relation", "source"] _VALID_BACKENDS = ("auto", "hybrid", "embedding", "fts5", "substring") +_RERANKER_CACHE: Any | None = None def _configured_backend(store: KBStore) -> str: @@ -74,6 +75,90 @@ def _configured_backend(store: KBStore) -> str: return "auto" +def _configured_rerank(store: KBStore, *, limit: int) -> tuple[bool, int]: + """Resolve the optional context rerank stage from config.yaml. + + Defaults to disabled so existing KBs keep byte-identical ordering unless + they opt in with ``retrieval.rerank.enabled: true``. ``top_k`` is the + window to reorder; by default it is the caller's context limit. + """ + try: + loaded = yaml.safe_load(store.config_path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError): + return False, limit + if not isinstance(loaded, dict): + return False, limit + retrieval = loaded.get("retrieval") + if not isinstance(retrieval, dict): + return False, limit + rerank = retrieval.get("rerank") + if not isinstance(rerank, dict): + return False, limit + + enabled = rerank.get("enabled", False) + enabled = enabled if isinstance(enabled, bool) else False + + top_k = rerank.get("top_k", limit) + top_k = ( + top_k + if isinstance(top_k, int) and not isinstance(top_k, bool) and top_k > 0 + else limit + ) + return enabled, top_k + + +def _default_reranker_cached() -> Any: + global _RERANKER_CACHE + if _RERANKER_CACHE is None: + from .embeddings.rerank import default_reranker + + _RERANKER_CACHE = default_reranker() + return _RERANKER_CACHE + + +def _maybe_rerank( + store: KBStore, + *, + query: str, + hits: list[tuple[str, str, str, float]], + limit: int, +) -> list[tuple[str, str, str, float]]: + enabled, top_k = _configured_rerank(store, limit=limit) + if not enabled or not hits or top_k <= 0: + return hits + + window_size = min(top_k, len(hits)) + window = hits[:window_size] + try: + from .embeddings.rerank import rerank as do_rerank + + reranked = do_rerank( + query=query, + hits=window, + reranker=_default_reranker_cached(), + top_k=window_size, + ) + except ImportError: + return hits + + # Keep reranking as an ordering-only stage: the configured window may move, + # but it must not add/drop artifacts from the already-scoped result set. + original_by_key = {(hit[0], hit[1]): hit for hit in window} + seen: set[tuple[str, str]] = set() + ordered: list[tuple[str, str, str, float]] = [] + for hit in reranked: + key = (hit[0], hit[1]) + if key in original_by_key and key not in seen: + ordered.append(original_by_key[key]) + seen.add(key) + for hit in window: + key = (hit[0], hit[1]) + if key not in seen: + ordered.append(hit) + seen.add(key) + return ordered + hits[window_size:] + + def _retrieve( store: KBStore, query: str, @@ -101,6 +186,7 @@ def _retrieve( fused = rrf_fuse(sem, lex, limit=fetch_limit) if fused: filtered = filter_hits(store, fused, viewer, limit=limit) + filtered = _maybe_rerank(store, query=query, hits=filtered, limit=limit) return [(k, i, s, sc, "hybrid") for k, i, s, sc in filtered] # both retrievers empty -> fall through to the substring scan below. diff --git a/tests/test_retrieval_backend.py b/tests/test_retrieval_backend.py index dcff5372..a3ccd181 100644 --- a/tests/test_retrieval_backend.py +++ b/tests/test_retrieval_backend.py @@ -32,6 +32,15 @@ def _set_backend(store: KBStore, backend: str) -> None: store.config_path.write_text(yaml.safe_dump(cfg)) +def _set_rerank(store: KBStore, *, enabled: bool, top_k: int | None = None) -> None: + cfg = yaml.safe_load(store.config_path.read_text()) + rerank_cfg = {"enabled": enabled} + if top_k is not None: + rerank_cfg["top_k"] = top_k + cfg.setdefault("retrieval", {})["rerank"] = rerank_cfg + store.config_path.write_text(yaml.safe_dump(cfg)) + + def _force_semantic_hit(monkeypatch: pytest.MonkeyPatch) -> None: """Make the embedding path always return a hit, so a backend label of "embedding" appears iff `_retrieve` actually consulted semantic search.""" @@ -124,6 +133,178 @@ def test_backend_hybrid_merges_semantic_and_lexical( assert _backends(pack) == {"hybrid"} +def test_context_rerank_disabled_preserves_hybrid_order( + store: KBStore, monkeypatch: pytest.MonkeyPatch +) -> None: + src = store.put_source(b"e2") + store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id])) + health.rebuild_index(store) + monkeypatch.setattr( + context.index_db, + "search_semantic", + lambda *a, **k: [ + ("claim", "c1", "JWT token rotation", 0.90), + ("claim", "c2", "OAuth refresh flow", 0.80), + ], + ) + monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + _set_backend(store, "hybrid") + _set_rerank(store, enabled=False) + + pack = context.build_context_pack(store, query="auth", limit=2) + + assert [item["id"] for item in pack["items"]] == ["c1", "c2"] + + +def test_context_rerank_enabled_reorders_scoped_window( + store: KBStore, monkeypatch: pytest.MonkeyPatch +) -> None: + from vouch.embeddings import rerank as rerank_mod + + src = store.put_source(b"e2") + store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id])) + store.put_claim(Claim(id="c3", text="SAML login flow", evidence=[src.id])) + health.rebuild_index(store) + monkeypatch.setattr( + context.index_db, + "search_semantic", + lambda *a, **k: [ + ("claim", "c1", "JWT token rotation", 0.90), + ("claim", "c2", "OAuth refresh flow", 0.80), + ("claim", "c3", "SAML login flow", 0.70), + ], + ) + monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + monkeypatch.setattr(context, "_RERANKER_CACHE", None) + monkeypatch.setattr(rerank_mod, "default_reranker", lambda: object()) + monkeypatch.setattr( + rerank_mod, + "rerank", + lambda *, query, hits, reranker, top_k: [ + (hits[1][0], hits[1][1], hits[1][2], 99.0), + (hits[0][0], hits[0][1], hits[0][2], 88.0), + ][:top_k], + ) + _set_backend(store, "hybrid") + _set_rerank(store, enabled=False) + before = context.build_context_pack(store, query="auth", limit=3) + scores_by_id = {item["id"]: item["score"] for item in before["items"]} + _set_rerank(store, enabled=True, top_k=2) + + pack = context.build_context_pack(store, query="auth", limit=3) + + assert [item["id"] for item in pack["items"]] == ["c2", "c1", "c3"] + assert [item["score"] for item in pack["items"]] == [ + scores_by_id["c2"], + scores_by_id["c1"], + scores_by_id["c3"], + ] + + +def test_context_rerank_bool_top_k_falls_back_to_limit( + store: KBStore, monkeypatch: pytest.MonkeyPatch +) -> None: + from vouch.embeddings import rerank as rerank_mod + + src = store.put_source(b"e2") + store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id])) + store.put_claim(Claim(id="c3", text="SAML login flow", evidence=[src.id])) + health.rebuild_index(store) + monkeypatch.setattr( + context.index_db, + "search_semantic", + lambda *a, **k: [ + ("claim", "c1", "JWT token rotation", 0.90), + ("claim", "c2", "OAuth refresh flow", 0.80), + ("claim", "c3", "SAML login flow", 0.70), + ], + ) + monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + seen_top_k: list[int] = [] + monkeypatch.setattr(context, "_RERANKER_CACHE", None) + monkeypatch.setattr(rerank_mod, "default_reranker", lambda: object()) + + def fake_rerank(*, query, hits, reranker, top_k): + seen_top_k.append(top_k) + return list(reversed(hits)) + + monkeypatch.setattr(rerank_mod, "rerank", fake_rerank) + _set_backend(store, "hybrid") + _set_rerank(store, enabled=True, top_k=True) + + pack = context.build_context_pack(store, query="auth", limit=3) + + assert seen_top_k == [3] + assert [item["id"] for item in pack["items"]] == ["c3", "c2", "c1"] + + +def test_context_rerank_reuses_default_reranker( + store: KBStore, monkeypatch: pytest.MonkeyPatch +) -> None: + from vouch.embeddings import rerank as rerank_mod + + src = store.put_source(b"e2") + store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id])) + health.rebuild_index(store) + monkeypatch.setattr( + context.index_db, + "search_semantic", + lambda *a, **k: [ + ("claim", "c1", "JWT token rotation", 0.90), + ("claim", "c2", "OAuth refresh flow", 0.80), + ], + ) + monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + calls = 0 + monkeypatch.setattr(context, "_RERANKER_CACHE", None) + + def fake_default_reranker(): + nonlocal calls + calls += 1 + return object() + + monkeypatch.setattr(rerank_mod, "default_reranker", fake_default_reranker) + monkeypatch.setattr(rerank_mod, "rerank", lambda *, query, hits, reranker, top_k: hits) + _set_backend(store, "hybrid") + _set_rerank(store, enabled=True) + + context.build_context_pack(store, query="auth", limit=2) + context.build_context_pack(store, query="auth", limit=2) + + assert calls == 1 + + +def test_context_rerank_missing_extra_degrades_to_fused_order( + store: KBStore, monkeypatch: pytest.MonkeyPatch +) -> None: + from vouch.embeddings import rerank as rerank_mod + + src = store.put_source(b"e2") + store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id])) + health.rebuild_index(store) + monkeypatch.setattr( + context.index_db, + "search_semantic", + lambda *a, **k: [ + ("claim", "c1", "JWT token rotation", 0.90), + ("claim", "c2", "OAuth refresh flow", 0.80), + ], + ) + monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + monkeypatch.setattr(context, "_RERANKER_CACHE", None) + monkeypatch.setattr( + rerank_mod, + "default_reranker", + lambda: (_ for _ in ()).throw(ImportError("missing optional extra")), + ) + _set_backend(store, "hybrid") + _set_rerank(store, enabled=True) + + pack = context.build_context_pack(store, query="auth", limit=2) + + assert [item["id"] for item in pack["items"]] == ["c1", "c2"] + + def test_near_duplicate_summaries_are_dropped( store: KBStore, monkeypatch: pytest.MonkeyPatch ) -> None: