From 7ab4f5a00a52c8442cf68698379bd7bb798fc0c8 Mon Sep 17 00:00:00 2001
From: Nick M <274344962+nickmopen@users.noreply.github.com>
Date: Wed, 8 Jul 2026 04:37:11 +0300
Subject: [PATCH] feat(retrieval): wire the cross-encoder reranker into
 kb.context

vouch already fuses lexical + semantic hits by reciprocal-rank fusion
and already ships a cross-encoder reranker, but the reranker was only
reachable from `vouch search --rerank`. the context-pack path that
actually feeds agents (build_context_pack -> _retrieve -> kb.context)
never reranked, and there was no config to turn it on.

add retrieval.rerank.enabled (default false) and retrieval.rerank.top_k
(default: the query's context limit) in config.yaml. when enabled,
fused hybrid hits are reordered by embeddings.rerank.rerank before
scoping filters run, mirroring the existing --rerank cli path.
degrades to the fused order if the reranker extra isn't installed.
off by default, so existing rankings are unaffected.

Fixes #429
---
 CHANGELOG.md                    |  10 +++
 src/vouch/context.py            |  43 +++++++++++++
 tests/test_retrieval_backend.py | 105 ++++++++++++++++++++++++++++++++
 3 files changed, 158 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 986dcb0e..8ef1262f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,16 @@ All notable changes to vouch are documented here. Format follows
 
 ## [Unreleased]
 
+### Added
+- `retrieval.rerank.enabled` (default false) wires the existing
+  cross-encoder reranker (`embeddings/rerank.py`, already used by
+  `vouch search --rerank`) into the `kb.context` / `kb.search` hybrid read
+  path. When enabled, fused RRF hits are reordered by the reranker before
+  scoping filters run; `retrieval.rerank.top_k` controls the rerank window
+  (default: the query's context limit). Off by default, so existing
+  rankings are byte-identical until a kb opts in; degrades to the fused
+  order if the optional reranker extra isn't installed.
+
 ## [1.2.2] — 2026-07-07
 
 ### Packaging
diff --git a/src/vouch/context.py b/src/vouch/context.py
index fa7e3f16..7679dfd1 100644
--- a/src/vouch/context.py
+++ b/src/vouch/context.py
@@ -74,6 +74,31 @@ def _configured_backend(store: KBStore) -> str:
     return "auto"
 
 
+def _rerank_cfg(store: KBStore) -> tuple[bool, int | None]:
+    """Read ``retrieval.rerank`` defensively. Returns (enabled, top_k).
+
+    Mirrors ``salience.reflex_cfg``'s config-reading shape. ``top_k=None``
+    means "use the caller's context limit" — resolved at the `_retrieve`
+    call site, since the config itself has no notion of a query's limit.
+    """
+    try:
+        loaded = yaml.safe_load(store.config_path.read_text(encoding="utf-8"))
+    except (OSError, yaml.YAMLError):
+        return False, None
+    retrieval = loaded.get("retrieval") if isinstance(loaded, dict) else None
+    rerank = retrieval.get("rerank") if isinstance(retrieval, dict) else None
+    if not isinstance(rerank, dict):
+        rerank = {}
+
+    enabled = rerank.get("enabled", False)
+    enabled = enabled if isinstance(enabled, bool) else False
+
+    top_k = rerank.get("top_k")
+    top_k = top_k if isinstance(top_k, int) and top_k > 0 else None
+
+    return enabled, top_k
+
+
 def _retrieve(
     store: KBStore,
     query: str,
@@ -88,6 +113,12 @@ def _retrieve(
       - "embedding": semantic search only
       - "fts5": lexical FTS5 only
       - "substring": substring scan only
+
+    When `retrieval.rerank.enabled` is set, the fused hybrid hits are
+    reordered by the cross-encoder reranker (`embeddings/rerank.py`) before
+    scoping filters run — same reranker `vouch search --rerank` already
+    uses. Off by default, so existing rankings are unaffected; degrades to
+    the fused order if the optional reranker extra isn't installed.
     """
     backend = _configured_backend(store)
     fetch_limit = scoped_fetch_limit(limit, viewer)
@@ -100,6 +131,18 @@ def _retrieve(
             lex = []
         fused = rrf_fuse(sem, lex, limit=fetch_limit)
         if fused:
+            rerank_enabled, rerank_top_k = _rerank_cfg(store)
+            if rerank_enabled:
+                try:
+                    from .embeddings.rerank import default_reranker
+                    from .embeddings.rerank import rerank as do_rerank
+
+                    fused = do_rerank(
+                        query=query, hits=fused, reranker=default_reranker(),
+                        top_k=rerank_top_k or limit,
+                    )
+                except ImportError:
+                    pass  # reranker extra not installed; keep fused order
             filtered = filter_hits(store, fused, viewer, 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..96a807d1 100644
--- a/tests/test_retrieval_backend.py
+++ b/tests/test_retrieval_backend.py
@@ -45,6 +45,27 @@ def _backends(pack: dict) -> set[str]:
     return {item["backend"] for item in pack["items"]}
 
 
+def _set_rerank(store: KBStore, *, enabled: bool, top_k: int | None = None) -> None:
+    cfg = yaml.safe_load(store.config_path.read_text())
+    rerank_cfg = cfg.setdefault("retrieval", {}).setdefault("rerank", {})
+    rerank_cfg["enabled"] = enabled
+    if top_k is not None:
+        rerank_cfg["top_k"] = top_k
+    store.config_path.write_text(yaml.safe_dump(cfg))
+
+
+class _StubReranker:
+    """Deterministic stand-in for the cross-encoder: longer snippet wins.
+
+    Exercises the real `embeddings.rerank.rerank` scoring/sort path without
+    needing the optional sentence-transformers extra, so this runs under
+    the base CI install like the rest of this file (#92-style pattern).
+    """
+
+    def score(self, query: str, candidates: list[str]) -> list[float]:
+        return [float(len(c)) for c in candidates]
+
+
 def test_backend_fts5_skips_embedding(
     store: KBStore, monkeypatch: pytest.MonkeyPatch
 ) -> None:
@@ -164,6 +185,90 @@ def test_dedupe_keeps_highest_scored_regardless_of_input_order() -> None:
     assert [i.id for i in out] == ["hi"]
 
 
+def test_rerank_disabled_by_default_ordering_unchanged(
+    store: KBStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """No `retrieval.rerank` config at all (#429): the reranker must never
+    even be constructed, and hybrid ordering is exactly the RRF-fused order."""
+    def _boom() -> None:
+        raise AssertionError("reranker must not be constructed when rerank is off")
+
+    monkeypatch.setattr("vouch.embeddings.rerank.default_reranker", _boom)
+    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", "short", 0.9),
+            ("claim", "c2", "a much longer snippet of text", 0.8),
+        ],
+    )
+    monkeypatch.setattr(context.index_db, "search", lambda *a, **k: [])
+    _set_backend(store, "hybrid")
+
+    pack = context.build_context_pack(store, query="auth")
+    assert [i["id"] for i in pack["items"]] == ["c1", "c2"]
+
+
+def test_rerank_enabled_reorders_by_cross_encoder_score(
+    store: KBStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """With `retrieval.rerank.enabled: true`, hybrid hits are reordered by
+    the reranker's scores instead of the RRF fusion order."""
+    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", "short", 0.9),
+            ("claim", "c2", "a much longer snippet of text", 0.8),
+        ],
+    )
+    monkeypatch.setattr(context.index_db, "search", lambda *a, **k: [])
+    _set_backend(store, "hybrid")
+
+    # sanity: fused order (rerank off) is c1 then c2 by RRF score.
+    baseline = context.build_context_pack(store, query="auth")
+    assert [i["id"] for i in baseline["items"]] == ["c1", "c2"]
+
+    monkeypatch.setattr(
+        "vouch.embeddings.rerank.default_reranker", lambda: _StubReranker()
+    )
+    _set_rerank(store, enabled=True)
+    reranked = context.build_context_pack(store, query="auth")
+    assert [i["id"] for i in reranked["items"]] == ["c2", "c1"]
+    assert _backends(reranked) == {"hybrid"}
+
+
+def test_rerank_missing_extra_degrades_to_fused_order(
+    store: KBStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """`sentence-transformers` not installed must not break `kb.context` —
+    it degrades to the unreranked fused order instead of raising."""
+    def _raise() -> None:
+        raise ImportError("sentence-transformers not installed")
+
+    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", "short", 0.9),
+            ("claim", "c2", "a much longer snippet of text", 0.8),
+        ],
+    )
+    monkeypatch.setattr(context.index_db, "search", lambda *a, **k: [])
+    _set_backend(store, "hybrid")
+    monkeypatch.setattr("vouch.embeddings.rerank.default_reranker", _raise)
+    _set_rerank(store, enabled=True)
+
+    pack = context.build_context_pack(store, query="auth")
+    assert [i["id"] for i in pack["items"]] == ["c1", "c2"]
+
+
 def test_dedupe_preserves_input_order_not_score_order() -> None:
     """Survivors keep the caller's order (ranked hits first, appended
     neighbours last) even when a later distinct item outscores an earlier one,