Skip to content

corpus-router.md

Latest commit

 

History

History
235 lines (186 loc) · 9.17 KB

corpus-router.md

File metadata and controls

235 lines (186 loc) · 9.17 KB
Raw

Corpus router — rag-retrieval (Step 3.5)

The corpus router answers which corpora a query should touch, then drives retrieval across them. It sits one level above the Step 2.10 RetrievalRouter (which picks backends for one corpus).

Overview

A multi-tenant tenant can own many corpora (Legal, Support KB, Engineering docs…). CorpusRouter selects a relevant subset per query — via an explicit caller pin, config-driven static rules, a dependency-free learned classifier, or a fall-through to all visible corpora — and then either fans out (one retrieval per corpus, RRF-fused) or runs a single corpus-scoped retrieval. The selection is returned as a frozen CorpusRoutingDecision, surfaced on the gateway query/retrieve responses, an OTel span, and an audit event.

Public API lives in rag_retrieval:

from rag_retrieval import (
    CorpusRouter,
    CorpusRouterConfig,
    CorpusRule,
    StaticRuleClassifier,
    LearnedCorpusClassifier,
)
from rag_retrieval.corpus_router import CorpusClassifier  # Protocol

The wire types are in rag_core:

from rag_core.types import (
    CorpusRoutingStrategy,   # EXPLICIT / STATIC_RULES / LEARNED / ALL / UNCONSTRAINED
    CorpusRoutingDecision,   # frozen selection output
    CorpusScore,             # per-corpus score + reason
)

Usage

Construct a router

router = CorpusRouter(
    corpus_store=corpus_store,        # CorpusStore SPI — the visibility source
    retrieval_router=retrieval_router,  # Step 2.10 RetrievalRouter
    static_classifier=StaticRuleClassifier([
        CorpusRule(corpus_ids=("legal",), match_any=("contract", "gdpr"), name="legal-kw"),
    ]),
    learned_classifier=LearnedCorpusClassifier(temperature=1.0),
    config=CorpusRouterConfig(fan_out=True, max_corpora=3, score_threshold=0.0),
    audit_writer=audit_writer,        # optional — enables corpus.route audit events
)

In production you rarely build this by hand — rag_gateway.wiring (build_corpus_router_from_config) constructs it from rag.yaml. See guides/corpus-routing.md.

Decide (selection only)

decision = await router.decide(ctx, text="how do I reset my password")
decision.strategy          # CorpusRoutingStrategy.STATIC_RULES
decision.selected_corpora  # (CorpusId('support'),)
decision.fan_out           # False (single corpus selected)
decision.candidate_n       # number of visible corpora considered
for s in decision.scores:  # CorpusScore(corpus_id, score, selected, reason)
    ...

decide is pure selection — no retrieval backends run. An explicit corpus_ids= that names only invisible corpora raises CorpusRouterError.

Route (selection + retrieval)

corpus_decision, routing_decision, refs = await router.route(
    ctx,
    text="review this contract",
    top_k=10,
    corpus_ids=None,      # or pin: ["legal"] → EXPLICIT
)

Returns the CorpusRoutingDecision (corpus selection), the backend-level RoutingDecision from the underlying RetrievalRouter, and the fused list[ChunkRef]. top_k <= 0 raises ValueError.

CorpusRoutingDecision

Field Type Meaning
strategy CorpusRoutingStrategy The effective strategy that drove this decision
selected_corpora tuple[CorpusId, ...] Corpora that will be queried
scores tuple[CorpusScore, ...] Per-corpus score + selected + advisory reason
fan_out bool Federated (one retrieval per corpus, RRF-fused) vs single-pass
candidate_n int Visible corpora considered before selection
reason str Advisory summary ("explicit", "static_rules", "no_visible_corpora", …)

CorpusScore.reason is advisory only (rule name, "learned", "explicit", "below_threshold", "excluded:<rule>") — never parse it in production.

Strategies

Strategy When it fires
EXPLICIT Caller passed corpus_ids; intersected with visibility, no scoring
STATIC_RULES A static rule boosted or excluded a corpus
LEARNED No rule fired but the learned classifier differentiated ≥2 corpora
ALL Nothing differentiated the visible corpora — query every one
UNCONSTRAINED Tenant has no visible corpora — single unconstrained retrieval (runtime-only; not a config choice)

