railtracks.retrieval

25@dataclass
26class BatchIngested:
27    """A batch of chunks that finished embedding and was written to the store.
28
29    ``batch_index`` is **per-document**: it starts at 0 for each document and
30    counts that document's batches (both successful and failed) in order. It
31    is not a run-global counter — to track overall progress, count events or
32    read ``IngestionStats``.
33
34    ``metrics`` carries the per-batch usage and timing reported by the
35    embedder (tokens, dollar cost, latency, vector count). Use it to track
36    per-ingest cost without having to wrap the embedder.
37    """
38
39    document_id: UUID
40    embedded_chunks: list[EmbeddedChunk]
41    batch_index: int
42    metrics: EmbeddingMetrics | None = None

74@dataclass
75class Chunk:
76    content: str
77    document_id: UUID
78    id: UUID = field(default_factory=uuid4)
79    index: int = 0
80    parent_chunk_id: UUID | None = None
81    offsets: tuple[int, int] | None = None
82    metadata: dict[str, Any] = field(default_factory=dict)

24@dataclass
25class Document:
26    """A unit of source content produced by a loader.
27
28    Attributes:
29        content: The decoded textual content of the document. Always a string;
30            binary loaders are responsible for their own decoding.
31        type: A :class:`DocumentType` describing the content format. Cloud
32            loaders infer it from the object's file extension; structured
33            loaders (CSV, SQL) set it explicitly.
34        id: Unique identifier. If not provided and ``source`` is set, derived
35            deterministically from ``source`` via UUID5 (RFC 4122 URL
36            namespace) so the same source yields the same id across processes
37            — required for the runtime's upsert (``delete_where`` on
38            ``document_id``) to find and clear the prior version when content
39            changes. Sourceless documents get a random UUID4 (no stable
40            identity → no upsert semantics).
41        source: The natural identifier of where this document came from —
42            a URI (``s3://bucket/key``, ``gs://bucket/name``, ``https://...``),
43            a file path, or a relational id. Cloud loaders always set this;
44            user-constructed documents may leave it ``None``.
45
46            Writers that derive a storage key (when no ``key_fn`` is supplied)
47            look here first; the cloud writers also strip their own URI prefix
48            so that "load from S3, write back to S3" produces a clean key
49            rather than a nested URI.
50        content_hash: SHA-256 of ``content``. Computed by the runtime at
51            ingest time; loaders should leave this ``None``. Used by
52            staleness-detection to skip re-embedding unchanged documents.
53        metadata: Arbitrary provider-specific or user-attached key-value data.
54            Loaders use this to expose details like ``bucket``, ``key``,
55            ``page``, ``row_index``, etc.
56    """
57
58    content: str
59    type: DocumentType = DocumentType.TEXT
60    id: UUID = _UNSET_DOCUMENT_ID
61    source: str | None = None
62    content_hash: str | None = None
63    metadata: dict[str, Any] = field(default_factory=dict)
64
65    def __post_init__(self) -> None:
66        if self.id == _UNSET_DOCUMENT_ID:
67            self.id = (
68                uuid5(NAMESPACE_URL, self.source)
69                if self.source is not None
70                else uuid4()
71            )

45@dataclass
46class DocumentFailed:
47    """A document that had at least one failed embedding batch.
48
49    Note: successful batches for the same document *are* written to the
50    store. ``DocumentFailed`` is an informational signal that the document
51    is now partial — callers may want to retry, delete, or accept the
52    partial state.
53    """
54
55    document_id: UUID
56    source: str | None
57    errors: list[Exception]

60@dataclass
61class DocumentSkipped:
62    """A document skipped during ingest because the store already has an
63    entry with the same ``source_path`` and ``content_hash``."""
64
65    document_id: UUID
66    source: str | None
67    reason: str = "unchanged"

15class DocumentType(str, Enum):
16    TEXT = "text"
17    MARKDOWN = "markdown"
18    PDF = "pdf"
19    CSV = "csv"
20    JSON = "json"
21    JSONL = "jsonl"

85@dataclass
86class EmbeddedChunk:
87    chunk: Chunk
88    vector: list[float]
89    embedding_model: str
90    embedding_version: str | None = None

