Skip to content

Commit 7ab4f5a

Browse files
committed
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
1 parent 661b59c commit 7ab4f5a

3 files changed

Lines changed: 158 additions & 0 deletions

File tree

CHANGELOG.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,16 @@ All notable changes to vouch are documented here. Format follows
66

77
## [Unreleased]
88

9+
### Added
10+
- `retrieval.rerank.enabled` (default false) wires the existing
11+
cross-encoder reranker (`embeddings/rerank.py`, already used by
12+
`vouch search --rerank`) into the `kb.context` / `kb.search` hybrid read
13+
path. When enabled, fused RRF hits are reordered by the reranker before
14+
scoping filters run; `retrieval.rerank.top_k` controls the rerank window
15+
(default: the query's context limit). Off by default, so existing
16+
rankings are byte-identical until a kb opts in; degrades to the fused
17+
order if the optional reranker extra isn't installed.
18+
919
## [1.2.2] — 2026-07-07
1020

1121
### Packaging

src/vouch/context.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,31 @@ def _configured_backend(store: KBStore) -> str:
7474
return "auto"
7575

7676

77+
def _rerank_cfg(store: KBStore) -> tuple[bool, int | None]:
78+
"""Read ``retrieval.rerank`` defensively. Returns (enabled, top_k).
79+
80+
Mirrors ``salience.reflex_cfg``'s config-reading shape. ``top_k=None``
81+
means "use the caller's context limit" — resolved at the `_retrieve`
82+
call site, since the config itself has no notion of a query's limit.
83+
"""
84+
try:
85+
loaded = yaml.safe_load(store.config_path.read_text(encoding="utf-8"))
86+
except (OSError, yaml.YAMLError):
87+
return False, None
88+
retrieval = loaded.get("retrieval") if isinstance(loaded, dict) else None
89+
rerank = retrieval.get("rerank") if isinstance(retrieval, dict) else None
90+
if not isinstance(rerank, dict):
91+
rerank = {}
92+
93+
enabled = rerank.get("enabled", False)
94+
enabled = enabled if isinstance(enabled, bool) else False
95+
96+
top_k = rerank.get("top_k")
97+
top_k = top_k if isinstance(top_k, int) and top_k > 0 else None
98+
99+
return enabled, top_k
100+
101+
77102
def _retrieve(
78103
store: KBStore,
79104
query: str,
@@ -88,6 +113,12 @@ def _retrieve(
88113
- "embedding": semantic search only
89114
- "fts5": lexical FTS5 only
90115
- "substring": substring scan only
116+
117+
When `retrieval.rerank.enabled` is set, the fused hybrid hits are
118+
reordered by the cross-encoder reranker (`embeddings/rerank.py`) before
119+
scoping filters run — same reranker `vouch search --rerank` already
120+
uses. Off by default, so existing rankings are unaffected; degrades to
121+
the fused order if the optional reranker extra isn't installed.
91122
"""
92123
backend = _configured_backend(store)
93124
fetch_limit = scoped_fetch_limit(limit, viewer)
@@ -100,6 +131,18 @@ def _retrieve(
100131
lex = []
101132
fused = rrf_fuse(sem, lex, limit=fetch_limit)
102133
if fused:
134+
rerank_enabled, rerank_top_k = _rerank_cfg(store)
135+
if rerank_enabled:
136+
try:
137+
from .embeddings.rerank import default_reranker
138+
from .embeddings.rerank import rerank as do_rerank
139+
140+
fused = do_rerank(
141+
query=query, hits=fused, reranker=default_reranker(),
142+
top_k=rerank_top_k or limit,
143+
)
144+
except ImportError:
145+
pass # reranker extra not installed; keep fused order
103146
filtered = filter_hits(store, fused, viewer, limit=limit)
104147
return [(k, i, s, sc, "hybrid") for k, i, s, sc in filtered]
105148
# both retrievers empty -> fall through to the substring scan below.

tests/test_retrieval_backend.py

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,27 @@ def _backends(pack: dict) -> set[str]:
4545
return {item["backend"] for item in pack["items"]}
4646

4747

48+
def _set_rerank(store: KBStore, *, enabled: bool, top_k: int | None = None) -> None:
49+
cfg = yaml.safe_load(store.config_path.read_text())
50+
rerank_cfg = cfg.setdefault("retrieval", {}).setdefault("rerank", {})
51+
rerank_cfg["enabled"] = enabled
52+
if top_k is not None:
53+
rerank_cfg["top_k"] = top_k
54+
store.config_path.write_text(yaml.safe_dump(cfg))
55+
56+
57+
class _StubReranker:
58+
"""Deterministic stand-in for the cross-encoder: longer snippet wins.
59+
60+
Exercises the real `embeddings.rerank.rerank` scoring/sort path without
61+
needing the optional sentence-transformers extra, so this runs under
62+
the base CI install like the rest of this file (#92-style pattern).
63+
"""
64+
65+
def score(self, query: str, candidates: list[str]) -> list[float]:
66+
return [float(len(c)) for c in candidates]
67+
68+
4869
def test_backend_fts5_skips_embedding(
4970
store: KBStore, monkeypatch: pytest.MonkeyPatch
5071
) -> None:
@@ -164,6 +185,90 @@ def test_dedupe_keeps_highest_scored_regardless_of_input_order() -> None:
164185
assert [i.id for i in out] == ["hi"]
165186

166187

188+
def test_rerank_disabled_by_default_ordering_unchanged(
189+
store: KBStore, monkeypatch: pytest.MonkeyPatch
190+
) -> None:
191+
"""No `retrieval.rerank` config at all (#429): the reranker must never
192+
even be constructed, and hybrid ordering is exactly the RRF-fused order."""
193+
def _boom() -> None:
194+
raise AssertionError("reranker must not be constructed when rerank is off")
195+
196+
monkeypatch.setattr("vouch.embeddings.rerank.default_reranker", _boom)
197+
src = store.put_source(b"e2")
198+
store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id]))
199+
health.rebuild_index(store)
200+
monkeypatch.setattr(
201+
context.index_db, "search_semantic",
202+
lambda *a, **k: [
203+
("claim", "c1", "short", 0.9),
204+
("claim", "c2", "a much longer snippet of text", 0.8),
205+
],
206+
)
207+
monkeypatch.setattr(context.index_db, "search", lambda *a, **k: [])
208+
_set_backend(store, "hybrid")
209+
210+
pack = context.build_context_pack(store, query="auth")
211+
assert [i["id"] for i in pack["items"]] == ["c1", "c2"]
212+
213+
214+
def test_rerank_enabled_reorders_by_cross_encoder_score(
215+
store: KBStore, monkeypatch: pytest.MonkeyPatch
216+
) -> None:
217+
"""With `retrieval.rerank.enabled: true`, hybrid hits are reordered by
218+
the reranker's scores instead of the RRF fusion order."""
219+
src = store.put_source(b"e2")
220+
store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id]))
221+
health.rebuild_index(store)
222+
monkeypatch.setattr(
223+
context.index_db, "search_semantic",
224+
lambda *a, **k: [
225+
("claim", "c1", "short", 0.9),
226+
("claim", "c2", "a much longer snippet of text", 0.8),
227+
],
228+
)
229+
monkeypatch.setattr(context.index_db, "search", lambda *a, **k: [])
230+
_set_backend(store, "hybrid")
231+
232+
# sanity: fused order (rerank off) is c1 then c2 by RRF score.
233+
baseline = context.build_context_pack(store, query="auth")
234+
assert [i["id"] for i in baseline["items"]] == ["c1", "c2"]
235+
236+
monkeypatch.setattr(
237+
"vouch.embeddings.rerank.default_reranker", lambda: _StubReranker()
238+
)
239+
_set_rerank(store, enabled=True)
240+
reranked = context.build_context_pack(store, query="auth")
241+
assert [i["id"] for i in reranked["items"]] == ["c2", "c1"]
242+
assert _backends(reranked) == {"hybrid"}
243+
244+
245+
def test_rerank_missing_extra_degrades_to_fused_order(
246+
store: KBStore, monkeypatch: pytest.MonkeyPatch
247+
) -> None:
248+
"""`sentence-transformers` not installed must not break `kb.context` —
249+
it degrades to the unreranked fused order instead of raising."""
250+
def _raise() -> None:
251+
raise ImportError("sentence-transformers not installed")
252+
253+
src = store.put_source(b"e2")
254+
store.put_claim(Claim(id="c2", text="OAuth refresh flow", evidence=[src.id]))
255+
health.rebuild_index(store)
256+
monkeypatch.setattr(
257+
context.index_db, "search_semantic",
258+
lambda *a, **k: [
259+
("claim", "c1", "short", 0.9),
260+
("claim", "c2", "a much longer snippet of text", 0.8),
261+
],
262+
)
263+
monkeypatch.setattr(context.index_db, "search", lambda *a, **k: [])
264+
_set_backend(store, "hybrid")
265+
monkeypatch.setattr("vouch.embeddings.rerank.default_reranker", _raise)
266+
_set_rerank(store, enabled=True)
267+
268+
pack = context.build_context_pack(store, query="auth")
269+
assert [i["id"] for i in pack["items"]] == ["c1", "c2"]
270+
271+
167272
def test_dedupe_preserves_input_order_not_score_order() -> None:
168273
"""Survivors keep the caller's order (ranked hits first, appended
169274
neighbours last) even when a later distinct item outscores an earlier one,

0 commit comments

Comments
 (0)