Skip to content

router.md

Latest commit

 

History

History
287 lines (227 loc) · 9.89 KB

router.md

File metadata and controls

287 lines (227 loc) · 9.89 KB
Raw

rag-retrieval — Retrieval router (Step 2.10)

rag-retrieval 0.2 ships a router that sits in front of the existing HybridRetriever and answers the two questions the Step 2.5 orchestrator deliberately left to its caller:

  1. Which backends should run for this query?
  2. Is any wired backend currently degraded?

When the answer to (2) leaves only the keyword backend healthy, the router falls back to BM25-only retrieval and surfaces a bm25_fallback=true signal so operators can dashboard fallback rates.

Overview

Symbol Purpose
RetrievalRouter Async orchestrator: classify shape → gate on health → call HybridRetriever with the right inputs
RouterConfig Per-router knobs (base weights, health-probe-on-decide)
BackendHealthTracker In-memory rolling-window failure tracker
HealthConfig Threshold + window + degrade-duration knobs
classify_shape Pure function: query text → QueryShape
RoutingDecision Frozen Pydantic record of the routing decision (re-exported from rag-core)
QueryShape Enum: KEYWORD_HEAVY / SEMANTIC / ENTITY_RICH / MIXED

RoutingDecision, QueryShape, and RouterError live in rag-core so they can be referenced from any package without dragging rag-retrieval into the import graph.

Usage

RetrievalRouter

Wire it on top of an already-constructed HybridRetriever:

from rag_retrieval import HybridRetriever, RetrievalRouter

hybrid = HybridRetriever(
    vector_backend=qdrant,
    keyword_backend=elasticsearch,
    graph_backend=neo4j,
)
router = RetrievalRouter(
    hybrid=hybrid,
    embedder=openai_embedder,   # optional; supplied vector / HyDE win when present
)

Per-call API mirrors HybridRetriever.retrieve but resolves per-backend inputs automatically:

decision, chunks = await router.route(
    ctx,
    text=user_query,
    hyde_vector=understood.hyde_vector,   # from Step 2.6 pipeline
    graph_seeds=key_entities,             # entity-extraction output
    top_k=10,
)

decision is a RoutingDecision carrying:

  • shape: QueryShape — what the classifier decided
  • use_vector / use_keyword / use_graph: bool — which backends fired
  • weights: tuple[float, float, float] — RRF weights actually applied
  • degraded_backends: tuple[str, ...] — labels skipped due to health
  • bm25_fallback: bool — true when keyword was the only survivor
  • reason: str — short human-readable summary

chunks is the same list[ChunkRef] shape HybridRetriever returns.

Pre-classify without executing

decide() returns the RoutingDecision without running retrieval — use this when you want to log / dashboard / replay routing without paying for backend calls:

decision = await router.decide(
    ctx,
    text=user_query,
    expansion_terms=understood.expansion_terms,
    hyde_vector=understood.hyde_vector,
    caller_has_graph_seeds=bool(key_entities),
)
if decision.bm25_fallback:
    metrics.increment("router.bm25_fallback")

Execute a pre-built decision (Step 2.11 G-02)

execute(ctx, decision, ...) runs a routed retrieval against an already-made decision — useful for sticky-mode agent loops where the iter-0 decision should be reused without paying the (small but non-zero) decide cost on every iteration:

# iter 0
decision = await router.decide(ctx, text=user_query, ...)
results = await router.execute(ctx, decision, text=user_query, top_k=10)

# iter 1+ (sticky: reuse the iter-0 decision)
for sub_query in followups:
    more = await router.execute(ctx, decision, text=sub_query, top_k=10)

The canonical retrieve.route OTel span is emitted by execute, so dashboards see one span per executed retrieval regardless of whether decide was bundled (route()) or pre-computed. execute raises RouterError if the decision has no active backends — defence-in-depth against silently returning zero results.

classify_shape

A pure helper exposed for offline analysis and tests:

from rag_retrieval import classify_shape
from rag_core import QueryShape

assert classify_shape(text="acme 2024 earnings") == QueryShape.KEYWORD_HEAVY
assert classify_shape(text="What is RAG?") == QueryShape.SEMANTIC
assert classify_shape(text="Alice Smith and Acme Corp") == QueryShape.ENTITY_RICH

BackendHealthTracker

In-memory per-label failure window. Construct your own to share across routers (rare) or to inject a fake clock for tests:

from rag_retrieval import BackendHealthTracker, HealthConfig

