Skip to content

Latest commit

 

History

History
295 lines (230 loc) · 11.3 KB

File metadata and controls

295 lines (230 loc) · 11.3 KB

Reference — rag-core

The rag-core package owns the domain types, the plugin SPI ABCs, the noop in-memory implementations, the AuditWriter facade, and the error hierarchy.

This page focuses on the public type surface introduced and modified in Step 1.1a. For the SPI ABCs themselves, read the module docstrings in rag_core.spi.*; for runtime behavior, see the conformance tests in tests/contract/.


Overview

All public types live in rag_core.types and are re-exported from rag_core. Every model is a frozen Pydantic v2 value object — create new instances, don't mutate.

Type Purpose Introduced / changed
RequestContext Per-request envelope threaded through every SPI Step 1.1a (new)
Principal User/service identity + ACL labels Step 0.2 (added acl_labels in Step 1.1a)
PiiPolicy Per-tenant PII handling rules Step 1.1a (new)
Budget Per-request resource envelope (tokens, dollars, wall_ms, iter) Step 1.1a (new)
Chunk Document piece Step 0.2 (added acl_labels, trust_level, content_ref in 1.1a)
Embedding Dense vector for a chunk Step 0.2 (added tenant_id, acl_labels, dtype in 1.1a)
BlobRef Lazy pointer to chunk text in Storage Step 1.1a (new) — ADR-0007
ChunkRef ID-only retrieval result Step 1.1a (new)
QueryPlan / PlanNode / Cost Planner output, cost-aware dispatch Step 1.1a (new) — ADR-0008
StageEvent Typed cross-stage observation Step 1.1a (new)
TrustLevel Chunk provenance for prompt-injection defense Step 1.1a (new)
EmbeddingDtype float32 / int8 / binary quantization Step 1.1a (new) — ADR-0009
IndexHint / WriteVolume Scale-tier hint passed to IndexBackend writes Step 1.1b (new) — ADR-0009
PiiAction block / redact / mask / encrypt / tag_only / allow Step 1.1a (new)

Usage

Constructing a RequestContext at the gateway boundary

from rag_core.types import (
    Budget,
    PiiAction,
    PiiPolicy,
    Principal,
    PrincipalId,
    PrincipalKind,
    RequestContext,
    TenantId,
)

ctx = RequestContext(
    tenant_id=TenantId("acme"),
    principal=Principal(
        id=PrincipalId("alice@acme"),
        kind=PrincipalKind.user,
        display_name="Alice",
        tenant_id=TenantId("acme"),
        acl_labels=frozenset({"engineering", "public"}),
    ),
    pii_policy=PiiPolicy(action=PiiAction.redact),
    budget=Budget(max_tokens=4000, max_dollars=0.05, max_wall_ms=2000),
)

RequestContext validates that tenant_id == principal.tenant_id and is frozen. Derive a per-turn variant with ctx.model_copy(update=...).

Passing ctx through SPIs

Every public SPI method takes ctx as its first argument. Tenant scoping is derived from ctx.tenant_id:

await vector_store.bulk_index(ctx, embeddings, hint=IndexHint(estimated_size=50_000))
refs = await vector_store.retrieve_ids(ctx, query_vec, top_k=10, corpus_ids=[])

The tests/contract/spi_signature.py linter fails CI if an SPI method is added without ctx: RequestContext as its first argument.

Read / write SPI split (Step 1.1b)

VectorStore, KeywordStore, and GraphStore are now thin composite ABCs over two narrower roles:

Role Methods Used by
*RetrievalBackend retrieve_ids, hydrate (where applicable), query (graph) Gateway query path, rerankers
*IndexBackend bulk_index, stream_index, bulk_delete (+ graph node/edge bulk variants) Ingest pipeline

Backends serving both sides keep inheriting from the composite class (VectorStore, KeywordStore, GraphStore) — nothing changes for implementers. Consumers should depend on the narrower role for clarity.

# Hot-path discipline: retrieve cheap refs, rerank, hydrate the survivors
refs = await vector_backend.retrieve_ids(ctx, qvec, top_k=200, corpus_ids=[])
top  = await reranker.rerank(ctx, refs, query_text)
hot  = await keyword_backend.hydrate(ctx, top[:20])

Embedder similarly exposes bulk_embed(ctx, texts, chunk_ids) (canonical batch entry) and a thin embed(ctx, text, chunk_id) for single-item callers. See performance.md for the p99 budgets per method.

Choosing an IndexHint

The hint is consulted at write time (and at initialize() time for backends that need to pick an index implementation). Reasonable starting values:

IndexHint(
    estimated_size=50_000,          # current + projected vector count
    recall_target=0.95,             # default; raise to 0.99 for legal/medical
    latency_target_ms=50.0,         # p99 query budget
    write_volume=WriteVolume.low,   # bump to medium/high for streaming ingest
)

PgVector and Qdrant currently honour hint only at initialize() time; the per-call writes ignore it. Per ADR-0009, future backends will use it to pick between flat / ivfflat / HNSW / IVF-PQ index types.

Lazy chunk text via BlobRef

