Skip to content

Latest commit

 

History

History
228 lines (183 loc) · 10.5 KB

File metadata and controls

228 lines (183 loc) · 10.5 KB

Architecture — Reranking (Step 2.7)

Two-stage cross-encoder reranking with optional MMR diversification, implemented in the rag-reranker package. Anchored in ADR-0006.

Overview

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]

Internals — design rationale

Why two stages (ADR-0006)

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.

Why MMR (and not learned diversification)

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:

  1. 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.
  2. Fall back to a wired Embedder — when metadata is absent, the pipeline calls Embedder.bulk_embed over the stage-2 candidates only (≤ stage2_top_n), keeping the extra cost bounded.
  3. 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.

Failure mode — propagate, don't degrade

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 is ctx-less

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.

Hydration as a callable, not a backend handle

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.

Extension points

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 ChunkRefChunk 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.

rerank OTel span

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.

Composition with HybridRetriever

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.

PolicyEngine boundary

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.

Reviewer checklist

When reviewing a Step 2.7 PR or a new reranker backend:

  • Backend subclasses Reranker, implements precise_rerank, and sets Chunk.score (Step 1.1c added score: 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 in asyncio.to_thread so the event loop stays responsive.
  • If the backend pulls embeddings (for MMR fallback), it accepts Chunk.metadata['embedding'] as a free hit before calling the Embedder.
  • 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.modules so CI without the extra still runs them.

Deliberately deferred

  • 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_exit threshold sweep — Step 4.6 latency tuning.
  • Per-stage child spans inside rerank — Step 5.1 (per-stage observability owns that surface).

See also