From ab134afaaff1a02871db16a0fa31bd1ff3ba4957 Mon Sep 17 00:00:00 2001 From: Deep Kumar Singh Kushwah Date: Sun, 24 May 2026 03:21:10 +0530 Subject: [PATCH] =?UTF-8?q?refactor(core):=20SPI=20split=20=E2=80=94=20Ret?= =?UTF-8?q?rievalBackend=20/=20IndexBackend=20(Step=201.1b)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split VectorStore / KeywordStore / GraphStore into narrow read+write roles (*RetrievalBackend / *IndexBackend) so consumers depend on the role they actually use; legacy composite ABCs preserved for backends serving both sides. Adds bulk + streaming + ID-only retrieval to the SPI surface, introduces IndexHint + WriteVolume for scale-tier writes (ADR-0009), and splits Embedder into single embed + canonical bulk_embed. - Retrieval side: retrieve_ids → list[ChunkRef]; hydrate on keyword stores loads full Chunks (vector stores pass-through, content lives elsewhere). - Index side: bulk_index / bulk_delete; stream_index default batches an AsyncIterator into bulk_index calls. Graph adds bulk_upsert_nodes/edges and bulk_delete_nodes + stream_upsert_nodes. - Backends migrated: NoopVectorStore / NoopKeywordStore / NoopGraphStore / NoopEmbedder + PgVectorStore + QdrantVectorStore. - tests/contract/spi_signature.py extended to enforce ctx-first on all new role ABCs. - Schemas regenerated (IndexHint.json now produced by gen_schemas). - Docs: rag-core.md adds "Read/write SPI split" + IndexHint section; storage-backends.md points to the new method names. Co-Authored-By: Claude Opus 4.7 (1M context) --- TRACKER.md | 8 +- docs/architecture/storage-backends.md | 2 +- docs/reference/rag-core.md | 49 +++++- .../src/rag_backends/vector/pgvector.py | 44 ++++-- .../src/rag_backends/vector/qdrant.py | 52 +++++-- packages/core/src/rag_core/__init__.py | 16 ++ packages/core/src/rag_core/gen_schemas.py | 2 + packages/core/src/rag_core/spi/__init__.py | 29 +++- packages/core/src/rag_core/spi/embedder.py | 33 +++- packages/core/src/rag_core/spi/graph_store.py | 114 ++++++++++++-- .../core/src/rag_core/spi/keyword_store.py | 116 +++++++++++--- .../core/src/rag_core/spi/noop/embedder.py | 2 +- .../src/rag_core/spi/noop/keyword_store.py | 62 ++++++-- .../src/rag_core/spi/noop/vector_store.py | 46 ++++-- .../core/src/rag_core/spi/vector_store.py | 143 +++++++++++++++--- packages/core/src/rag_core/types.py | 44 ++++++ tests/contract/spi_signature.py | 18 +++ tests/contract/test_embedder.py | 41 +++-- tests/contract/test_graph_store.py | 32 ++++ tests/contract/test_keyword_store.py | 67 +++++--- tests/contract/test_vector_store.py | 78 ++++++---- tests/integration/test_pgvector.py | 70 +++++---- tests/integration/test_qdrant.py | 69 +++++---- uv.lock | 2 +- 24 files changed, 907 insertions(+), 232 deletions(-) diff --git a/TRACKER.md b/TRACKER.md index 13211d0..9485bc8 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.1b — SPI split (Retrieval/Index, bulk + streaming + ID-only methods, IndexHint) +**Next action:** Phase 1 Step 1.1c — PolicyEngine package (`rag-policy`): `PolicyEngine` SPI + noop impl as single PDP for ACL/PII/quotas/redaction > **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 | **2** | 14 | +| 1 | Ingestion + Knowledge Store | 16 | **3** | 13 | | 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** | **15** | **69** | +| **Total** | | **84** | **16** | **68** | --- @@ -69,7 +69,7 @@ |------|-------|--------|--------|----|-----------------| | 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 | ✅ | `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.1b | SPI split — Retrieval/Index, bulk + streaming + ID-only | ✅ | `build/phase-1/step-1.1b-spi-split-retrieval-index` | _pending_ | Split `VectorStore`/`KeywordStore`/`GraphStore` into `*RetrievalBackend` (read) + `*IndexBackend` (write) composite ABCs. `retrieve_ids` returns `list[ChunkRef]`; `hydrate` lives on the retrieval side (keyword full-Chunk; vector pass-through). Bulk: `bulk_index`/`bulk_delete` (+ graph bulk node/edge variants); streaming: `stream_index` async-iterator default that batches into `bulk_index`. `Embedder` split into single `embed` + canonical `bulk_embed`. New `IndexHint` + `WriteVolume` types passed to writes (per ADR-0009). Noop impls, `PgVectorStore`, `QdrantVectorStore` migrated; conformance + integration tests updated; `spi_signature.py` extended to enforce the split. Schemas regenerated (`IndexHint.json`). | | 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. | | 1.1e | Cache SPI split + perf discipline + async telemetry | ⏳ | 1.1a | — | Split `Cache` into `EmbeddingCache` (key model_id+text_hash), `RetrievalCache` (plan_hash+corpus_version), `AnswerCache` (plan_hash+corpus_version+policy_version). Each has distinct invalidation. Hot-path convention doc (Pydantic at SPI boundary, `model_construct`/msgspec inside). Async telemetry path with bounded buffer + drop-on-overflow counter. | diff --git a/docs/architecture/storage-backends.md b/docs/architecture/storage-backends.md index a9c654e..c676751 100644 --- a/docs/architecture/storage-backends.md +++ b/docs/architecture/storage-backends.md @@ -86,7 +86,7 @@ stack. To add a new VectorStore backend: 1. Create `packages/backends/src/rag_backends/vector/.py` -2. Implement `VectorStore` ABC (upsert, query, delete, health) +2. Implement `VectorStore` ABC — `bulk_index`, `retrieve_ids`, `bulk_delete`, `health` (the inherited `stream_index` and `hydrate` defaults are usually fine). See Step 1.1b notes in [docs/reference/rag-core.md](../reference/rag-core.md#read--write-spi-split-step-11b). 3. Add an `initialize()` method for schema/collection setup 4. Add an integration test in `tests/integration/test_.py` 5. Export from `rag_backends.vector.__init__` and `rag_backends.__init__` diff --git a/docs/reference/rag-core.md b/docs/reference/rag-core.md index 7d1e46b..1ddc1e2 100644 --- a/docs/reference/rag-core.md +++ b/docs/reference/rag-core.md @@ -30,6 +30,7 @@ don't mutate. | `StageEvent` | Typed cross-stage observation | Step 1.1a (new) | | `TrustLevel` | Chunk provenance for prompt-injection defense | Step 1.1a (new) | | `EmbeddingDtype` | float32 / int8 / binary quantization | Step 1.1a (new) — ADR-0009 | +| `IndexHint` / `WriteVolume` | Scale-tier hint passed to `IndexBackend` writes | Step 1.1b (new) — ADR-0009 | | `PiiAction` | block / redact / mask / encrypt / tag_only / allow | Step 1.1a (new) | --- @@ -73,13 +74,57 @@ 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=[]) +await vector_store.bulk_index(ctx, embeddings, hint=IndexHint(estimated_size=50_000)) +refs = await vector_store.retrieve_ids(ctx, query_vec, top_k=10, corpus_ids=[]) ``` The `tests/contract/spi_signature.py` linter fails CI if an SPI method is added without `ctx: RequestContext` as its first argument. +### Read / write SPI split (Step 1.1b) + +`VectorStore`, `KeywordStore`, and `GraphStore` are now thin composite ABCs +over two narrower roles: + +| Role | Methods | Used by | +|---|---|---| +| `*RetrievalBackend` | `retrieve_ids`, `hydrate` (where applicable), `query` (graph) | Gateway query path, rerankers | +| `*IndexBackend` | `bulk_index`, `stream_index`, `bulk_delete` (+ graph node/edge bulk variants) | Ingest pipeline | + +Backends serving both sides keep inheriting from the composite class +(`VectorStore`, `KeywordStore`, `GraphStore`) — nothing changes for +implementers. Consumers should depend on the narrower role for clarity. + +```python +# Hot-path discipline: retrieve cheap refs, rerank, hydrate the survivors +refs = await vector_backend.retrieve_ids(ctx, qvec, top_k=200, corpus_ids=[]) +top = await reranker.rerank(ctx, refs, query_text) +hot = await keyword_backend.hydrate(ctx, top[:20]) +``` + +`Embedder` similarly exposes `bulk_embed(ctx, texts, chunk_ids)` (canonical +batch entry) and a thin `embed(ctx, text, chunk_id)` for single-item callers. +See [performance.md](../architecture/performance.md) for the p99 budgets per +method. + +### Choosing an `IndexHint` + +The hint is consulted at write time (and at `initialize()` time for backends +that need to pick an index implementation). Reasonable starting values: + +```python +IndexHint( + estimated_size=50_000, # current + projected vector count + recall_target=0.95, # default; raise to 0.99 for legal/medical + latency_target_ms=50.0, # p99 query budget + write_volume=WriteVolume.low, # bump to medium/high for streaming ingest +) +``` + +PgVector and Qdrant currently honour `hint` only at `initialize()` time; the +per-call writes ignore it. Per ADR-0009, future backends will use it to pick +between flat / ivfflat / HNSW / IVF-PQ index types. + ### Lazy chunk text via `BlobRef` `Chunk.content` and `Chunk.content_ref` are mutually exclusive at the diff --git a/packages/backends/src/rag_backends/vector/pgvector.py b/packages/backends/src/rag_backends/vector/pgvector.py index 1fdcbcd..5c59a88 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(ctx, embeddings) - results = await store.query(ctx, query_vector, top_k=10, corpus_ids=[]) + await store.bulk_index(ctx, embeddings) + refs = await store.retrieve_ids(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,15 @@ 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, RequestContext +from rag_core.types import ( + ChunkId, + ChunkRef, + CorpusId, + Embedding, + IndexHint, + RequestContext, + TenantId, +) from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -65,6 +73,7 @@ _SQL_QUERY_BASE = """ SELECT chunk_id, + corpus_id, 1.0 - (vector <=> $2) AS score FROM {table} WHERE tenant_id = $1 @@ -145,10 +154,18 @@ async def close(self) -> None: self._pool = None # ------------------------------------------------------------------ - # VectorStore SPI + # VectorStore SPI (Step 1.1b split: bulk_index / retrieve_ids / bulk_delete) # ------------------------------------------------------------------ - async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: + async def bulk_index( + self, + ctx: RequestContext, + embeddings: list[Embedding], + *, + hint: IndexHint | None = None, + ) -> None: + # ``hint`` is consulted only at ``initialize()`` time (index choice); + # per-call writes ignore it. See ADR-0009. if not embeddings: return pool = await self._get_pool() @@ -167,14 +184,14 @@ async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None async with pool.acquire() as conn: await conn.executemany(_SQL_UPSERT.format(table=self._table), rows) - async def query( + async def retrieve_ids( self, ctx: RequestContext, vector: list[float], top_k: int, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, - ) -> list[tuple[ChunkId, float]]: + ) -> list[ChunkRef]: pool = await self._get_pool() qvec = _to_np(vector) @@ -195,9 +212,18 @@ async def query( async with pool.acquire() as conn: rows = await conn.fetch(sql, *args) - return [(ChunkId(row["chunk_id"]), float(row["score"])) for row in rows] + tenant = TenantId(ctx.tenant_id) + return [ + ChunkRef( + chunk_id=ChunkId(row["chunk_id"]), + tenant_id=tenant, + score=float(row["score"]), + corpus_id=CorpusId(row["corpus_id"]) if row.get("corpus_id") else None, + ) + for row in rows + ] - async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: + async def bulk_delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: if not chunk_ids: return pool = await self._get_pool() diff --git a/packages/backends/src/rag_backends/vector/qdrant.py b/packages/backends/src/rag_backends/vector/qdrant.py index 361a728..f7cdfb5 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(ctx, embeddings) - results = await store.query(ctx, vector, top_k=10, corpus_ids=[]) + await store.bulk_index(ctx, embeddings) + refs = await store.retrieve_ids(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,15 @@ VectorParams, ) from rag_core.spi.vector_store import VectorStore -from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext +from rag_core.types import ( + ChunkId, + ChunkRef, + CorpusId, + Embedding, + IndexHint, + RequestContext, + TenantId, +) from rag_observability.logging import get_logger _log = get_logger(__name__) @@ -101,10 +109,18 @@ async def close(self) -> None: await self._client.close() # ------------------------------------------------------------------ - # VectorStore SPI + # VectorStore SPI (Step 1.1b split: bulk_index / retrieve_ids / bulk_delete) # ------------------------------------------------------------------ - async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: + async def bulk_index( + self, + ctx: RequestContext, + embeddings: list[Embedding], + *, + hint: IndexHint | None = None, + ) -> None: + # Qdrant collection-level HNSW parameters are configured at + # ``initialize()`` time; ``hint`` is reserved for that path (ADR-0009). if not embeddings: return @@ -124,14 +140,14 @@ async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None ] await self._client.upsert(collection_name=self._collection, points=points) - async def query( + async def retrieve_ids( self, ctx: RequestContext, vector: list[float], top_k: int, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, - ) -> list[tuple[ChunkId, float]]: + ) -> list[ChunkRef]: must: list[FieldCondition] = [ FieldCondition(key=_FIELD_TENANT, match=MatchValue(value=str(ctx.tenant_id))) ] @@ -151,13 +167,23 @@ async def query( with_payload=True, ) - return [ - (ChunkId(str(hit.payload[_FIELD_CHUNK])), float(hit.score)) - for hit in response.points - if hit.payload and _FIELD_CHUNK in hit.payload - ] + tenant = TenantId(ctx.tenant_id) + refs: list[ChunkRef] = [] + for hit in response.points: + if not hit.payload or _FIELD_CHUNK not in hit.payload: + continue + corpus_payload = hit.payload.get(_FIELD_CORPUS) or None + refs.append( + ChunkRef( + chunk_id=ChunkId(str(hit.payload[_FIELD_CHUNK])), + tenant_id=tenant, + score=float(hit.score), + corpus_id=CorpusId(str(corpus_payload)) if corpus_payload else None, + ) + ) + return refs - async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: + async def bulk_delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: if not chunk_ids: return diff --git a/packages/core/src/rag_core/__init__.py b/packages/core/src/rag_core/__init__.py index af7f042..1c0d0fc 100644 --- a/packages/core/src/rag_core/__init__.py +++ b/packages/core/src/rag_core/__init__.py @@ -53,7 +53,11 @@ Cache, Connector, Embedder, + GraphIndexBackend, + GraphRetrievalBackend, GraphStore, + KeywordIndexBackend, + KeywordRetrievalBackend, KeywordStore, LLMMessage, LLMResponse, @@ -67,6 +71,8 @@ Secrets, Storage, Telemetry, + VectorIndexBackend, + VectorRetrievalBackend, VectorStore, ) from rag_core.telemetry import ( @@ -95,6 +101,7 @@ DocumentStatus, Embedding, EmbeddingDtype, + IndexHint, PiiAction, PiiPolicy, PlanNode, @@ -112,6 +119,7 @@ TenantId, TraceContext, TrustLevel, + WriteVolume, ) __version__ = "0.5.0" @@ -165,6 +173,7 @@ "DocumentStatus", "Embedding", "EmbeddingDtype", + "IndexHint", "PiiAction", "PiiPolicy", "PlanNode", @@ -182,6 +191,7 @@ "TenantId", "TraceContext", "TrustLevel", + "WriteVolume", # errors "ACLDeniedError", "AuthError", @@ -211,7 +221,11 @@ "Cache", "Connector", "Embedder", + "GraphIndexBackend", + "GraphRetrievalBackend", "GraphStore", + "KeywordIndexBackend", + "KeywordRetrievalBackend", "KeywordStore", "LLM", "LLMMessage", @@ -227,5 +241,7 @@ "Secrets", "Storage", "Telemetry", + "VectorIndexBackend", + "VectorRetrievalBackend", "VectorStore", ] diff --git a/packages/core/src/rag_core/gen_schemas.py b/packages/core/src/rag_core/gen_schemas.py index 96227f5..d4112ca 100644 --- a/packages/core/src/rag_core/gen_schemas.py +++ b/packages/core/src/rag_core/gen_schemas.py @@ -26,6 +26,7 @@ Cost, Document, Embedding, + IndexHint, PiiPolicy, PlanNode, Principal, @@ -54,6 +55,7 @@ Citation, ChunkRef, Cost, + IndexHint, PlanNode, QueryPlan, StageEvent, diff --git a/packages/core/src/rag_core/spi/__init__.py b/packages/core/src/rag_core/spi/__init__.py index 320394c..fcad201 100644 --- a/packages/core/src/rag_core/spi/__init__.py +++ b/packages/core/src/rag_core/spi/__init__.py @@ -3,6 +3,11 @@ Import from this package to depend on an interface without coupling to any concrete implementation. Noop (in-memory) implementations live in ``rag_core.spi.noop`` and are suitable for tests and local development. + +Step 1.1b split the storage SPIs into narrow read/write roles +(``*RetrievalBackend`` / ``*IndexBackend``) while keeping the composite +``VectorStore`` / ``KeywordStore`` / ``GraphStore`` ABCs for backends that +serve both sides. """ from rag_core.spi.audit_store import AuditStore @@ -10,8 +15,16 @@ from rag_core.spi.cache import Cache from rag_core.spi.connector import Connector from rag_core.spi.embedder import Embedder -from rag_core.spi.graph_store import GraphStore -from rag_core.spi.keyword_store import KeywordStore +from rag_core.spi.graph_store import ( + GraphIndexBackend, + GraphRetrievalBackend, + GraphStore, +) +from rag_core.spi.keyword_store import ( + KeywordIndexBackend, + KeywordRetrievalBackend, + KeywordStore, +) from rag_core.spi.llm import LLM, LLMMessage, LLMResponse from rag_core.spi.ocr import OCR, OCRResult from rag_core.spi.parser import Parser @@ -21,7 +34,11 @@ from rag_core.spi.secrets import Secrets from rag_core.spi.storage import Storage from rag_core.spi.telemetry import Telemetry -from rag_core.spi.vector_store import VectorStore +from rag_core.spi.vector_store import ( + VectorIndexBackend, + VectorRetrievalBackend, + VectorStore, +) __all__ = [ "AuditStore", @@ -29,7 +46,11 @@ "Cache", "Connector", "Embedder", + "GraphIndexBackend", + "GraphRetrievalBackend", "GraphStore", + "KeywordIndexBackend", + "KeywordRetrievalBackend", "KeywordStore", "LLM", "LLMMessage", @@ -45,5 +66,7 @@ "Secrets", "Storage", "Telemetry", + "VectorIndexBackend", + "VectorRetrievalBackend", "VectorStore", ] diff --git a/packages/core/src/rag_core/spi/embedder.py b/packages/core/src/rag_core/spi/embedder.py index 1c44e95..57ae564 100644 --- a/packages/core/src/rag_core/spi/embedder.py +++ b/packages/core/src/rag_core/spi/embedder.py @@ -1,4 +1,11 @@ -"""Embedder SPI — text-to-vector embedding.""" +"""Embedder SPI — text-to-vector embedding. + +Step 1.1b split ``embed`` into a single-item ``embed`` plus a batch +``bulk_embed``. Both are part of the SPI surface to give callers explicit +control over batching at the call site (the agent loop wants single-item; +the ingest pipeline wants batched). See [docs/architecture/performance.md] +for the p99 budgets. +""" from __future__ import annotations @@ -27,7 +34,7 @@ def dimension(self) -> int: """Output vector dimension.""" @abc.abstractmethod - async def embed( + async def bulk_embed( self, ctx: RequestContext, texts: list[str], @@ -37,5 +44,25 @@ async def embed( ``texts`` and ``chunk_ids`` must have the same length. The returned list preserves input order. Implementations may split into sub-batches - for API limits but must return all results in one call. + to fit provider API limits but must return all results in one call. + + This is the canonical batched entry point — providers nearly always + accept batched requests, and the per-call overhead is the same for 1 + item or 100. + """ + + async def embed( + self, + ctx: RequestContext, + text: str, + chunk_id: ChunkId, + ) -> Embedding: + """Embed a single text. + + Default implementation calls :meth:`bulk_embed` with a list of one. + Backends with a single-item fast path (avoiding batch-coordination + overhead) may override. """ + + result = await self.bulk_embed(ctx, [text], [chunk_id]) + return result[0] diff --git a/packages/core/src/rag_core/spi/graph_store.py b/packages/core/src/rag_core/spi/graph_store.py index 2c6f342..ba6a2ad 100644 --- a/packages/core/src/rag_core/spi/graph_store.py +++ b/packages/core/src/rag_core/spi/graph_store.py @@ -1,21 +1,45 @@ -"""GraphStore SPI — knowledge-graph node/edge storage and traversal.""" +"""GraphStore SPI — knowledge-graph node/edge storage and traversal. + +Step 1.1b split this SPI into :class:`GraphRetrievalBackend` (read / +traversal) and :class:`GraphIndexBackend` (write). The legacy +:class:`GraphStore` composite ABC inherits both. + +Unlike vector and keyword stores, graph queries do not map cleanly onto a +single ``retrieve_ids → hydrate`` flow — a Cypher / Gremlin / SPARQL query +returns arbitrary rows — so this SPI does **not** force a +``retrieve_ids``/``hydrate`` pair. Phase 2 GraphRAG (Step 2.9) will layer a +chunk-shaped wrapper on top. +""" from __future__ import annotations import abc +from collections.abc import AsyncIterator from typing import Any from rag_core.spi._base import HealthCheckMixin -from rag_core.types import RequestContext +from rag_core.types import IndexHint, RequestContext -class GraphStore(HealthCheckMixin, abc.ABC): - """Abstract knowledge-graph backend (Neo4j, Kuzu, Amazon Neptune, …). +class GraphRetrievalBackend(HealthCheckMixin, abc.ABC): + """Read side of a knowledge-graph backend.""" - 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 query( + self, + ctx: RequestContext, + statement: str, + parameters: dict[str, Any], + ) -> list[dict[str, Any]]: + """Execute a graph query (Cypher, Gremlin, or SPARQL — backend-specific). + + Returns a list of row dicts. The caller is responsible for knowing the + query language accepted by the concrete implementation. + """ + + +class GraphIndexBackend(HealthCheckMixin, abc.ABC): + """Write side of a knowledge-graph backend.""" @abc.abstractmethod async def upsert_node( @@ -42,15 +66,73 @@ async def upsert_edge( 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( + async def bulk_upsert_nodes( self, ctx: RequestContext, - statement: str, - parameters: dict[str, Any], - ) -> list[dict[str, Any]]: - """Execute a graph query (Cypher, Gremlin, or SPARQL — backend-specific). + nodes: list[tuple[str, list[str], dict[str, Any]]], + *, + hint: IndexHint | None = None, + ) -> None: + """Upsert a batch of ``(node_id, labels, properties)`` tuples. - Returns a list of row dicts. The caller is responsible for knowing the - query language accepted by the concrete implementation. + Default implementation calls :meth:`upsert_node` per item; backends + with native UNWIND / batch APIs (Neo4j, Memgraph) should override. """ + + for node_id, labels, properties in nodes: + await self.upsert_node(ctx, node_id, labels, properties) + + async def bulk_upsert_edges( + self, + ctx: RequestContext, + edges: list[tuple[str, str, str, dict[str, Any]]], + *, + hint: IndexHint | None = None, + ) -> None: + """Upsert a batch of ``(from_id, to_id, rel_type, properties)`` tuples.""" + + for from_id, to_id, rel_type, properties in edges: + await self.upsert_edge(ctx, from_id, to_id, rel_type, properties) + + async def bulk_delete_nodes( + self, + ctx: RequestContext, + node_ids: list[str], + ) -> None: + """Delete a batch of nodes; unknown IDs are silently ignored.""" + + for node_id in node_ids: + await self.delete_node(ctx, node_id) + + async def stream_upsert_nodes( + self, + ctx: RequestContext, + nodes: AsyncIterator[tuple[str, list[str], dict[str, Any]]], + *, + hint: IndexHint | None = None, + batch_size: int = 256, + ) -> int: + """Consume an async stream of nodes and write them in batches. + + Returns the total number of nodes written. + """ + + total = 0 + batch: list[tuple[str, list[str], dict[str, Any]]] = [] + async for item in nodes: + batch.append(item) + if len(batch) >= batch_size: + await self.bulk_upsert_nodes(ctx, batch, hint=hint) + total += len(batch) + batch = [] + if batch: + await self.bulk_upsert_nodes(ctx, batch, hint=hint) + total += len(batch) + return total + + +class GraphStore(GraphRetrievalBackend, GraphIndexBackend, abc.ABC): + """Backend that handles both read and write of a knowledge graph. + + Composite ABC for the common full-service case (Neo4j, Memgraph, Neptune). + """ diff --git a/packages/core/src/rag_core/spi/keyword_store.py b/packages/core/src/rag_core/spi/keyword_store.py index c9b5459..508006b 100644 --- a/packages/core/src/rag_core/spi/keyword_store.py +++ b/packages/core/src/rag_core/spi/keyword_store.py @@ -1,36 +1,44 @@ -"""KeywordStore SPI — BM25 / full-text keyword search.""" +"""KeywordStore SPI — BM25 / full-text keyword search. + +Step 1.1b split this SPI into :class:`KeywordRetrievalBackend` (read) and +:class:`KeywordIndexBackend` (write). The legacy :class:`KeywordStore` is a +composite ABC inheriting both so backends that handle the full surface +(Elasticsearch, OpenSearch, Tantivy, Postgres FTS) keep a single class to +implement. +""" from __future__ import annotations import abc +from collections.abc import AsyncIterator from rag_core.spi._base import HealthCheckMixin -from rag_core.types import Chunk, ChunkId, CorpusId, RequestContext - +from rag_core.types import ( + Chunk, + ChunkId, + ChunkRef, + CorpusId, + IndexHint, + 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. Tenant isolation comes from ``ctx.tenant_id``. - """ +class KeywordRetrievalBackend(HealthCheckMixin, abc.ABC): + """Read side of a keyword (BM25-style) index.""" @abc.abstractmethod - 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. - """ - - @abc.abstractmethod - async def search( + async def retrieve_ids( self, ctx: RequestContext, query_text: str, top_k: int, corpus_ids: list[CorpusId], - ) -> list[tuple[ChunkId, float]]: - """Return ``(chunk_id, bm25_score)`` pairs ordered by descending score. + ) -> list[ChunkRef]: + """Return up to ``top_k`` ``ChunkRef`` ordered by descending BM25 score. + + ID-only retrieval keeps the hot path light. Callers may then call + :meth:`hydrate` on the survivors of reranking / policy checks to load + full chunk content. Args: ctx: Per-request envelope (tenant, principal, budget, trace). @@ -40,5 +48,75 @@ async def search( """ @abc.abstractmethod - async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: + async def hydrate( + self, + ctx: RequestContext, + chunk_refs: list[ChunkRef], + ) -> list[Chunk]: + """Load full :class:`Chunk` payloads for the given refs. + + Keyword stores typically retain chunk text (it is what BM25 indexes + over), so they can serve hydration directly. Order of the result must + match the order of ``chunk_refs``; missing chunks are skipped. + """ + + +class KeywordIndexBackend(HealthCheckMixin, abc.ABC): + """Write side of a keyword index.""" + + @abc.abstractmethod + async def bulk_index( + self, + ctx: RequestContext, + chunks: list[Chunk], + *, + hint: IndexHint | None = None, + ) -> None: + """Add or update a batch of chunks in the keyword index. + + Existing entries for the same ``chunk_id`` are overwritten. ``hint`` + is advisory and may be used by backends to size analyzer pools or + refresh intervals. + """ + + async def stream_index( + self, + ctx: RequestContext, + chunks: AsyncIterator[Chunk], + *, + hint: IndexHint | None = None, + batch_size: int = 256, + ) -> int: + """Consume an async stream of chunks and write them in batches. + + Default implementation buffers up to ``batch_size`` then forwards to + :meth:`bulk_index`. Returns the total number of chunks written. + """ + + total = 0 + batch: list[Chunk] = [] + async for chunk in chunks: + batch.append(chunk) + if len(batch) >= batch_size: + await self.bulk_index(ctx, batch, hint=hint) + total += len(batch) + batch = [] + if batch: + await self.bulk_index(ctx, batch, hint=hint) + total += len(batch) + return total + + @abc.abstractmethod + async def bulk_delete( + self, + ctx: RequestContext, + chunk_ids: list[ChunkId], + ) -> None: """Remove index entries for the given chunk IDs. Unknown IDs ignored.""" + + +class KeywordStore(KeywordRetrievalBackend, KeywordIndexBackend, abc.ABC): + """Backend that handles both read and write of a keyword index. + + Composite ABC for the common full-service case. + """ diff --git a/packages/core/src/rag_core/spi/noop/embedder.py b/packages/core/src/rag_core/spi/noop/embedder.py index 6debb14..46e0963 100644 --- a/packages/core/src/rag_core/spi/noop/embedder.py +++ b/packages/core/src/rag_core/spi/noop/embedder.py @@ -20,7 +20,7 @@ def model(self) -> str: def dimension(self) -> int: return self._dimension - async def embed( + async def bulk_embed( self, ctx: RequestContext, texts: list[str], 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 644e180..ba75627 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,15 @@ from __future__ import annotations from rag_core.spi.keyword_store import KeywordStore -from rag_core.types import Chunk, ChunkId, CorpusId, RequestContext +from rag_core.types import ( + Chunk, + ChunkId, + ChunkRef, + CorpusId, + IndexHint, + RequestContext, + TenantId, +) class NoopKeywordStore(KeywordStore): @@ -13,20 +21,26 @@ def __init__(self) -> None: # (tenant_id, chunk_id) -> Chunk self._index: dict[tuple[str, str], Chunk] = {} - async def index(self, ctx: RequestContext, chunks: list[Chunk]) -> None: + async def bulk_index( + self, + ctx: RequestContext, + chunks: list[Chunk], + *, + hint: IndexHint | None = None, + ) -> None: for chunk in chunks: self._index[(ctx.tenant_id, chunk.id)] = chunk - async def search( + async def retrieve_ids( self, ctx: RequestContext, query_text: str, top_k: int, corpus_ids: list[CorpusId], - ) -> list[tuple[ChunkId, float]]: + ) -> list[ChunkRef]: query_lower = query_text.lower() - results: list[tuple[ChunkId, float]] = [] - for (tid, cid), chunk in self._index.items(): + scored: list[tuple[float, Chunk]] = [] + for (tid, _cid), chunk in self._index.items(): if tid != ctx.tenant_id: continue if corpus_ids and chunk.corpus_id not in corpus_ids: @@ -38,12 +52,38 @@ async def search( tokens = query_lower.split() hits = sum(1 for t in tokens if t in content_lower) if hits: - score = hits / len(tokens) - results.append((ChunkId(cid), score)) - results.sort(key=lambda x: x[1], reverse=True) - return results[:top_k] + scored.append((hits / len(tokens), chunk)) + scored.sort(key=lambda x: x[0], reverse=True) + return [ + ChunkRef( + chunk_id=chunk.id, + tenant_id=TenantId(ctx.tenant_id), + score=score, + acl_labels=chunk.acl_labels, + corpus_id=chunk.corpus_id, + ) + for score, chunk in scored[:top_k] + ] + + async def hydrate( + self, + ctx: RequestContext, + chunk_refs: list[ChunkRef], + ) -> list[Chunk]: + chunks: list[Chunk] = [] + for ref in chunk_refs: + if ref.tenant_id != ctx.tenant_id: + continue + chunk = self._index.get((ctx.tenant_id, ref.chunk_id)) + if chunk is not None: + chunks.append(chunk) + return chunks - async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: + async def bulk_delete( + self, + ctx: RequestContext, + chunk_ids: list[ChunkId], + ) -> None: for cid in chunk_ids: self._index.pop((ctx.tenant_id, cid), None) 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 844fc8c..8320ad6 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,15 @@ from typing import Any from rag_core.spi.vector_store import VectorStore -from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext +from rag_core.types import ( + ChunkId, + ChunkRef, + CorpusId, + Embedding, + IndexHint, + RequestContext, + TenantId, +) def _cosine(a: list[float], b: list[float]) -> float: @@ -23,28 +31,46 @@ def __init__(self) -> None: # (tenant_id, chunk_id) -> Embedding self._store: dict[tuple[str, str], Embedding] = {} - async def upsert(self, ctx: RequestContext, embeddings: list[Embedding]) -> None: + async def bulk_index( + self, + ctx: RequestContext, + embeddings: list[Embedding], + *, + hint: IndexHint | None = None, + ) -> None: for emb in embeddings: self._store[(ctx.tenant_id, emb.chunk_id)] = emb - async def query( + async def retrieve_ids( self, ctx: RequestContext, vector: list[float], top_k: int, 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(): + ) -> list[ChunkRef]: + scored: list[tuple[float, Embedding]] = [] + for (tid, _cid), emb in self._store.items(): 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] + scored.append((score, emb)) + scored.sort(key=lambda x: x[0], reverse=True) + return [ + ChunkRef( + chunk_id=emb.chunk_id, + tenant_id=TenantId(ctx.tenant_id), + score=score, + acl_labels=emb.acl_labels, + ) + for score, emb in scored[:top_k] + ] - async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: + async def bulk_delete( + self, + ctx: RequestContext, + chunk_ids: list[ChunkId], + ) -> None: for cid in chunk_ids: self._store.pop((ctx.tenant_id, cid), None) diff --git a/packages/core/src/rag_core/spi/vector_store.py b/packages/core/src/rag_core/spi/vector_store.py index 885e92c..727ccb8 100644 --- a/packages/core/src/rag_core/spi/vector_store.py +++ b/packages/core/src/rag_core/spi/vector_store.py @@ -1,42 +1,57 @@ -"""VectorStore SPI — dense-vector storage and ANN search.""" +"""VectorStore SPI — dense-vector storage and ANN search. + +Step 1.1b split the SPI into two narrow ABCs: + +* :class:`VectorRetrievalBackend` — read path: ``retrieve_ids`` + ``hydrate``. +* :class:`VectorIndexBackend` — write path: ``bulk_index`` / + ``stream_index`` / ``bulk_delete``. + +The legacy :class:`VectorStore` ABC inherits both, so backends that handle the +full read+write surface (PgVector, Qdrant, Weaviate, …) keep a single class to +implement. Pure-read consumers (gateway query path) and pure-write consumers +(ingest pipeline) should depend on the narrower role for clarity. + +See [docs/adr/ADR-0009-vector-index-strategy.md] for the index-selection +strategy and [docs/architecture/performance.md] for p99 budgets. +""" from __future__ import annotations import abc +from collections.abc import AsyncIterator from typing import Any from rag_core.spi._base import HealthCheckMixin -from rag_core.types import ChunkId, CorpusId, Embedding, RequestContext - +from rag_core.types import ( + ChunkId, + ChunkRef, + CorpusId, + Embedding, + IndexHint, + RequestContext, +) -class VectorStore(HealthCheckMixin, abc.ABC): - """Abstract store for dense embeddings with ANN (approximate nearest-neighbour) search. - Tenant isolation is enforced via ``ctx.tenant_id`` on every call. See - docs/architecture/request-context.md. +class VectorRetrievalBackend(HealthCheckMixin, abc.ABC): + """Read side of a dense-vector backend. - Step 1.1b will split this into ``RetrievalBackend`` (read) + - ``IndexBackend`` (write) and introduce ID-only retrieval. + All methods take ``ctx: RequestContext`` first; ``ctx.tenant_id`` namespaces + every operation (see [docs/architecture/request-context.md]). """ @abc.abstractmethod - 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( + async def retrieve_ids( self, ctx: RequestContext, vector: list[float], top_k: int, corpus_ids: list[CorpusId], filters: dict[str, Any] | None = None, - ) -> list[tuple[ChunkId, float]]: - """Return ``(chunk_id, score)`` pairs ordered by descending similarity. + ) -> list[ChunkRef]: + """Return up to ``top_k`` ``ChunkRef`` ordered by descending similarity. + + ID-only retrieval keeps the hot path light: the caller reranks / + policy-checks the refs, then calls :meth:`hydrate` on the survivors. Args: ctx: Per-request envelope (tenant, principal, budget, trace). @@ -44,8 +59,94 @@ async def query( top_k: Maximum number of results to return. corpus_ids: Restrict results to these corpora (empty = all corpora). filters: Optional metadata key-value equality filters. + + Returns: + List of :class:`ChunkRef` with ``chunk_id``, ``score``, ``tenant_id``, + ``corpus_id`` and any cheap-to-include metadata (acl_labels) needed + for downstream filtering. Content is **not** hydrated. + """ + + async def hydrate( + self, + ctx: RequestContext, + chunk_refs: list[ChunkRef], + ) -> list[ChunkRef]: + """Optionally enrich ``chunk_refs`` with backend-side metadata. + + Vector backends typically do not store chunk content (that lives in a + :class:`KeywordRetrievalBackend` or via :class:`BlobRef` in storage), so + the default implementation is a pass-through. Subclasses that *do* + carry side-data in their payload may override to attach it. + + Callers wanting full :class:`Chunk` content should instead call + :meth:`KeywordRetrievalBackend.hydrate` or load via the Storage SPI. """ + return list(chunk_refs) + + +class VectorIndexBackend(HealthCheckMixin, abc.ABC): + """Write side of a dense-vector backend.""" + @abc.abstractmethod - async def delete(self, ctx: RequestContext, chunk_ids: list[ChunkId]) -> None: + async def bulk_index( + self, + ctx: RequestContext, + embeddings: list[Embedding], + *, + hint: IndexHint | None = None, + ) -> None: + """Insert or overwrite a batch of embeddings. + + If an embedding for a given ``chunk_id`` already exists it is replaced. + ``ctx.tenant_id`` namespaces the data. Backends may consult ``hint`` + to choose / tune the index implementation (see ADR-0009). + """ + + async def stream_index( + self, + ctx: RequestContext, + embeddings: AsyncIterator[Embedding], + *, + hint: IndexHint | None = None, + batch_size: int = 256, + ) -> int: + """Consume an async stream of embeddings and write them in batches. + + Default implementation buffers up to ``batch_size`` then forwards to + :meth:`bulk_index`. Backends with native streaming bulk APIs (e.g. + Qdrant's gRPC stream) may override. + + Returns: + Total number of embeddings written. + """ + + total = 0 + batch: list[Embedding] = [] + async for emb in embeddings: + batch.append(emb) + if len(batch) >= batch_size: + await self.bulk_index(ctx, batch, hint=hint) + total += len(batch) + batch = [] + if batch: + await self.bulk_index(ctx, batch, hint=hint) + total += len(batch) + return total + + @abc.abstractmethod + async def bulk_delete( + self, + ctx: RequestContext, + chunk_ids: list[ChunkId], + ) -> None: """Remove embeddings by chunk ID. Unknown IDs are silently ignored.""" + + +class VectorStore(VectorRetrievalBackend, VectorIndexBackend, abc.ABC): + """Backend that handles both read and write of dense vectors. + + Composite ABC for the common case where a single backend (PgVector, Qdrant, + Weaviate, …) serves both sides. Inherit from the narrower roles directly + when a backend is read-only or write-only. + """ diff --git a/packages/core/src/rag_core/types.py b/packages/core/src/rag_core/types.py index 41b5477..00a8bc2 100644 --- a/packages/core/src/rag_core/types.py +++ b/packages/core/src/rag_core/types.py @@ -103,6 +103,23 @@ class EmbeddingDtype(StrEnum): binary = "binary" +class WriteVolume(StrEnum): + """Expected write traffic class for an index, used by ``IndexHint``. + + Backends use this to balance build-time cost against query latency: + + ``low`` — mostly read; rebuild on schedule is fine. + ``medium`` — incremental updates expected; favour incremental indexes. + ``high`` — heavy concurrent writes; favour lock-free / append structures. + + See ADR-0009. + """ + + low = "low" + medium = "medium" + high = "high" + + class PiiAction(StrEnum): """Per-tenant PII handling policy.""" @@ -471,6 +488,33 @@ class ChunkRef(BaseModel): metadata: dict[str, Any] = Field(default_factory=dict) +# --------------------------------------------------------------------------- +# IndexHint — scale-tier hint passed to IndexBackend writes +# --------------------------------------------------------------------------- +class IndexHint(BaseModel): + """Hint passed by callers to an ``IndexBackend`` at write / initialise time. + + Backends pick the index implementation (flat / ivfflat / HNSW / IVF-PQ / + DiskANN, etc.) based on these fields. See ADR-0009 for the default mapping + per backend. Callers may omit the hint — backends fall back to a sensible + default for the current corpus size. + """ + + model_config = {"frozen": True} + + estimated_size: int = 0 + """Current + projected vector / chunk count for this corpus.""" + + recall_target: float = 0.95 + """Recall floor in ``[0.0, 1.0]``. Higher → exact-er index, slower queries.""" + + latency_target_ms: float = 50.0 + """p99 query latency target the backend should optimise for.""" + + write_volume: WriteVolume = WriteVolume.low + """Expected write traffic — see :class:`WriteVolume`.""" + + # --------------------------------------------------------------------------- # QueryPlan — output of the planner, input to retrieval # --------------------------------------------------------------------------- diff --git a/tests/contract/spi_signature.py b/tests/contract/spi_signature.py index ae9cea1..04ed826 100644 --- a/tests/contract/spi_signature.py +++ b/tests/contract/spi_signature.py @@ -36,7 +36,11 @@ Cache, Connector, Embedder, + GraphIndexBackend, + GraphRetrievalBackend, GraphStore, + KeywordIndexBackend, + KeywordRetrievalBackend, KeywordStore, Parser, PIIDetector, @@ -44,6 +48,8 @@ Reranker, Secrets, Storage, + VectorIndexBackend, + VectorRetrievalBackend, VectorStore, ) from rag_core.types import RequestContext @@ -57,7 +63,11 @@ Cache: {"health"}, Connector: {"health"}, Embedder: {"health", "model", "dimension"}, + GraphRetrievalBackend: {"health"}, + GraphIndexBackend: {"health"}, GraphStore: {"health"}, + KeywordRetrievalBackend: {"health"}, + KeywordIndexBackend: {"health"}, KeywordStore: {"health"}, LLM: {"health"}, OCR: {"health"}, @@ -67,6 +77,8 @@ Reranker: {"health"}, Secrets: {"health"}, Storage: {"health"}, + VectorRetrievalBackend: {"health"}, + VectorIndexBackend: {"health"}, VectorStore: {"health"}, } @@ -76,7 +88,11 @@ def _ctx_threaded_classes() -> list[type]: Cache, Connector, Embedder, + GraphRetrievalBackend, + GraphIndexBackend, GraphStore, + KeywordRetrievalBackend, + KeywordIndexBackend, KeywordStore, LLM, OCR, @@ -86,6 +102,8 @@ def _ctx_threaded_classes() -> list[type]: Reranker, Secrets, Storage, + VectorRetrievalBackend, + VectorIndexBackend, VectorStore, ] diff --git a/tests/contract/test_embedder.py b/tests/contract/test_embedder.py index 8558a7f..b2613b6 100644 --- a/tests/contract/test_embedder.py +++ b/tests/contract/test_embedder.py @@ -1,8 +1,8 @@ -"""Conformance tests for Embedder SPI.""" +"""Conformance tests for Embedder SPI (Step 1.1b: bulk_embed + single embed).""" import pytest from rag_core.spi.noop import NoopEmbedder -from rag_core.types import ChunkId, RequestContext +from rag_core.types import ChunkId, Embedding, RequestContext pytestmark = pytest.mark.contract @@ -11,35 +11,48 @@ async def test_health(embedder: NoopEmbedder) -> None: assert await embedder.health() is True -async def test_embed_returns_correct_count(embedder: NoopEmbedder, ctx: RequestContext) -> None: +async def test_bulk_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(ctx, texts, ids) + results = await embedder.bulk_embed(ctx, texts, ids) assert len(results) == 3 -async def test_embed_dimension_matches(embedder: NoopEmbedder, ctx: RequestContext) -> None: - results = await embedder.embed(ctx, ["test"], [ChunkId("c1")]) +async def test_bulk_embed_dimension_matches(embedder: NoopEmbedder, ctx: RequestContext) -> None: + results = await embedder.bulk_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, ctx: RequestContext) -> None: +async def test_bulk_embed_preserves_chunk_ids(embedder: NoopEmbedder, ctx: RequestContext) -> None: ids = [ChunkId("alpha"), ChunkId("beta")] - results = await embedder.embed(ctx, ["a", "b"], ids) + results = await embedder.bulk_embed(ctx, ["a", "b"], ids) assert [r.chunk_id for r in results] == ids -async def test_embed_sets_tenant_id(embedder: NoopEmbedder, ctx: RequestContext) -> None: - results = await embedder.embed(ctx, ["x"], [ChunkId("c1")]) +async def test_bulk_embed_sets_tenant_id(embedder: NoopEmbedder, ctx: RequestContext) -> None: + results = await embedder.bulk_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")]) +async def test_bulk_embed_model_name(embedder: NoopEmbedder, ctx: RequestContext) -> None: + results = await embedder.bulk_embed(ctx, ["x"], [ChunkId("c1")]) assert results[0].model == embedder.model -async def test_embed_mismatched_lengths_raises(embedder: NoopEmbedder, ctx: RequestContext) -> None: +async def test_bulk_embed_mismatched_lengths_raises( + embedder: NoopEmbedder, ctx: RequestContext +) -> None: with pytest.raises(ValueError): - await embedder.embed(ctx, ["a", "b"], [ChunkId("c1")]) + await embedder.bulk_embed(ctx, ["a", "b"], [ChunkId("c1")]) + + +async def test_embed_single(embedder: NoopEmbedder, ctx: RequestContext) -> None: + """Default single-item path delegates to bulk_embed and unwraps.""" + result = await embedder.embed(ctx, "single text", ChunkId("c-single")) + assert isinstance(result, Embedding) + assert result.chunk_id == ChunkId("c-single") + assert result.dimension == embedder.dimension + assert len(result.vector) == embedder.dimension diff --git a/tests/contract/test_graph_store.py b/tests/contract/test_graph_store.py index 774c2c8..7e345c6 100644 --- a/tests/contract/test_graph_store.py +++ b/tests/contract/test_graph_store.py @@ -41,3 +41,35 @@ async def test_delete_removes_edges(graph_store: NoopGraphStore, ctx: RequestCon 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) + + +async def test_bulk_upsert_nodes(graph_store: NoopGraphStore, ctx: RequestContext) -> None: + await graph_store.bulk_upsert_nodes( + ctx, + [ + ("a", ["X"], {"k": 1}), + ("b", ["X"], {"k": 2}), + ("c", ["Y"], {"k": 3}), + ], + ) + rows = await graph_store.query(ctx, "", {}) + assert {r["node_id"] for r in rows} == {"a", "b", "c"} + + +async def test_bulk_delete_nodes(graph_store: NoopGraphStore, ctx: RequestContext) -> None: + for n in ("a", "b", "c"): + await graph_store.upsert_node(ctx, n, [], {}) + await graph_store.bulk_delete_nodes(ctx, ["a", "b"]) + rows = await graph_store.query(ctx, "", {}) + assert {r["node_id"] for r in rows} == {"c"} + + +async def test_stream_upsert_nodes(graph_store: NoopGraphStore, ctx: RequestContext) -> None: + async def gen(): + for i in range(5): + yield (f"n{i}", [], {}) + + written = await graph_store.stream_upsert_nodes(ctx, gen(), batch_size=2) + assert written == 5 + rows = await graph_store.query(ctx, "", {}) + assert {r["node_id"] for r in rows} == {f"n{i}" for i in range(5)} diff --git a/tests/contract/test_keyword_store.py b/tests/contract/test_keyword_store.py index a23e4b4..2884347 100644 --- a/tests/contract/test_keyword_store.py +++ b/tests/contract/test_keyword_store.py @@ -1,8 +1,8 @@ -"""Conformance tests for KeywordStore SPI.""" +"""Conformance tests for KeywordStore SPI (Step 1.1b split).""" import pytest from rag_core.spi.noop import NoopKeywordStore -from rag_core.types import ChunkId, CorpusId, RequestContext +from rag_core.types import ChunkId, ChunkRef, CorpusId, RequestContext from tests.contract.conftest import make_chunk, make_ctx @@ -13,49 +13,72 @@ async def test_health(keyword_store: NoopKeywordStore) -> None: assert await keyword_store.health() is True -async def test_index_and_search( +async def test_bulk_index_and_retrieve_ids( keyword_store: NoopKeywordStore, ctx: RequestContext, corpus_id: CorpusId ) -> None: 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) + await keyword_store.bulk_index(ctx, [chunk]) + refs = await keyword_store.retrieve_ids(ctx, "quick fox", top_k=5, corpus_ids=[corpus_id]) + assert any(r.chunk_id == ChunkId("c1") for r in refs) + assert all(isinstance(r, ChunkRef) for r in refs) -async def test_search_tenant_isolation(keyword_store: NoopKeywordStore) -> None: +async def test_retrieve_ids_tenant_isolation(keyword_store: NoopKeywordStore) -> None: 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(ctx_t1, [chunk]) - results = await keyword_store.search(ctx_t2, "secret", top_k=5, corpus_ids=[]) - assert results == [] + await keyword_store.bulk_index(ctx_t1, [chunk]) + refs = await keyword_store.retrieve_ids(ctx_t2, "secret", top_k=5, corpus_ids=[]) + assert refs == [] -async def test_delete(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: +async def test_bulk_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] + await keyword_store.bulk_index(ctx, chunks) + await keyword_store.bulk_delete(ctx, [ChunkId("c1")]) + refs = await keyword_store.retrieve_ids(ctx, "apple", top_k=5, corpus_ids=[]) + ids = [r.chunk_id for r in refs] assert ChunkId("c1") not in ids async def test_no_match_returns_empty(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: - await keyword_store.index( + await keyword_store.bulk_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 == [] + refs = await keyword_store.retrieve_ids(ctx, "zebra", top_k=5, corpus_ids=[]) + assert refs == [] async def test_scores_in_range(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: - await keyword_store.index( + await keyword_store.bulk_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 + refs = await keyword_store.retrieve_ids(ctx, "cat mat", top_k=5, corpus_ids=[]) + for ref in refs: + assert 0.0 <= ref.score <= 1.0 + + +async def test_hydrate(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: + chunk = make_chunk("c1", content="hello world", tenant_id=str(ctx.tenant_id)) + await keyword_store.bulk_index(ctx, [chunk]) + refs = await keyword_store.retrieve_ids(ctx, "hello", top_k=5, corpus_ids=[]) + assert refs + hydrated = await keyword_store.hydrate(ctx, refs) + assert len(hydrated) == 1 + assert hydrated[0].id == "c1" + assert hydrated[0].content == "hello world" + + +async def test_stream_index(keyword_store: NoopKeywordStore, ctx: RequestContext) -> None: + async def gen(): + for i in range(5): + yield make_chunk(f"s{i}", content=f"text {i}", tenant_id=str(ctx.tenant_id)) + + written = await keyword_store.stream_index(ctx, gen(), batch_size=2) + assert written == 5 + refs = await keyword_store.retrieve_ids(ctx, "text", top_k=10, corpus_ids=[]) + assert {r.chunk_id for r in refs} == {ChunkId(f"s{i}") for i in range(5)} diff --git a/tests/contract/test_vector_store.py b/tests/contract/test_vector_store.py index 17c7c2e..110dc42 100644 --- a/tests/contract/test_vector_store.py +++ b/tests/contract/test_vector_store.py @@ -1,8 +1,8 @@ -"""Conformance tests for VectorStore SPI.""" +"""Conformance tests for VectorStore SPI (Step 1.1b split).""" import pytest from rag_core.spi.noop import NoopVectorStore -from rag_core.types import ChunkId, RequestContext +from rag_core.types import ChunkId, ChunkRef, RequestContext, TenantId from tests.contract.conftest import make_ctx, make_embedding @@ -13,53 +13,57 @@ async def test_health(vector_store: NoopVectorStore) -> None: assert await vector_store.health() is True -async def test_upsert_and_query(vector_store: NoopVectorStore, ctx: RequestContext) -> None: +async def test_bulk_index_and_retrieve_ids( + 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") - assert score == pytest.approx(1.0, abs=1e-6) + await vector_store.bulk_index(ctx, [emb]) + refs = await vector_store.retrieve_ids(ctx, emb.vector, top_k=5, corpus_ids=[]) + assert len(refs) == 1 + ref = refs[0] + assert isinstance(ref, ChunkRef) + assert ref.chunk_id == ChunkId("c1") + assert ref.tenant_id == ctx.tenant_id + assert ref.score == pytest.approx(1.0, abs=1e-6) -async def test_query_tenant_isolation(vector_store: NoopVectorStore) -> None: +async def test_retrieve_ids_tenant_isolation(vector_store: NoopVectorStore) -> None: 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 == [] + await vector_store.bulk_index(ctx_t1, [make_embedding("c1", tenant_id="t1")]) + refs = await vector_store.retrieve_ids(ctx_t2, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + assert refs == [] -async def test_delete(vector_store: NoopVectorStore, ctx: RequestContext) -> None: - await vector_store.upsert( +async def test_bulk_delete(vector_store: NoopVectorStore, ctx: RequestContext) -> None: + await vector_store.bulk_index( 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] + await vector_store.bulk_delete(ctx, [ChunkId("c1")]) + refs = await vector_store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + ids = [r.chunk_id for r in refs] assert ChunkId("c1") not in ids assert ChunkId("c2") in ids -async def test_delete_unknown_id_is_noop( +async def test_bulk_delete_unknown_id_is_noop( vector_store: NoopVectorStore, ctx: RequestContext ) -> None: - await vector_store.delete(ctx, [ChunkId("nonexistent")]) # must not raise + await vector_store.bulk_delete(ctx, [ChunkId("nonexistent")]) # must not raise async def test_top_k_respected(vector_store: NoopVectorStore, ctx: RequestContext) -> None: for i in range(10): - 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 + await vector_store.bulk_index(ctx, [make_embedding(f"c{i}", tenant_id=str(ctx.tenant_id))]) + refs = await vector_store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=3, corpus_ids=[]) + assert len(refs) <= 3 -async def test_upsert_overwrites(vector_store: NoopVectorStore, ctx: RequestContext) -> None: +async def test_bulk_index_overwrites(vector_store: NoopVectorStore, ctx: RequestContext) -> None: from rag_core.types import Embedding emb1 = make_embedding("c1", tenant_id=str(ctx.tenant_id)) @@ -70,7 +74,25 @@ async def test_upsert_overwrites(vector_store: NoopVectorStore, ctx: RequestCont 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 + await vector_store.bulk_index(ctx, [emb1]) + await vector_store.bulk_index(ctx, [emb2]) + refs = await vector_store.retrieve_ids(ctx, [0.9, 0.9, 0.9, 0.9], top_k=1, corpus_ids=[]) + assert len(refs) == 1 + + +async def test_stream_index(vector_store: NoopVectorStore, ctx: RequestContext) -> None: + async def gen(): + for i in range(7): + yield make_embedding(f"s{i}", tenant_id=str(ctx.tenant_id)) + + written = await vector_store.stream_index(ctx, gen(), batch_size=3) + assert written == 7 + refs = await vector_store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=100, corpus_ids=[]) + assert {r.chunk_id for r in refs} == {ChunkId(f"s{i}") for i in range(7)} + + +async def test_hydrate_passthrough(vector_store: NoopVectorStore, ctx: RequestContext) -> None: + # Vector backends do not own chunk content; default hydrate is a no-op echo. + refs = [ChunkRef(chunk_id=ChunkId("c1"), tenant_id=TenantId(str(ctx.tenant_id)), score=0.5)] + out = await vector_store.hydrate(ctx, refs) + assert out == refs diff --git a/tests/integration/test_pgvector.py b/tests/integration/test_pgvector.py index 52b0ed1..241b1fc 100644 --- a/tests/integration/test_pgvector.py +++ b/tests/integration/test_pgvector.py @@ -1,4 +1,4 @@ -"""Integration tests for PgVectorStore. +"""Integration tests for PgVectorStore (Step 1.1b SPI split). Requires: PostgreSQL with pgvector extension (task dev). Skipped automatically when Postgres is not reachable. @@ -46,58 +46,74 @@ async def test_health(store: PgVectorStore) -> None: assert await store.health() is True -async def test_upsert_and_query(store: PgVectorStore) -> None: +async def test_bulk_index_and_retrieve_ids(store: PgVectorStore) -> None: 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] + await store.bulk_index(ctx, [_emb("c1", ctx.tenant_id)]) + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + assert len(refs) >= 1 + ids = [r.chunk_id for r in refs] assert ChunkId("c1") in ids # Score for identical vectors should be ~1.0 - score = next(s for cid, s in results if cid == ChunkId("c1")) + score = next(r.score for r in refs if r.chunk_id == ChunkId("c1")) assert score == pytest.approx(1.0, abs=1e-4) + # tenant_id is populated on the returned ref + assert all(r.tenant_id == ctx.tenant_id for r in refs) async def test_tenant_isolation(store: PgVectorStore) -> None: 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] + await store.bulk_index(ctx_t1, [_emb("c-iso", ctx_t1.tenant_id)]) + refs = await store.retrieve_ids(ctx_t2, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + ids = [r.chunk_id for r in refs] assert ChunkId("c-iso") not in ids -async def test_delete(store: PgVectorStore) -> None: +async def test_bulk_delete(store: PgVectorStore) -> None: ctx = make_ctx(tenant_id="t-pg-del") - await store.upsert( + await store.bulk_index( 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] + await store.bulk_delete(ctx, [ChunkId("c-del-1")]) + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + ids = [r.chunk_id for r in refs] 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: +async def test_bulk_delete_unknown_is_noop(store: PgVectorStore) -> None: ctx = make_ctx(tenant_id="t-pg-ghost") - await store.delete(ctx, [ChunkId("ghost")]) + await store.bulk_delete(ctx, [ChunkId("ghost")]) -async def test_upsert_overwrites(store: PgVectorStore) -> None: +async def test_bulk_index_overwrites(store: PgVectorStore) -> None: 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) + await store.bulk_index(ctx, [_emb("c-ow", ctx.tenant_id, [0.1, 0.2, 0.3, 0.4])]) + await store.bulk_index(ctx, [_emb("c-ow", ctx.tenant_id, [0.9, 0.9, 0.9, 0.9])]) + refs = await store.retrieve_ids(ctx, [0.9, 0.9, 0.9, 0.9], top_k=1, corpus_ids=[]) + assert len(refs) == 1 + assert refs[0].chunk_id == ChunkId("c-ow") + assert refs[0].score == pytest.approx(1.0, abs=1e-4) async def test_top_k_respected(store: PgVectorStore) -> None: ctx = make_ctx(tenant_id="t-pg-topk") for i in range(8): - 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 + await store.bulk_index(ctx, [_emb(f"c-topk-{i}", ctx.tenant_id)]) + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=3, corpus_ids=[]) + assert len(refs) <= 3 + + +async def test_stream_index(store: PgVectorStore) -> None: + ctx = make_ctx(tenant_id="t-pg-stream") + + async def gen(): + for i in range(7): + yield _emb(f"c-stream-{i}", ctx.tenant_id) + + written = await store.stream_index(ctx, gen(), batch_size=3) + assert written == 7 + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[]) + ids = {r.chunk_id for r in refs} + assert {ChunkId(f"c-stream-{i}") for i in range(7)} <= ids diff --git a/tests/integration/test_qdrant.py b/tests/integration/test_qdrant.py index dac306f..eceb6d8 100644 --- a/tests/integration/test_qdrant.py +++ b/tests/integration/test_qdrant.py @@ -1,4 +1,4 @@ -"""Integration tests for QdrantVectorStore. +"""Integration tests for QdrantVectorStore (Step 1.1b SPI split). Requires: Qdrant service (task dev). Skipped automatically when Qdrant is not reachable. @@ -41,57 +41,72 @@ async def test_health(store: QdrantVectorStore) -> None: assert await store.health() is True -async def test_upsert_and_query(store: QdrantVectorStore) -> None: +async def test_bulk_index_and_retrieve_ids(store: QdrantVectorStore) -> None: 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] + await store.bulk_index(ctx, [_emb("c1", ctx.tenant_id)]) + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + assert len(refs) >= 1 + ids = [r.chunk_id for r in refs] assert ChunkId("c1") in ids - score = next(s for cid, s in results if cid == ChunkId("c1")) + score = next(r.score for r in refs if r.chunk_id == ChunkId("c1")) assert score == pytest.approx(1.0, abs=1e-4) + assert all(r.tenant_id == ctx.tenant_id for r in refs) async def test_tenant_isolation(store: QdrantVectorStore) -> None: 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] + await store.bulk_index(ctx_t1, [_emb("c-qd-iso", ctx_t1.tenant_id)]) + refs = await store.retrieve_ids(ctx_t2, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + ids = [r.chunk_id for r in refs] assert ChunkId("c-qd-iso") not in ids -async def test_delete(store: QdrantVectorStore) -> None: +async def test_bulk_delete(store: QdrantVectorStore) -> None: ctx = make_ctx(tenant_id="t-qd-del") - await store.upsert( + await store.bulk_index( 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] + await store.bulk_delete(ctx, [ChunkId("c-qd-del-1")]) + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=5, corpus_ids=[]) + ids = [r.chunk_id for r in refs] 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: +async def test_bulk_delete_unknown_is_noop(store: QdrantVectorStore) -> None: ctx = make_ctx(tenant_id="t-qd-ghost") - await store.delete(ctx, [ChunkId("ghost")]) + await store.bulk_delete(ctx, [ChunkId("ghost")]) -async def test_upsert_overwrites(store: QdrantVectorStore) -> None: +async def test_bulk_index_overwrites(store: QdrantVectorStore) -> None: 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) + await store.bulk_index(ctx, [_emb("c-qd-ow", ctx.tenant_id, [0.1, 0.2, 0.3, 0.4])]) + await store.bulk_index(ctx, [_emb("c-qd-ow", ctx.tenant_id, [0.9, 0.9, 0.9, 0.9])]) + refs = await store.retrieve_ids(ctx, [0.9, 0.9, 0.9, 0.9], top_k=1, corpus_ids=[]) + assert len(refs) == 1 + assert refs[0].chunk_id == ChunkId("c-qd-ow") + assert refs[0].score == pytest.approx(1.0, abs=1e-4) async def test_top_k_respected(store: QdrantVectorStore) -> None: ctx = make_ctx(tenant_id="t-qd-topk") for i in range(8): - 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 + await store.bulk_index(ctx, [_emb(f"c-qd-topk-{i}", ctx.tenant_id)]) + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=3, corpus_ids=[]) + assert len(refs) <= 3 + + +async def test_stream_index(store: QdrantVectorStore) -> None: + ctx = make_ctx(tenant_id="t-qd-stream") + + async def gen(): + for i in range(7): + yield _emb(f"c-qd-stream-{i}", ctx.tenant_id) + + written = await store.stream_index(ctx, gen(), batch_size=3) + assert written == 7 + refs = await store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[]) + ids = {r.chunk_id for r in refs} + assert {ChunkId(f"c-qd-stream-{i}") for i in range(7)} <= ids diff --git a/uv.lock b/uv.lock index b5287eb..840bd21 100644 --- a/uv.lock +++ b/uv.lock @@ -3246,7 +3246,7 @@ provides-extras = ["eval"] [[package]] name = "rag-core" -version = "0.2.0" +version = "0.3.0" source = { editable = "packages/core" } dependencies = [ { name = "opentelemetry-api" },