72@dataclass
73class EmbeddingFailure:
74    """A failed batch embedding attempt.
75
76    Attributes:
77        chunks: Source chunks that could not be embedded.
78        errors: Exceptions raised during embedding.
79    """
80
81    chunks: list[Chunk]
82    errors: list[Exception]

 7class EmbeddingModelMismatchError(RuntimeError):
 8    """Raised when the runtime's embedder model differs from the store's.
 9
10    Mixing vectors from different embedding models silently produces
11    meaningless similarity scores, so the runtime fails loudly before
12    issuing the search.
13    """

70@dataclass
71class IngestionStats:
72    """Summary of a complete ingest run.
73
74    ``total_metrics`` accumulates per-batch ``EmbeddingMetrics`` (tokens,
75    dollar cost, latency, vector count) across every successful batch in
76    the run, so callers can read a single total for billing/observability.
77    """
78
79    documents_loaded: int = 0
80    documents_failed: int = 0
81    documents_skipped: int = 0
82    chunks_created: int = 0
83    chunks_embedded: int = 0
84    batches_failed: int = 0
85    batch_failures: list[EmbeddingFailure] = field(default_factory=list)
86    failed_documents: list[DocumentFailed] = field(default_factory=list)
87    total_metrics: EmbeddingMetrics = field(default_factory=EmbeddingMetrics)

102@dataclass
103class RetrievalResult:
104    query: str
105    chunks: list[RetrievedChunk]
106    total_candidates: int | None = None
107    metadata: dict[str, Any] = field(default_factory=dict)

