diff --git a/dist/schemas/Embedding.json b/dist/schemas/Embedding.json index b32a934..c2e464a 100644 --- a/dist/schemas/Embedding.json +++ b/dist/schemas/Embedding.json @@ -11,7 +11,7 @@ "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.", + "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.\n\n``corpus_id`` carries the owning corpus so vector backends can honour the\n``corpus_ids`` filter on ``retrieve_ids`` (ADR-0004 \u00a73 \u2014 the embedder SPI\nhas no corpus, so the ingest pipeline stamps it from the source ``Chunk``).\nOptional because not every embedding has a corpus: query-side embeddings\n(HyDE) carry synthetic chunk ids and no corpus. A ``None`` corpus never\nmatches a non-empty ``corpus_ids`` filter \u2014 the safe (no-leak) direction.", "properties": { "chunk_id": { "title": "Chunk Id", @@ -21,6 +21,18 @@ "title": "Tenant Id", "type": "string" }, + "corpus_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Corpus Id" + }, "model": { "title": "Model", "type": "string" diff --git a/docs/adr/ADR-0004-storage-backends.md b/docs/adr/ADR-0004-storage-backends.md index ffbb5d5..d40920d 100644 --- a/docs/adr/ADR-0004-storage-backends.md +++ b/docs/adr/ADR-0004-storage-backends.md @@ -42,13 +42,18 @@ Both have an `initialize()` method (idempotent DDL / collection creation). ### 3. Corpus ID gap in Embedding -`rag_core.types.Embedding` does not carry `corpus_id`. The VectorStore SPI -takes `corpus_ids` as a filter parameter, but backends store `corpus_id = ""` -for now and skip corpus filtering when the list is empty. - -**Resolution:** Phase 1.8 (embedder pipeline) will extend `Embedding` with -`corpus_id`. Backends will be updated at that point. The gap is tracked in -[TRACKER.md](../../TRACKER.md). +`rag_core.types.Embedding` did not carry `corpus_id`. The VectorStore SPI +takes `corpus_ids` as a filter parameter, but backends stored `corpus_id = ""` +and could not filter by corpus — `NoopVectorStore` leaked chunks across corpora +in the same tenant, and the real backends (which filtered on the always-empty +value) returned nothing for a corpus-scoped query. + +**Resolution (closed):** `Embedding` now has an optional `corpus_id`. Because +the embedder SPI (`bulk_embed(ctx, texts, chunk_ids)`) has no corpus, the ingest +pipeline stamps each embedding from its source `Chunk` before indexing. All +VectorStore backends persist `corpus_id` and honour the `corpus_ids` filter, and +each returned `ChunkRef` carries its `corpus_id`. A shared contract test in +`tests/contract/test_vector_store.py` guards corpus isolation. ### 4. Cache backend: RedisCache @@ -79,7 +84,8 @@ against the local stack. - **Positive:** Real persistence available from Phase 1.1; all tests self-skip without services; SPI boundary unchanged. -- **Negative:** `corpus_id` filtering in vector stores is a stub until Phase - 1.8 — callers must pass `corpus_ids=[]` or accept unfiltered results. +- **Negative:** ~~`corpus_id` filtering in vector stores is a stub until Phase + 1.8~~ — *resolved (see §3): `Embedding.corpus_id` is stamped by the ingest + pipeline and every backend filters on it.* - **Neutral:** numpy added as a dependency for pgvector's asyncpg codec; acceptable given it is already ubiquitous in the ML stack. diff --git a/docs/reference/backends.md b/docs/reference/backends.md index ce8b34d..3f0405a 100644 --- a/docs/reference/backends.md +++ b/docs/reference/backends.md @@ -492,11 +492,14 @@ vector = embedding.vector payload = {tenant_id, chunk_id, corpus_id, model} ``` -### Corpus ID gap - -`rag_core.types.Embedding` does not carry `corpus_id`. Both VectorStore -implementations store `corpus_id = ""` and skip corpus filtering until Phase 1.8 -extends the `Embedding` type. See ADR-0004 for details. +### Corpus scoping + +`rag_core.types.Embedding` carries an optional `corpus_id` (the ingest pipeline +stamps it from the source `Chunk`, since the embedder SPI has no corpus). Every +VectorStore persists it (`""` when unset) and honours the `corpus_ids` argument +to `retrieve_ids`: a non-empty list restricts results to those corpora, an empty +list means all corpora. Each returned `ChunkRef` carries its `corpus_id`. This +closes the gap originally deferred in ADR-0004 §3. --- diff --git a/packages/backends/src/rag_backends/vector/elasticsearch.py b/packages/backends/src/rag_backends/vector/elasticsearch.py index 4b94fe5..7f05010 100644 --- a/packages/backends/src/rag_backends/vector/elasticsearch.py +++ b/packages/backends/src/rag_backends/vector/elasticsearch.py @@ -201,7 +201,7 @@ async def bulk_index( { _FIELD_TENANT: tenant_id, _FIELD_CHUNK: str(emb.chunk_id), - _FIELD_CORPUS: "", + _FIELD_CORPUS: str(emb.corpus_id) if emb.corpus_id else "", # ADR-0004 §3 _FIELD_MODEL: emb.model, _FIELD_ACL: sorted(emb.acl_labels), _FIELD_VECTOR: list(emb.vector), diff --git a/packages/backends/src/rag_backends/vector/pgvector.py b/packages/backends/src/rag_backends/vector/pgvector.py index fee9de8..8f8fd1e 100644 --- a/packages/backends/src/rag_backends/vector/pgvector.py +++ b/packages/backends/src/rag_backends/vector/pgvector.py @@ -237,7 +237,7 @@ async def bulk_index( ( str(ctx.tenant_id), str(emb.chunk_id), - "", # corpus_id not available in Embedding — see ADR-0004 + str(emb.corpus_id) if emb.corpus_id else "", # ADR-0004 §3 emb.model, _to_np(emb.vector), emb.dimension, diff --git a/packages/backends/src/rag_backends/vector/pinecone.py b/packages/backends/src/rag_backends/vector/pinecone.py index 134c531..f76d9ba 100644 --- a/packages/backends/src/rag_backends/vector/pinecone.py +++ b/packages/backends/src/rag_backends/vector/pinecone.py @@ -203,7 +203,7 @@ async def bulk_index( "metadata": { _META_TENANT: tenant_id, _META_CHUNK: str(emb.chunk_id), - _META_CORPUS: "", + _META_CORPUS: str(emb.corpus_id) if emb.corpus_id else "", # ADR-0004 §3 _META_MODEL: emb.model, _META_ACL: sorted(emb.acl_labels), }, diff --git a/packages/backends/src/rag_backends/vector/qdrant.py b/packages/backends/src/rag_backends/vector/qdrant.py index 66ad71f..c81e2db 100644 --- a/packages/backends/src/rag_backends/vector/qdrant.py +++ b/packages/backends/src/rag_backends/vector/qdrant.py @@ -228,7 +228,7 @@ async def bulk_index( payload={ _FIELD_TENANT: tenant_id, _FIELD_CHUNK: str(emb.chunk_id), - _FIELD_CORPUS: "", # corpus_id not in Embedding — see ADR-0004 + _FIELD_CORPUS: str(emb.corpus_id) if emb.corpus_id else "", # ADR-0004 §3 _FIELD_MODEL: emb.model, # Step 2.1: persist ACL labels so FilterExpr push-down works. _FIELD_ACL: sorted(emb.acl_labels), diff --git a/packages/backends/src/rag_backends/vector/weaviate.py b/packages/backends/src/rag_backends/vector/weaviate.py index 194db7d..29ef125 100644 --- a/packages/backends/src/rag_backends/vector/weaviate.py +++ b/packages/backends/src/rag_backends/vector/weaviate.py @@ -208,7 +208,7 @@ async def bulk_index( properties={ _PROP_TENANT: tenant_id, _PROP_CHUNK: str(emb.chunk_id), - _PROP_CORPUS: "", + _PROP_CORPUS: str(emb.corpus_id) if emb.corpus_id else "", # ADR-0004 §3 _PROP_MODEL: emb.model, _PROP_ACL: sorted(emb.acl_labels), }, 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 536cdc6..faa23a8 100644 --- a/packages/core/src/rag_core/spi/noop/vector_store.py +++ b/packages/core/src/rag_core/spi/noop/vector_store.py @@ -58,6 +58,7 @@ def _embedding_attrs(emb: Embedding) -> dict[str, Any]: return { "tenant_id": str(emb.tenant_id), "chunk_id": str(emb.chunk_id), + "corpus_id": str(emb.corpus_id) if emb.corpus_id is not None else None, "acl_labels": emb.acl_labels, "model": emb.model, } @@ -83,11 +84,18 @@ async def retrieve_ids( ) -> list[ChunkRef]: index = self._index_key(ctx) scored: list[tuple[float, Embedding]] = [] - for (idx, tid, _cid), emb in self._store.items(): + for (idx, tid, _chunk_id), emb in self._store.items(): # Physical isolation (Step 6.2) first, then the tenant filter — a read # never scans another partition, independent of the tenant filter. if idx != index or tid != ctx.tenant_id: continue + # Corpus scoping — mirror NoopKeywordStore. A non-empty ``corpus_ids`` + # restricts to those corpora; an embedding with no corpus_id never + # matches (the no-leak direction). ``corpus_id`` lives on the + # Embedding (ADR-0004 §3), not the store key (whose third element is + # the chunk id). + if corpus_ids and emb.corpus_id not in corpus_ids: + continue # Filter push-down (Step 2.1) — reference semantics. if filters is not None and not evaluate(filters, self._embedding_attrs(emb)): continue @@ -100,6 +108,7 @@ async def retrieve_ids( tenant_id=TenantId(ctx.tenant_id), score=score, acl_labels=emb.acl_labels, + corpus_id=emb.corpus_id, ) for score, emb in scored[:top_k] ] diff --git a/packages/core/src/rag_core/types.py b/packages/core/src/rag_core/types.py index 81a1fcd..e25e2b4 100644 --- a/packages/core/src/rag_core/types.py +++ b/packages/core/src/rag_core/types.py @@ -729,12 +729,20 @@ class Embedding(BaseModel): Step 1.1a promoted ``tenant_id`` and ``acl_labels`` to typed required fields and added ``dtype`` to capture int8 / binary quantization. + + ``corpus_id`` carries the owning corpus so vector backends can honour the + ``corpus_ids`` filter on ``retrieve_ids`` (ADR-0004 §3 — the embedder SPI + has no corpus, so the ingest pipeline stamps it from the source ``Chunk``). + Optional because not every embedding has a corpus: query-side embeddings + (HyDE) carry synthetic chunk ids and no corpus. A ``None`` corpus never + matches a non-empty ``corpus_ids`` filter — the safe (no-leak) direction. """ model_config = {"frozen": True} chunk_id: ChunkId tenant_id: TenantId + corpus_id: CorpusId | None = None model: str vector: list[float] dimension: int diff --git a/packages/ingest/src/rag_ingest/pipeline.py b/packages/ingest/src/rag_ingest/pipeline.py index 4b7d949..7b6b6d1 100644 --- a/packages/ingest/src/rag_ingest/pipeline.py +++ b/packages/ingest/src/rag_ingest/pipeline.py @@ -42,6 +42,7 @@ from rag_core.types import ( Chunk, Document, + Embedding, IngestResult, IngestRunSummary, IngestStatus, @@ -311,12 +312,20 @@ async def _embed( backends: _Backends, ctx: RequestContext, chunks: list[Chunk], -) -> list[Any]: +) -> list[Embedding]: if not chunks: return [] texts = [c.content or "" for c in chunks] chunk_ids = [c.id for c in chunks] - return await backends.embedder.bulk_embed(ctx, texts, chunk_ids) + embeddings = await backends.embedder.bulk_embed(ctx, texts, chunk_ids) + # The embedder SPI carries no corpus, so stamp each embedding with its + # source chunk's corpus_id — vector backends need it to honour the + # corpus_ids filter on retrieve_ids (ADR-0004 §3). + corpus_by_chunk = {c.id: c.corpus_id for c in chunks} + return [ + emb.model_copy(update={"corpus_id": corpus_by_chunk.get(emb.chunk_id)}) + for emb in embeddings + ] def _summarise(results: list[IngestResult]) -> IngestRunSummary: diff --git a/packages/ragctl/src/ragctl/main.py b/packages/ragctl/src/ragctl/main.py index 9338a17..d827335 100644 --- a/packages/ragctl/src/ragctl/main.py +++ b/packages/ragctl/src/ragctl/main.py @@ -1102,6 +1102,7 @@ async def _run() -> None: Embedding( chunk_id=c.id, tenant_id=tenant_id, + corpus_id=c.corpus_id, model="noop", vector=[1.0 if j == i else 0.0 for j in range(len(documents))], dimension=len(documents), @@ -1875,6 +1876,7 @@ async def _run() -> None: Embedding( chunk_id=c.id, tenant_id=tenant_id, + corpus_id=c.corpus_id, model="noop", vector=[1.0 if j == i else 0.0 for j in range(len(documents))], dimension=len(documents), @@ -2130,6 +2132,7 @@ async def _run() -> None: Embedding( chunk_id=c.id, tenant_id=tenant_id, + corpus_id=c.corpus_id, model="noop", vector=[1.0 if j == i else 0.0 for j in range(len(documents))], dimension=len(documents), @@ -3340,6 +3343,7 @@ async def _run() -> None: Embedding( chunk_id=c.id, tenant_id=tenant_id, + corpus_id=c.corpus_id, model="noop", vector=[1.0 if j == i else 0.0 for j in range(len(documents))], dimension=len(documents), diff --git a/proto/core.proto b/proto/core.proto index 550c32b..6ce4477 100644 --- a/proto/core.proto +++ b/proto/core.proto @@ -157,6 +157,10 @@ message Embedding { repeated float vector = 3; int32 dimension = 4; google.protobuf.Timestamp created_at = 5; + // Owning corpus, so vector backends can honour the corpus_ids filter on + // retrieve_ids (ADR-0004 §3). Empty string when unset (query-side / HyDE + // embeddings carry no corpus). + string corpus_id = 6; } // --------------------------------------------------------------------------- diff --git a/tests/contract/conftest.py b/tests/contract/conftest.py index f598393..f87ebc8 100644 --- a/tests/contract/conftest.py +++ b/tests/contract/conftest.py @@ -224,10 +224,12 @@ def make_embedding( chunk_id: str = "chunk-1", dim: int = 4, tenant_id: str = "tenant-test", + corpus_id: str | None = None, ) -> Embedding: return Embedding( chunk_id=ChunkId(chunk_id), tenant_id=TenantId(tenant_id), + corpus_id=CorpusId(corpus_id) if corpus_id is not None else None, model="noop-embedder", vector=[0.1, 0.2, 0.3, 0.4][:dim], dimension=dim, diff --git a/tests/contract/test_vector_store.py b/tests/contract/test_vector_store.py index f9c5775..52b9128 100644 --- a/tests/contract/test_vector_store.py +++ b/tests/contract/test_vector_store.py @@ -2,7 +2,7 @@ import pytest from rag_core.spi.noop import NoopVectorStore -from rag_core.types import ChunkId, ChunkRef, RequestContext, TenantId +from rag_core.types import ChunkId, ChunkRef, CorpusId, RequestContext, TenantId from tests.contract.conftest import make_ctx, make_embedding @@ -223,3 +223,70 @@ async def test_retrieve_ids_filter_true_is_noop( ) refs_none = await vector_store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[]) assert {r.chunk_id for r in refs_true} == {r.chunk_id for r in refs_none} + + +# --------------------------------------------------------------------------- +# Corpus scoping — ADR-0004 §3. A retrieval scoped to one corpus must NOT +# return chunks from other corpora in the same tenant (the bug that leaked +# cross-corpus chunks via the vector arm). +# --------------------------------------------------------------------------- + + +async def test_retrieve_ids_corpus_isolation( + vector_store: NoopVectorStore, ctx: RequestContext +) -> None: + """A non-empty ``corpus_ids`` restricts results to those corpora.""" + tenant = str(ctx.tenant_id) + await vector_store.bulk_index( + ctx, + [ + make_embedding("a", tenant_id=tenant, corpus_id="x"), + make_embedding("b", tenant_id=tenant, corpus_id="y"), + ], + ) + refs = await vector_store.retrieve_ids( + ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[CorpusId("x")] + ) + assert [r.chunk_id for r in refs] == [ChunkId("a")] + + +async def test_retrieve_ids_surfaces_corpus_id( + vector_store: NoopVectorStore, ctx: RequestContext +) -> None: + """Each ``ChunkRef`` carries the embedding's ``corpus_id`` (SPI contract).""" + await vector_store.bulk_index( + ctx, [make_embedding("a", tenant_id=str(ctx.tenant_id), corpus_id="x")] + ) + refs = await vector_store.retrieve_ids( + ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[CorpusId("x")] + ) + assert [r.corpus_id for r in refs] == [CorpusId("x")] + + +async def test_retrieve_ids_empty_corpus_ids_returns_all( + vector_store: NoopVectorStore, ctx: RequestContext +) -> None: + """Empty ``corpus_ids`` means no corpus filter — every corpus is returned.""" + tenant = str(ctx.tenant_id) + await vector_store.bulk_index( + ctx, + [ + make_embedding("a", tenant_id=tenant, corpus_id="x"), + make_embedding("b", tenant_id=tenant, corpus_id="y"), + ], + ) + refs = await vector_store.retrieve_ids(ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[]) + assert {r.chunk_id for r in refs} == {ChunkId("a"), ChunkId("b")} + + +async def test_retrieve_ids_none_corpus_never_matches_scoped_query( + vector_store: NoopVectorStore, ctx: RequestContext +) -> None: + """An embedding with no corpus_id must not leak into a corpus-scoped query.""" + await vector_store.bulk_index( + ctx, [make_embedding("a", tenant_id=str(ctx.tenant_id), corpus_id=None)] + ) + refs = await vector_store.retrieve_ids( + ctx, [0.1, 0.2, 0.3, 0.4], top_k=10, corpus_ids=[CorpusId("x")] + ) + assert refs == [] diff --git a/tests/ingest/test_pipeline_document.py b/tests/ingest/test_pipeline_document.py index 00a2ae4..d31d504 100644 --- a/tests/ingest/test_pipeline_document.py +++ b/tests/ingest/test_pipeline_document.py @@ -7,6 +7,7 @@ from rag_chunker import HeadingAwareChunker from rag_core.spi.noop import NoopEmbedder, NoopVectorStore from rag_core.types import ( + CorpusId, IngestResult, IngestStatus, RequestContext, @@ -125,6 +126,40 @@ async def test_embeddings_actually_indexed( assert len(vector_store._store) == result.embedding_count +async def test_indexed_embeddings_are_corpus_scoped( + ctx: RequestContext, + make_document, + token_counter: Any, +) -> None: + """End-to-end: the pipeline stamps each embedding with its document's + corpus, so a corpus-scoped vector retrieval does not leak other corpora + (ADR-0004 §3). The embedder SPI has no corpus — this proves the pipeline + supplies it.""" + vector_store = NoopVectorStore() + pipeline = IngestPipeline( + chunker=HeadingAwareChunker(max_tokens=256, overlap_tokens=0, token_counter=token_counter), + enricher=DefaultEnricher(), + embedder=NoopEmbedder(dimension=8), + index_backend=vector_store, + parser_registry=default_registry().select, + ) + doc, content = make_document(content=b"Hello world. " * 20, mime_type="text/plain") + result = await pipeline.ingest_document(ctx, doc, content) + assert result.status == IngestStatus.ingested + + query_vec = [0.0] * 8 + hits = await vector_store.retrieve_ids( + ctx, query_vec, top_k=100, corpus_ids=[doc.corpus_id] + ) + assert len(hits) == result.embedding_count + assert all(h.corpus_id == doc.corpus_id for h in hits) + + leaked = await vector_store.retrieve_ids( + ctx, query_vec, top_k=100, corpus_ids=[CorpusId("some-other-corpus")] + ) + assert leaked == [] + + async def test_blocked_by_policy_short_circuits_pipeline( ctx: RequestContext, make_document,