diff --git a/TRACKER.md b/TRACKER.md index 0404dca..13211d0 100644 --- a/TRACKER.md +++ b/TRACKER.md @@ -10,7 +10,7 @@ **Last updated:** 2026-05-24 **Current phase:** Phase 1 — Ingestion + Knowledge Store -**Next action:** Phase 1 Step 1.1a — Core type & SPI refactor (RequestContext, typed ACL, dtype, BlobRef, QueryPlan, ChunkRef, StageEvent) +**Next action:** Phase 1 Step 1.1b — SPI split (Retrieval/Index, bulk + streaming + ID-only methods, IndexHint) > **Refactor window (Steps 1.1a–1.1f):** Before resuming the connectors framework (1.2), we insert a six-step refactor that locks in architecture + optimization decisions which are very expensive to retrofit later (PolicyEngine PDP, RequestContext-threaded SPIs, split Retrieval/Index backends, bulk + streaming + ID-only methods, Pipeline + Batcher primitives, three-way cache split, hot-path discipline). See [docs/adr/ADR-0005…0009] and [docs/architecture/policy-engine.md], [request-context.md], [caching.md], [performance.md]. @@ -32,14 +32,14 @@ | Phase | Title | Steps | ✅ Done | Remaining | |-------|-------|------:|-------:|----------:| | 0 | Foundation | 13 | **13** | 0 | -| 1 | Ingestion + Knowledge Store | 16 | **1** | 15 | +| 1 | Ingestion + Knowledge Store | 16 | **2** | 14 | | 2 | Retrieval Engine | 11 | 0 | 11 | | 3 | Gateway & Agent Runtime | 11 | 0 | 11 | | 4 | Reliability | 6 | 0 | 6 | | 5 | Eval & Observability | 7 | 0 | 7 | | 6 | Governance & Tenancy | 10 | 0 | 10 | | 7 | Pilot, Harden, GA | 10 | 0 | 10 | -| **Total** | | **84** | **14** | **70** | +| **Total** | | **84** | **15** | **69** | --- @@ -68,7 +68,7 @@ | Step | Title | Status | Branch | PR | Key Deliverables | |------|-------|--------|--------|----|-----------------| | 1.1 | Storage backends | ✅ | `build/phase-1/step-1.1-storage-backends` | [#40](https://github.com/officialCodeWork/AgentContextOS/pull/40) | `rag-backends` package: `PgVectorStore` (asyncpg + pgvector, ivfflat), `QdrantVectorStore` (query_points API), `RedisCache`, `S3Storage` (aioboto3, MinIO-compatible), `LocalFileStorage`; integration tests (skip-if-no-service); MinIO added to dev stack; `task test-integration` + `task test-backends`; ADR-0004 | -| 1.1a | Core type & SPI refactor | ⏳ | — | — | `RequestContext` frozen model threaded through every SPI; `tenant_id` + `acl_labels` typed required on `Chunk`/`Embedding` (not metadata dict); `trust_level` on `Chunk` for prompt-injection defense; `dtype` on `Embedding` (float32/int8/binary); `BlobRef` for lazy chunk text; `QueryPlan` + `ChunkRef` types; typed `StageEvent`. Conformance tests updated. ADR-0005 (PolicyEngine) + ADR-0008 (cost-aware planner) referenced. | +| 1.1a | Core type & SPI refactor | ✅ | `build/phase-1/step-1.1a-core-type-spi-refactor` | _pending_ | `RequestContext` frozen model threaded through every SPI; `tenant_id` + `acl_labels` typed required on `Chunk`/`Embedding` (not metadata dict); `trust_level` on `Chunk` for prompt-injection defense; `dtype` on `Embedding` (float32/int8/binary); `BlobRef` for lazy chunk text; `QueryPlan` + `ChunkRef` + `Cost` + `PlanNode` types; typed `StageEvent`. `tests/contract/spi_signature.py` linter (RequestContext-first); rag-backends (`PgVectorStore`, `QdrantVectorStore`, `RedisCache`, `S3Storage`, `LocalFileStorage`) migrated; conformance + integration tests updated; `Budget.spend()` for agent-loop sub-turn budgets; schemas regenerated. ADR-0005 / ADR-0007 / ADR-0008 / ADR-0009 referenced. | | 1.1b | SPI split — Retrieval/Index, bulk + streaming + ID-only | ⏳ | 1.1a | — | Split `VectorStore`/`KeywordStore`/`GraphStore` into `RetrievalBackend` (read) + `IndexBackend` (write); add `bulk_embed`/`bulk_index`/`bulk_delete`; add async-iterator ingest variants; add `retrieve_ids` + `hydrate` pair; add `IndexHint` parameter for scale-tier index selection. ADR-0009 (vector index strategy) authored. | | 1.1c | PolicyEngine package | ⏳ | 1.1a | — | New `packages/policy/` (`rag-policy`): `PolicyEngine` SPI + noop impl; single decision point consulted by every retrieval/ingest path for ACL, PII, quotas, redaction. Replaces scattered checks across Steps 1.7, 4.5, 6.3, 6.5. ADR-0005 finalized. | | 1.1d | Pipeline + Batcher primitives | ⏳ | 1.1a | — | `Pipeline` primitive in `rag-core`: async DAG with bounded queues, per-stage worker counts, backpressure (used by Step 1.10 write path). `Batcher[Req, Resp]` middleware (DataLoader pattern) coalescing concurrent SPI calls into batched provider calls; sits under Embedder/Reranker SPIs. | diff --git a/dist/schemas/Chunk.json b/dist/schemas/Chunk.json index 54c97c1..be19b56 100644 --- a/dist/schemas/Chunk.json +++ b/dist/schemas/Chunk.json @@ -1,5 +1,53 @@ { - "description": "A contiguous piece of a Document produced by the chunking pipeline.", + "$defs": { + "BlobRef": { + "description": "Reference to a blob stored in the ``Storage`` SPI rather than inline.\n\nUsed by chunks whose text exceeds the inline-storage threshold (see\nADR-0007 tiered storage). Callers must hydrate via ``Storage.get(uri)``\nonly when the text is actually needed.", + "properties": { + "uri": { + "title": "Uri", + "type": "string" + }, + "size_bytes": { + "title": "Size Bytes", + "type": "integer" + }, + "content_type": { + "default": "text/plain; charset=utf-8", + "title": "Content Type", + "type": "string" + }, + "sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Sha256" + } + }, + "required": [ + "uri", + "size_bytes" + ], + "title": "BlobRef", + "type": "object" + }, + "TrustLevel": { + "description": "Provenance of a chunk's text \u2014 used by the prompt-injection defense.\n\n``trusted`` \u2014 first-party content authored under tenant control.\n``ingested`` \u2014 content fetched from a known external source (vetted feed,\n enterprise SharePoint, etc.).\n``user_supplied`` \u2014 content directly contributed by an end-user channel\n (web upload, chat-attached file, \u2026) which may contain\n adversarial instructions.", + "enum": [ + "trusted", + "ingested", + "user_supplied" + ], + "title": "TrustLevel", + "type": "string" + } + }, + "description": "A contiguous piece of a Document produced by the chunking pipeline.\n\nStep 1.1a promoted three governance-relevant fields out of the\n``metadata`` dict into typed required fields:\n\n- ``acl_labels`` \u2014 set of ACL labels the PolicyEngine compares against\n the requesting principal's ``acl_labels``.\n- ``trust_level`` \u2014 provenance, used by the prompt-injection defense.\n- ``content_ref`` \u2014 optional ``BlobRef`` for tiered text storage; when set,\n ``content`` may be ``None`` and callers must hydrate via Storage.", "properties": { "id": { "title": "Id", @@ -18,8 +66,27 @@ "type": "string" }, "content": { - "title": "Content", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Content" + }, + "content_ref": { + "anyOf": [ + { + "$ref": "#/$defs/BlobRef" + }, + { + "type": "null" + } + ], + "default": null }, "position": { "title": "Position", @@ -49,6 +116,18 @@ "default": null, "title": "Token Count" }, + "acl_labels": { + "items": { + "type": "string" + }, + "title": "Acl Labels", + "type": "array", + "uniqueItems": true + }, + "trust_level": { + "$ref": "#/$defs/TrustLevel", + "default": "ingested" + }, "metadata": { "additionalProperties": true, "title": "Metadata", @@ -64,7 +143,6 @@ "document_id", "tenant_id", "corpus_id", - "content", "position" ], "title": "Chunk", diff --git a/dist/schemas/Embedding.json b/dist/schemas/Embedding.json index 742b752..b32a934 100644 --- a/dist/schemas/Embedding.json +++ b/dist/schemas/Embedding.json @@ -1,10 +1,26 @@ { - "description": "Dense vector representation of a Chunk.", + "$defs": { + "EmbeddingDtype": { + "description": "Numeric representation of an embedding vector at rest.\n\n``float32`` \u2014 full-precision (default).\n``int8`` \u2014 quantized 8-bit integers (\u22484\u00d7 smaller, slight recall loss).\n``binary`` \u2014 1-bit per dim packed (\u224832\u00d7 smaller, larger recall loss,\n used as a coarse first-stage in two-stage retrieval).\n\nSee ADR-0009 for the vector index + quantization strategy.", + "enum": [ + "float32", + "int8", + "binary" + ], + "title": "EmbeddingDtype", + "type": "string" + } + }, + "description": "Dense vector representation of a Chunk.\n\nStep 1.1a promoted ``tenant_id`` and ``acl_labels`` to typed required\nfields and added ``dtype`` to capture int8 / binary quantization.", "properties": { "chunk_id": { "title": "Chunk Id", "type": "string" }, + "tenant_id": { + "title": "Tenant Id", + "type": "string" + }, "model": { "title": "Model", "type": "string" @@ -20,6 +36,18 @@ "title": "Dimension", "type": "integer" }, + "dtype": { + "$ref": "#/$defs/EmbeddingDtype", + "default": "float32" + }, + "acl_labels": { + "items": { + "type": "string" + }, + "title": "Acl Labels", + "type": "array", + "uniqueItems": true + }, "created_at": { "format": "date-time", "title": "Created At", @@ -28,6 +56,7 @@ }, "required": [ "chunk_id", + "tenant_id", "model", "vector", "dimension" diff --git a/dist/schemas/Principal.json b/dist/schemas/Principal.json index fd28891..5256dcd 100644 --- a/dist/schemas/Principal.json +++ b/dist/schemas/Principal.json @@ -10,7 +10,7 @@ "type": "string" } }, - "description": "Authenticated identity \u2014 user, service account, or group.", + "description": "Authenticated identity \u2014 user, service account, or group.\n\n``acl_labels`` is the set of ACL labels the principal carries; it is the\n*typed* counterpart of the ``acl_labels`` field on ``Chunk`` /\n``Embedding``. The PolicyEngine compares the two during retrieval.", "properties": { "id": { "title": "Id", @@ -46,6 +46,14 @@ }, "title": "Roles", "type": "array" + }, + "acl_labels": { + "items": { + "type": "string" + }, + "title": "Acl Labels", + "type": "array", + "uniqueItems": true } }, "required": [ diff --git a/docs/README.md b/docs/README.md index c004557..bff953e 100644 --- a/docs/README.md +++ b/docs/README.md @@ -20,6 +20,7 @@ |------|-------------| | [ragctl.md](reference/ragctl.md) | Full `ragctl` command reference — public usage, internals, extension points | | [backends.md](reference/backends.md) | `rag-backends` reference — PgVectorStore, QdrantVectorStore, RedisCache, S3Storage, LocalFileStorage | +| [rag-core.md](reference/rag-core.md) | `rag-core` type surface — `RequestContext`, `Budget`, `BlobRef`, `QueryPlan`, `ChunkRef`, `StageEvent` | ## guides/ diff --git a/docs/architecture/request-context.md b/docs/architecture/request-context.md index 69fcc54..b1b7126 100644 --- a/docs/architecture/request-context.md +++ b/docs/architecture/request-context.md @@ -1,5 +1,8 @@ # RequestContext — the per-request envelope +**Status:** Implemented in Step 1.1a (PR landing the `RequestContext` type, +ctx-first SPI signatures, and the `tests/contract/spi_signature.py` linter). + ## Overview `RequestContext` is a frozen Pydantic v2 model that travels with every SPI call. It carries everything downstream layers need to make correct, governed, budgeted decisions without reaching back to global state or guessing. @@ -22,14 +25,15 @@ A single typed envelope on every call signature eliminates all three. class RequestContext(BaseModel): model_config = {"frozen": True} - request_id: str # ULID; round-trips with the trace - tenant_id: str # required; non-empty - principal: Principal # user/service identity + ACL set - pii_policy: PiiPolicy # per-tenant: redact | mask | encrypt | tag-only | block + request_id: RequestId # ULID-like ID; round-trips with the trace + tenant_id: TenantId # required; matched against principal.tenant_id + principal: Principal # user/service identity + ACL labels + pii_policy: PiiPolicy # per-tenant: redact | mask | encrypt | tag_only | block | allow trace: TraceContext # OTel span + correlation IDs - budget: Budget # tokens, cost (dollars), wall_ms, max_iter + budget: Budget # tokens, dollars, wall_ms, max_iter feature_flags: frozenset[str] # for shadow mode, A/B, killswitches corpus_routing_hint: str | None # optional sticky routing across agent-loop turns + created_at: datetime # autopopulated at construction ``` `Principal`, `PiiPolicy`, `TraceContext`, `Budget` are all frozen Pydantic models in `rag_core.types`. diff --git a/docs/reference/rag-core.md b/docs/reference/rag-core.md new file mode 100644 index 0000000..7d1e46b --- /dev/null +++ b/docs/reference/rag-core.md @@ -0,0 +1,185 @@ +# 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 | +| `PiiAction` | block / redact / mask / encrypt / tag_only / allow | Step 1.1a (new) | + +--- + +## Usage + +### Constructing a `RequestContext` at the gateway boundary + +```python +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`: + +```python +await vector_store.upsert(ctx, embeddings) +results = await vector_store.query(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. + +### 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. + +```python +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 + +```python +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 + +```python +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](../architecture/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 +`...` 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. + +--- + +## Related + +- [docs/architecture/request-context.md](../architecture/request-context.md) — design rationale. +- [docs/architecture/policy-engine.md](../architecture/policy-engine.md) — primary `ctx` consumer (Step 1.1c). +- [docs/architecture/performance.md](../architecture/performance.md) — hot-path / validation discipline. +- [ADR-0005](../adr/ADR-0005-policy-engine.md), [ADR-0007](../adr/ADR-0007-tiered-storage.md), [ADR-0008](../adr/ADR-0008-cost-aware-planner.md), [ADR-0009](../adr/ADR-0009-vector-index-strategy.md). +- `tests/contract/spi_signature.py` — the signature linter. diff --git a/packages/backends/src/rag_backends/cache/redis.py b/packages/backends/src/rag_backends/cache/redis.py index f45098b..b73546e 100644 --- a/packages/backends/src/rag_backends/cache/redis.py +++ b/packages/backends/src/rag_backends/cache/redis.py @@ -15,6 +15,7 @@ import redis.asyncio as aioredis from rag_core.spi.cache import Cache +from rag_core.types import RequestContext from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -60,25 +61,32 @@ async def close(self) -> None: # Cache SPI # ------------------------------------------------------------------ - def _k(self, key: str) -> str: - return self._prefix + key + def _k(self, ctx: RequestContext, key: str) -> str: + # Namespace every key by tenant to enforce isolation at the cache layer. + return f"{self._prefix}{ctx.tenant_id}:{key}" - async def get(self, key: str) -> bytes | None: - value: bytes | None = await self._redis.get(self._k(key)) + async def get(self, ctx: RequestContext, key: str) -> bytes | None: + value: bytes | None = await self._redis.get(self._k(ctx, key)) return value - async def set(self, key: str, value: bytes, ttl_seconds: int | None = None) -> None: + async def set( + self, + ctx: RequestContext, + key: str, + value: bytes, + ttl_seconds: int | None = None, + ) -> None: ttl = ttl_seconds if ttl_seconds is not None else self._default_ttl if ttl is not None: - await self._redis.setex(self._k(key), ttl, value) + await self._redis.setex(self._k(ctx, key), ttl, value) else: - await self._redis.set(self._k(key), value) + await self._redis.set(self._k(ctx, key), value) - async def delete(self, key: str) -> None: - await self._redis.delete(self._k(key)) + async def delete(self, ctx: RequestContext, key: str) -> None: + await self._redis.delete(self._k(ctx, key)) - async def exists(self, key: str) -> bool: - count: int = await self._redis.exists(self._k(key)) + async def exists(self, ctx: RequestContext, key: str) -> bool: + count: int = await self._redis.exists(self._k(ctx, key)) return count > 0 async def health(self) -> bool: diff --git a/packages/backends/src/rag_backends/storage/local.py b/packages/backends/src/rag_backends/storage/local.py index 4e227bc..5db7821 100644 --- a/packages/backends/src/rag_backends/storage/local.py +++ b/packages/backends/src/rag_backends/storage/local.py @@ -6,8 +6,8 @@ Usage:: store = LocalFileStorage(root="/tmp/rag-dev") - await store.put("docs/a.txt", b"hello") - data = await store.get("docs/a.txt") + await store.put(ctx, "docs/a.txt", b"hello") + data = await store.get(ctx, "docs/a.txt") """ from __future__ import annotations @@ -19,6 +19,7 @@ import aiofiles import aiofiles.os from rag_core.spi.storage import Storage +from rag_core.types import RequestContext from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -27,8 +28,9 @@ class LocalFileStorage(Storage): """Storage backend that writes objects to a local directory tree. - Keys are mapped to filesystem paths relative to ``root``. Directory - separators in keys (``/``) become OS path separators. + Keys are mapped to filesystem paths under ``//``. + The tenant prefix is derived from ``ctx.tenant_id`` so callers cannot + accidentally read another tenant's blobs. Args: root: Absolute path to the root directory. Created on first use. @@ -37,11 +39,15 @@ class LocalFileStorage(Storage): def __init__(self, root: str | Path) -> None: self._root = Path(root) - def _path(self, key: str) -> Path: + def _tenant_root(self, ctx: RequestContext) -> Path: + return self._root / str(ctx.tenant_id) + + def _path(self, ctx: RequestContext, key: str) -> Path: # Resolve to prevent path-traversal attacks - resolved = (self._root / key).resolve() - if not str(resolved).startswith(str(self._root.resolve())): - raise ValueError(f"Key {key!r} resolves outside storage root") + tenant_root = self._tenant_root(ctx) + resolved = (tenant_root / key).resolve() + if not str(resolved).startswith(str(tenant_root.resolve())): + raise ValueError(f"Key {key!r} resolves outside tenant storage root") return resolved async def _ensure_parent(self, path: Path) -> None: @@ -51,43 +57,49 @@ async def _ensure_parent(self, path: Path) -> None: # Storage SPI # ------------------------------------------------------------------ - async def put(self, key: str, data: bytes, content_type: str | None = None) -> None: - path = self._path(key) + async def put( + self, + ctx: RequestContext, + key: str, + data: bytes, + content_type: str | None = None, + ) -> None: + path = self._path(ctx, key) await self._ensure_parent(path) async with aiofiles.open(path, "wb") as f: await f.write(data) - async def get(self, key: str) -> bytes: - path = self._path(key) + async def get(self, ctx: RequestContext, key: str) -> bytes: + path = self._path(ctx, key) if not path.exists(): raise KeyError(key) async with aiofiles.open(path, "rb") as f: data: bytes = await f.read() return data - async def delete(self, key: str) -> None: - path = self._path(key) + async def delete(self, ctx: RequestContext, key: str) -> None: + path = self._path(ctx, key) try: await aiofiles.os.remove(str(path)) except FileNotFoundError: pass # idempotent - async def exists(self, key: str) -> bool: - return self._path(key).exists() + async def exists(self, ctx: RequestContext, key: str) -> bool: + return self._path(ctx, key).exists() - def list_keys(self, prefix: str = "") -> AsyncIterator[str]: - root = self._root - prefix_path = root / prefix + def list_keys(self, ctx: RequestContext, prefix: str = "") -> AsyncIterator[str]: + tenant_root = self._tenant_root(ctx) + prefix_path = tenant_root / prefix async def _gen() -> AsyncIterator[str]: - base = str(root) + base = str(tenant_root) start = str(prefix_path) if not prefix_path.exists(): return for dirpath, _dirs, files in os.walk(start): for fname in files: full = os.path.join(dirpath, fname) - # Return the key relative to root (using forward slashes) + # Return the key relative to tenant root (using forward slashes) rel = os.path.relpath(full, base).replace(os.sep, "/") yield rel diff --git a/packages/backends/src/rag_backends/storage/s3.py b/packages/backends/src/rag_backends/storage/s3.py index 8c4b20d..08957c3 100644 --- a/packages/backends/src/rag_backends/storage/s3.py +++ b/packages/backends/src/rag_backends/storage/s3.py @@ -16,9 +16,9 @@ aws_secret_access_key="minioadmin", ) - await store.put("docs/a.txt", b"hello", content_type="text/plain") - data = await store.get("docs/a.txt") - await store.delete("docs/a.txt") + await store.put(ctx, "docs/a.txt", b"hello", content_type="text/plain") + data = await store.get(ctx, "docs/a.txt") + await store.delete(ctx, "docs/a.txt") """ from __future__ import annotations @@ -29,6 +29,7 @@ import aioboto3 from botocore.exceptions import ClientError from rag_core.spi.storage import Storage +from rag_core.types import RequestContext from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -37,6 +38,9 @@ class S3Storage(Storage): """Storage backend backed by Amazon S3 or any S3-compatible store. + Object keys are namespaced as ``/`` so tenants + cannot accidentally read each other's blobs. + Args: bucket: S3 bucket name. Must already exist. region: AWS region (ignored for MinIO/custom endpoints). @@ -45,7 +49,7 @@ class S3Storage(Storage): aws_access_key_id: Explicit credentials (optional; falls back to the default boto3 credential chain: env vars, ~/.aws, instance role). aws_secret_access_key: Matching secret key. - prefix: Optional key prefix prepended to every object key. + prefix: Optional global key prefix prepended before the tenant segment. """ def __init__( @@ -69,14 +73,23 @@ def __init__( if endpoint_url: self._s3_kwargs["endpoint_url"] = endpoint_url - def _key(self, key: str) -> str: - return self._prefix + key + def _tenant_prefix(self, ctx: RequestContext) -> str: + return f"{self._prefix}{ctx.tenant_id}/" + + def _key(self, ctx: RequestContext, key: str) -> str: + return f"{self._tenant_prefix(ctx)}{key}" # ------------------------------------------------------------------ # Storage SPI # ------------------------------------------------------------------ - async def put(self, key: str, data: bytes, content_type: str | None = None) -> None: + async def put( + self, + ctx: RequestContext, + key: str, + data: bytes, + content_type: str | None = None, + ) -> None: extra: dict[str, str] = {} if content_type: extra["ContentType"] = content_type @@ -84,15 +97,15 @@ async def put(self, key: str, data: bytes, content_type: str | None = None) -> N async with self._session.client("s3", **self._s3_kwargs) as s3: await s3.put_object( Bucket=self._bucket, - Key=self._key(key), + Key=self._key(ctx, key), Body=data, **extra, ) - async def get(self, key: str) -> bytes: + async def get(self, ctx: RequestContext, key: str) -> bytes: async with self._session.client("s3", **self._s3_kwargs) as s3: try: - response = await s3.get_object(Bucket=self._bucket, Key=self._key(key)) + response = await s3.get_object(Bucket=self._bucket, Key=self._key(ctx, key)) body: bytes = await response["Body"].read() return body except ClientError as exc: @@ -100,25 +113,25 @@ async def get(self, key: str) -> bytes: raise KeyError(key) from exc raise - async def delete(self, key: str) -> None: + async def delete(self, ctx: RequestContext, key: str) -> None: async with self._session.client("s3", **self._s3_kwargs) as s3: - await s3.delete_object(Bucket=self._bucket, Key=self._key(key)) + await s3.delete_object(Bucket=self._bucket, Key=self._key(ctx, key)) - async def exists(self, key: str) -> bool: + async def exists(self, ctx: RequestContext, key: str) -> bool: async with self._session.client("s3", **self._s3_kwargs) as s3: try: - await s3.head_object(Bucket=self._bucket, Key=self._key(key)) + await s3.head_object(Bucket=self._bucket, Key=self._key(ctx, key)) return True except ClientError as exc: if exc.response["Error"]["Code"] in ("NoSuchKey", "404", "403"): return False raise - def list_keys(self, prefix: str = "") -> AsyncIterator[str]: - full_prefix = self._prefix + prefix + def list_keys(self, ctx: RequestContext, prefix: str = "") -> AsyncIterator[str]: + tenant_prefix = self._tenant_prefix(ctx) + full_prefix = tenant_prefix + prefix bucket = self._bucket s3_kwargs = self._s3_kwargs - store_prefix = self._prefix session = self._session async def _gen() -> AsyncIterator[str]: @@ -127,8 +140,8 @@ async def _gen() -> AsyncIterator[str]: async for page in paginator.paginate(Bucket=bucket, Prefix=full_prefix): for obj in page.get("Contents", []): raw_key: str = obj["Key"] - # Strip the store-level prefix before returning - yield raw_key[len(store_prefix) :] + # Strip the tenant prefix before returning + yield raw_key[len(tenant_prefix) :] return _gen() diff --git a/packages/backends/src/rag_backends/vector/pgvector.py b/packages/backends/src/rag_backends/vector/pgvector.py index 7e6a476..1fdcbcd 100644 --- a/packages/backends/src/rag_backends/vector/pgvector.py +++ b/packages/backends/src/rag_backends/vector/pgvector.py @@ -4,8 +4,8 @@ store = PgVectorStore(dsn="postgresql://rag:rag@localhost:5432/rag") await store.initialize() # creates table + index once - await store.upsert(embeddings, tenant_id) - results = await store.query(query_vector, top_k=10, tenant_id=..., corpus_ids=[]) + await store.upsert(ctx, embeddings) + results = await store.query(ctx, query_vector, top_k=10, corpus_ids=[]) await store.close() The table ``rag_vector_store`` is created on ``initialize()`` and is separate @@ -22,7 +22,7 @@ import numpy as np from pgvector.asyncpg import register_vector from rag_core.spi.vector_store import VectorStore -from rag_core.types import ChunkId, CorpusId, Embedding, TenantId +from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -148,13 +148,13 @@ async def close(self) -> None: # VectorStore SPI # ------------------------------------------------------------------ - async def upsert(self, embeddings: list[Embedding], tenant_id: TenantId) -> None: + async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: if not embeddings: return pool = await self._get_pool() rows = [ ( - str(tenant_id), + str(ctx.tenant_id), str(emb.chunk_id), "", # corpus_id not available in Embedding — see ADR-0004 emb.model, @@ -169,9 +169,9 @@ async def upsert(self, embeddings: list[Embedding], tenant_id: TenantId) -> None async def query( self, + ctx: RequestContext, vector: list[float], top_k: int, - tenant_id: TenantId, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, ) -> list[tuple[ChunkId, float]]: @@ -180,7 +180,7 @@ async def query( # Build WHERE clause additions extra_where = "" - args: list[Any] = [str(tenant_id), qvec, top_k] + args: list[Any] = [str(ctx.tenant_id), qvec, top_k] if corpus_ids: extra_where = f" AND corpus_id = ANY(${len(args) + 1})" @@ -197,14 +197,14 @@ async def query( return [(ChunkId(row["chunk_id"]), float(row["score"])) for row in rows] - async def delete(self, chunk_ids: list[ChunkId], tenant_id: TenantId) -> None: + async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: if not chunk_ids: return pool = await self._get_pool() async with pool.acquire() as conn: await conn.execute( _SQL_DELETE.format(table=self._table), - str(tenant_id), + str(ctx.tenant_id), [str(c) for c in chunk_ids], ) diff --git a/packages/backends/src/rag_backends/vector/qdrant.py b/packages/backends/src/rag_backends/vector/qdrant.py index 8ac1c24..361a728 100644 --- a/packages/backends/src/rag_backends/vector/qdrant.py +++ b/packages/backends/src/rag_backends/vector/qdrant.py @@ -4,8 +4,8 @@ store = QdrantVectorStore(url="http://localhost:6333") await store.initialize(dimension=1536) # create collection once - await store.upsert(embeddings, tenant_id) - results = await store.query(vector, top_k=10, tenant_id=..., corpus_ids=[]) + await store.upsert(ctx, embeddings) + results = await store.query(ctx, vector, top_k=10, corpus_ids=[]) A single Qdrant collection (``rag_embeddings`` by default) stores all tenants. Tenant isolation is enforced via a ``tenant_id`` payload filter on every query @@ -29,7 +29,7 @@ VectorParams, ) from rag_core.spi.vector_store import VectorStore -from rag_core.types import ChunkId, CorpusId, Embedding, TenantId +from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -104,16 +104,17 @@ async def close(self) -> None: # VectorStore SPI # ------------------------------------------------------------------ - async def upsert(self, embeddings: list[Embedding], tenant_id: TenantId) -> None: + async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: if not embeddings: return + tenant_id = str(ctx.tenant_id) points = [ PointStruct( - id=_chunk_uuid(str(tenant_id), str(emb.chunk_id)), + id=_chunk_uuid(tenant_id, str(emb.chunk_id)), vector=emb.vector, payload={ - _FIELD_TENANT: str(tenant_id), + _FIELD_TENANT: tenant_id, _FIELD_CHUNK: str(emb.chunk_id), _FIELD_CORPUS: "", # corpus_id not in Embedding — see ADR-0004 _FIELD_MODEL: emb.model, @@ -125,14 +126,14 @@ async def upsert(self, embeddings: list[Embedding], tenant_id: TenantId) -> None async def query( self, + ctx: RequestContext, vector: list[float], top_k: int, - tenant_id: TenantId, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, ) -> list[tuple[ChunkId, float]]: must: list[FieldCondition] = [ - FieldCondition(key=_FIELD_TENANT, match=MatchValue(value=str(tenant_id))) + FieldCondition(key=_FIELD_TENANT, match=MatchValue(value=str(ctx.tenant_id))) ] if corpus_ids: must.append( @@ -156,11 +157,11 @@ async def query( if hit.payload and _FIELD_CHUNK in hit.payload ] - async def delete(self, chunk_ids: list[ChunkId], tenant_id: TenantId) -> None: + async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: if not chunk_ids: return - point_ids = [_chunk_uuid(str(tenant_id), str(cid)) for cid in chunk_ids] + point_ids = [_chunk_uuid(str(ctx.tenant_id), str(cid)) for cid in chunk_ids] await self._client.delete( collection_name=self._collection, points_selector=point_ids, # type: ignore[arg-type] diff --git a/packages/core/pyproject.toml b/packages/core/pyproject.toml index 1086be8..7e75c5f 100644 --- a/packages/core/pyproject.toml +++ b/packages/core/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "rag-core" -version = "0.2.0" +version = "0.3.0" description = "AgentContextOS — core domain types, error hierarchy, and plugin SPIs" readme = "README.md" requires-python = ">=3.12" diff --git a/packages/core/src/rag_core/__init__.py b/packages/core/src/rag_core/__init__.py index 1dc6b43..af7f042 100644 --- a/packages/core/src/rag_core/__init__.py +++ b/packages/core/src/rag_core/__init__.py @@ -82,25 +82,39 @@ ACLAction, AuditEvent, AuditOutcome, + BlobRef, + Budget, Chunk, ChunkId, + ChunkRef, Citation, CorpusId, + Cost, Document, DocumentId, DocumentStatus, Embedding, + EmbeddingDtype, + PiiAction, + PiiPolicy, + PlanNode, Principal, PrincipalId, PrincipalKind, Query, + QueryPlan, + RequestContext, + RequestId, RetrievalResult, + StageEvent, + StageEventKind, Tenant, TenantId, TraceContext, + TrustLevel, ) -__version__ = "0.4.0" +__version__ = "0.5.0" __all__ = [ # eval @@ -138,22 +152,36 @@ "ACLAction", "AuditEvent", "AuditOutcome", + "BlobRef", + "Budget", "Chunk", "ChunkId", + "ChunkRef", "Citation", "CorpusId", + "Cost", "Document", "DocumentId", "DocumentStatus", "Embedding", + "EmbeddingDtype", + "PiiAction", + "PiiPolicy", + "PlanNode", "Principal", "PrincipalId", "PrincipalKind", "Query", + "QueryPlan", + "RequestContext", + "RequestId", "RetrievalResult", + "StageEvent", + "StageEventKind", "Tenant", "TenantId", "TraceContext", + "TrustLevel", # errors "ACLDeniedError", "AuthError", diff --git a/packages/core/src/rag_core/gen_schemas.py b/packages/core/src/rag_core/gen_schemas.py index e02bd7c..96227f5 100644 --- a/packages/core/src/rag_core/gen_schemas.py +++ b/packages/core/src/rag_core/gen_schemas.py @@ -18,13 +18,22 @@ from rag_core.types import ( ACL, AuditEvent, + BlobRef, + Budget, Chunk, + ChunkRef, Citation, + Cost, Document, Embedding, + PiiPolicy, + PlanNode, Principal, Query, + QueryPlan, + RequestContext, RetrievalResult, + StageEvent, Tenant, TraceContext, ) @@ -33,12 +42,21 @@ TraceContext, ACL, Principal, + PiiPolicy, + Budget, + RequestContext, Tenant, + BlobRef, Document, Chunk, Embedding, Query, Citation, + ChunkRef, + Cost, + PlanNode, + QueryPlan, + StageEvent, RetrievalResult, AuditEvent, ] diff --git a/packages/core/src/rag_core/spi/auth.py b/packages/core/src/rag_core/spi/auth.py index 3a023cb..f78711c 100644 --- a/packages/core/src/rag_core/spi/auth.py +++ b/packages/core/src/rag_core/spi/auth.py @@ -1,4 +1,13 @@ -"""Auth SPI — authentication and authorisation.""" +"""Auth SPI — authentication and authorisation. + +``Auth`` runs at the gateway boundary, *before* a ``RequestContext`` exists — +the gateway calls ``authenticate`` to obtain a ``Principal``, then constructs +``RequestContext`` from it. Therefore, the methods on ``Auth`` are the one +documented exception to the "every SPI method takes ctx first" rule. + +Once a request is in flight, governance checks go through ``PolicyEngine`` +(Step 1.1c) which does take ``ctx``. +""" from __future__ import annotations diff --git a/packages/core/src/rag_core/spi/cache.py b/packages/core/src/rag_core/spi/cache.py index 49773af..139dc58 100644 --- a/packages/core/src/rag_core/spi/cache.py +++ b/packages/core/src/rag_core/spi/cache.py @@ -1,35 +1,51 @@ -"""Cache SPI — key-value byte cache with optional TTL.""" +"""Cache SPI — key-value byte cache with optional TTL. + +Step 1.1e will split this into separate ``EmbeddingCache`` / +``RetrievalCache`` / ``AnswerCache`` SPIs. This base interface remains for +generic byte-level caching needs. +""" from __future__ import annotations import abc from rag_core.spi._base import HealthCheckMixin +from rag_core.types import RequestContext class Cache(HealthCheckMixin, abc.ABC): """Abstract cache backend (Redis, Memcached, in-process, …). Keys are plain strings; values are raw bytes. Callers handle - serialisation (JSON, msgpack, pickle, …). + serialisation (JSON, msgpack, pickle, …). Tenant scoping (per-tenant + quotas, key namespacing) comes from ``ctx``. """ @abc.abstractmethod - async def get(self, key: str) -> bytes | None: + async def get(self, ctx: RequestContext, key: str) -> bytes | None: """Return cached bytes or ``None`` on a miss.""" @abc.abstractmethod - async def set(self, key: str, value: bytes, ttl_seconds: int | None = None) -> None: + async def set( + self, + ctx: RequestContext, + key: str, + value: bytes, + ttl_seconds: int | None = None, + ) -> None: """Store ``value`` under ``key``. Args: + ctx: Per-request envelope. + key: Cache key. + value: Bytes to store. ttl_seconds: Seconds until the entry expires. ``None`` = no expiry. """ @abc.abstractmethod - async def delete(self, key: str) -> None: + async def delete(self, ctx: RequestContext, key: str) -> None: """Remove ``key``. No-op if the key does not exist.""" @abc.abstractmethod - async def exists(self, key: str) -> bool: + async def exists(self, ctx: RequestContext, key: str) -> bool: """Return True if ``key`` is present and not expired.""" diff --git a/packages/core/src/rag_core/spi/connector.py b/packages/core/src/rag_core/spi/connector.py index b9ca80d..765cff4 100644 --- a/packages/core/src/rag_core/spi/connector.py +++ b/packages/core/src/rag_core/spi/connector.py @@ -7,7 +7,7 @@ from datetime import datetime from rag_core.spi._base import HealthCheckMixin -from rag_core.types import Document, TenantId +from rag_core.types import Document, RequestContext class Connector(HealthCheckMixin, abc.ABC): @@ -15,19 +15,19 @@ class Connector(HealthCheckMixin, abc.ABC): A connector is responsible for discovering documents in an external source and fetching their raw bytes. The ingestion pipeline handles parsing, - chunking, and embedding downstream. + chunking, and embedding downstream. Tenant scoping comes from ``ctx``. """ @abc.abstractmethod def list_documents( self, - tenant_id: TenantId, + ctx: RequestContext, since: datetime | None = None, ) -> AsyncIterator[Document]: """Async-iterate document stubs discovered in the source. Args: - tenant_id: Tenant on whose behalf the crawl is running. + ctx: Per-request envelope (tenant, trace, budget). since: If set, only yield documents modified after this timestamp (incremental sync). None = full crawl. @@ -37,7 +37,7 @@ def list_documents( """ @abc.abstractmethod - async def fetch(self, document: Document) -> bytes: + async def fetch(self, ctx: RequestContext, document: Document) -> bytes: """Download and return the raw bytes for ``document``. Raises: diff --git a/packages/core/src/rag_core/spi/embedder.py b/packages/core/src/rag_core/spi/embedder.py index 1979777..1c44e95 100644 --- a/packages/core/src/rag_core/spi/embedder.py +++ b/packages/core/src/rag_core/spi/embedder.py @@ -5,7 +5,7 @@ import abc from rag_core.spi._base import HealthCheckMixin -from rag_core.types import ChunkId, Embedding, TenantId +from rag_core.types import ChunkId, Embedding, RequestContext class Embedder(HealthCheckMixin, abc.ABC): @@ -13,6 +13,7 @@ class Embedder(HealthCheckMixin, abc.ABC): Implementations handle batching, retries, and rate-limit back-off internally. Callers pass raw texts and receive ``Embedding`` objects. + Tenant-scoped rate limiting / quota enforcement comes from ``ctx``. """ @property @@ -28,9 +29,9 @@ def dimension(self) -> int: @abc.abstractmethod async def embed( self, + ctx: RequestContext, texts: list[str], chunk_ids: list[ChunkId], - tenant_id: TenantId, ) -> list[Embedding]: """Embed a batch of texts. diff --git a/packages/core/src/rag_core/spi/graph_store.py b/packages/core/src/rag_core/spi/graph_store.py index 2715cef..2c6f342 100644 --- a/packages/core/src/rag_core/spi/graph_store.py +++ b/packages/core/src/rag_core/spi/graph_store.py @@ -6,48 +6,48 @@ from typing import Any from rag_core.spi._base import HealthCheckMixin -from rag_core.types import TenantId +from rag_core.types import RequestContext class GraphStore(HealthCheckMixin, abc.ABC): """Abstract knowledge-graph backend (Neo4j, Kuzu, Amazon Neptune, …). - All operations are scoped to a ``tenant_id`` so implementations can - isolate graph data per tenant via label namespacing, separate databases, - or row-level predicates. + Tenant isolation comes from ``ctx.tenant_id`` — implementations namespace + graph data per tenant via label prefixing, separate databases, or + row-level predicates. """ @abc.abstractmethod async def upsert_node( self, + ctx: RequestContext, node_id: str, labels: list[str], properties: dict[str, Any], - tenant_id: TenantId, ) -> None: """Create or update a node. ``node_id`` is the stable external key.""" @abc.abstractmethod async def upsert_edge( self, + ctx: RequestContext, from_id: str, to_id: str, rel_type: str, properties: dict[str, Any], - tenant_id: TenantId, ) -> None: """Create or update a directed edge ``from_id --[rel_type]--> to_id``.""" @abc.abstractmethod - async def delete_node(self, node_id: str, tenant_id: TenantId) -> None: + async def delete_node(self, ctx: RequestContext, node_id: str) -> None: """Remove a node and all its incident edges. Unknown node_id ignored.""" @abc.abstractmethod async def query( self, + ctx: RequestContext, statement: str, parameters: dict[str, Any], - tenant_id: TenantId, ) -> list[dict[str, Any]]: """Execute a graph query (Cypher, Gremlin, or SPARQL — backend-specific). diff --git a/packages/core/src/rag_core/spi/keyword_store.py b/packages/core/src/rag_core/spi/keyword_store.py index 90d3169..c9b5459 100644 --- a/packages/core/src/rag_core/spi/keyword_store.py +++ b/packages/core/src/rag_core/spi/keyword_store.py @@ -5,18 +5,18 @@ import abc from rag_core.spi._base import HealthCheckMixin -from rag_core.types import Chunk, ChunkId, CorpusId, TenantId +from rag_core.types import Chunk, ChunkId, CorpusId, RequestContext class KeywordStore(HealthCheckMixin, abc.ABC): """Abstract full-text index for keyword (BM25-style) retrieval. Implementations may delegate to Elasticsearch, OpenSearch, Postgres FTS, - Tantivy, etc. The contract is always stateless from the caller's view. + Tantivy, etc. Tenant isolation comes from ``ctx.tenant_id``. """ @abc.abstractmethod - async def index(self, chunks: list[Chunk], tenant_id: TenantId) -> None: + async def index(self, ctx: RequestContext, chunks: list[Chunk]) -> None: """Add or update chunks in the keyword index. Existing entries for the same ``chunk_id`` are overwritten. @@ -25,20 +25,20 @@ async def index(self, chunks: list[Chunk], tenant_id: TenantId) -> None: @abc.abstractmethod async def search( self, + ctx: RequestContext, query_text: str, top_k: int, - tenant_id: TenantId, corpus_ids: list[CorpusId], ) -> list[tuple[ChunkId, float]]: """Return ``(chunk_id, bm25_score)`` pairs ordered by descending score. Args: + ctx: Per-request envelope (tenant, principal, budget, trace). query_text: Raw query string (tokenisation is backend-specific). top_k: Maximum number of results. - tenant_id: Restrict results to this tenant. corpus_ids: Restrict results to these corpora (empty = all corpora). """ @abc.abstractmethod - async def delete(self, chunk_ids: list[ChunkId], tenant_id: TenantId) -> None: + async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: """Remove index entries for the given chunk IDs. Unknown IDs ignored.""" diff --git a/packages/core/src/rag_core/spi/llm.py b/packages/core/src/rag_core/spi/llm.py index b66e9e5..71cd0b6 100644 --- a/packages/core/src/rag_core/spi/llm.py +++ b/packages/core/src/rag_core/spi/llm.py @@ -8,7 +8,7 @@ from typing import Any from rag_core.spi._base import HealthCheckMixin -from rag_core.types import TenantId +from rag_core.types import RequestContext @dataclass(frozen=True) @@ -33,15 +33,14 @@ class LLMResponse: class LLM(HealthCheckMixin, abc.ABC): """Abstract LLM backend (OpenAI, Anthropic, Mistral, local vLLM, …). - All calls are tenant-scoped so implementations can apply per-tenant - rate limits, model routing, or audit logging. + Tenant scoping (rate limits, model routing, audit) comes from ``ctx``. """ @abc.abstractmethod async def complete( self, + ctx: RequestContext, messages: list[LLMMessage], - tenant_id: TenantId, max_tokens: int = 1024, temperature: float = 0.0, stop: list[str] | None = None, @@ -51,8 +50,8 @@ async def complete( @abc.abstractmethod async def stream( self, + ctx: RequestContext, messages: list[LLMMessage], - tenant_id: TenantId, max_tokens: int = 1024, temperature: float = 0.0, stop: list[str] | None = None, diff --git a/packages/core/src/rag_core/spi/noop/cache.py b/packages/core/src/rag_core/spi/noop/cache.py index 3faf8d2..efc69f6 100644 --- a/packages/core/src/rag_core/spi/noop/cache.py +++ b/packages/core/src/rag_core/spi/noop/cache.py @@ -3,22 +3,29 @@ from __future__ import annotations from rag_core.spi.cache import Cache +from rag_core.types import RequestContext class NoopCache(Cache): def __init__(self) -> None: self._store: dict[str, bytes] = {} - async def get(self, key: str) -> bytes | None: + async def get(self, ctx: RequestContext, key: str) -> bytes | None: return self._store.get(key) - async def set(self, key: str, value: bytes, ttl_seconds: int | None = None) -> None: + async def set( + self, + ctx: RequestContext, + key: str, + value: bytes, + ttl_seconds: int | None = None, + ) -> None: self._store[key] = value # TTL ignored in noop - async def delete(self, key: str) -> None: + async def delete(self, ctx: RequestContext, key: str) -> None: self._store.pop(key, None) - async def exists(self, key: str) -> bool: + async def exists(self, ctx: RequestContext, key: str) -> bool: return key in self._store async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/noop/connector.py b/packages/core/src/rag_core/spi/noop/connector.py index b322d5c..27738ac 100644 --- a/packages/core/src/rag_core/spi/noop/connector.py +++ b/packages/core/src/rag_core/spi/noop/connector.py @@ -6,7 +6,7 @@ from datetime import datetime from rag_core.spi.connector import Connector -from rag_core.types import Document, TenantId +from rag_core.types import Document, RequestContext class NoopConnector(Connector): @@ -17,10 +17,10 @@ def __init__(self, documents: list[Document] | None = None) -> None: def list_documents( self, - tenant_id: TenantId, + ctx: RequestContext, since: datetime | None = None, ) -> AsyncIterator[Document]: - docs = [d for d in self._documents if d.tenant_id == tenant_id] + docs = [d for d in self._documents if d.tenant_id == ctx.tenant_id] if since is not None: docs = [d for d in docs if d.updated_at >= since] @@ -30,7 +30,7 @@ async def _gen() -> AsyncIterator[Document]: return _gen() - async def fetch(self, document: Document) -> bytes: + async def fetch(self, ctx: RequestContext, document: Document) -> bytes: return b"" async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/noop/embedder.py b/packages/core/src/rag_core/spi/noop/embedder.py index 4cec448..6debb14 100644 --- a/packages/core/src/rag_core/spi/noop/embedder.py +++ b/packages/core/src/rag_core/spi/noop/embedder.py @@ -3,7 +3,7 @@ from __future__ import annotations from rag_core.spi.embedder import Embedder -from rag_core.types import ChunkId, Embedding, TenantId +from rag_core.types import ChunkId, Embedding, RequestContext class NoopEmbedder(Embedder): @@ -22,15 +22,16 @@ def dimension(self) -> int: async def embed( self, + ctx: RequestContext, texts: list[str], chunk_ids: list[ChunkId], - tenant_id: TenantId, ) -> list[Embedding]: if len(texts) != len(chunk_ids): raise ValueError(f"texts length {len(texts)} != chunk_ids length {len(chunk_ids)}") return [ Embedding( chunk_id=cid, + tenant_id=ctx.tenant_id, model=self.model, vector=[0.0] * self._dimension, dimension=self._dimension, diff --git a/packages/core/src/rag_core/spi/noop/graph_store.py b/packages/core/src/rag_core/spi/noop/graph_store.py index 6de1b3c..d5f106f 100644 --- a/packages/core/src/rag_core/spi/noop/graph_store.py +++ b/packages/core/src/rag_core/spi/noop/graph_store.py @@ -5,7 +5,7 @@ from typing import Any from rag_core.spi.graph_store import GraphStore -from rag_core.types import TenantId +from rag_core.types import RequestContext class NoopGraphStore(GraphStore): @@ -17,42 +17,42 @@ def __init__(self) -> None: async def upsert_node( self, + ctx: RequestContext, node_id: str, labels: list[str], properties: dict[str, Any], - tenant_id: TenantId, ) -> None: - self._nodes[(tenant_id, node_id)] = {"labels": labels, **properties} + self._nodes[(ctx.tenant_id, node_id)] = {"labels": labels, **properties} async def upsert_edge( self, + ctx: RequestContext, from_id: str, to_id: str, rel_type: str, properties: dict[str, Any], - tenant_id: TenantId, ) -> None: - self._edges.append((tenant_id, from_id, to_id, rel_type, properties)) + self._edges.append((ctx.tenant_id, from_id, to_id, rel_type, properties)) - async def delete_node(self, node_id: str, tenant_id: TenantId) -> None: - self._nodes.pop((tenant_id, node_id), None) + async def delete_node(self, ctx: RequestContext, node_id: str) -> None: + self._nodes.pop((ctx.tenant_id, node_id), None) self._edges = [ e for e in self._edges - if not (e[0] == tenant_id and (e[1] == node_id or e[2] == node_id)) + if not (e[0] == ctx.tenant_id and (e[1] == node_id or e[2] == node_id)) ] async def query( self, + ctx: RequestContext, statement: str, parameters: dict[str, Any], - tenant_id: TenantId, ) -> list[dict[str, Any]]: # Noop: return all nodes for this tenant as rows return [ {"node_id": nid, **props} for (tid, nid), props in self._nodes.items() - if tid == tenant_id + if tid == ctx.tenant_id ] async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/noop/keyword_store.py b/packages/core/src/rag_core/spi/noop/keyword_store.py index b8f020c..644e180 100644 --- a/packages/core/src/rag_core/spi/noop/keyword_store.py +++ b/packages/core/src/rag_core/spi/noop/keyword_store.py @@ -3,7 +3,7 @@ from __future__ import annotations from rag_core.spi.keyword_store import KeywordStore -from rag_core.types import Chunk, ChunkId, CorpusId, TenantId +from rag_core.types import Chunk, ChunkId, CorpusId, RequestContext class NoopKeywordStore(KeywordStore): @@ -13,26 +13,28 @@ def __init__(self) -> None: # (tenant_id, chunk_id) -> Chunk self._index: dict[tuple[str, str], Chunk] = {} - async def index(self, chunks: list[Chunk], tenant_id: TenantId) -> None: + async def index(self, ctx: RequestContext, chunks: list[Chunk]) -> None: for chunk in chunks: - self._index[(tenant_id, chunk.id)] = chunk + self._index[(ctx.tenant_id, chunk.id)] = chunk async def search( self, + ctx: RequestContext, query_text: str, top_k: int, - tenant_id: TenantId, corpus_ids: list[CorpusId], ) -> list[tuple[ChunkId, float]]: query_lower = query_text.lower() results: list[tuple[ChunkId, float]] = [] for (tid, cid), chunk in self._index.items(): - if tid != tenant_id: + if tid != ctx.tenant_id: continue if corpus_ids and chunk.corpus_id not in corpus_ids: continue + if chunk.content is None: + # Noop store does not hydrate BlobRef-backed chunks. + continue content_lower = chunk.content.lower() - # Score = fraction of query tokens present in content tokens = query_lower.split() hits = sum(1 for t in tokens if t in content_lower) if hits: @@ -41,9 +43,9 @@ async def search( results.sort(key=lambda x: x[1], reverse=True) return results[:top_k] - async def delete(self, chunk_ids: list[ChunkId], tenant_id: TenantId) -> None: + async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: for cid in chunk_ids: - self._index.pop((tenant_id, cid), None) + self._index.pop((ctx.tenant_id, cid), None) async def health(self) -> bool: return True diff --git a/packages/core/src/rag_core/spi/noop/llm.py b/packages/core/src/rag_core/spi/noop/llm.py index dd27ecf..94fd800 100644 --- a/packages/core/src/rag_core/spi/noop/llm.py +++ b/packages/core/src/rag_core/spi/noop/llm.py @@ -5,7 +5,7 @@ from collections.abc import AsyncIterator from rag_core.spi.llm import LLM, LLMMessage, LLMResponse -from rag_core.types import TenantId +from rag_core.types import RequestContext class NoopLLM(LLM): @@ -13,8 +13,8 @@ class NoopLLM(LLM): async def complete( self, + ctx: RequestContext, messages: list[LLMMessage], - tenant_id: TenantId, max_tokens: int = 1024, temperature: float = 0.0, stop: list[str] | None = None, @@ -30,13 +30,13 @@ async def complete( async def stream( self, + ctx: RequestContext, messages: list[LLMMessage], - tenant_id: TenantId, max_tokens: int = 1024, temperature: float = 0.0, stop: list[str] | None = None, ) -> AsyncIterator[str]: - response = await self.complete(messages, tenant_id, max_tokens, temperature, stop) + response = await self.complete(ctx, messages, max_tokens, temperature, stop) async def _gen() -> AsyncIterator[str]: for word in response.content.split(): diff --git a/packages/core/src/rag_core/spi/noop/ocr.py b/packages/core/src/rag_core/spi/noop/ocr.py index d31da24..82a0adf 100644 --- a/packages/core/src/rag_core/spi/noop/ocr.py +++ b/packages/core/src/rag_core/spi/noop/ocr.py @@ -3,10 +3,16 @@ from __future__ import annotations from rag_core.spi.ocr import OCR, OCRResult +from rag_core.types import RequestContext class NoopOCR(OCR): - async def extract(self, image_bytes: bytes, mime_type: str) -> OCRResult: + async def extract( + self, + ctx: RequestContext, + image_bytes: bytes, + mime_type: str, + ) -> OCRResult: return OCRResult(text="", confidence=1.0) async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/noop/parser.py b/packages/core/src/rag_core/spi/noop/parser.py index 262cf50..5eecd83 100644 --- a/packages/core/src/rag_core/spi/noop/parser.py +++ b/packages/core/src/rag_core/spi/noop/parser.py @@ -5,7 +5,7 @@ import hashlib from rag_core.spi.parser import Parser -from rag_core.types import CorpusId, Document, DocumentId, DocumentStatus, TenantId +from rag_core.types import CorpusId, Document, DocumentId, DocumentStatus, RequestContext class NoopParser(Parser): @@ -16,17 +16,17 @@ def supports(self, mime_type: str) -> bool: async def parse( self, + ctx: RequestContext, data: bytes, mime_type: str, document_id: DocumentId, - tenant_id: TenantId, corpus_id: CorpusId, source_uri: str, ) -> Document: content_hash = hashlib.sha256(data).hexdigest() return Document( id=document_id, - tenant_id=tenant_id, + tenant_id=ctx.tenant_id, corpus_id=corpus_id, source_uri=source_uri, content_hash=content_hash, diff --git a/packages/core/src/rag_core/spi/noop/pii_detector.py b/packages/core/src/rag_core/spi/noop/pii_detector.py index 58619c3..04f7552 100644 --- a/packages/core/src/rag_core/spi/noop/pii_detector.py +++ b/packages/core/src/rag_core/spi/noop/pii_detector.py @@ -3,16 +3,16 @@ from __future__ import annotations from rag_core.spi.pii_detector import PIIDetector, PIISpan -from rag_core.types import TenantId +from rag_core.types import RequestContext class NoopPIIDetector(PIIDetector): """Reports no PII in any text. Safe for dev; never use in production.""" - async def detect(self, text: str, tenant_id: TenantId) -> list[PIISpan]: + async def detect(self, ctx: RequestContext, text: str) -> list[PIISpan]: return [] - async def redact(self, text: str, tenant_id: TenantId) -> str: + async def redact(self, ctx: RequestContext, text: str) -> str: return text # nothing to redact async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/noop/queue.py b/packages/core/src/rag_core/spi/noop/queue.py index 308d663..4c1e8bd 100644 --- a/packages/core/src/rag_core/spi/noop/queue.py +++ b/packages/core/src/rag_core/spi/noop/queue.py @@ -6,6 +6,7 @@ from collections.abc import AsyncIterator from rag_core.spi.queue import Queue, QueueMessage +from rag_core.types import RequestContext class NoopQueue(Queue): @@ -16,6 +17,7 @@ def __init__(self) -> None: async def publish( self, + ctx: RequestContext, topic: str, payload: bytes, attributes: dict[str, str] | None = None, @@ -29,10 +31,15 @@ async def publish( self.published.append(msg) return msg.message_id - def subscribe(self, topic: str, subscription: str) -> AsyncIterator[QueueMessage]: + def subscribe( + self, + ctx: RequestContext, + topic: str, + subscription: str, + ) -> AsyncIterator[QueueMessage]: raise NotImplementedError("NoopQueue.subscribe is not supported; use NoopQueue.published") - async def ack(self, topic: str, message_id: str) -> None: + async def ack(self, ctx: RequestContext, topic: str, message_id: str) -> None: pass # auto-ack: nothing to do async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/noop/reranker.py b/packages/core/src/rag_core/spi/noop/reranker.py index 20618de..3488612 100644 --- a/packages/core/src/rag_core/spi/noop/reranker.py +++ b/packages/core/src/rag_core/spi/noop/reranker.py @@ -3,14 +3,18 @@ from __future__ import annotations from rag_core.spi.reranker import Reranker -from rag_core.types import Chunk +from rag_core.types import Chunk, RequestContext class NoopReranker(Reranker): """Returns the first ``top_k`` chunks unchanged with score 1.0.""" async def rerank( - self, query: str, chunks: list[Chunk], top_k: int + self, + ctx: RequestContext, + query: str, + chunks: list[Chunk], + top_k: int, ) -> list[tuple[Chunk, float]]: return [(c, 1.0) for c in chunks[:top_k]] diff --git a/packages/core/src/rag_core/spi/noop/secrets.py b/packages/core/src/rag_core/spi/noop/secrets.py index b9d08c3..9744af1 100644 --- a/packages/core/src/rag_core/spi/noop/secrets.py +++ b/packages/core/src/rag_core/spi/noop/secrets.py @@ -3,6 +3,7 @@ from __future__ import annotations from rag_core.spi.secrets import Secrets +from rag_core.types import RequestContext class NoopSecrets(Secrets): @@ -11,13 +12,13 @@ class NoopSecrets(Secrets): def __init__(self, secrets: dict[str, str] | None = None) -> None: self._secrets: dict[str, str] = secrets or {} - async def get(self, name: str) -> str: + async def get(self, ctx: RequestContext, name: str) -> str: if name not in self._secrets: raise KeyError(name) return self._secrets[name] - async def get_bytes(self, name: str) -> bytes: - return (await self.get(name)).encode() + async def get_bytes(self, ctx: RequestContext, name: str) -> bytes: + return (await self.get(ctx, name)).encode() async def health(self) -> bool: return True diff --git a/packages/core/src/rag_core/spi/noop/storage.py b/packages/core/src/rag_core/spi/noop/storage.py index 6c8a4b8..8876f7c 100644 --- a/packages/core/src/rag_core/spi/noop/storage.py +++ b/packages/core/src/rag_core/spi/noop/storage.py @@ -5,27 +5,34 @@ from collections.abc import AsyncIterator from rag_core.spi.storage import Storage +from rag_core.types import RequestContext class NoopStorage(Storage): def __init__(self) -> None: self._blobs: dict[str, bytes] = {} - async def put(self, key: str, data: bytes, content_type: str | None = None) -> None: + async def put( + self, + ctx: RequestContext, + key: str, + data: bytes, + content_type: str | None = None, + ) -> None: self._blobs[key] = data - async def get(self, key: str) -> bytes: + async def get(self, ctx: RequestContext, key: str) -> bytes: if key not in self._blobs: raise KeyError(key) return self._blobs[key] - async def delete(self, key: str) -> None: + async def delete(self, ctx: RequestContext, key: str) -> None: self._blobs.pop(key, None) - async def exists(self, key: str) -> bool: + async def exists(self, ctx: RequestContext, key: str) -> bool: return key in self._blobs - def list_keys(self, prefix: str = "") -> AsyncIterator[str]: + def list_keys(self, ctx: RequestContext, prefix: str = "") -> AsyncIterator[str]: async def _gen() -> AsyncIterator[str]: for k in self._blobs: if k.startswith(prefix): diff --git a/packages/core/src/rag_core/spi/noop/vector_store.py b/packages/core/src/rag_core/spi/noop/vector_store.py index d90349f..844fc8c 100644 --- a/packages/core/src/rag_core/spi/noop/vector_store.py +++ b/packages/core/src/rag_core/spi/noop/vector_store.py @@ -6,7 +6,7 @@ from typing import Any from rag_core.spi.vector_store import VectorStore -from rag_core.types import ChunkId, CorpusId, Embedding, TenantId +from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext def _cosine(a: list[float], b: list[float]) -> float: @@ -23,30 +23,30 @@ def __init__(self) -> None: # (tenant_id, chunk_id) -> Embedding self._store: dict[tuple[str, str], Embedding] = {} - async def upsert(self, embeddings: list[Embedding], tenant_id: TenantId) -> None: + async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: for emb in embeddings: - self._store[(tenant_id, emb.chunk_id)] = emb + self._store[(ctx.tenant_id, emb.chunk_id)] = emb async def query( self, + ctx: RequestContext, vector: list[float], top_k: int, - tenant_id: TenantId, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, ) -> list[tuple[ChunkId, float]]: results: list[tuple[ChunkId, float]] = [] for (tid, cid), emb in self._store.items(): - if tid != tenant_id: + if tid != ctx.tenant_id: continue score = _cosine(vector, emb.vector) results.append((ChunkId(cid), score)) results.sort(key=lambda x: x[1], reverse=True) return results[:top_k] - async def delete(self, chunk_ids: list[ChunkId], tenant_id: TenantId) -> None: + async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: for cid in chunk_ids: - self._store.pop((tenant_id, cid), None) + self._store.pop((ctx.tenant_id, cid), None) async def health(self) -> bool: return True diff --git a/packages/core/src/rag_core/spi/ocr.py b/packages/core/src/rag_core/spi/ocr.py index 0ba140d..d011aff 100644 --- a/packages/core/src/rag_core/spi/ocr.py +++ b/packages/core/src/rag_core/spi/ocr.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from rag_core.spi._base import HealthCheckMixin +from rag_core.types import RequestContext @dataclass(frozen=True) @@ -25,10 +26,16 @@ class OCR(HealthCheckMixin, abc.ABC): """ @abc.abstractmethod - async def extract(self, image_bytes: bytes, mime_type: str) -> OCRResult: + async def extract( + self, + ctx: RequestContext, + image_bytes: bytes, + mime_type: str, + ) -> OCRResult: """Run OCR on a single image and return extracted text. Args: + ctx: Per-request envelope (tenant, budget, trace). image_bytes: Raw image bytes (PNG, JPEG, TIFF, …). mime_type: MIME type of ``image_bytes``, e.g. ``"image/png"``. diff --git a/packages/core/src/rag_core/spi/parser.py b/packages/core/src/rag_core/spi/parser.py index 0b2b7cd..f36dd5c 100644 --- a/packages/core/src/rag_core/spi/parser.py +++ b/packages/core/src/rag_core/spi/parser.py @@ -5,7 +5,7 @@ import abc from rag_core.spi._base import HealthCheckMixin -from rag_core.types import CorpusId, Document, DocumentId, TenantId +from rag_core.types import CorpusId, Document, DocumentId, RequestContext class Parser(HealthCheckMixin, abc.ABC): @@ -22,16 +22,17 @@ def supports(self, mime_type: str) -> bool: @abc.abstractmethod async def parse( self, + ctx: RequestContext, data: bytes, mime_type: str, document_id: DocumentId, - tenant_id: TenantId, corpus_id: CorpusId, source_uri: str, ) -> Document: """Parse ``data`` and return a ``Document`` with ``content_hash`` set. The returned ``Document.status`` should be ``DocumentStatus.ready``. + ``tenant_id`` is taken from ``ctx``. Raises: rag_core.errors.ParseError: If the bytes cannot be parsed. diff --git a/packages/core/src/rag_core/spi/pii_detector.py b/packages/core/src/rag_core/spi/pii_detector.py index 6413731..d6e69f6 100644 --- a/packages/core/src/rag_core/spi/pii_detector.py +++ b/packages/core/src/rag_core/spi/pii_detector.py @@ -6,7 +6,7 @@ from dataclasses import dataclass from rag_core.spi._base import HealthCheckMixin -from rag_core.types import TenantId +from rag_core.types import RequestContext @dataclass(frozen=True) @@ -23,12 +23,13 @@ class PIISpan: class PIIDetector(HealthCheckMixin, abc.ABC): """Abstract PII detection backend (Presidio, AWS Comprehend, …). - Per-tenant policy (block / redact / allow) is enforced by the pipeline - layer above this SPI — the detector only identifies and redacts spans. + Per-tenant policy (block / redact / allow) is enforced by the + ``PolicyEngine`` (Step 1.1c) above this SPI — the detector only + identifies and redacts spans. """ @abc.abstractmethod - async def detect(self, text: str, tenant_id: TenantId) -> list[PIISpan]: + async def detect(self, ctx: RequestContext, text: str) -> list[PIISpan]: """Identify PII spans in ``text``. Returns an empty list if no PII is found. Spans do not overlap. @@ -36,5 +37,5 @@ async def detect(self, text: str, tenant_id: TenantId) -> list[PIISpan]: """ @abc.abstractmethod - async def redact(self, text: str, tenant_id: TenantId) -> str: + async def redact(self, ctx: RequestContext, text: str) -> str: """Return ``text`` with all detected PII replaced by ``PIISpan.replacement`` strings.""" diff --git a/packages/core/src/rag_core/spi/queue.py b/packages/core/src/rag_core/spi/queue.py index 2f99ead..7385850 100644 --- a/packages/core/src/rag_core/spi/queue.py +++ b/packages/core/src/rag_core/spi/queue.py @@ -7,6 +7,8 @@ from dataclasses import dataclass, field from datetime import UTC, datetime +from rag_core.types import RequestContext + @dataclass(frozen=True) class QueueMessage: @@ -25,11 +27,19 @@ class Queue(abc.ABC): Producers call ``publish``; consumers iterate ``subscribe``. Acknowledgement semantics are implementation-specific — callers that need at-least-once delivery must call ``ack`` if it is supported. + + ``ctx`` carries tenant scoping, audit context, and budget. Consumers + typically synthesize a fresh ``RequestContext`` per message before calling + ``ack``. """ @abc.abstractmethod async def publish( - self, topic: str, payload: bytes, attributes: dict[str, str] | None = None + self, + ctx: RequestContext, + topic: str, + payload: bytes, + attributes: dict[str, str] | None = None, ) -> str: """Publish a message to ``topic``. @@ -38,7 +48,12 @@ async def publish( """ @abc.abstractmethod - def subscribe(self, topic: str, subscription: str) -> AsyncIterator[QueueMessage]: + def subscribe( + self, + ctx: RequestContext, + topic: str, + subscription: str, + ) -> AsyncIterator[QueueMessage]: """Async-iterate messages from ``topic`` via ``subscription``. The iterator does not return until cancelled. Callers are responsible @@ -46,7 +61,7 @@ def subscribe(self, topic: str, subscription: str) -> AsyncIterator[QueueMessage """ @abc.abstractmethod - async def ack(self, topic: str, message_id: str) -> None: + async def ack(self, ctx: RequestContext, topic: str, message_id: str) -> None: """Acknowledge a message. No-op for backends with auto-ack.""" async def health(self) -> bool: diff --git a/packages/core/src/rag_core/spi/reranker.py b/packages/core/src/rag_core/spi/reranker.py index 9a85531..a8f5135 100644 --- a/packages/core/src/rag_core/spi/reranker.py +++ b/packages/core/src/rag_core/spi/reranker.py @@ -5,7 +5,7 @@ import abc from rag_core.spi._base import HealthCheckMixin -from rag_core.types import Chunk +from rag_core.types import Chunk, RequestContext class Reranker(HealthCheckMixin, abc.ABC): @@ -19,6 +19,7 @@ class Reranker(HealthCheckMixin, abc.ABC): @abc.abstractmethod async def rerank( self, + ctx: RequestContext, query: str, chunks: list[Chunk], top_k: int, @@ -26,6 +27,7 @@ async def rerank( """Return ``(chunk, relevance_score)`` pairs ordered by descending score. Args: + ctx: Per-request envelope. query: The user query string. chunks: Candidate chunks from the first-stage retrieval. top_k: Maximum number of chunks to return. diff --git a/packages/core/src/rag_core/spi/secrets.py b/packages/core/src/rag_core/spi/secrets.py index f4d9aa1..13e86fa 100644 --- a/packages/core/src/rag_core/spi/secrets.py +++ b/packages/core/src/rag_core/spi/secrets.py @@ -5,17 +5,19 @@ import abc from rag_core.spi._base import HealthCheckMixin +from rag_core.types import RequestContext class Secrets(HealthCheckMixin, abc.ABC): """Abstract secret store (AWS Secrets Manager, Vault, GCP Secret Manager, …). Callers request secrets by logical name; implementations resolve the - physical path and handle caching / rotation transparently. + physical path (often per-tenant via ``ctx.tenant_id``) and handle caching + / rotation transparently. """ @abc.abstractmethod - async def get(self, name: str) -> str: + async def get(self, ctx: RequestContext, name: str) -> str: """Return the current plaintext value of ``name``. Raises: @@ -23,7 +25,7 @@ async def get(self, name: str) -> str: """ @abc.abstractmethod - async def get_bytes(self, name: str) -> bytes: + async def get_bytes(self, ctx: RequestContext, name: str) -> bytes: """Return the raw binary value of ``name`` (e.g. TLS certs, keys). Raises: diff --git a/packages/core/src/rag_core/spi/storage.py b/packages/core/src/rag_core/spi/storage.py index bd141d0..ceeeb4b 100644 --- a/packages/core/src/rag_core/spi/storage.py +++ b/packages/core/src/rag_core/spi/storage.py @@ -6,22 +6,29 @@ from collections.abc import AsyncIterator from rag_core.spi._base import HealthCheckMixin +from rag_core.types import RequestContext class Storage(HealthCheckMixin, abc.ABC): """Abstract object store (S3, GCS, Azure Blob, local filesystem, …). Keys are arbitrary slash-separated strings (like S3 object keys). - Implementations may add a tenant prefix automatically; callers should not - rely on the physical key format. + Implementations may add a tenant prefix automatically (derived from + ``ctx.tenant_id``); callers should not rely on the physical key format. """ @abc.abstractmethod - async def put(self, key: str, data: bytes, content_type: str | None = None) -> None: + async def put( + self, + ctx: RequestContext, + key: str, + data: bytes, + content_type: str | None = None, + ) -> None: """Write ``data`` at ``key``. Overwrites any existing object.""" @abc.abstractmethod - async def get(self, key: str) -> bytes: + async def get(self, ctx: RequestContext, key: str) -> bytes: """Read the object at ``key``. Raises: @@ -29,13 +36,13 @@ async def get(self, key: str) -> bytes: """ @abc.abstractmethod - async def delete(self, key: str) -> None: + async def delete(self, ctx: RequestContext, key: str) -> None: """Delete the object at ``key``. No-op if ``key`` does not exist.""" @abc.abstractmethod - async def exists(self, key: str) -> bool: + async def exists(self, ctx: RequestContext, key: str) -> bool: """Return True if ``key`` exists.""" @abc.abstractmethod - def list_keys(self, prefix: str = "") -> AsyncIterator[str]: + def list_keys(self, ctx: RequestContext, prefix: str = "") -> AsyncIterator[str]: """Async-iterate all keys that start with ``prefix``.""" diff --git a/packages/core/src/rag_core/spi/vector_store.py b/packages/core/src/rag_core/spi/vector_store.py index 55dce68..885e92c 100644 --- a/packages/core/src/rag_core/spi/vector_store.py +++ b/packages/core/src/rag_core/spi/vector_store.py @@ -6,43 +6,46 @@ from typing import Any from rag_core.spi._base import HealthCheckMixin -from rag_core.types import ChunkId, CorpusId, Embedding, TenantId +from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext class VectorStore(HealthCheckMixin, abc.ABC): """Abstract store for dense embeddings with ANN (approximate nearest-neighbour) search. - Tenant isolation is enforced at the SPI boundary — every mutating call - and every query carries ``tenant_id`` so implementations can namespace - data (separate collections, schema-per-tenant, row-level security, …). + Tenant isolation is enforced via ``ctx.tenant_id`` on every call. See + docs/architecture/request-context.md. + + Step 1.1b will split this into ``RetrievalBackend`` (read) + + ``IndexBackend`` (write) and introduce ID-only retrieval. """ @abc.abstractmethod - async def upsert(self, embeddings: list[Embedding], tenant_id: TenantId) -> None: + async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: """Insert or overwrite embeddings. If an embedding for a given ``chunk_id`` already exists it is replaced. + ``ctx.tenant_id`` namespaces the data. """ @abc.abstractmethod async def query( self, + ctx: RequestContext, vector: list[float], top_k: int, - tenant_id: TenantId, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, ) -> list[tuple[ChunkId, float]]: """Return ``(chunk_id, score)`` pairs ordered by descending similarity. Args: + ctx: Per-request envelope (tenant, principal, budget, trace). vector: Query vector; must match the dimension of stored embeddings. top_k: Maximum number of results to return. - tenant_id: Restrict results to this tenant. corpus_ids: Restrict results to these corpora (empty = all corpora). filters: Optional metadata key-value equality filters. """ @abc.abstractmethod - async def delete(self, chunk_ids: list[ChunkId], tenant_id: TenantId) -> None: + async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: """Remove embeddings by chunk ID. Unknown IDs are silently ignored.""" diff --git a/packages/core/src/rag_core/types.py b/packages/core/src/rag_core/types.py index e0af189..41b5477 100644 --- a/packages/core/src/rag_core/types.py +++ b/packages/core/src/rag_core/types.py @@ -4,6 +4,14 @@ objects — callers create new instances rather than mutating existing ones. JSON Schema generation is handled by gen_schemas.py; proto/core.proto mirrors these definitions as the cross-language canonical schema. + +Step 1.1a introduces ``RequestContext`` and its substructures (``Principal``, +``PiiPolicy``, ``Budget``) as the per-request envelope threaded through every +SPI call. It also adds tiered-storage and planner types (``BlobRef``, +``QueryPlan``, ``ChunkRef``, ``StageEvent``) and promotes ``tenant_id`` / +``acl_labels`` / ``trust_level`` / ``dtype`` to typed required fields on +``Chunk`` / ``Embedding``. See docs/architecture/request-context.md and +ADR-0005 / ADR-0007 / ADR-0008. """ from __future__ import annotations @@ -23,6 +31,7 @@ CorpusId = NewType("CorpusId", str) DocumentId = NewType("DocumentId", str) ChunkId = NewType("ChunkId", str) +RequestId = NewType("RequestId", str) def _new_id() -> str: @@ -62,6 +71,58 @@ class DocumentStatus(StrEnum): failed = "failed" +class TrustLevel(StrEnum): + """Provenance of a chunk's text — used by the prompt-injection defense. + + ``trusted`` — first-party content authored under tenant control. + ``ingested`` — content fetched from a known external source (vetted feed, + enterprise SharePoint, etc.). + ``user_supplied`` — content directly contributed by an end-user channel + (web upload, chat-attached file, …) which may contain + adversarial instructions. + """ + + trusted = "trusted" + ingested = "ingested" + user_supplied = "user_supplied" + + +class EmbeddingDtype(StrEnum): + """Numeric representation of an embedding vector at rest. + + ``float32`` — full-precision (default). + ``int8`` — quantized 8-bit integers (≈4× smaller, slight recall loss). + ``binary`` — 1-bit per dim packed (≈32× smaller, larger recall loss, + used as a coarse first-stage in two-stage retrieval). + + See ADR-0009 for the vector index + quantization strategy. + """ + + float32 = "float32" + int8 = "int8" + binary = "binary" + + +class PiiAction(StrEnum): + """Per-tenant PII handling policy.""" + + block = "block" + redact = "redact" + mask = "mask" + encrypt = "encrypt" + tag_only = "tag_only" + allow = "allow" + + +class StageEventKind(StrEnum): + """Kind of a `StageEvent` emitted by a pipeline stage.""" + + started = "started" + completed = "completed" + failed = "failed" + skipped = "skipped" + + # --------------------------------------------------------------------------- # TraceContext # --------------------------------------------------------------------------- @@ -94,7 +155,12 @@ class ACL(BaseModel): # Principal # --------------------------------------------------------------------------- class Principal(BaseModel): - """Authenticated identity — user, service account, or group.""" + """Authenticated identity — user, service account, or group. + + ``acl_labels`` is the set of ACL labels the principal carries; it is the + *typed* counterpart of the ``acl_labels`` field on ``Chunk`` / + ``Embedding``. The PolicyEngine compares the two during retrieval. + """ model_config = {"frozen": True} @@ -104,6 +170,102 @@ class Principal(BaseModel): email: str | None = None tenant_id: TenantId roles: list[str] = Field(default_factory=list) + acl_labels: frozenset[str] = Field(default_factory=frozenset) + + +# --------------------------------------------------------------------------- +# PiiPolicy +# --------------------------------------------------------------------------- +class PiiPolicy(BaseModel): + """Per-tenant PII enforcement contract attached to a `RequestContext`.""" + + model_config = {"frozen": True} + + action: PiiAction = PiiAction.redact + # Entity types this policy applies to (e.g. {"EMAIL", "PERSON"}); empty = all. + entities: frozenset[str] = Field(default_factory=frozenset) + # Score threshold below which detections are ignored. + min_score: float = 0.5 + + +# --------------------------------------------------------------------------- +# Budget +# --------------------------------------------------------------------------- +class Budget(BaseModel): + """Per-request resource envelope consulted by the cost-aware planner.""" + + model_config = {"frozen": True} + + max_tokens: int | None = None + max_dollars: float | None = None + max_wall_ms: int | None = None + max_iter: int | None = None + + def spend( + self, + *, + tokens: int = 0, + dollars: float = 0.0, + wall_ms: int = 0, + iter_: int = 0, + ) -> Budget: + """Return a new ``Budget`` reduced by the supplied spend amounts. + + Negative results clamp to ``0`` (signalling exhaustion). See ADR-0008 + for how the planner uses this between agent-loop turns. + """ + + def _sub(cur: int | None, used: int) -> int | None: + if cur is None: + return None + return max(cur - used, 0) + + def _sub_f(cur: float | None, used: float) -> float | None: + if cur is None: + return None + return max(cur - used, 0.0) + + return Budget( + max_tokens=_sub(self.max_tokens, tokens), + max_dollars=_sub_f(self.max_dollars, dollars), + max_wall_ms=_sub(self.max_wall_ms, wall_ms), + max_iter=_sub(self.max_iter, iter_), + ) + + +# --------------------------------------------------------------------------- +# RequestContext — the per-request envelope +# --------------------------------------------------------------------------- +class RequestContext(BaseModel): + """Per-request envelope threaded through every SPI call. + + Constructed exactly once at the gateway boundary. Downstream code treats + it as a trusted, immutable object: validation runs at construction time + only — hot paths MUST NOT re-validate. + + See [docs/architecture/request-context.md](../../../../docs/architecture/request-context.md). + """ + + model_config = {"frozen": True} + + request_id: RequestId = Field(default_factory=lambda: RequestId(_new_id())) + tenant_id: TenantId + principal: Principal + pii_policy: PiiPolicy = Field(default_factory=PiiPolicy) + trace: TraceContext = Field(default_factory=TraceContext) + budget: Budget = Field(default_factory=Budget) + feature_flags: frozenset[str] = Field(default_factory=frozenset) + corpus_routing_hint: str | None = None + created_at: datetime = Field(default_factory=_utcnow) + + @model_validator(mode="after") + def _check_tenant_match(self) -> RequestContext: + if self.principal.tenant_id != self.tenant_id: + raise ValueError( + "RequestContext.tenant_id must match principal.tenant_id " + f"(ctx={self.tenant_id!r}, principal={self.principal.tenant_id!r})" + ) + return self # --------------------------------------------------------------------------- @@ -121,6 +283,25 @@ class Tenant(BaseModel): created_at: datetime = Field(default_factory=_utcnow) +# --------------------------------------------------------------------------- +# BlobRef — lazy pointer to chunk text or other large blobs +# --------------------------------------------------------------------------- +class BlobRef(BaseModel): + """Reference to a blob stored in the ``Storage`` SPI rather than inline. + + Used by chunks whose text exceeds the inline-storage threshold (see + ADR-0007 tiered storage). Callers must hydrate via ``Storage.get(uri)`` + only when the text is actually needed. + """ + + model_config = {"frozen": True} + + uri: str + size_bytes: int + content_type: str = "text/plain; charset=utf-8" + sha256: str | None = None + + # --------------------------------------------------------------------------- # Document # --------------------------------------------------------------------------- @@ -148,7 +329,17 @@ class Document(BaseModel): # Chunk # --------------------------------------------------------------------------- class Chunk(BaseModel): - """A contiguous piece of a Document produced by the chunking pipeline.""" + """A contiguous piece of a Document produced by the chunking pipeline. + + Step 1.1a promoted three governance-relevant fields out of the + ``metadata`` dict into typed required fields: + + - ``acl_labels`` — set of ACL labels the PolicyEngine compares against + the requesting principal's ``acl_labels``. + - ``trust_level`` — provenance, used by the prompt-injection defense. + - ``content_ref`` — optional ``BlobRef`` for tiered text storage; when set, + ``content`` may be ``None`` and callers must hydrate via Storage. + """ model_config = {"frozen": True} @@ -156,28 +347,44 @@ class Chunk(BaseModel): document_id: DocumentId tenant_id: TenantId corpus_id: CorpusId - content: str + content: str | None = None + content_ref: BlobRef | None = None # 0-based index of this chunk within its parent document position: int # set by hierarchical chunkers to link a child chunk to its parent section parent_id: ChunkId | None = None token_count: int | None = None + acl_labels: frozenset[str] = Field(default_factory=frozenset) + trust_level: TrustLevel = TrustLevel.ingested metadata: dict[str, Any] = Field(default_factory=dict) created_at: datetime = Field(default_factory=_utcnow) + @model_validator(mode="after") + def _check_content_present(self) -> Chunk: + if self.content is None and self.content_ref is None: + raise ValueError("Chunk must have either content or content_ref") + return self + # --------------------------------------------------------------------------- # Embedding # --------------------------------------------------------------------------- class Embedding(BaseModel): - """Dense vector representation of a Chunk.""" + """Dense vector representation of a Chunk. + + Step 1.1a promoted ``tenant_id`` and ``acl_labels`` to typed required + fields and added ``dtype`` to capture int8 / binary quantization. + """ model_config = {"frozen": True} chunk_id: ChunkId + tenant_id: TenantId model: str vector: list[float] dimension: int + dtype: EmbeddingDtype = EmbeddingDtype.float32 + acl_labels: frozenset[str] = Field(default_factory=frozenset) created_at: datetime = Field(default_factory=_utcnow) @model_validator(mode="after") @@ -242,6 +449,100 @@ class RetrievalResult(BaseModel): created_at: datetime = Field(default_factory=_utcnow) +# --------------------------------------------------------------------------- +# ChunkRef — ID-only retrieval result +# --------------------------------------------------------------------------- +class ChunkRef(BaseModel): + """Lightweight reference to a chunk returned by ID-only retrieval paths. + + Used by `RetrievalBackend.retrieve_ids` (Step 1.1b) to avoid hydrating + full chunk content for results that may be rejected by the reranker or + PolicyEngine. Carries only what downstream stages need to fuse, rerank, + and policy-check before hydration. + """ + + model_config = {"frozen": True} + + chunk_id: ChunkId + tenant_id: TenantId + score: float + acl_labels: frozenset[str] = Field(default_factory=frozenset) + corpus_id: CorpusId | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# QueryPlan — output of the planner, input to retrieval +# --------------------------------------------------------------------------- +class Cost(BaseModel): + """Best-effort cost envelope for a plan node (ADR-0008).""" + + model_config = {"frozen": True} + + ms_estimate: float = 0.0 + tokens_estimate: int = 0 + dollars_estimate: float = 0.0 + + +class PlanNode(BaseModel): + """Single step in a `QueryPlan`.""" + + model_config = {"frozen": True} + + op: str # e.g. "vector.retrieve", "keyword.retrieve", "rerank.precise" + backend: str # implementation identifier, e.g. "qdrant", "pgvector" + params: dict[str, Any] = Field(default_factory=dict) + estimated_cost: Cost = Field(default_factory=Cost) + + +class QueryPlan(BaseModel): + """Planner output describing how a query will be executed. + + See ADR-0008 (cost-aware planner) for the role this type plays in the + pre-dispatch budget check and plan mutation. + """ + + model_config = {"frozen": True} + + request_id: RequestId + nodes: list[PlanNode] + # Version bumped whenever the planner mutates the plan (e.g. budget cut). + version: int = 1 + corpus_version: str | None = None # used as a cache-key component + created_at: datetime = Field(default_factory=_utcnow) + + @property + def total_estimated_cost(self) -> Cost: + return Cost( + ms_estimate=sum(n.estimated_cost.ms_estimate for n in self.nodes), + tokens_estimate=sum(n.estimated_cost.tokens_estimate for n in self.nodes), + dollars_estimate=sum(n.estimated_cost.dollars_estimate for n in self.nodes), + ) + + +# --------------------------------------------------------------------------- +# StageEvent — typed cross-stage observation +# --------------------------------------------------------------------------- +class StageEvent(BaseModel): + """Typed event emitted by a pipeline stage. + + Used both for tracing (the gateway publishes them as SSE) and to feed the + cost-aware planner's online cost estimator (ADR-0008). + """ + + model_config = {"frozen": True} + + request_id: RequestId + stage: str # e.g. "embed", "vector.retrieve", "rerank.precise" + kind: StageEventKind + latency_ms: float | None = None + tokens: int | None = None + dollars: float | None = None + error: str | None = None + attributes: dict[str, Any] = Field(default_factory=dict) + occurred_at: datetime = Field(default_factory=_utcnow) + + # --------------------------------------------------------------------------- # AuditEvent # --------------------------------------------------------------------------- diff --git a/packages/core/tests/test_types.py b/packages/core/tests/test_types.py index db89295..acfa0f5 100644 --- a/packages/core/tests/test_types.py +++ b/packages/core/tests/test_types.py @@ -7,22 +7,33 @@ ACLAction, AuditEvent, AuditOutcome, + BlobRef, + Budget, Chunk, ChunkId, + ChunkRef, Citation, CorpusId, Document, DocumentId, DocumentStatus, Embedding, + EmbeddingDtype, + PiiAction, + PiiPolicy, Principal, PrincipalId, PrincipalKind, Query, + QueryPlan, + RequestContext, RetrievalResult, + StageEvent, + StageEventKind, Tenant, TenantId, TraceContext, + TrustLevel, ) @@ -203,18 +214,245 @@ def test_hierarchical_parent(self) -> None: class TestEmbedding: def test_valid(self) -> None: vec = [0.1, 0.2, 0.3] - emb = Embedding(chunk_id=_chunk_id(), model="bge-large", vector=vec, dimension=3) + emb = Embedding( + chunk_id=_chunk_id(), + tenant_id=_tid(), + model="bge-large", + vector=vec, + dimension=3, + ) assert emb.dimension == 3 + assert emb.tenant_id == _tid() + assert emb.dtype == EmbeddingDtype.float32 + assert emb.acl_labels == frozenset() def test_dimension_mismatch_raises(self) -> None: with pytest.raises(ValidationError, match="does not match dimension"): - Embedding(chunk_id=_chunk_id(), model="bge-large", vector=[0.1, 0.2], dimension=3) + Embedding( + chunk_id=_chunk_id(), + tenant_id=_tid(), + model="bge-large", + vector=[0.1, 0.2], + dimension=3, + ) def test_frozen(self) -> None: - emb = Embedding(chunk_id=_chunk_id(), model="bge-large", vector=[0.1], dimension=1) + emb = Embedding( + chunk_id=_chunk_id(), + tenant_id=_tid(), + model="bge-large", + vector=[0.1], + dimension=1, + ) with pytest.raises(ValidationError): emb.model = "other" # type: ignore[misc] + def test_int8_dtype(self) -> None: + emb = Embedding( + chunk_id=_chunk_id(), + tenant_id=_tid(), + model="bge-int8", + vector=[0.1], + dimension=1, + dtype=EmbeddingDtype.int8, + ) + assert emb.dtype == EmbeddingDtype.int8 + + def test_acl_labels(self) -> None: + emb = Embedding( + chunk_id=_chunk_id(), + tenant_id=_tid(), + model="bge", + vector=[0.1], + dimension=1, + acl_labels=frozenset({"pii", "internal"}), + ) + assert "pii" in emb.acl_labels + + +# --------------------------------------------------------------------------- +# Chunk extras (Step 1.1a additions) +# --------------------------------------------------------------------------- +class TestChunkStep11a: + def test_acl_labels_typed(self) -> None: + chunk = Chunk( + document_id=_doc_id(), + tenant_id=_tid(), + corpus_id=_cid(), + content="hi", + position=0, + acl_labels=frozenset({"public"}), + ) + assert chunk.acl_labels == frozenset({"public"}) + + def test_default_trust_level(self) -> None: + chunk = Chunk( + document_id=_doc_id(), + tenant_id=_tid(), + corpus_id=_cid(), + content="hi", + position=0, + ) + assert chunk.trust_level == TrustLevel.ingested + + def test_user_supplied_trust_level(self) -> None: + chunk = Chunk( + document_id=_doc_id(), + tenant_id=_tid(), + corpus_id=_cid(), + content="from a webhook", + position=0, + trust_level=TrustLevel.user_supplied, + ) + assert chunk.trust_level == TrustLevel.user_supplied + + def test_blob_ref_instead_of_content(self) -> None: + ref = BlobRef(uri="s3://bucket/chunks/c1.txt", size_bytes=4096) + chunk = Chunk( + document_id=_doc_id(), + tenant_id=_tid(), + corpus_id=_cid(), + content_ref=ref, + position=0, + ) + assert chunk.content is None + assert chunk.content_ref is not None + assert chunk.content_ref.size_bytes == 4096 + + def test_missing_content_and_ref_raises(self) -> None: + with pytest.raises(ValidationError, match="content or content_ref"): + Chunk( + document_id=_doc_id(), + tenant_id=_tid(), + corpus_id=_cid(), + position=0, + ) + + +# --------------------------------------------------------------------------- +# Step 1.1a — new envelope types +# --------------------------------------------------------------------------- +class TestRequestContext: + def _ctx(self) -> RequestContext: + return RequestContext( + tenant_id=_tid(), + principal=Principal( + id=_pid(), + kind=PrincipalKind.user, + display_name="alice", + tenant_id=_tid(), + ), + ) + + def test_defaults(self) -> None: + ctx = self._ctx() + assert ctx.tenant_id == _tid() + assert ctx.principal.tenant_id == _tid() + assert isinstance(ctx.trace, TraceContext) + assert ctx.budget.max_tokens is None + + def test_frozen(self) -> None: + ctx = self._ctx() + with pytest.raises(ValidationError): + ctx.tenant_id = TenantId("other") # type: ignore[misc] + + def test_tenant_mismatch_raises(self) -> None: + with pytest.raises(ValidationError, match="principal.tenant_id"): + RequestContext( + tenant_id=_tid(), + principal=Principal( + display_name="x", + tenant_id=TenantId("other-tenant"), + ), + ) + + +class TestBudget: + def test_spend_reduces_remaining(self) -> None: + b = Budget(max_tokens=1000, max_dollars=1.0, max_wall_ms=5000, max_iter=4) + b2 = b.spend(tokens=200, dollars=0.25, wall_ms=400, iter_=1) + assert b2.max_tokens == 800 + assert b2.max_dollars == pytest.approx(0.75) + assert b2.max_wall_ms == 4600 + assert b2.max_iter == 3 + + def test_spend_clamps_at_zero(self) -> None: + b = Budget(max_tokens=100) + assert b.spend(tokens=200).max_tokens == 0 + + def test_spend_ignores_unbounded(self) -> None: + b = Budget() # all None + b2 = b.spend(tokens=10, dollars=1.0) + assert b2.max_tokens is None + assert b2.max_dollars is None + + +class TestPiiPolicy: + def test_defaults(self) -> None: + p = PiiPolicy() + assert p.action == PiiAction.redact + assert p.min_score == pytest.approx(0.5) + assert p.entities == frozenset() + + +class TestBlobRef: + def test_basic(self) -> None: + ref = BlobRef(uri="s3://b/k", size_bytes=10) + assert ref.uri == "s3://b/k" + assert ref.content_type.startswith("text/plain") + + +class TestChunkRef: + def test_basic(self) -> None: + ref = ChunkRef( + chunk_id=_chunk_id(), + tenant_id=_tid(), + score=0.87, + acl_labels=frozenset({"public"}), + ) + assert ref.score == pytest.approx(0.87) + assert ref.acl_labels == frozenset({"public"}) + + +class TestQueryPlan: + def test_total_cost_sums(self) -> None: + from rag_core.types import Cost, PlanNode, RequestId + + plan = QueryPlan( + request_id=RequestId("r1"), + nodes=[ + PlanNode( + op="vector.retrieve", + backend="qdrant", + estimated_cost=Cost(ms_estimate=10, tokens_estimate=0, dollars_estimate=0.001), + ), + PlanNode( + op="rerank.precise", + backend="cohere", + estimated_cost=Cost(ms_estimate=50, tokens_estimate=400, dollars_estimate=0.01), + ), + ], + ) + total = plan.total_estimated_cost + assert total.ms_estimate == pytest.approx(60) + assert total.tokens_estimate == 400 + assert total.dollars_estimate == pytest.approx(0.011) + + +class TestStageEvent: + def test_basic(self) -> None: + from rag_core.types import RequestId + + ev = StageEvent( + request_id=RequestId("r1"), + stage="vector.retrieve", + kind=StageEventKind.completed, + latency_ms=42.0, + tokens=0, + ) + assert ev.kind == StageEventKind.completed + assert ev.latency_ms == pytest.approx(42.0) + # --------------------------------------------------------------------------- # Query diff --git a/pyproject.toml b/pyproject.toml index 15965b0..6870f10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,9 @@ ignore_missing_imports = true asyncio_mode = "auto" testpaths = ["tests", "packages"] pythonpath = ["packages/core/src", "packages/config/src", "packages/observability/src", "packages/ragctl/src", "packages/backends/src", "apps/gateway/src"] +# spi_signature.py is the SPI signature linter (Step 1.1a); referenced by name in +# docs/architecture/request-context.md. Collected alongside test_*.py files. +python_files = ["test_*.py", "*_test.py", "spi_signature.py"] addopts = "-x -q --tb=short --import-mode=importlib" markers = [ "contract: SPI conformance tests", diff --git a/tests/contract/conftest.py b/tests/contract/conftest.py index 0e399ab..baced0d 100644 --- a/tests/contract/conftest.py +++ b/tests/contract/conftest.py @@ -28,7 +28,10 @@ CorpusId, DocumentId, Embedding, + Principal, PrincipalId, + PrincipalKind, + RequestContext, TenantId, ) @@ -42,6 +45,21 @@ PRINCIPAL = PrincipalId("alice") +def make_ctx(tenant_id: str = "tenant-test", principal_id: str = "alice") -> RequestContext: + """Build a `RequestContext` for tests with sensible defaults.""" + + tid = TenantId(tenant_id) + return RequestContext( + tenant_id=tid, + principal=Principal( + id=PrincipalId(principal_id), + kind=PrincipalKind.user, + display_name=principal_id, + tenant_id=tid, + ), + ) + + @pytest.fixture() def tenant_id() -> TenantId: return TENANT @@ -52,6 +70,17 @@ def corpus_id() -> CorpusId: return CORPUS +@pytest.fixture() +def ctx() -> RequestContext: + return make_ctx(tenant_id=str(TENANT), principal_id=str(PRINCIPAL)) + + +@pytest.fixture() +def other_ctx() -> RequestContext: + """A request context for a different tenant — used by isolation tests.""" + return make_ctx(tenant_id="tenant-other", principal_id="mallory") + + # --------------------------------------------------------------------------- # SPI fixtures # --------------------------------------------------------------------------- @@ -143,9 +172,14 @@ def pii_detector() -> NoopPIIDetector: # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -def make_embedding(chunk_id: str = "chunk-1", dim: int = 4) -> Embedding: +def make_embedding( + chunk_id: str = "chunk-1", + dim: int = 4, + tenant_id: str = "tenant-test", +) -> Embedding: return Embedding( chunk_id=ChunkId(chunk_id), + tenant_id=TenantId(tenant_id), model="noop-embedder", vector=[0.1, 0.2, 0.3, 0.4][:dim], dimension=dim, @@ -158,8 +192,6 @@ def make_chunk( tenant_id: str = "tenant-test", corpus_id: str = "corpus-a", ) -> Chunk: - from rag_core.types import DocumentId - return Chunk( id=ChunkId(chunk_id), document_id=DocumentId("doc-1"), diff --git a/tests/contract/spi_signature.py b/tests/contract/spi_signature.py new file mode 100644 index 0000000..ae9cea1 --- /dev/null +++ b/tests/contract/spi_signature.py @@ -0,0 +1,161 @@ +"""Linter: every public SPI abstract method takes ``ctx: RequestContext`` first. + +This is the type-level enforcement promised in +docs/architecture/request-context.md (Step 1.1a). If a new SPI method is added +that omits ``ctx`` — or a refactor silently drops it — this test fails CI. + +Exceptions +---------- + +A small set of methods are documented exceptions: + +* ``HealthCheckMixin.health`` — pre-ctx liveness probe. +* ``Auth.*`` — runs at the gateway boundary *before* a + ``RequestContext`` exists. Produces the ``Principal`` that the gateway then + uses to build ``ctx``. +* ``Telemetry.*`` — cross-cutting infrastructure; owns + ``TraceContext`` directly (which is itself a field of ``RequestContext``). +* ``AuditStore.*`` — consumes already-built ``AuditEvent`` records + via the ``AuditWriter`` facade (which is the caller that takes ``ctx``). +* ``Embedder.model`` / ``Embedder.dimension`` — properties, not callable. +* ``Parser.supports`` — pure capability check, no I/O. + +Add exceptions only with a docstring comment justifying *why* the method +operates below the ctx layer. +""" + +from __future__ import annotations + +import inspect +from typing import Any + +import pytest +from rag_core.spi import ( + LLM, + OCR, + Cache, + Connector, + Embedder, + GraphStore, + KeywordStore, + Parser, + PIIDetector, + Queue, + Reranker, + Secrets, + Storage, + VectorStore, +) +from rag_core.types import RequestContext + +pytestmark = pytest.mark.contract + + +# Per-class set of methods that legitimately do *not* take ctx as first arg. +_PER_CLASS_EXEMPT: dict[type, set[str]] = { + # health() lives on HealthCheckMixin and is universally pre-ctx. + Cache: {"health"}, + Connector: {"health"}, + Embedder: {"health", "model", "dimension"}, + GraphStore: {"health"}, + KeywordStore: {"health"}, + LLM: {"health"}, + OCR: {"health"}, + Parser: {"health", "supports"}, + PIIDetector: {"health"}, + Queue: {"health"}, + Reranker: {"health"}, + Secrets: {"health"}, + Storage: {"health"}, + VectorStore: {"health"}, +} + + +def _ctx_threaded_classes() -> list[type]: + return [ + Cache, + Connector, + Embedder, + GraphStore, + KeywordStore, + LLM, + OCR, + Parser, + PIIDetector, + Queue, + Reranker, + Secrets, + Storage, + VectorStore, + ] + + +def _public_methods(cls: type) -> dict[str, Any]: + """Return public callables declared directly on ``cls`` (not inherited).""" + + members: dict[str, Any] = {} + for name in cls.__dict__: + if name.startswith("_"): + continue + attr = inspect.getattr_static(cls, name) + members[name] = attr + return members + + +def test_every_spi_method_takes_ctx_first() -> None: + """Every non-exempt public method on an SPI ABC has ``ctx: RequestContext`` first.""" + + failures: list[str] = [] + + for cls in _ctx_threaded_classes(): + exempt = _PER_CLASS_EXEMPT.get(cls, set()) + for name, attr in _public_methods(cls).items(): + if name in exempt: + continue + # Skip non-function descriptors (e.g. @property) + if isinstance(attr, property): + continue + + # `attr` may be a function, async function, or classmethod/staticmethod. + func = attr + if isinstance(attr, (classmethod, staticmethod)): + func = attr.__func__ # type: ignore[assignment] + + if not callable(func): + continue + + try: + sig = inspect.signature(func) + except (TypeError, ValueError): + continue + + params = list(sig.parameters.values()) + # Drop `self` for instance methods + if params and params[0].name == "self": + params = params[1:] + + if not params: + failures.append(f"{cls.__name__}.{name}(): missing ctx parameter (no args)") + continue + + first = params[0] + if first.name != "ctx": + failures.append( + f"{cls.__name__}.{name}(): first parameter is {first.name!r}, expected 'ctx'" + ) + continue + + ann = first.annotation + # Accept either the class or the string forward-reference form. + ok = ann is RequestContext or ( + isinstance(ann, str) and ann.split(".")[-1] == "RequestContext" + ) + if not ok: + failures.append( + f"{cls.__name__}.{name}(): ctx annotation is {ann!r}, expected RequestContext" + ) + + assert not failures, ( + "SPI signature linter failed — every public SPI method must take " + "`ctx: RequestContext` as its first argument:\n - " + "\n - ".join(failures) + ) diff --git a/tests/contract/test_cache.py b/tests/contract/test_cache.py index 9066397..94144c5 100644 --- a/tests/contract/test_cache.py +++ b/tests/contract/test_cache.py @@ -2,6 +2,7 @@ import pytest from rag_core.spi.noop import NoopCache +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -10,33 +11,33 @@ async def test_health(cache: NoopCache) -> None: assert await cache.health() is True -async def test_set_and_get(cache: NoopCache) -> None: - await cache.set("key1", b"value1") - assert await cache.get("key1") == b"value1" +async def test_set_and_get(cache: NoopCache, ctx: RequestContext) -> None: + await cache.set(ctx, "key1", b"value1") + assert await cache.get(ctx, "key1") == b"value1" -async def test_get_miss_returns_none(cache: NoopCache) -> None: - assert await cache.get("nonexistent") is None +async def test_get_miss_returns_none(cache: NoopCache, ctx: RequestContext) -> None: + assert await cache.get(ctx, "nonexistent") is None -async def test_exists(cache: NoopCache) -> None: - assert not await cache.exists("k") - await cache.set("k", b"v") - assert await cache.exists("k") +async def test_exists(cache: NoopCache, ctx: RequestContext) -> None: + assert not await cache.exists(ctx, "k") + await cache.set(ctx, "k", b"v") + assert await cache.exists(ctx, "k") -async def test_delete(cache: NoopCache) -> None: - await cache.set("k", b"v") - await cache.delete("k") - assert await cache.get("k") is None - assert not await cache.exists("k") +async def test_delete(cache: NoopCache, ctx: RequestContext) -> None: + await cache.set(ctx, "k", b"v") + await cache.delete(ctx, "k") + assert await cache.get(ctx, "k") is None + assert not await cache.exists(ctx, "k") -async def test_delete_nonexistent_is_noop(cache: NoopCache) -> None: - await cache.delete("ghost") # must not raise +async def test_delete_nonexistent_is_noop(cache: NoopCache, ctx: RequestContext) -> None: + await cache.delete(ctx, "ghost") # must not raise -async def test_overwrite(cache: NoopCache) -> None: - await cache.set("k", b"first") - await cache.set("k", b"second") - assert await cache.get("k") == b"second" +async def test_overwrite(cache: NoopCache, ctx: RequestContext) -> None: + await cache.set(ctx, "k", b"first") + await cache.set(ctx, "k", b"second") + assert await cache.get(ctx, "k") == b"second" diff --git a/tests/contract/test_connector.py b/tests/contract/test_connector.py index 5b533a9..e94b122 100644 --- a/tests/contract/test_connector.py +++ b/tests/contract/test_connector.py @@ -2,14 +2,21 @@ import pytest from rag_core.spi.noop import NoopConnector -from rag_core.types import CorpusId, DocumentId, DocumentStatus, TenantId +from rag_core.types import ( + CorpusId, + Document, + DocumentId, + DocumentStatus, + RequestContext, + TenantId, +) -pytestmark = pytest.mark.contract +from tests.contract.conftest import make_ctx +pytestmark = pytest.mark.contract -def _make_doc(tenant_id: TenantId, corpus_id: CorpusId) -> object: - from rag_core.types import Document +def _make_doc(tenant_id: TenantId, corpus_id: CorpusId) -> Document: return Document( id=DocumentId("d1"), tenant_id=tenant_id, @@ -24,30 +31,31 @@ async def test_health(connector: NoopConnector) -> None: assert await connector.health() is True -async def test_list_empty(connector: NoopConnector, tenant_id: TenantId) -> None: - docs = [d async for d in connector.list_documents(tenant_id)] +async def test_list_empty(connector: NoopConnector, ctx: RequestContext) -> None: + docs = [d async for d in connector.list_documents(ctx)] assert docs == [] -async def test_list_with_documents(tenant_id: TenantId, corpus_id: CorpusId) -> None: - doc = _make_doc(tenant_id, corpus_id) - c = NoopConnector(documents=[doc]) # type: ignore[arg-type] - docs = [d async for d in c.list_documents(tenant_id)] +async def test_list_with_documents(ctx: RequestContext, corpus_id: CorpusId) -> None: + doc = _make_doc(ctx.tenant_id, corpus_id) + c = NoopConnector(documents=[doc]) + docs = [d async for d in c.list_documents(ctx)] assert len(docs) == 1 assert docs[0].id == DocumentId("d1") async def test_tenant_isolation(corpus_id: CorpusId) -> None: - t1, t2 = TenantId("t1"), TenantId("t2") - doc = _make_doc(t1, corpus_id) - c = NoopConnector(documents=[doc]) # type: ignore[arg-type] - docs = [d async for d in c.list_documents(t2)] + ctx_t1 = make_ctx(tenant_id="t1") + ctx_t2 = make_ctx(tenant_id="t2") + doc = _make_doc(ctx_t1.tenant_id, corpus_id) + c = NoopConnector(documents=[doc]) + docs = [d async for d in c.list_documents(ctx_t2)] assert docs == [] async def test_fetch_returns_bytes( - connector: NoopConnector, tenant_id: TenantId, corpus_id: CorpusId + connector: NoopConnector, ctx: RequestContext, corpus_id: CorpusId ) -> None: - doc = _make_doc(tenant_id, corpus_id) - data = await connector.fetch(doc) # type: ignore[arg-type] + doc = _make_doc(ctx.tenant_id, corpus_id) + data = await connector.fetch(ctx, doc) assert isinstance(data, bytes) diff --git a/tests/contract/test_embedder.py b/tests/contract/test_embedder.py index a9e22ef..8558a7f 100644 --- a/tests/contract/test_embedder.py +++ b/tests/contract/test_embedder.py @@ -2,7 +2,7 @@ import pytest from rag_core.spi.noop import NoopEmbedder -from rag_core.types import ChunkId, TenantId +from rag_core.types import ChunkId, RequestContext pytestmark = pytest.mark.contract @@ -11,30 +11,35 @@ async def test_health(embedder: NoopEmbedder) -> None: assert await embedder.health() is True -async def test_embed_returns_correct_count(embedder: NoopEmbedder, tenant_id: TenantId) -> None: +async def test_embed_returns_correct_count(embedder: NoopEmbedder, ctx: RequestContext) -> None: texts = ["hello", "world", "foo"] ids = [ChunkId(f"c{i}") for i in range(len(texts))] - results = await embedder.embed(texts, ids, tenant_id) + results = await embedder.embed(ctx, texts, ids) assert len(results) == 3 -async def test_embed_dimension_matches(embedder: NoopEmbedder, tenant_id: TenantId) -> None: - results = await embedder.embed(["test"], [ChunkId("c1")], tenant_id) +async def test_embed_dimension_matches(embedder: NoopEmbedder, ctx: RequestContext) -> None: + results = await embedder.embed(ctx, ["test"], [ChunkId("c1")]) assert results[0].dimension == embedder.dimension assert len(results[0].vector) == embedder.dimension -async def test_embed_preserves_chunk_ids(embedder: NoopEmbedder, tenant_id: TenantId) -> None: +async def test_embed_preserves_chunk_ids(embedder: NoopEmbedder, ctx: RequestContext) -> None: ids = [ChunkId("alpha"), ChunkId("beta")] - results = await embedder.embed(["a", "b"], ids, tenant_id) + results = await embedder.embed(ctx, ["a", "b"], ids) assert [r.chunk_id for r in results] == ids -async def test_embed_model_name(embedder: NoopEmbedder, tenant_id: TenantId) -> None: - results = await embedder.embed(["x"], [ChunkId("c1")], tenant_id) +async def test_embed_sets_tenant_id(embedder: NoopEmbedder, ctx: RequestContext) -> None: + results = await embedder.embed(ctx, ["x"], [ChunkId("c1")]) + assert results[0].tenant_id == ctx.tenant_id + + +async def test_embed_model_name(embedder: NoopEmbedder, ctx: RequestContext) -> None: + results = await embedder.embed(ctx, ["x"], [ChunkId("c1")]) assert results[0].model == embedder.model -async def test_embed_mismatched_lengths_raises(embedder: NoopEmbedder, tenant_id: TenantId) -> None: +async def test_embed_mismatched_lengths_raises(embedder: NoopEmbedder, ctx: RequestContext) -> None: with pytest.raises(ValueError): - await embedder.embed(["a", "b"], [ChunkId("c1")], tenant_id) + await embedder.embed(ctx, ["a", "b"], [ChunkId("c1")]) diff --git a/tests/contract/test_graph_store.py b/tests/contract/test_graph_store.py index 8442161..774c2c8 100644 --- a/tests/contract/test_graph_store.py +++ b/tests/contract/test_graph_store.py @@ -2,7 +2,9 @@ import pytest from rag_core.spi.noop import NoopGraphStore -from rag_core.types import TenantId +from rag_core.types import RequestContext + +from tests.contract.conftest import make_ctx pytestmark = pytest.mark.contract @@ -11,30 +13,31 @@ async def test_health(graph_store: NoopGraphStore) -> None: assert await graph_store.health() is True -async def test_upsert_and_query_node(graph_store: NoopGraphStore, tenant_id: TenantId) -> None: - await graph_store.upsert_node("n1", ["Document"], {"title": "Hello"}, tenant_id) - rows = await graph_store.query("MATCH (n) RETURN n", {}, tenant_id) +async def test_upsert_and_query_node(graph_store: NoopGraphStore, ctx: RequestContext) -> None: + await graph_store.upsert_node(ctx, "n1", ["Document"], {"title": "Hello"}) + rows = await graph_store.query(ctx, "MATCH (n) RETURN n", {}) node_ids = [r["node_id"] for r in rows] assert "n1" in node_ids async def test_tenant_isolation(graph_store: NoopGraphStore) -> None: - t1, t2 = TenantId("t1"), TenantId("t2") - await graph_store.upsert_node("n1", ["X"], {}, t1) - rows = await graph_store.query("MATCH (n) RETURN n", {}, t2) + ctx_t1 = make_ctx(tenant_id="t1") + ctx_t2 = make_ctx(tenant_id="t2") + await graph_store.upsert_node(ctx_t1, "n1", ["X"], {}) + rows = await graph_store.query(ctx_t2, "MATCH (n) RETURN n", {}) assert rows == [] -async def test_delete_node(graph_store: NoopGraphStore, tenant_id: TenantId) -> None: - await graph_store.upsert_node("n1", ["X"], {}, tenant_id) - await graph_store.delete_node("n1", tenant_id) - rows = await graph_store.query("", {}, tenant_id) +async def test_delete_node(graph_store: NoopGraphStore, ctx: RequestContext) -> None: + await graph_store.upsert_node(ctx, "n1", ["X"], {}) + await graph_store.delete_node(ctx, "n1") + rows = await graph_store.query(ctx, "", {}) assert all(r["node_id"] != "n1" for r in rows) -async def test_delete_removes_edges(graph_store: NoopGraphStore, tenant_id: TenantId) -> None: - await graph_store.upsert_node("a", [], {}, tenant_id) - await graph_store.upsert_node("b", [], {}, tenant_id) - await graph_store.upsert_edge("a", "b", "LINKS_TO", {}, tenant_id) - await graph_store.delete_node("a", tenant_id) +async def test_delete_removes_edges(graph_store: NoopGraphStore, ctx: RequestContext) -> None: + await graph_store.upsert_node(ctx, "a", [], {}) + await graph_store.upsert_node(ctx, "b", [], {}) + await graph_store.upsert_edge(ctx, "a", "b", "LINKS_TO", {}) + await graph_store.delete_node(ctx, "a") assert not any(e[1] == "a" or e[2] == "a" for e in graph_store._edges) diff --git a/tests/contract/test_keyword_store.py b/tests/contract/test_keyword_store.py index 2a22393..a23e4b4 100644 --- a/tests/contract/test_keyword_store.py +++ b/tests/contract/test_keyword_store.py @@ -2,9 +2,9 @@ import pytest from rag_core.spi.noop import NoopKeywordStore -from rag_core.types import ChunkId, CorpusId, TenantId +from rag_core.types import ChunkId, CorpusId, RequestContext -from tests.contract.conftest import make_chunk +from tests.contract.conftest import make_chunk, make_ctx pytestmark = pytest.mark.contract @@ -14,41 +14,48 @@ async def test_health(keyword_store: NoopKeywordStore) -> None: async def test_index_and_search( - keyword_store: NoopKeywordStore, tenant_id: TenantId, corpus_id: CorpusId + keyword_store: NoopKeywordStore, ctx: RequestContext, corpus_id: CorpusId ) -> None: - chunk = make_chunk("c1", content="the quick brown fox") - await keyword_store.index([chunk], tenant_id) - results = await keyword_store.search( - "quick fox", top_k=5, tenant_id=tenant_id, corpus_ids=[corpus_id] - ) + chunk = make_chunk("c1", content="the quick brown fox", tenant_id=str(ctx.tenant_id)) + await keyword_store.index(ctx, [chunk]) + results = await keyword_store.search(ctx, "quick fox", top_k=5, corpus_ids=[corpus_id]) assert any(cid == ChunkId("c1") for cid, _ in results) async def test_search_tenant_isolation(keyword_store: NoopKeywordStore) -> None: - t1, t2 = TenantId("t1"), TenantId("t2") + ctx_t1 = make_ctx(tenant_id="t1") + ctx_t2 = make_ctx(tenant_id="t2") chunk = make_chunk("c1", content="secret data", tenant_id="t1") - await keyword_store.index([chunk], t1) - results = await keyword_store.search("secret", top_k=5, tenant_id=t2, corpus_ids=[]) + await keyword_store.index(ctx_t1, [chunk]) + results = await keyword_store.search(ctx_t2, "secret", top_k=5, corpus_ids=[]) assert results == [] -async def test_delete(keyword_store: NoopKeywordStore, tenant_id: TenantId) -> None: - chunks = [make_chunk("c1", content="apple"), make_chunk("c2", content="apple")] - await keyword_store.index(chunks, tenant_id) - await keyword_store.delete([ChunkId("c1")], tenant_id) - results = await keyword_store.search("apple", top_k=5, tenant_id=tenant_id, corpus_ids=[]) +async def test_delete(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: + chunks = [ + make_chunk("c1", content="apple", tenant_id=str(ctx.tenant_id)), + make_chunk("c2", content="apple", tenant_id=str(ctx.tenant_id)), + ] + await keyword_store.index(ctx, chunks) + await keyword_store.delete(ctx, [ChunkId("c1")]) + results = await keyword_store.search(ctx, "apple", top_k=5, corpus_ids=[]) ids = [r[0] for r in results] assert ChunkId("c1") not in ids -async def test_no_match_returns_empty(keyword_store: NoopKeywordStore, tenant_id: TenantId) -> None: - await keyword_store.index([make_chunk("c1", content="apple pie")], tenant_id) - results = await keyword_store.search("zebra", top_k=5, tenant_id=tenant_id, corpus_ids=[]) +async def test_no_match_returns_empty(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: + await keyword_store.index( + ctx, [make_chunk("c1", content="apple pie", tenant_id=str(ctx.tenant_id))] + ) + results = await keyword_store.search(ctx, "zebra", top_k=5, corpus_ids=[]) assert results == [] -async def test_scores_in_range(keyword_store: NoopKeywordStore, tenant_id: TenantId) -> None: - await keyword_store.index([make_chunk("c1", content="cat sat on the mat")], tenant_id) - results = await keyword_store.search("cat mat", top_k=5, tenant_id=tenant_id, corpus_ids=[]) +async def test_scores_in_range(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: + await keyword_store.index( + ctx, + [make_chunk("c1", content="cat sat on the mat", tenant_id=str(ctx.tenant_id))], + ) + results = await keyword_store.search(ctx, "cat mat", top_k=5, corpus_ids=[]) for _, score in results: assert 0.0 <= score <= 1.0 diff --git a/tests/contract/test_llm.py b/tests/contract/test_llm.py index 062f119..4d14b63 100644 --- a/tests/contract/test_llm.py +++ b/tests/contract/test_llm.py @@ -3,7 +3,7 @@ import pytest from rag_core.spi.llm import LLMMessage from rag_core.spi.noop import NoopLLM -from rag_core.types import TenantId +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -12,24 +12,24 @@ async def test_health(llm: NoopLLM) -> None: assert await llm.health() is True -async def test_complete_returns_response(llm: NoopLLM, tenant_id: TenantId) -> None: +async def test_complete_returns_response(llm: NoopLLM, ctx: RequestContext) -> None: msgs = [LLMMessage(role="user", content="Hello")] - resp = await llm.complete(msgs, tenant_id) + resp = await llm.complete(ctx, msgs) assert isinstance(resp.content, str) assert len(resp.content) > 0 assert resp.model == "noop-llm" -async def test_complete_token_counts_positive(llm: NoopLLM, tenant_id: TenantId) -> None: +async def test_complete_token_counts_positive(llm: NoopLLM, ctx: RequestContext) -> None: msgs = [LLMMessage(role="user", content="count tokens")] - resp = await llm.complete(msgs, tenant_id) + resp = await llm.complete(ctx, msgs) assert resp.input_tokens >= 0 assert resp.output_tokens >= 0 -async def test_stream_yields_strings(llm: NoopLLM, tenant_id: TenantId) -> None: +async def test_stream_yields_strings(llm: NoopLLM, ctx: RequestContext) -> None: msgs = [LLMMessage(role="user", content="stream test")] - tokens = [tok async for tok in await llm.stream(msgs, tenant_id)] + tokens = [tok async for tok in await llm.stream(ctx, msgs)] assert len(tokens) > 0 assert all(isinstance(t, str) for t in tokens) joined = "".join(tokens) diff --git a/tests/contract/test_ocr.py b/tests/contract/test_ocr.py index 6bec459..ed80f91 100644 --- a/tests/contract/test_ocr.py +++ b/tests/contract/test_ocr.py @@ -2,6 +2,7 @@ import pytest from rag_core.spi.noop import NoopOCR +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -10,12 +11,12 @@ async def test_health(ocr: NoopOCR) -> None: assert await ocr.health() is True -async def test_extract_returns_result(ocr: NoopOCR) -> None: - result = await ocr.extract(b"\x89PNG\r\n", "image/png") +async def test_extract_returns_result(ocr: NoopOCR, ctx: RequestContext) -> None: + result = await ocr.extract(ctx, b"\x89PNG\r\n", "image/png") assert isinstance(result.text, str) assert 0.0 <= result.confidence <= 1.0 -async def test_extract_empty_bytes(ocr: NoopOCR) -> None: - result = await ocr.extract(b"", "image/jpeg") +async def test_extract_empty_bytes(ocr: NoopOCR, ctx: RequestContext) -> None: + result = await ocr.extract(ctx, b"", "image/jpeg") assert result.text == "" diff --git a/tests/contract/test_parser.py b/tests/contract/test_parser.py index 6b0d5e5..f0157e8 100644 --- a/tests/contract/test_parser.py +++ b/tests/contract/test_parser.py @@ -2,7 +2,7 @@ import pytest from rag_core.spi.noop import NoopParser -from rag_core.types import CorpusId, DocumentId, DocumentStatus, TenantId +from rag_core.types import CorpusId, DocumentId, DocumentStatus, RequestContext pytestmark = pytest.mark.contract @@ -17,32 +17,31 @@ async def test_supports_any_mime(parser: NoopParser) -> None: async def test_parse_returns_document( - parser: NoopParser, tenant_id: TenantId, corpus_id: CorpusId + parser: NoopParser, ctx: RequestContext, corpus_id: CorpusId ) -> None: doc = await parser.parse( + ctx, b"hello world", "text/plain", DocumentId("doc-1"), - tenant_id, corpus_id, "file:///test.txt", ) assert doc.status == DocumentStatus.ready assert doc.content_hash != "" - assert doc.tenant_id == tenant_id + assert doc.tenant_id == ctx.tenant_id assert doc.corpus_id == corpus_id async def test_parse_content_hash_deterministic( - parser: NoopParser, tenant_id: TenantId, corpus_id: CorpusId + parser: NoopParser, ctx: RequestContext, corpus_id: CorpusId ) -> None: kwargs = dict( mime_type="text/plain", document_id=DocumentId("d1"), - tenant_id=tenant_id, corpus_id=corpus_id, source_uri="file:///a.txt", ) - doc1 = await parser.parse(b"same content", **kwargs) - doc2 = await parser.parse(b"same content", **kwargs) + doc1 = await parser.parse(ctx, b"same content", **kwargs) + doc2 = await parser.parse(ctx, b"same content", **kwargs) assert doc1.content_hash == doc2.content_hash diff --git a/tests/contract/test_pii_detector.py b/tests/contract/test_pii_detector.py index ed9a1ef..02b40a4 100644 --- a/tests/contract/test_pii_detector.py +++ b/tests/contract/test_pii_detector.py @@ -2,7 +2,7 @@ import pytest from rag_core.spi.noop import NoopPIIDetector -from rag_core.types import TenantId +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -11,21 +11,21 @@ async def test_health(pii_detector: NoopPIIDetector) -> None: assert await pii_detector.health() is True -async def test_detect_returns_list(pii_detector: NoopPIIDetector, tenant_id: TenantId) -> None: - spans = await pii_detector.detect("My name is Alice", tenant_id) +async def test_detect_returns_list(pii_detector: NoopPIIDetector, ctx: RequestContext) -> None: + spans = await pii_detector.detect(ctx, "My name is Alice") assert isinstance(spans, list) -async def test_redact_returns_string(pii_detector: NoopPIIDetector, tenant_id: TenantId) -> None: - result = await pii_detector.redact("My name is Alice", tenant_id) +async def test_redact_returns_string(pii_detector: NoopPIIDetector, ctx: RequestContext) -> None: + result = await pii_detector.redact(ctx, "My name is Alice") assert isinstance(result, str) -async def test_noop_detects_nothing(pii_detector: NoopPIIDetector, tenant_id: TenantId) -> None: - spans = await pii_detector.detect("alice@example.com", tenant_id) +async def test_noop_detects_nothing(pii_detector: NoopPIIDetector, ctx: RequestContext) -> None: + spans = await pii_detector.detect(ctx, "alice@example.com") assert spans == [] -async def test_noop_redact_unchanged(pii_detector: NoopPIIDetector, tenant_id: TenantId) -> None: +async def test_noop_redact_unchanged(pii_detector: NoopPIIDetector, ctx: RequestContext) -> None: text = "call me at 555-1234" - assert await pii_detector.redact(text, tenant_id) == text + assert await pii_detector.redact(ctx, text) == text diff --git a/tests/contract/test_queue.py b/tests/contract/test_queue.py index 850fe38..d48db25 100644 --- a/tests/contract/test_queue.py +++ b/tests/contract/test_queue.py @@ -2,6 +2,7 @@ import pytest from rag_core.spi.noop import NoopQueue +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -10,20 +11,20 @@ async def test_health(queue: NoopQueue) -> None: assert await queue.health() is True -async def test_publish_returns_message_id(queue: NoopQueue) -> None: - msg_id = await queue.publish("topic.ingest", b'{"doc_id": "1"}') +async def test_publish_returns_message_id(queue: NoopQueue, ctx: RequestContext) -> None: + msg_id = await queue.publish(ctx, "topic.ingest", b'{"doc_id": "1"}') assert isinstance(msg_id, str) assert len(msg_id) > 0 -async def test_published_messages_recorded(queue: NoopQueue) -> None: - await queue.publish("topic.a", b"payload1") - await queue.publish("topic.a", b"payload2") +async def test_published_messages_recorded(queue: NoopQueue, ctx: RequestContext) -> None: + await queue.publish(ctx, "topic.a", b"payload1") + await queue.publish(ctx, "topic.a", b"payload2") assert len(queue.published) == 2 assert queue.published[0].payload == b"payload1" assert queue.published[1].payload == b"payload2" -async def test_message_attributes(queue: NoopQueue) -> None: - await queue.publish("topic.x", b"data", attributes={"source": "test"}) +async def test_message_attributes(queue: NoopQueue, ctx: RequestContext) -> None: + await queue.publish(ctx, "topic.x", b"data", attributes={"source": "test"}) assert queue.published[0].attributes["source"] == "test" diff --git a/tests/contract/test_reranker.py b/tests/contract/test_reranker.py index 50886b1..4bdc50e 100644 --- a/tests/contract/test_reranker.py +++ b/tests/contract/test_reranker.py @@ -2,6 +2,7 @@ import pytest from rag_core.spi.noop import NoopReranker +from rag_core.types import RequestContext from tests.contract.conftest import make_chunk @@ -12,20 +13,20 @@ async def test_health(reranker: NoopReranker) -> None: assert await reranker.health() is True -async def test_rerank_returns_pairs(reranker: NoopReranker) -> None: +async def test_rerank_returns_pairs(reranker: NoopReranker, ctx: RequestContext) -> None: chunks = [make_chunk(f"c{i}", content=f"doc {i}") for i in range(5)] - results = await reranker.rerank("query", chunks, top_k=3) + results = await reranker.rerank(ctx, "query", chunks, top_k=3) assert len(results) <= 3 for _chunk, score in results: assert 0.0 <= score <= 1.0 -async def test_rerank_empty_input(reranker: NoopReranker) -> None: - results = await reranker.rerank("query", [], top_k=5) +async def test_rerank_empty_input(reranker: NoopReranker, ctx: RequestContext) -> None: + results = await reranker.rerank(ctx, "query", [], top_k=5) assert results == [] -async def test_rerank_top_k_respected(reranker: NoopReranker) -> None: +async def test_rerank_top_k_respected(reranker: NoopReranker, ctx: RequestContext) -> None: chunks = [make_chunk(f"c{i}") for i in range(10)] - results = await reranker.rerank("q", chunks, top_k=4) + results = await reranker.rerank(ctx, "q", chunks, top_k=4) assert len(results) <= 4 diff --git a/tests/contract/test_secrets.py b/tests/contract/test_secrets.py index 3ef7483..53d6fc7 100644 --- a/tests/contract/test_secrets.py +++ b/tests/contract/test_secrets.py @@ -2,6 +2,7 @@ import pytest from rag_core.spi.noop import NoopSecrets +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -10,16 +11,16 @@ async def test_health(secrets: NoopSecrets) -> None: assert await secrets.health() is True -async def test_get_known_secret(secrets: NoopSecrets) -> None: - value = await secrets.get("db_password") +async def test_get_known_secret(secrets: NoopSecrets, ctx: RequestContext) -> None: + value = await secrets.get(ctx, "db_password") assert value == "s3cr3t" -async def test_get_unknown_raises(secrets: NoopSecrets) -> None: +async def test_get_unknown_raises(secrets: NoopSecrets, ctx: RequestContext) -> None: with pytest.raises(KeyError): - await secrets.get("nonexistent") + await secrets.get(ctx, "nonexistent") -async def test_get_bytes(secrets: NoopSecrets) -> None: - data = await secrets.get_bytes("db_password") +async def test_get_bytes(secrets: NoopSecrets, ctx: RequestContext) -> None: + data = await secrets.get_bytes(ctx, "db_password") assert data == b"s3cr3t" diff --git a/tests/contract/test_storage.py b/tests/contract/test_storage.py index a78c76b..b84244f 100644 --- a/tests/contract/test_storage.py +++ b/tests/contract/test_storage.py @@ -2,6 +2,7 @@ import pytest from rag_core.spi.noop import NoopStorage +from rag_core.types import RequestContext pytestmark = pytest.mark.contract @@ -10,35 +11,35 @@ async def test_health(storage: NoopStorage) -> None: assert await storage.health() is True -async def test_put_and_get(storage: NoopStorage) -> None: - await storage.put("docs/a.txt", b"hello") - assert await storage.get("docs/a.txt") == b"hello" +async def test_put_and_get(storage: NoopStorage, ctx: RequestContext) -> None: + await storage.put(ctx, "docs/a.txt", b"hello") + assert await storage.get(ctx, "docs/a.txt") == b"hello" -async def test_get_missing_raises(storage: NoopStorage) -> None: +async def test_get_missing_raises(storage: NoopStorage, ctx: RequestContext) -> None: with pytest.raises(KeyError): - await storage.get("missing/key") + await storage.get(ctx, "missing/key") -async def test_exists(storage: NoopStorage) -> None: - assert not await storage.exists("k") - await storage.put("k", b"v") - assert await storage.exists("k") +async def test_exists(storage: NoopStorage, ctx: RequestContext) -> None: + assert not await storage.exists(ctx, "k") + await storage.put(ctx, "k", b"v") + assert await storage.exists(ctx, "k") -async def test_delete(storage: NoopStorage) -> None: - await storage.put("k", b"v") - await storage.delete("k") - assert not await storage.exists("k") +async def test_delete(storage: NoopStorage, ctx: RequestContext) -> None: + await storage.put(ctx, "k", b"v") + await storage.delete(ctx, "k") + assert not await storage.exists(ctx, "k") -async def test_delete_nonexistent_is_noop(storage: NoopStorage) -> None: - await storage.delete("ghost") +async def test_delete_nonexistent_is_noop(storage: NoopStorage, ctx: RequestContext) -> None: + await storage.delete(ctx, "ghost") -async def test_list_keys(storage: NoopStorage) -> None: - await storage.put("a/1", b"x") - await storage.put("a/2", b"y") - await storage.put("b/1", b"z") - keys = [k async for k in storage.list_keys("a/")] +async def test_list_keys(storage: NoopStorage, ctx: RequestContext) -> None: + await storage.put(ctx, "a/1", b"x") + await storage.put(ctx, "a/2", b"y") + await storage.put(ctx, "b/1", b"z") + keys = [k async for k in storage.list_keys(ctx, "a/")] assert sorted(keys) == ["a/1", "a/2"] diff --git a/tests/contract/test_vector_store.py b/tests/contract/test_vector_store.py index df0aad2..17c7c2e 100644 --- a/tests/contract/test_vector_store.py +++ b/tests/contract/test_vector_store.py @@ -2,9 +2,9 @@ import pytest from rag_core.spi.noop import NoopVectorStore -from rag_core.types import ChunkId, TenantId +from rag_core.types import ChunkId, RequestContext -from tests.contract.conftest import make_embedding +from tests.contract.conftest import make_ctx, make_embedding pytestmark = pytest.mark.contract @@ -13,10 +13,10 @@ async def test_health(vector_store: NoopVectorStore) -> None: assert await vector_store.health() is True -async def test_upsert_and_query(vector_store: NoopVectorStore, tenant_id: TenantId) -> None: - emb = make_embedding("c1") - await vector_store.upsert([emb], tenant_id) - results = await vector_store.query(emb.vector, top_k=5, tenant_id=tenant_id, corpus_ids=[]) +async def test_upsert_and_query(vector_store: NoopVectorStore, ctx: RequestContext) -> None: + emb = make_embedding("c1", tenant_id=str(ctx.tenant_id)) + await vector_store.upsert(ctx, [emb]) + results = await vector_store.query(ctx, emb.vector, top_k=5, corpus_ids=[]) assert len(results) == 1 chunk_id, score = results[0] assert chunk_id == ChunkId("c1") @@ -24,46 +24,53 @@ async def test_upsert_and_query(vector_store: NoopVectorStore, tenant_id: Tenant async def test_query_tenant_isolation(vector_store: NoopVectorStore) -> None: - t1, t2 = TenantId("t1"), TenantId("t2") - await vector_store.upsert([make_embedding("c1")], t1) - results = await vector_store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=t2, corpus_ids=[]) + ctx_t1 = make_ctx(tenant_id="t1") + ctx_t2 = make_ctx(tenant_id="t2") + await vector_store.upsert(ctx_t1, [make_embedding("c1", tenant_id="t1")]) + results = await vector_store.query(ctx_t2, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) assert results == [] -async def test_delete(vector_store: NoopVectorStore, tenant_id: TenantId) -> None: - await vector_store.upsert([make_embedding("c1"), make_embedding("c2")], tenant_id) - await vector_store.delete([ChunkId("c1")], tenant_id) - results = await vector_store.query( - [0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=tenant_id, corpus_ids=[] +async def test_delete(vector_store: NoopVectorStore, ctx: RequestContext) -> None: + await vector_store.upsert( + ctx, + [ + make_embedding("c1", tenant_id=str(ctx.tenant_id)), + make_embedding("c2", tenant_id=str(ctx.tenant_id)), + ], ) + await vector_store.delete(ctx, [ChunkId("c1")]) + results = await vector_store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) ids = [r[0] for r in results] assert ChunkId("c1") not in ids assert ChunkId("c2") in ids async def test_delete_unknown_id_is_noop( - vector_store: NoopVectorStore, tenant_id: TenantId + vector_store: NoopVectorStore, ctx: RequestContext ) -> None: - await vector_store.delete([ChunkId("nonexistent")], tenant_id) # must not raise + await vector_store.delete(ctx, [ChunkId("nonexistent")]) # must not raise -async def test_top_k_respected(vector_store: NoopVectorStore, tenant_id: TenantId) -> None: +async def test_top_k_respected(vector_store: NoopVectorStore, ctx: RequestContext) -> None: for i in range(10): - await vector_store.upsert([make_embedding(f"c{i}")], tenant_id) - results = await vector_store.query( - [0.1, 0.2, 0.3, 0.4], top_k=3, tenant_id=tenant_id, corpus_ids=[] - ) + await vector_store.upsert(ctx, [make_embedding(f"c{i}", tenant_id=str(ctx.tenant_id))]) + results = await vector_store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=3, corpus_ids=[]) assert len(results) <= 3 -async def test_upsert_overwrites(vector_store: NoopVectorStore, tenant_id: TenantId) -> None: +async def test_upsert_overwrites(vector_store: NoopVectorStore, ctx: RequestContext) -> None: from rag_core.types import Embedding - emb1 = make_embedding("c1") - emb2 = Embedding(chunk_id=ChunkId("c1"), model="noop", vector=[0.9, 0.9, 0.9, 0.9], dimension=4) - await vector_store.upsert([emb1], tenant_id) - await vector_store.upsert([emb2], tenant_id) - results = await vector_store.query( - [0.9, 0.9, 0.9, 0.9], top_k=1, tenant_id=tenant_id, corpus_ids=[] + emb1 = make_embedding("c1", tenant_id=str(ctx.tenant_id)) + emb2 = Embedding( + chunk_id=ChunkId("c1"), + tenant_id=ctx.tenant_id, + model="noop", + vector=[0.9, 0.9, 0.9, 0.9], + dimension=4, ) + await vector_store.upsert(ctx, [emb1]) + await vector_store.upsert(ctx, [emb2]) + results = await vector_store.query(ctx, [0.9, 0.9, 0.9, 0.9], top_k=1, corpus_ids=[]) assert len(results) == 1 diff --git a/tests/integration/test_local_storage.py b/tests/integration/test_local_storage.py index 301ccf8..415940a 100644 --- a/tests/integration/test_local_storage.py +++ b/tests/integration/test_local_storage.py @@ -9,6 +9,10 @@ import pytest from rag_backends.storage.local import LocalFileStorage +from tests.contract.conftest import make_ctx + +CTX = make_ctx(tenant_id="local-store-test") + @pytest.fixture() def store(tmp_path: Path) -> LocalFileStorage: @@ -20,45 +24,54 @@ async def test_health(store: LocalFileStorage) -> None: async def test_put_and_get(store: LocalFileStorage) -> None: - await store.put("docs/hello.txt", b"hello") - assert await store.get("docs/hello.txt") == b"hello" + await store.put(CTX, "docs/hello.txt", b"hello") + assert await store.get(CTX, "docs/hello.txt") == b"hello" async def test_get_missing_raises(store: LocalFileStorage) -> None: with pytest.raises(KeyError): - await store.get("no/such/file.txt") + await store.get(CTX, "no/such/file.txt") async def test_exists(store: LocalFileStorage) -> None: - assert not await store.exists("a.bin") - await store.put("a.bin", b"x") - assert await store.exists("a.bin") + assert not await store.exists(CTX, "a.bin") + await store.put(CTX, "a.bin", b"x") + assert await store.exists(CTX, "a.bin") async def test_delete(store: LocalFileStorage) -> None: - await store.put("b.bin", b"v") - await store.delete("b.bin") - assert not await store.exists("b.bin") + await store.put(CTX, "b.bin", b"v") + await store.delete(CTX, "b.bin") + assert not await store.exists(CTX, "b.bin") async def test_delete_nonexistent_is_noop(store: LocalFileStorage) -> None: - await store.delete("ghost.bin") + await store.delete(CTX, "ghost.bin") async def test_list_keys(store: LocalFileStorage) -> None: - await store.put("a/1.txt", b"x") - await store.put("a/2.txt", b"y") - await store.put("b/3.txt", b"z") - keys = sorted([k async for k in store.list_keys("a/")]) + await store.put(CTX, "a/1.txt", b"x") + await store.put(CTX, "a/2.txt", b"y") + await store.put(CTX, "b/3.txt", b"z") + keys = sorted([k async for k in store.list_keys(CTX, "a/")]) assert keys == ["a/1.txt", "a/2.txt"] async def test_overwrite(store: LocalFileStorage) -> None: - await store.put("c.txt", b"first") - await store.put("c.txt", b"second") - assert await store.get("c.txt") == b"second" + await store.put(CTX, "c.txt", b"first") + await store.put(CTX, "c.txt", b"second") + assert await store.get(CTX, "c.txt") == b"second" async def test_path_traversal_blocked(store: LocalFileStorage) -> None: with pytest.raises(ValueError): - await store.put("../escape.txt", b"evil") + await store.put(CTX, "../escape.txt", b"evil") + + +async def test_tenant_isolation(store: LocalFileStorage) -> None: + ctx_a = make_ctx(tenant_id="tenant-a") + ctx_b = make_ctx(tenant_id="tenant-b") + await store.put(ctx_a, "iso.txt", b"from-a") + assert not await store.exists(ctx_b, "iso.txt") + with pytest.raises(KeyError): + await store.get(ctx_b, "iso.txt") diff --git a/tests/integration/test_pgvector.py b/tests/integration/test_pgvector.py index 8b54cfa..52b0ed1 100644 --- a/tests/integration/test_pgvector.py +++ b/tests/integration/test_pgvector.py @@ -8,6 +8,8 @@ from rag_backends.vector.pgvector import PgVectorStore from rag_core.types import ChunkId, Embedding, TenantId +from tests.contract.conftest import make_ctx + pytestmark = pytest.mark.integration # Unique table name per test run avoids cross-test pollution @@ -29,9 +31,15 @@ async def store(pg_dsn: str) -> PgVectorStore: # type: ignore[misc] await s.close() -def _emb(chunk_id: str, vec: list[float] | None = None) -> Embedding: +def _emb(chunk_id: str, tenant_id: TenantId, vec: list[float] | None = None) -> Embedding: v = vec or [0.1, 0.2, 0.3, 0.4] - return Embedding(chunk_id=ChunkId(chunk_id), model="test", vector=v, dimension=_DIM) + return Embedding( + chunk_id=ChunkId(chunk_id), + tenant_id=tenant_id, + model="test", + vector=v, + dimension=_DIM, + ) async def test_health(store: PgVectorStore) -> None: @@ -39,9 +47,9 @@ async def test_health(store: PgVectorStore) -> None: async def test_upsert_and_query(store: PgVectorStore) -> None: - tid = TenantId("t-pg-1") - await store.upsert([_emb("c1")], tid) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=tid, corpus_ids=[]) + ctx = make_ctx(tenant_id="t-pg-1") + await store.upsert(ctx, [_emb("c1", ctx.tenant_id)]) + results = await store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) assert len(results) >= 1 ids = [r[0] for r in results] assert ChunkId("c1") in ids @@ -51,40 +59,45 @@ async def test_upsert_and_query(store: PgVectorStore) -> None: async def test_tenant_isolation(store: PgVectorStore) -> None: - t1, t2 = TenantId("t-pg-iso-1"), TenantId("t-pg-iso-2") - await store.upsert([_emb("c-iso")], t1) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=t2, corpus_ids=[]) + ctx_t1 = make_ctx(tenant_id="t-pg-iso-1") + ctx_t2 = make_ctx(tenant_id="t-pg-iso-2") + await store.upsert(ctx_t1, [_emb("c-iso", ctx_t1.tenant_id)]) + results = await store.query(ctx_t2, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) ids = [r[0] for r in results] assert ChunkId("c-iso") not in ids async def test_delete(store: PgVectorStore) -> None: - tid = TenantId("t-pg-del") - await store.upsert([_emb("c-del-1"), _emb("c-del-2")], tid) - await store.delete([ChunkId("c-del-1")], tid) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=tid, corpus_ids=[]) + ctx = make_ctx(tenant_id="t-pg-del") + await store.upsert( + ctx, + [_emb("c-del-1", ctx.tenant_id), _emb("c-del-2", ctx.tenant_id)], + ) + await store.delete(ctx, [ChunkId("c-del-1")]) + results = await store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) ids = [r[0] for r in results] assert ChunkId("c-del-1") not in ids assert ChunkId("c-del-2") in ids async def test_delete_unknown_is_noop(store: PgVectorStore) -> None: - await store.delete([ChunkId("ghost")], TenantId("t-pg-ghost")) + ctx = make_ctx(tenant_id="t-pg-ghost") + await store.delete(ctx, [ChunkId("ghost")]) async def test_upsert_overwrites(store: PgVectorStore) -> None: - tid = TenantId("t-pg-overwrite") - await store.upsert([_emb("c-ow", [0.1, 0.2, 0.3, 0.4])], tid) - await store.upsert([_emb("c-ow", [0.9, 0.9, 0.9, 0.9])], tid) - results = await store.query([0.9, 0.9, 0.9, 0.9], top_k=1, tenant_id=tid, corpus_ids=[]) + ctx = make_ctx(tenant_id="t-pg-overwrite") + await store.upsert(ctx, [_emb("c-ow", ctx.tenant_id, [0.1, 0.2, 0.3, 0.4])]) + await store.upsert(ctx, [_emb("c-ow", ctx.tenant_id, [0.9, 0.9, 0.9, 0.9])]) + results = await store.query(ctx, [0.9, 0.9, 0.9, 0.9], top_k=1, corpus_ids=[]) assert len(results) == 1 assert results[0][0] == ChunkId("c-ow") assert results[0][1] == pytest.approx(1.0, abs=1e-4) async def test_top_k_respected(store: PgVectorStore) -> None: - tid = TenantId("t-pg-topk") + ctx = make_ctx(tenant_id="t-pg-topk") for i in range(8): - await store.upsert([_emb(f"c-topk-{i}")], tid) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=3, tenant_id=tid, corpus_ids=[]) + await store.upsert(ctx, [_emb(f"c-topk-{i}", ctx.tenant_id)]) + results = await store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=3, corpus_ids=[]) assert len(results) <= 3 diff --git a/tests/integration/test_qdrant.py b/tests/integration/test_qdrant.py index efaa086..dac306f 100644 --- a/tests/integration/test_qdrant.py +++ b/tests/integration/test_qdrant.py @@ -8,6 +8,8 @@ from rag_backends.vector.qdrant import QdrantVectorStore from rag_core.types import ChunkId, Embedding, TenantId +from tests.contract.conftest import make_ctx + pytestmark = pytest.mark.integration _COLLECTION = "rag_inttest_qdrant" @@ -24,9 +26,15 @@ async def store(qdrant_url: str) -> QdrantVectorStore: # type: ignore[misc] await s.close() -def _emb(chunk_id: str, vec: list[float] | None = None) -> Embedding: +def _emb(chunk_id: str, tenant_id: TenantId, vec: list[float] | None = None) -> Embedding: v = vec or [0.1, 0.2, 0.3, 0.4] - return Embedding(chunk_id=ChunkId(chunk_id), model="test", vector=v, dimension=_DIM) + return Embedding( + chunk_id=ChunkId(chunk_id), + tenant_id=tenant_id, + model="test", + vector=v, + dimension=_DIM, + ) async def test_health(store: QdrantVectorStore) -> None: @@ -34,9 +42,9 @@ async def test_health(store: QdrantVectorStore) -> None: async def test_upsert_and_query(store: QdrantVectorStore) -> None: - tid = TenantId("t-qd-1") - await store.upsert([_emb("c1")], tid) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=tid, corpus_ids=[]) + ctx = make_ctx(tenant_id="t-qd-1") + await store.upsert(ctx, [_emb("c1", ctx.tenant_id)]) + results = await store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) assert len(results) >= 1 ids = [r[0] for r in results] assert ChunkId("c1") in ids @@ -45,40 +53,45 @@ async def test_upsert_and_query(store: QdrantVectorStore) -> None: async def test_tenant_isolation(store: QdrantVectorStore) -> None: - t1, t2 = TenantId("t-qd-iso-1"), TenantId("t-qd-iso-2") - await store.upsert([_emb("c-qd-iso")], t1) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=t2, corpus_ids=[]) + ctx_t1 = make_ctx(tenant_id="t-qd-iso-1") + ctx_t2 = make_ctx(tenant_id="t-qd-iso-2") + await store.upsert(ctx_t1, [_emb("c-qd-iso", ctx_t1.tenant_id)]) + results = await store.query(ctx_t2, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) ids = [r[0] for r in results] assert ChunkId("c-qd-iso") not in ids async def test_delete(store: QdrantVectorStore) -> None: - tid = TenantId("t-qd-del") - await store.upsert([_emb("c-qd-del-1"), _emb("c-qd-del-2")], tid) - await store.delete([ChunkId("c-qd-del-1")], tid) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=5, tenant_id=tid, corpus_ids=[]) + ctx = make_ctx(tenant_id="t-qd-del") + await store.upsert( + ctx, + [_emb("c-qd-del-1", ctx.tenant_id), _emb("c-qd-del-2", ctx.tenant_id)], + ) + await store.delete(ctx, [ChunkId("c-qd-del-1")]) + results = await store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) ids = [r[0] for r in results] assert ChunkId("c-qd-del-1") not in ids assert ChunkId("c-qd-del-2") in ids async def test_delete_unknown_is_noop(store: QdrantVectorStore) -> None: - await store.delete([ChunkId("ghost")], TenantId("t-qd-ghost")) + ctx = make_ctx(tenant_id="t-qd-ghost") + await store.delete(ctx, [ChunkId("ghost")]) async def test_upsert_overwrites(store: QdrantVectorStore) -> None: - tid = TenantId("t-qd-overwrite") - await store.upsert([_emb("c-qd-ow", [0.1, 0.2, 0.3, 0.4])], tid) - await store.upsert([_emb("c-qd-ow", [0.9, 0.9, 0.9, 0.9])], tid) - results = await store.query([0.9, 0.9, 0.9, 0.9], top_k=1, tenant_id=tid, corpus_ids=[]) + ctx = make_ctx(tenant_id="t-qd-overwrite") + await store.upsert(ctx, [_emb("c-qd-ow", ctx.tenant_id, [0.1, 0.2, 0.3, 0.4])]) + await store.upsert(ctx, [_emb("c-qd-ow", ctx.tenant_id, [0.9, 0.9, 0.9, 0.9])]) + results = await store.query(ctx, [0.9, 0.9, 0.9, 0.9], top_k=1, corpus_ids=[]) assert len(results) == 1 assert results[0][0] == ChunkId("c-qd-ow") assert results[0][1] == pytest.approx(1.0, abs=1e-4) async def test_top_k_respected(store: QdrantVectorStore) -> None: - tid = TenantId("t-qd-topk") + ctx = make_ctx(tenant_id="t-qd-topk") for i in range(8): - await store.upsert([_emb(f"c-qd-topk-{i}")], tid) - results = await store.query([0.1, 0.2, 0.3, 0.4], top_k=3, tenant_id=tid, corpus_ids=[]) + await store.upsert(ctx, [_emb(f"c-qd-topk-{i}", ctx.tenant_id)]) + results = await store.query(ctx, [0.1, 0.2, 0.3, 0.4], top_k=3, corpus_ids=[]) assert len(results) <= 3 diff --git a/tests/integration/test_redis.py b/tests/integration/test_redis.py index ca74990..7bf1b98 100644 --- a/tests/integration/test_redis.py +++ b/tests/integration/test_redis.py @@ -9,8 +9,12 @@ import pytest from rag_backends.cache.redis import RedisCache +from tests.contract.conftest import make_ctx + pytestmark = pytest.mark.integration +CTX = make_ctx(tenant_id="inttest-redis") + @pytest.fixture() async def cache(redis_url: str) -> RedisCache: # type: ignore[misc] @@ -24,50 +28,50 @@ async def test_health(cache: RedisCache) -> None: async def test_set_and_get(cache: RedisCache) -> None: - await cache.set("k1", b"hello") - assert await cache.get("k1") == b"hello" + await cache.set(CTX, "k1", b"hello") + assert await cache.get(CTX, "k1") == b"hello" async def test_get_missing_returns_none(cache: RedisCache) -> None: - assert await cache.get("missing-key-xyz") is None + assert await cache.get(CTX, "missing-key-xyz") is None async def test_exists(cache: RedisCache) -> None: - assert not await cache.exists("k-exists-check") - await cache.set("k-exists-check", b"v") - assert await cache.exists("k-exists-check") + assert not await cache.exists(CTX, "k-exists-check") + await cache.set(CTX, "k-exists-check", b"v") + assert await cache.exists(CTX, "k-exists-check") async def test_delete(cache: RedisCache) -> None: - await cache.set("k-del", b"v") - await cache.delete("k-del") - assert not await cache.exists("k-del") + await cache.set(CTX, "k-del", b"v") + await cache.delete(CTX, "k-del") + assert not await cache.exists(CTX, "k-del") async def test_delete_nonexistent_is_noop(cache: RedisCache) -> None: - await cache.delete("ghost-key") + await cache.delete(CTX, "ghost-key") async def test_ttl_expiry(cache: RedisCache) -> None: - await cache.set("k-ttl", b"expires", ttl_seconds=1) - assert await cache.get("k-ttl") == b"expires" + await cache.set(CTX, "k-ttl", b"expires", ttl_seconds=1) + assert await cache.get(CTX, "k-ttl") == b"expires" await asyncio.sleep(1.1) - assert await cache.get("k-ttl") is None + assert await cache.get(CTX, "k-ttl") is None async def test_overwrite(cache: RedisCache) -> None: - await cache.set("k-ow", b"first") - await cache.set("k-ow", b"second") - assert await cache.get("k-ow") == b"second" + await cache.set(CTX, "k-ow", b"first") + await cache.set(CTX, "k-ow", b"second") + assert await cache.get(CTX, "k-ow") == b"second" -async def test_prefix_isolation(redis_url: str) -> None: - """Two RedisCache instances with different prefixes must not collide.""" - c1 = RedisCache(url=redis_url, prefix="ns1:") - c2 = RedisCache(url=redis_url, prefix="ns2:") +async def test_tenant_isolation(redis_url: str) -> None: + """Two RequestContexts with different tenant_ids must not collide.""" + c = RedisCache(url=redis_url, prefix="ns:") + ctx_a = make_ctx(tenant_id="tenant-a") + ctx_b = make_ctx(tenant_id="tenant-b") try: - await c1.set("shared-key", b"from-ns1") - assert await c2.get("shared-key") is None + await c.set(ctx_a, "shared-key", b"from-a") + assert await c.get(ctx_b, "shared-key") is None finally: - await c1.close() - await c2.close() + await c.close() diff --git a/tests/integration/test_s3.py b/tests/integration/test_s3.py index 1d833e4..c39a686 100644 --- a/tests/integration/test_s3.py +++ b/tests/integration/test_s3.py @@ -7,8 +7,12 @@ import pytest from rag_backends.storage.s3 import S3Storage +from tests.contract.conftest import make_ctx + pytestmark = pytest.mark.integration +CTX = make_ctx(tenant_id="inttest-s3") + @pytest.fixture() async def store(minio_config: dict[str, str]) -> S3Storage: # type: ignore[misc] @@ -26,43 +30,52 @@ async def test_health(store: S3Storage) -> None: async def test_put_and_get(store: S3Storage) -> None: - await store.put("test/hello.txt", b"hello world", content_type="text/plain") - data = await store.get("test/hello.txt") + await store.put(CTX, "test/hello.txt", b"hello world", content_type="text/plain") + data = await store.get(CTX, "test/hello.txt") assert data == b"hello world" async def test_get_missing_raises(store: S3Storage) -> None: with pytest.raises(KeyError): - await store.get("test/does-not-exist.bin") + await store.get(CTX, "test/does-not-exist.bin") async def test_exists(store: S3Storage) -> None: - assert not await store.exists("test/exists-check.bin") - await store.put("test/exists-check.bin", b"x") - assert await store.exists("test/exists-check.bin") + assert not await store.exists(CTX, "test/exists-check.bin") + await store.put(CTX, "test/exists-check.bin", b"x") + assert await store.exists(CTX, "test/exists-check.bin") async def test_delete(store: S3Storage) -> None: - await store.put("test/del.txt", b"bye") - await store.delete("test/del.txt") - assert not await store.exists("test/del.txt") + await store.put(CTX, "test/del.txt", b"bye") + await store.delete(CTX, "test/del.txt") + assert not await store.exists(CTX, "test/del.txt") async def test_delete_nonexistent_is_noop(store: S3Storage) -> None: - await store.delete("test/ghost.bin") + await store.delete(CTX, "test/ghost.bin") async def test_list_keys(store: S3Storage) -> None: - await store.put("list/a.txt", b"a") - await store.put("list/b.txt", b"b") - await store.put("other/c.txt", b"c") - keys = [k async for k in store.list_keys("list/")] + await store.put(CTX, "list/a.txt", b"a") + await store.put(CTX, "list/b.txt", b"b") + await store.put(CTX, "other/c.txt", b"c") + keys = [k async for k in store.list_keys(CTX, "list/")] assert "list/a.txt" in keys assert "list/b.txt" in keys assert "other/c.txt" not in keys async def test_overwrite(store: S3Storage) -> None: - await store.put("test/ow.txt", b"v1") - await store.put("test/ow.txt", b"v2") - assert await store.get("test/ow.txt") == b"v2" + await store.put(CTX, "test/ow.txt", b"v1") + await store.put(CTX, "test/ow.txt", b"v2") + assert await store.get(CTX, "test/ow.txt") == b"v2" + + +async def test_tenant_isolation(store: S3Storage) -> None: + ctx_a = make_ctx(tenant_id="tenant-a-s3iso") + ctx_b = make_ctx(tenant_id="tenant-b-s3iso") + await store.put(ctx_a, "iso.txt", b"from-a") + assert not await store.exists(ctx_b, "iso.txt") + with pytest.raises(KeyError): + await store.get(ctx_b, "iso.txt")