106class RetrievalRuntime:
107    """Orchestrates loading, chunking, embedding, storage, and retrieval.
108
109    The runtime captures *how* to process documents (chunker + embedder +
110    store); the loader passed to :meth:`ingest` decides *what* to
111    process. A single runtime can ingest from multiple sources, mix
112    chunking strategies via separate runtimes against the same store,
113    and update existing documents by re-ingesting them. Multi-tenant
114    callers share one runtime and pass ``scope`` per :meth:`ingest` or
115    :meth:`retrieve` call.
116
117    Args:
118        chunker: Splits documents into chunks.
119        embedder: Embeds chunk text into vectors.
120        store: Receives written ``StoreEntry``s and serves similarity search.
121        batch_size: Items per embedding batch. Falls back to
122            ``embedder.default_batch_size`` when omitted; raises
123            ``ValueError`` at construction if neither is set.
124        on_ingest: Synchronous callback invoked with each ``IngestionEvent``
125            as it is yielded. Wrap in ``asyncio.create_task`` for async logging.
126        on_retrieve: Synchronous callback invoked with the query string and
127            the ``RetrievalResult`` after each retrieve call.
128        max_tokens: When set, chunks whose token count exceeds this limit
129            are dropped before embedding and reported via
130            ``EmbeddingFailure`` rather than being sent to the provider.
131            Requires ``tokenizer`` (defaults to ``TiktokenTokenizer``).
132        tokenizer: Tokenizer used to enforce ``max_tokens``. Defaults to
133            ``TiktokenTokenizer`` lazily when ``max_tokens`` is set.
134    """
135
136    def __init__(
137        self,
138        chunker: Chunker,
139        embedder: Embedding,
140        store: Store,
141        *,
142        batch_size: int | None = None,
143        on_ingest: Callable[
144            [BatchIngested | EmbeddingFailure | DocumentFailed | DocumentSkipped], None
145        ]
146        | None = None,
147        on_retrieve: Callable[[str, RetrievalResult], None] | None = None,
148        max_tokens: int | None = None,
149        tokenizer: Tokenizer | None = None,
150    ) -> None:
151        self._chunker = chunker
152        self._embedder = embedder
153        self._store = store
154        self._batch_size = self._resolve_batch_size(batch_size, embedder)
155        self._on_ingest = on_ingest
156        self._on_retrieve = on_retrieve
157        self._max_tokens = max_tokens
158        if max_tokens is not None and tokenizer is None:
159            from .chunking.tokenization import TiktokenTokenizer
160
161            tokenizer = TiktokenTokenizer()
162        self._tokenizer = tokenizer
163        # Captured on the first successful embedded batch and checked at
164        # retrieve time; survives process restarts by lazy-seeding from
165        # an existing store entry on the first ingest/retrieve.
166        self._captured_model: str | None = None
167        self._seed_attempted: bool = False
168
169    @property
170    def store(self) -> Store:
171        return self._store
172
173    @property
174    def embedder(self) -> Embedding:
175        return self._embedder
176
177    @property
178    def chunker(self) -> Chunker:
179        return self._chunker
180
181    @property
182    def batch_size(self) -> int:
183        return self._batch_size
184
185    @property
186    def max_tokens(self) -> int | None:
187        return self._max_tokens
188
189    @staticmethod
190    def _resolve_batch_size(batch_size: int | None, embedder: Embedding) -> int:
191        bs = batch_size if batch_size is not None else embedder.default_batch_size
192        if bs is None:
193            raise ValueError(
194                f"{type(embedder).__name__} does not declare a "
195                "default_batch_size. Pass batch_size= to RetrievalRuntime "
196                "or set default_batch_size on the embedder class."
197            )
198        return bs
199
200    async def ingest(
201        self,
202        loader: BaseDocumentLoader,
203        *,
204        scope: StoreScope | None = None,
205    ) -> AsyncGenerator[
206        BatchIngested | EmbeddingFailure | DocumentFailed | DocumentSkipped, None
207    ]:
208        """Stream loader → chunker → embedder → store, yielding per-batch events.
209
210        Args:
211            loader: Source of ``Document`` objects to ingest.
212            scope: Tag written onto every ``StoreEntry`` produced by this
213                call. Single-tenant callers can leave this ``None``.
214
215        Yields:
216            ``BatchIngested`` after each successful batch finishes writing,
217            ``EmbeddingFailure`` for any failed batch, and ``DocumentFailed``
218            once at end-of-document for each document that had any failed
219            batch. Successful batches for a partially-failed document are
220            still written; ``DocumentFailed`` signals the partial state.
221        """
222        stats = IngestionStats()
223        async for event in self._ingest_with_stats(loader, stats, scope):
224            if self._on_ingest is not None:
225                self._on_ingest(event)
226            yield event
227
228    async def ingest_all(
229        self,
230        loader: BaseDocumentLoader,
231        *,
232        scope: StoreScope | None = None,
233    ) -> IngestionStats:
234        """Drain `ingest` and return aggregate counts."""
235        stats = IngestionStats()
236        async for event in self._ingest_with_stats(loader, stats, scope):
237            if self._on_ingest is not None:
238                self._on_ingest(event)
239        return stats
240
241    async def _ingest_with_stats(
242        self,
243        loader: BaseDocumentLoader,
244        stats: IngestionStats,
245        scope: StoreScope | None,
246    ) -> AsyncGenerator[
247        BatchIngested | EmbeddingFailure | DocumentFailed | DocumentSkipped, None
248    ]:
249        async for doc in loader.astream():
250            async for event in self._ingest_document(doc, stats, scope):
251                yield event
252
253    async def _ingest_document(
254        self, doc: Document, stats: IngestionStats, scope: StoreScope | None
255    ) -> AsyncGenerator[
256        BatchIngested | EmbeddingFailure | DocumentFailed | DocumentSkipped, None
257    ]:
258        stats.documents_loaded += 1
259        doc.content_hash = _content_hash(doc.content)
260
261        await self._ensure_captured_model_seeded()
262
263        if await self._is_complete_duplicate(doc):
264            stats.documents_skipped += 1
265            yield DocumentSkipped(document_id=doc.id, source=doc.source)
266            return
267
268        chunks = await self._chunker.achunk(doc)
269        stats.chunks_created += len(chunks)
270        if not chunks:
271            return
272
273        self._stamp_staleness_metadata(doc, chunks)
274
275        # Token-size guard: drop oversized chunks before embedding to avoid
276        # provider 4xx errors. Each oversize chunk surfaces as an
277        # EmbeddingFailure carried into the document's accumulated errors.
278        doc_errors: list[Exception] = []
279        chunks, failures = self._split_oversized(chunks, stats)
280        for failure in failures:
281            doc_errors.extend(failure.errors)
282            yield failure
283        if not chunks:
284            if doc_errors:
285                yield self._record_document_failed(doc, doc_errors, stats)
286            return
287
288        # Stamp the final (post-token-guard) chunk count onto every chunk so a
289        # later staleness check can tell a complete document from a
290        # partially-written one. Every chunk carries the same total, so reading
291        # any one persisted chunk reveals how many were expected.
292        for chunk in chunks:
293            chunk.metadata["doc_chunk_count"] = len(chunks)
294
295        async for event in self._embed_and_store(doc, chunks, stats, doc_errors, scope):
296            yield event
297
298        if doc_errors:
299            yield self._record_document_failed(doc, doc_errors, stats)
300
301    async def _is_complete_duplicate(self, doc: Document) -> bool:
302        """Whether the store already holds a *complete* copy of ``doc``.
303
304        Skip re-embedding only when as many chunks are present as the last
305        write expected. A partially-written document (some chunks present after
306        an interrupted ingest) has fewer than expected and is re-ingested rather
307        than left broken. find() is metadata-only (no vector search) and only
308        fetches a single entry; the presence check is a count() so no payloads
309        are transferred and large documents never require a single oversized
310        read (Chroma Cloud caps the per-get ``limit`` value).
311        """
312        if doc.source is None:
313            return False
314        stale_filters = {
315            "source_path": doc.source,
316            "content_hash": doc.content_hash,
317        }
318        existing = await self._store.find(stale_filters, limit=1)
319        if not existing:
320            return False
321        expected = existing[0].chunk_metadata.get("doc_chunk_count")
322        if expected is None:
323            # Not written by this runtime (every ingest stamps the count):
324            # completeness can't be verified, so re-ingest.
325            return False
326        return await self._store.count(stale_filters) >= expected
327
328    @staticmethod
329    def _stamp_staleness_metadata(doc: Document, chunks: list[Chunk]) -> None:
330        """Inject staleness-detection metadata into every chunk so future
331        `find` calls can identify whether this document has changed."""
332        for chunk in chunks:
333            if doc.source is not None:
334                chunk.metadata.setdefault("source_path", doc.source)
335            if doc.content_hash is not None:
336                chunk.metadata.setdefault("content_hash", doc.content_hash)
337
338    def _split_oversized(
339        self, chunks: list[Chunk], stats: IngestionStats
340    ) -> tuple[list[Chunk], list[EmbeddingFailure]]:
341        """Partition chunks into embeddable ones and per-chunk failures.
342
343        Returns ``(ok_chunks, failures)``; each oversize chunk becomes a
344        single-chunk ``EmbeddingFailure`` and is recorded in ``stats``.
345        """
346        if self._max_tokens is None or self._tokenizer is None:
347            return chunks, []
348        ok_chunks: list[Chunk] = []
349        failures: list[EmbeddingFailure] = []
350        for chunk in chunks:
351            tokens = self._tokenizer.count(chunk.content)
352            if tokens > self._max_tokens:
353                err = ValueError(
354                    f"chunk {chunk.id} has {tokens} tokens "
355                    f"(>{self._max_tokens}); dropped before embedding"
356                )
357                stats.batches_failed += 1
358                failure = EmbeddingFailure(chunks=[chunk], errors=[err])
359                stats.batch_failures.append(failure)
360                failures.append(failure)
361            else:
362                ok_chunks.append(chunk)
363        return ok_chunks, failures
364
365    async def _embed_and_store(
366        self,
367        doc: Document,
368        chunks: list[Chunk],
369        stats: IngestionStats,
370        doc_errors: list[Exception],
371        scope: StoreScope | None,
372    ) -> AsyncGenerator[
373        BatchIngested | EmbeddingFailure | DocumentFailed | DocumentSkipped, None
374    ]:
375        # batch_index is per-document: it counts batches (successful and
376        # failed) within this document and resets for the next one.
377        batch_index = 0
378        delete_done = False
379        async for batch in self._embedder.astream_batches(
380            chunks, batch_size=self._batch_size
381        ):
382            if isinstance(batch, EmbeddingResult):
383                # Check model BEFORE delete_where / write — a mismatch here
384                # must not corrupt the store by clearing prior chunks first.
385                self._check_model(batch.metrics.model)
386                if not delete_done:
387                    await self._store.delete_where({"document_id": str(doc.id)})
388                    delete_done = True
389                for embedded in batch.chunks:
390                    self._capture_model(embedded)
391                    entry = StoreEntry.from_chunk(embedded, scope=scope)
392                    await self._store.write(entry)
393                stats.chunks_embedded += len(batch.chunks)
394                stats.total_metrics = stats.total_metrics + batch.metrics
395                yield BatchIngested(
396                    document_id=doc.id,
397                    embedded_chunks=batch.chunks,
398                    batch_index=batch_index,
399                    metrics=batch.metrics,
400                )
401            else:
402                doc_errors.extend(batch.errors)
403                stats.batches_failed += 1
404                stats.batch_failures.append(batch)
405                yield batch
406            batch_index += 1
407
408    def _capture_model(self, embedded: EmbeddedChunk) -> None:
409        """Record the embedding model from the first successful chunk so later
410        retrieve() calls can enforce model consistency."""
411        if self._captured_model is None and embedded.embedding_model:
412            self._captured_model = embedded.embedding_model
413            logger.info(
414                "RetrievalRuntime captured embedding model %r "
415                "from first successful batch; subsequent retrieve() "
416                "calls will enforce this model.",
417                self._captured_model,
418            )
419
420    async def _ensure_captured_model_seeded(self) -> None:
421        """Lazily seed ``_captured_model`` from an existing store entry so the
422        guard survives across process restarts. ``StoreEntry.embedding_model``
423        is recorded on every persisted entry, so a single ``find`` call is
424        enough — no schema change required. Runs at most once per runtime;
425        a miss against an empty store sets ``_seed_attempted`` so we don't
426        re-query on every doc."""
427        if self._captured_model is not None or self._seed_attempted:
428            return
429        self._seed_attempted = True
430        existing = await self._store.find({}, limit=1)
431        if existing and existing[0].embedding_model:
432            self._captured_model = existing[0].embedding_model
433            logger.info(
434                "RetrievalRuntime seeded captured embedding model %r from "
435                "an existing store entry; mismatched embedders will raise.",
436                self._captured_model,
437            )
438
439    def _check_model(self, embed_model: str | None) -> None:
440        """Raise if ``embed_model`` disagrees with the captured model."""
441        if (
442            self._captured_model is not None
443            and embed_model
444            and embed_model != self._captured_model
445        ):
446            raise EmbeddingModelMismatchError(
447                f"Embedder produced vectors with model {embed_model!r} but "
448                f"store was built with {self._captured_model!r}. Similarity "
449                "scores across models are meaningless; rebuild the store "
450                "with the correct embedder or switch embedders."
451            )
452
453    @staticmethod
454    def _record_document_failed(
455        doc: Document, doc_errors: list[Exception], stats: IngestionStats
456    ) -> DocumentFailed:
457        failed = DocumentFailed(
458            document_id=doc.id,
459            source=doc.source,
460            errors=doc_errors,
461        )
462        stats.documents_failed += 1
463        stats.failed_documents.append(failed)
464        return failed
465
466    async def delete_document(self, document_id: UUID) -> None:
467        """Remove all chunks for a document from the store.
468
469        Convenience wrapper around ``store.delete_where({"document_id": ...})``
470        so callers don't need to know the metadata key.
471        """
472        await self._store.delete_where({"document_id": str(document_id)})
473
474    async def retrieve(
475        self,
476        query: str,
477        top_k: int = 5,
478        metadata_filters: dict[str, Any] | None = None,
479        scope: StoreScope | None = None,
480    ) -> RetrievalResult:
481        """Embed ``query`` and return the top ``top_k`` matches from the store.
482
483        Args:
484            query: The text to embed and search with.
485            top_k: Maximum number of results.
486            metadata_filters: Additional equality filters on chunk metadata.
487            scope: Restricts the search to entries written with the same
488                scope. Leave ``None`` to search across all scopes.
489
490        Raises:
491            EmbeddingModelMismatchError: When the embedder reports a model
492                different from the one captured on first ingest.
493        """
494        await self._ensure_captured_model_seeded()
495        text_result = await self._embedder.aembed([query])
496        self._check_model(text_result.metrics.model)
497
498        store_query = StoreQuery(
499            text=query,
500            scope=scope,
501            embedding=text_result.vectors[0],
502            top_k=top_k,
503            metadata_filters=metadata_filters,
504        )
505        store_hits = await self._store.read(store_query)
506        chunks = [
507            RetrievedChunk(
508                chunk=_entry_to_chunk(hit.entry),
509                score=hit.score,
510                rank=hit.rank,
511                source_retriever=hit.source_retriever,
512                rerank_score=hit.rerank_score,
513            )
514            for hit in store_hits
515        ]
516        result = RetrievalResult(query=query, chunks=chunks)
517        if self._on_retrieve is not None:
518            self._on_retrieve(query, result)
519        return result