Configuration

CorpusRouterConfig (the runtime knobs):

Field Default Meaning
min_corpora 1 Floor on selected set; always wins over score_threshold
max_corpora 0 Cap on selected set (0 = unbounded)
score_threshold 0.0 Drop corpora scoring below this (subject to min_corpora)
fan_out True Federated execution (per-corpus retrieval, RRF-fused) vs single-pass
rrf_k 60 RRF constant for fan-out fusion

Validated at construction: min_corpora >= 1, max_corpora >= 0, max_corpora >= min_corpora (when set), rrf_k > 0.

rag.yaml surface

Operators configure routing under retrieval.routing (CorpusRoutingConfig in rag-config) and register corpora under corpora (CorpusDefn):

backends:
  corpus_store:
    provider: in_memory          # none | in_memory | postgres
retrieval:
  routing:
    strategy: static_rules        # explicit | static_rules | learned | all
    fan_out: true
    max_corpora: 3
    score_threshold: 0.0
    learned_temperature: 1.0
    rules:
      - name: legal-keywords
        corpus_ids: [legal]
        match_any: [contract, liability, gdpr]
corpora:
  - id: legal
    tenant_id: acme
    display_name: Legal Documents
    term_weights: {contract: 2.0, gdpr: 1.8}   # learned-classifier profile

The config strategy decides which classifiers are wired; the router still auto-detects the effective strategy per query (configured vs effective). unconstrained is intentionally absent from the config enum — it is a runtime fallback only.

Classifiers

StaticRuleClassifier

Deterministic, config-driven. Each CorpusRule fires when the query matches match_any (case-insensitive token / substring) or match_regex; an empty match condition fires unconditionally (always-on default boost). A fired rule either boosts (exclude=False, additive) or removes (exclude=True) its corpus_ids. Exclude wins over boost for the same corpus.

LearnedCorpusClassifier

Dependency-free per-corpus term-weight scorer. Each corpus carries a profile (metadata["term_weights"], or an explicit profiles= map); the query bag-of-words is scored as a dot product against each profile, normalised with a temperature-scaled softmax. A uniform softmax (no profiles) reads as "no signal" and the router falls through. temperature must be > 0.

CorpusClassifier (extension point)

Any object implementing the runtime-checkable Protocol can be plugged in as the learned classifier:

class CorpusClassifier(Protocol):
    def score(self, *, text: str, corpora: Sequence[Corpus]) -> Mapping[CorpusId, float]: ...

Higher = more relevant; scores combine additively with rule boosts. To exclude a corpus, omit it or return a non-positive score — hard exclusion is the static-rule layer's job.

ragctl corpus

Config-driven smoke commands (no DB needed with in_memory):

ragctl corpus list  -f rag.yaml                     # show provider, tenant, seeded corpora
ragctl corpus seed  -f rag.yaml                     # upsert corpora (no-op for in_memory)
ragctl corpus route "reset my password" -f rag.yaml # print the routing decision
ragctl corpus route "anything" -f rag.yaml -c legal # pin → EXPLICIT

corpus route prints configured vs effective strategy, fan_out, candidate count, selected corpora, per-corpus scores + reasons. See the full ragctl reference.

Internals

  • Selection precedence (highest first): explicit pin → static rules → learned classifier → all visible. Effective strategy is auto-detected: a fired rule beats a learned signal, which beats all.
  • Fan-out fusion reuses rrf_fuse (Step 2.5), weighting each corpus's ranking by its relevance score (floored at a small positive value so a zero score does not erase a selected corpus).
  • Observability: route opens a corpus.route OTel span (rag.corpus_route.{strategy,candidate_n,selected_n,fan_out,selected,results_n,elapsed_ms}), emits a corpus.route_decision structured-log event, and appends a corpus.route audit event when an AuditWriter is wired.

Extension points

  • Custom scorer — implement CorpusClassifier and pass as learned_classifier.
  • Custom store — implement the CorpusStore SPI; provider: postgres ships a PgCorpusStore (see backends reference).
  • Single-pass vs fan-out — fan_out: false for homogeneous corpora with directly comparable scores.

See also