Two-stage cross-encoder reranking with optional MMR diversification,
implemented in the rag-reranker package.
Anchored in ADR-0006.
Reranking sits between retrieval (Step 2.5
HybridRetriever) and the context packer (Step
2.8). Its job is to re-score the top-N candidates returned by hybrid
retrieval and (optionally) diversify the final shortlist so the LLM
sees relevant and non-redundant context.
The pipeline is structured as:
ChunkRef[N] → fast_rerank → ChunkRef[M] → should_early_exit?
↓ ↓
yes-skip no
↓ ↓
hydrate hydrate
↓ ↓
Chunk[K] Chunk[M] → precise_rerank → Chunk[S]
↓
MMR (optional)
↓
Chunk[top_k]
The naive "score every candidate with a cross-encoder" approach blows the latency budget: at ~10 ms / pair on CPU, scoring 100 hybrid candidates costs ~1 s per query. Phase 2's p95 budget is ≤ 500 ms (dev) / ≤ 250 ms (SaaS).
The fix is a two-stage cascade. Stage 1 (fast_rerank) is a cheap
heuristic — typically pass-through truncation (for hosted-API rerankers
like Cohere, where the cheapest invocation is already the precise one)
or a bi-encoder dot-product (for local cross-encoder backends). Stage
2 (precise_rerank) is the expensive cross-encoder pass, but it only
ever sees the top ~50 candidates that stage 1 admitted.
ADR-0006 commits the SPI to the two-stage shape from day one because
refactoring an SPI after backends ship is far more expensive than the
upfront complexity. Single-stage backends ship as precise_rerank
only, inheriting the default pass-through fast_rerank.
MMR (Carbonell & Goldstein 1998) gives us a knob (λ) to trade
relevance against diversity with zero model dependencies. Learned
diversification (DPP, ListNet) needs a corpus-specific training set
that we don't have.
The pure :func:mmr_rerank primitive runs over (relevance_score, embedding) tuples; the orchestrator resolves embeddings as follows:
- Prefer
Chunk.metadata['embedding']— when the upstream retrieval path stashed the embedding on the chunk (most vector backends can do this cheaply since they already had it in hand), MMR pays nothing extra. - Fall back to a wired
Embedder— when metadata is absent, the pipeline callsEmbedder.bulk_embedover the stage-2 candidates only (≤stage2_top_n), keeping the extra cost bounded. - Refuse to silently drop diversity — if neither is available,
the pipeline raises
RerankerError(stage="mmr").
The first-pick convention is "argmax(relevance), no diversity
penalty". This makes the algorithm stable at the λ = 0 boundary
(where the multiplicative λ would zero out every relevance term and
fall back to index order) and matches the conventional formulation in
the literature.
Step 2.5's HybridRetriever and Step 2.6's QueryUnderstandingPipeline
both follow the "degrade gracefully on partial failure" pattern: one
component fails, the orchestrator logs and continues with reduced
output.
RerankPipeline deliberately breaks that pattern. A reranker failure
raises RerankerError. Rationale:
- A hybrid-retrieval source failing leaves the user with degraded but still useful results from the remaining sources.
- A reranker failing leaves the user with unranked candidates — the upstream order is dense-vs-keyword RRF, not "the right answer." Silently passing through that order would deceive the caller into thinking they got rerank quality.
- The Phase 3 gateway is the right place to decide between "fail the query" and "return upstream order with a banner." That decision needs the request's policy / SLO / user tier — context the reranker doesn't have.
should_early_exit(stage_1_scores) is a pure predicate over a list of
floats. It carries no RequestContext — listed in the SPI signature
linter's exempt set (tests/contract/spi_signature.py) for the same
reason Embedder.dimension is exempt: it's not an I/O method and
threading tenant plumbing through it would add cost without value.
Default returns False. Reranker subclasses override with
variance-based heuristics (e.g. "skip stage 2 when the top score is
above some τ and the gap between #1 and #2 is large"). Threshold
sweeps are deferred to Step 4.6 latency tuning.
The pipeline accepts a Hydrator callable (async (ctx, refs) → chunks) rather than a VectorRetrievalBackend handle. This keeps
rag-reranker decoupled from any specific backend package; callers
adapt their *RetrievalBackend.hydrate method into the protocol with
a one-line lambda. The Phase 3 gateway is expected to keep a hydrator
shared across HybridRetriever, RerankPipeline, and the context
packer.
| Hook | Purpose | Default |
|---|---|---|
Reranker.precise_rerank |
The actual cross-encoder pass. | Required override. |
Reranker.fast_rerank |
Cheap stage-1 narrowing. | Pass-through truncation. |
Reranker.should_early_exit |
Skip stage 2 on tight stage-1 score clusters. | Always False. |
MMRConfig |
Diversification knobs (λ, top_k). |
Not applied by default. |
Hydrator |
Convert ChunkRef → Chunk for stage 2. |
Caller-supplied. |
Embedder (optional) |
MMR embedding fallback when metadata is absent. | None — pipeline raises if MMR needs it. |
To add a new reranker backend, subclass
Reranker,
override precise_rerank, and (optionally) fast_rerank /
should_early_exit. See reference/reranker.md
for the public constructor shapes.
Every RerankPipeline.rerank() call opens a single rerank span via
span_from_trace_context with these attributes:
| Attribute | Type | Notes |
|---|---|---|
rag.tenant_id |
str | From ctx.tenant_id. |
rag.rerank.backend |
str | Reranker subclass name. |
rag.rerank.top_k_in |
int | Length of the input ChunkRef list. |
rag.rerank.top_k_out |
int | Length of the returned Chunk list. |
rag.rerank.stage1_n |
int | Output size of stage 1. 0 on empty / early-exit paths that bypass stage 1. |
rag.rerank.stage2_n |
int | Output size of stage 2. 0 when early-exit fired. |
rag.rerank.early_exit |
bool | Whether should_early_exit returned True. |
rag.rerank.mmr_applied |
bool | Whether MMR ran (config supplied AND > 1 candidate). |
rag.rerank.mmr_lambda |
float | Set when mmr_applied=True. |
rag.rerank.elapsed_ms |
float | Wall time, monotonic. |
The schema mirrors Step 2.6's query.understand span for visual
consistency on Grafana dashboards.
Step 2.5 HybridRetriever.retrieve() returns list[ChunkRef]. Step
2.7 keeps that interface untouched — the reranker is NOT wired inside
HybridRetriever. The expected wiring lives at the Phase 3 gateway:
refs = await hybrid.retrieve(ctx, ...)
chunks = await reranker_pipeline.rerank(ctx, query, refs, top_k=10)
packed = await packer.pack(ctx, chunks, budget=...)Rationale: the corpus router (Step 2.10) needs to decide per-query
whether reranking is worth its latency cost. Coupling reranking into
the retriever forces the router to choose a different HybridRetriever
per query, which is more plumbing than a separate component.
The reranker does not call PolicyEngine. Every ChunkRef arriving
at RerankPipeline.rerank() has already passed PolicyDecision.read_chunk
inside HybridRetriever (Step 2.5). Hydration is over the
already-policed ref set.
The policy-coverage linter
(tests/policy/coverage.py) confirms this by scanning the pipeline
for governed SPI calls; the only one is _hydrator(ctx, refs), which
doesn't match the \.hydrate\( pattern.
When reviewing a Step 2.7 PR or a new reranker backend:
- Backend subclasses
Reranker, implementsprecise_rerank, and setsChunk.score(Step 1.1c addedscore: float | None). - Scores are normalised into
[0, 1]— either natively (Cohere, Jina) or via sigmoid (cross-encoder logits). - SDK imports are lazy inside
__init__; the module imports cleanly without the matching[extra]installed. - Backend failure raises
RerankerError, not the underlying SDK exception — the pipeline wraps any wider exception too. - If the backend is sync (
sentence-transformers), heavy work runs inasyncio.to_threadso the event loop stays responsive. - If the backend pulls embeddings (for MMR fallback), it accepts
Chunk.metadata['embedding']as a free hit before calling theEmbedder. - Tenant isolation — backend code never assumes the caller hasn't already filtered by tenant; doesn't introduce cross-tenant similarity comparisons.
- Tests stub the SDK at
sys.modulesso CI without the extra still runs them.
- LLM-as-reranker (RankGPT, RankZephyr) — Phase 5 quality work.
- Learned reranker fine-tuning / distillation — Phase 7.
- Per-tenant model selection — Phase 6.
- Per-corpus rerank policy in
rag.yaml— Step 3.5 (corpus router). should_early_exitthreshold sweep — Step 4.6 latency tuning.- Per-stage child spans inside
rerank— Step 5.1 (per-stage observability owns that surface).
- ADR-0006 — the architectural anchor.
- hybrid-fusion.md — what feeds into stage 1.
- context-packing.md — what consumes the reranker output (Step 2.8).
- performance.md — hot-path discipline + per-SPI p99 budgets.
reference/reranker.md— public API surface.