tracker = BackendHealthTracker(
    config=HealthConfig(failure_threshold=5, window_seconds=120.0, degrade_seconds=60.0),
)
router = RetrievalRouter(hybrid=hybrid, health=tracker)

# Manual override — e.g. from a background health-probe loop:
await tracker.mark_degraded("vector")

The router calls record_failure itself when HybridRetriever.retrieve raises RetrievalError (every selected source failed) and record_success on the happy path. Manual record_* calls layer on top.

RouterConfig

@dataclass(frozen=True)
class RouterConfig:
    base_weights: HybridWeights = HybridWeights()
    probe_health_on_decide: bool = False
  • base_weights — multiplied by a shape-specific bias to compute the effective per-call weights. Default (1.0, 1.0, 1.0) is correct for most deployments.
  • probe_health_on_decide — when True, the router calls each wired backend's health() on every decide/route call and marks any failing probe as degraded. Default False because doubling per-request network round-trips for healthy backends is too expensive for production traffic — enable only on low-traffic test gateways.

HealthConfig

@dataclass(frozen=True)
class HealthConfig:
    failure_threshold: int = 3      # failures in window → degraded
    window_seconds: float = 60.0    # rolling-window width
    degrade_seconds: float = 30.0   # how long mark_degraded() sticks

Tune these off recorded p50/p99 latency once Phase 4's reliability work lands. Defaults are conservative for a dev gateway with low traffic.

CLI smoke

$ ragctl route "What is RAG?"
query:    'What is RAG?'
tenant:   ragctl-local
shape:    semantic
selected: vector=True keyword=True graph=True
weights:  v=1.50 k=1.00 g=0.80
degraded: none
bm25_fb:  False
reason:   shape=semantic
results:  5
  #1  score=0.053702  chunk_id=chunk-1
  ...

Force the BM25-only fallback:

$ ragctl route "acme 2024 earnings" --simulate-failure vector --no-graph
shape:    keyword_heavy
selected: vector=False keyword=True graph=False
weights:  v=0.30 k=1.50 g=0.50
degraded: ['vector']
bm25_fb:  True

Internals

Shape bias table

Multiplicative on top of RouterConfig.base_weights:

Shape Vector Keyword Graph
KEYWORD_HEAVY 0.3 1.5 0.5
SEMANTIC 1.5 1.0 0.8
ENTITY_RICH 1.0 1.2 1.8
MIXED 1.0 1.0 1.0

Down-weighting (rather than excluding) lets RRF still incorporate the weaker source when it surprises us with a relevant hit.

Decision flags

A backend is use_*=True if all of the following hold:

  1. It is wired on the underlying HybridRetriever.
  2. The shape-biased weight is > 0.
  3. It is not currently degraded in the health tracker.
  4. The router can supply its input for this call:
    • vector — caller-supplied vector OR HyDE vector OR an Embedder
    • keyword — always (uses the query text)
    • graph — caller passed non-empty graph_seeds

OTel span schema

Every route() call opens a retrieve.route span carrying:

Attribute Type Notes
rag.tenant_id string
rag.route.shape string QueryShape.value
rag.route.use_vector bool
rag.route.use_keyword bool
rag.route.use_graph bool
rag.route.bm25_fallback bool
rag.route.degraded_n int
rag.route.degraded string comma-separated; omitted when empty
rag.route.weight_vector float
rag.route.weight_keyword float
rag.route.weight_graph float
rag.route.results_n int final fused count
rag.route.elapsed_ms float monotonic clock
rag.route.error string only when an error propagated

A structured log event retrieval.route_decision mirrors the same data for log-only deployments.

Extension points

  • Replace the shape classifier. RetrievalRouter calls rag_retrieval._shape.classify_shape directly today. A learned or LLM-based classifier can land behind the same function signature (text + expansion_terms + hyde_vector → QueryShape) without API churn elsewhere — RoutingDecision.shape already abstracts the enum.
  • Replace the health tracker. Hand any object implementing the same record_success / record_failure / mark_degraded / is_degraded shape via the health constructor arg. A Redis-backed shared tracker is the obvious next step once Step 4.x's reliability work needs cross-worker degradation state.
  • Per-corpus weights. RouterConfig.base_weights is a single set today; per-corpus_id overrides are a one-call-site change once Phase 3's corpus router lands.
  • Sub-query fan-out. UnderstoodQuery.sub_queries is not consumed by the router — that fan-out is owned by Step 2.11's agent loop, which iterates the router rather than driving it from the inside. See docs/architecture/retrieval-router.md for the deferral rationale.

See also