From 199c6ce2ebd8bb0606c5ed10229920e20923fac7 Mon Sep 17 00:00:00 2001 From: Steve-too Date: Tue, 7 Jul 2026 15:12:51 +0000 Subject: [PATCH 1/3] feat(context): gate reranking behind retrieval config --- src/vouch/context.py | 73 +++++++++++++++++++++++++ tests/test_retrieval_backend.py | 95 +++++++++++++++++++++++++++++++++ 2 files changed, 168 insertions(+) diff --git a/src/vouch/context.py b/src/vouch/context.py index fa7e3f16..1388f349 100644 --- a/src/vouch/context.py +++ b/src/vouch/context.py @@ -74,6 +74,78 @@ 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 top_k > 0 else limit + return enabled, top_k + + +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 default_reranker + from .embeddings.rerank import rerank as do_rerank + + reranked = do_rerank( + query=query, + hits=window, + reranker=default_reranker(), + 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_keys = {(kind, artifact_id) for kind, artifact_id, _summary, _score 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_keys and key not in seen: + ordered.append(hit) + 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 +173,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..6c1d1ac0 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,92 @@ 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(rerank_mod, "default_reranker", lambda: object()) + monkeypatch.setattr( + rerank_mod, + "rerank", + lambda *, query, hits, reranker, top_k: [hits[1], hits[0]][:top_k], + ) + _set_backend(store, "hybrid") + _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"] + + +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( + 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: From 83091eb463d18d7e1d991ed94d58e30b440ed422 Mon Sep 17 00:00:00 2001 From: Steve-too Date: Tue, 7 Jul 2026 18:40:50 +0000 Subject: [PATCH 2/3] fix(context): address rerank review feedback --- src/vouch/context.py | 25 +++++++--- tests/test_retrieval_backend.py | 88 ++++++++++++++++++++++++++++++++- 2 files changed, 106 insertions(+), 7 deletions(-) diff --git a/src/vouch/context.py b/src/vouch/context.py index 1388f349..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: @@ -98,10 +99,23 @@ def _configured_rerank(store: KBStore, *, limit: int) -> tuple[bool, int]: 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 top_k > 0 else 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, *, @@ -116,13 +130,12 @@ def _maybe_rerank( window_size = min(top_k, len(hits)) window = hits[:window_size] try: - from .embeddings.rerank import default_reranker from .embeddings.rerank import rerank as do_rerank reranked = do_rerank( query=query, hits=window, - reranker=default_reranker(), + reranker=_default_reranker_cached(), top_k=window_size, ) except ImportError: @@ -130,13 +143,13 @@ def _maybe_rerank( # 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_keys = {(kind, artifact_id) for kind, artifact_id, _summary, _score in window} + 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_keys and key not in seen: - ordered.append(hit) + 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]) diff --git a/tests/test_retrieval_backend.py b/tests/test_retrieval_backend.py index 6c1d1ac0..58e6b1f2 100644 --- a/tests/test_retrieval_backend.py +++ b/tests/test_retrieval_backend.py @@ -175,18 +175,103 @@ def test_context_rerank_enabled_reorders_scoped_window( ], ) monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + 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], hits[0]][:top_k], + 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] = [] + 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 + 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( @@ -206,6 +291,7 @@ def test_context_rerank_missing_extra_degrades_to_fused_order( ], ) monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) + context._RERANKER_CACHE = None monkeypatch.setattr( rerank_mod, "default_reranker", From 6d236de7899785cf9c723321a2cea0a5877b5e08 Mon Sep 17 00:00:00 2001 From: Steve-too Date: Tue, 7 Jul 2026 18:46:13 +0000 Subject: [PATCH 3/3] test(context): isolate reranker cache reset --- tests/test_retrieval_backend.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_retrieval_backend.py b/tests/test_retrieval_backend.py index 58e6b1f2..a3ccd181 100644 --- a/tests/test_retrieval_backend.py +++ b/tests/test_retrieval_backend.py @@ -175,7 +175,7 @@ def test_context_rerank_enabled_reorders_scoped_window( ], ) monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) - context._RERANKER_CACHE = None + monkeypatch.setattr(context, "_RERANKER_CACHE", None) monkeypatch.setattr(rerank_mod, "default_reranker", lambda: object()) monkeypatch.setattr( rerank_mod, @@ -221,7 +221,7 @@ def test_context_rerank_bool_top_k_falls_back_to_limit( ) monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) seen_top_k: list[int] = [] - context._RERANKER_CACHE = None + monkeypatch.setattr(context, "_RERANKER_CACHE", None) monkeypatch.setattr(rerank_mod, "default_reranker", lambda: object()) def fake_rerank(*, query, hits, reranker, top_k): @@ -256,7 +256,7 @@ def test_context_rerank_reuses_default_reranker( ) monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) calls = 0 - context._RERANKER_CACHE = None + monkeypatch.setattr(context, "_RERANKER_CACHE", None) def fake_default_reranker(): nonlocal calls @@ -291,7 +291,7 @@ def test_context_rerank_missing_extra_degrades_to_fused_order( ], ) monkeypatch.setattr(context.index_db, "search", lambda *a, **k: []) - context._RERANKER_CACHE = None + monkeypatch.setattr(context, "_RERANKER_CACHE", None) monkeypatch.setattr( rerank_mod, "default_reranker",