93@dataclass
94class RetrievedChunk:
95    chunk: Chunk
96    score: float
97    rank: int
98    source_retriever: str | None = None
99    rerank_score: float | None = None

10@runtime_checkable
11class Store(Protocol):
12    async def write(self, entry: StoreEntry) -> str: ...
13    async def read(self, query: StoreQuery) -> list[RetrievedStoreEntry]: ...
14    async def delete(self, id: UUID) -> None: ...
15    async def clear(self, scope: StoreScope) -> None: ...
16    async def delete_where(self, filters: dict[str, Any]) -> None: ...
17    async def find(
18        self, filters: dict[str, Any], limit: int = 1
19    ) -> list[StoreEntry]: ...
20    async def count(self, filters: dict[str, Any] | None = None) -> int: ...

 44@dataclass
 45class StoreEntry:
 46    # Required fields
 47    id: UUID
 48    content: str
 49    vector: list[float] | None
 50    embedding_model: str
 51    chunk_id: UUID
 52    document_id: UUID
 53    # Optional enrichment fields
 54    abstract: str | None = None
 55    summary: str | None = None
 56    scope: StoreScope | None = None
 57    # Optional chunk provenance
 58    chunk_index: int = 0
 59    parent_chunk_id: UUID | None = None
 60    chunk_offsets: tuple[int, int] | None = None
 61    chunk_metadata: dict = field(default_factory=dict)
 62    # Optional embedding provenance
 63    embedding_version: str | None = None
 64    # Optional store metadata
 65    entities: list[Entity] | None = None
 66    valid_from: datetime | None = None
 67    valid_until: datetime | None = None
 68    created_at: datetime = field(default_factory=lambda: datetime.now(tz=timezone.utc))
 69
 70    @classmethod
 71    def from_chunk(
 72        cls,
 73        embedded_chunk: EmbeddedChunk,
 74        *,
 75        scope: StoreScope | None = None,
 76        abstract: str | None = None,
 77        summary: str | None = None,
 78        entities: list[Entity] | None = None,
 79        valid_from: datetime | None = None,
 80        valid_until: datetime | None = None,
 81    ) -> StoreEntry:
 82        chunk = embedded_chunk.chunk
 83        return cls(
 84            id=chunk.id,
 85            content=chunk.content,
 86            vector=embedded_chunk.vector,
 87            embedding_model=embedded_chunk.embedding_model,
 88            embedding_version=embedded_chunk.embedding_version,
 89            chunk_id=chunk.id,
 90            document_id=chunk.document_id,
 91            chunk_index=chunk.index,
 92            parent_chunk_id=chunk.parent_chunk_id,
 93            chunk_offsets=chunk.offsets,
 94            chunk_metadata=chunk.metadata,
 95            scope=scope,
 96            abstract=abstract,
 97            summary=summary,
 98            entities=entities,
 99            valid_from=valid_from,
100            valid_until=valid_until,
101        )