Skip to content

Latest commit

 

History

History
211 lines (159 loc) · 7.32 KB

File metadata and controls

211 lines (159 loc) · 7.32 KB

rag-retrieval — Hybrid retrieval orchestration

rag-retrieval (Step 2.5) is the first package that consumes the storage backends shipped in Steps 2.1 – 2.4. It combines dense vector, sparse keyword (BM25), and graph traversal results into a single ranked list via Reciprocal Rank Fusion.

Overview

Two public surfaces:

Symbol Purpose
rrf_fuse Pure function over already-ranked ChunkRef lists
HybridRetriever Async orchestrator: fan out to vector / keyword / graph, fuse, return top-k

The split lets you compose them independently — call rrf_fuse directly when you have ranked lists from somewhere else (an A/B shadow path, a cached ranking, a reranker output to merge with the live one), or use the orchestrator for the common case.

Usage

rrf_fuse

from rag_retrieval import rrf_fuse

# Three pre-ranked lists, each sorted by descending source relevance.
fused = rrf_fuse(
    [vector_results, keyword_results, graph_results],
    weights=[1.0, 1.0, 0.5],   # optional; default = all 1.0
    k=60,                      # smoothing constant (default 60)
    top_k=10,                  # truncate output (default = no truncate)
)

Returns list[ChunkRef] ordered by descending RRF score; the per-source score is replaced by the fused RRF score on each returned ChunkRef.

Formula

For a chunk d appearing in input lists L_1, …, L_n with weights w_1, …, w_n and smoothing constant k:

RRF(d) = Σ_i  w_i / (k + rank_i(d))

where rank_i(d) is the 1-indexed rank of d in L_i (chunk treated as absent — contribution 0 — if it does not appear in L_i).

Validation

  • k > 0 — required.
  • Weights ≥ 0; a weight of 0 disables that source's contribution.
  • Length of weights must equal length of rankings.
  • All input ChunkRef must share a single tenant_id — fusion raises RetrievalError on mixed tenants (defence-in-depth against an upstream wiring bug).

HybridRetriever

from rag_retrieval import HybridRetriever, HybridWeights

retriever = HybridRetriever(
    vector_backend=pgvector,        # any VectorRetrievalBackend
    keyword_backend=elasticsearch,  # any KeywordRetrievalBackend
    graph_backend=neo4j,            # any GraphRetrievalBackend
    weights=HybridWeights(vector=1.0, keyword=1.0, graph=0.5),
    k=60,
    policy_engine=policy,           # optional — see below
)

refs = await retriever.retrieve(
    ctx,
    vector=query_vector,            # skipped if None
    query_text="quick brown fox",   # skipped if None
    graph_seeds=["entity-42"],      # skipped if None or empty
    corpus_ids=[CorpusId("docs")],
    top_k=10,
    per_backend_top_k=30,           # default: max(top_k*3, 30)
    filters=acl_filter,             # FilterExpr passed to every backend
    graph_hops=1,
)

Per-call source routing

A backend wired at construction time is skipped on calls where its input is None (or, for graph_seeds, empty). This keeps query-shape routing (Step 2.10's job) a caller concern — you can wire all three backends once and decide per query which sources to exercise.

Per-backend top-k

per_backend_top_k is the limit passed to each backend's retrieve_ids. Default is max(top_k * 3, 30) — wider than the final top_k because RRF benefits from seeing low-rank candidates that consistently appear across sources. Reduce it for tighter latency.

Partial-failure tolerance

If one backend raises an Exception, the orchestrator logs a hybrid_retriever.source_failed event and drops that source from the fusion — the surviving backends still produce a fused result. If every active source fails, the orchestrator raises RetrievalError so the caller sees the degradation.

KeyboardInterrupt and SystemExit propagate out unchanged (asyncio's gather(return_exceptions=True) does this naturally).

PolicyEngine integration

When policy_engine is supplied, the retriever calls policy_engine.filter_pushdown(ctx, PolicyDecision.read_chunk) once per request and merges the result with any caller-supplied filters via And before passing them down to every backend. Callers without a policy engine should embed tenant / ACL push-down in filters themselves.

This is the canonical read_chunk policy call site for hybrid retrieval — see docs/architecture/policy-engine.md.

Why policy_engine is typed Any | None

rag-retrieval deliberately depends only on rag-core (see packages/retrieval/pyproject.toml) so the wheel stays composable without pulling sibling packages. To honour that, the policy_engine constructor parameter is annotated Any | None and rag_policy is imported lazily inside _merge_policy_filters. The use site uses cast(FilterExpr | None, ...) to restore the SPI's declared return type — PolicyEngine.filter_pushdown is typed -> FilterExpr at its source in rag_policy.engine, so the cast crosses the lazy-import boundary without weakening the public guarantee.

Internals

Default graph adapter

GraphRetrievalBackend.expand() returns NeighborResult (node-shaped); RRF fuses chunk-shaped lists. The default adapter assumes a chunk-graph projection where node_id == chunk_id:

ChunkRef(
    chunk_id=ChunkId(neighbor.node_id),
    tenant_id=ctx.tenant_id,
    score=1.0 / max(neighbor.distance, 1),
)

The score here is informational only — RRF uses rank, not score, so fusion treats the SPI-guaranteed ascending-distance ordering as the authoritative rank.

Metadata preservation

When a chunk appears in more than one source list, fusion sums their RRF contributions but keeps the first source's tenant_id, acl_labels, corpus_id, and metadata on the returned ChunkRef. That matches RRF's "rank is everything" framing — the metadata is opaque to the fusion decision.

Extension points

Custom graph adapter

GraphRAG (Step 2.9) will model entities as nodes, not chunks, so the default node_id == chunk_id assumption breaks. Pass your own adapter to the HybridRetriever constructor:

from rag_retrieval import GraphAdapter

async def gobble_adapter(ctx, neighbors):
    # Look up chunk_ids from a node-property table, attach ACL labels,
    # etc.
    ...

retriever = HybridRetriever(
    graph_backend=neo4j,
    graph_adapter=gobble_adapter,
    # ...
)

The adapter is async so real implementations can hydrate properties via I/O.

Per-query weighting

HybridWeights is frozen and constructor-only — recreating the retriever per request is cheap (no SPI state). Step 2.10's retrieval router will own per-query weighting via per-corpus / per-intent rules.

Pre-fusion filtering

The orchestrator does not transform per-source results before fusion — that's deliberately the reranker's job (Step 2.7). If you need to strip duplicates before RRF (e.g. when sources fan out across corpora that share chunks), do it in a wrapper around the backend's retrieve_ids rather than inside the orchestrator.

See also