Chunk.content and Chunk.content_ref are mutually exclusive at the storage level (at least one must be set). Tiered-storage backends (ADR-0007) return Chunk instances with content=None, content_ref=BlobRef(...); consumers must hydrate via Storage.get(ctx, content_ref.uri) before reading text.

ref = BlobRef(uri="s3://blobs/c1.txt", size_bytes=4096)
chunk = Chunk(
    document_id=..., tenant_id=..., corpus_id=...,
    position=0,
    content_ref=ref,    # content is None
)

Quantized embeddings

emb = Embedding(
    chunk_id=cid, tenant_id=tid, model="bge-large",
    vector=[...], dimension=1024,
    dtype=EmbeddingDtype.int8,   # ADR-0009 quantization tier
    acl_labels=frozenset({"public"}),
)

Cost-aware planner types

plan = QueryPlan(
    request_id=ctx.request_id,
    nodes=[
        PlanNode(
            op="vector.retrieve",
            backend="qdrant",
            estimated_cost=Cost(ms_estimate=12, tokens_estimate=0, dollars_estimate=0.0002),
        ),
        PlanNode(
            op="rerank.precise",
            backend="cohere-rerank-v3",
            estimated_cost=Cost(ms_estimate=180, tokens_estimate=400, dollars_estimate=0.012),
        ),
    ],
)
if plan.total_estimated_cost.dollars_estimate > ctx.budget.max_dollars:
    plan = mutate_for_budget(plan, ctx.budget)   # ADR-0008

Internals

Validation policy

RequestContext is validated once at construction time (the gateway boundary) and treated as trusted downstream. Hot-path code MUST NOT re-validate — see performance.md for the rule and the rationale.

Why acl_labels is now a typed field

Before Step 1.1a, ACL labels lived in Chunk.metadata["acl_labels"] — an untyped dict. Promoting them to a typed frozenset[str] on Chunk / Embedding / Principal:

  • Catches typos at type-check time.
  • Lets the PolicyEngine compare sets directly without dict lookups in hot paths.
  • Survives serialization round-trips with predictable types.

Why trust_level is on Chunk

The prompt-injection defense layer (future Phase 3 work) uses Chunk.trust_level to decide whether to wrap content in <untrusted_content>...</untrusted_content> markers when assembling the LLM prompt. Putting it on the chunk (not the document) lets per-chunk override — a user-uploaded section in an otherwise-trusted document remains flagged.


Extension points

Adding a field to RequestContext is a minor-version change:

  1. Add the field with a sensible default in rag_core.types.RequestContext.
  2. Update docs/architecture/request-context.md (Shape).
  3. Update the gateway's RequestContext construction site.
  4. Update downstream consumers that need it.

Removing or renaming a field is a major change and requires an ADR.


Concurrency primitives

rag-core also ships two cross-cutting concurrency primitives used by ingest and retrieval consumers (added in Step 1.1d):

Name Module Purpose
Pipeline, Stage rag_core.pipeline Async DAG with bounded queues + per-stage worker counts; drives the ingest write path
Batcher[Req, Resp] rag_core.batcher DataLoader-pattern middleware; coalesces concurrent SPI calls into batched provider calls

Full design notes, usage, and error semantics live in docs/architecture/pipeline-batcher.md. The reviewer checklist in docs/architecture/performance.md calls out where each must be used.

Cache SPIs (Step 1.1e)

Three distinct cache SPIs replace the single semantic-cache plan — each keyed differently, each with a different invalidation contract:

SPI Module Key Invalidation
EmbeddingCache rag_core.spi.embedding_cache (tenant, model_id, model_version, text_hash) Model / version change
RetrievalCache rag_core.spi.retrieval_cache (tenant, plan_hash, corpus_version) Corpus mutation (version bump)
AnswerCache rag_core.spi.answer_cache (tenant, plan_hash, corpus_version, policy_version) Any of the above + policy rotation

Each SPI takes ctx: RequestContext as first arg (linter-enforced) and ships with a noop in-memory reference impl in rag_core.spi.noop. See docs/architecture/caching.md for the full contract, hit-rate targets, and reference Redis-backed deployment shape.

FilterExpr mini-language (Step 2.1)

rag_core.filter ships the AST that retrieval backends accept on retrieve_ids — relocated here from rag-policy in Step 2.1 so the SPI boundary can use it without inverting the dependency graph.

from rag_core.filter import and_, any_in, eq, evaluate

expr = and_(
    eq("tenant_id", "acme"),
    any_in("acl_labels", ["public", "engineering"]),
)

# In-memory reference semantics — every backend translator MUST agree.
evaluate(expr, {"tenant_id": "acme", "acl_labels": {"engineering"}})  # True
Symbol Purpose
Eq, AnyIn, And, Or, Not, TrueExpr Discriminated-union node types
eq(), any_in(), and_(), or_(), not_(), true() Sugar constructors
FilterExpr The discriminated-union type alias
evaluate(expr, attrs) → bool Reference evaluator (used by noop stores, tests, and translator parity gates)

rag_policy.filter continues to re-export the same names for backwards compatibility — from rag_policy import FilterExpr, eq, ... still works. See docs/architecture/retrieval-read-layer.md for the full read-layer contract, allowed fields, and backend-translator rules.

Related