feat(context): gate reranking behind retrieval config

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.
 

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: