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)

10class ContentExtractor(Protocol):
11    """Convert one loader field value into searchable text."""
12
13    def __call__(self, value: Any, /) -> str: ...

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)

23class JsonExtractor:
24    """Render values as valid JSON while preserving Unicode text."""
25
26    def __call__(self, value: Any) -> str:
27        return json.dumps(value, ensure_ascii=False)

110@dataclass
111class OCRResult:
112    """Structured output from an OCR engine.
113
114    Attributes:
115        markdown: Full page text as markdown. Tables and headings are
116            represented in markdown syntax where the engine supports it.
117        json_result: Raw structured layout data from the engine, if available.
118            Schema is engine-specific; ``None`` when the engine does not
119            return structured layout data. For GLM-OCR this is the value of
120            ``PipelineResult.to_dict()["json_result"]``, which contains
121            bounding boxes, tables, and reading order.
122    """
123
124    markdown: str
125    json_result: dict[str, Any] | list[Any] | None = None
126
127    def to_text(self) -> str:
128        return self.markdown

30class ProseExtractor:
31    """Flatten nested JSON-like values into readable, labelled text."""
32
33    def __call__(self, value: Any) -> str:
34        return self._render(value, nested=False)
35
36    def _render(self, value: Any, *, nested: bool = True) -> str:
37        if isinstance(value, dict):
38            rendered = "; ".join(
39                f"{key}: {self._render(item)}" for key, item in value.items()
40            )
41            return f"{{{rendered}}}" if nested else rendered
42        if isinstance(value, (list, tuple)):
43            return f"[{', '.join(self._render(item) for item in value)}]"
44        if value is None:
45            return "null"
46        if isinstance(value, bool):
47            return "true" if value else "false"
48        return str(value)

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, scope):
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(
302        self, doc: Document, scope: StoreScope | None
303    ) -> bool:
304        """Whether the store already holds a *complete* copy of ``doc``.
305
306        Skip re-embedding only when as many chunks are present as the last
307        write expected. A partially-written document (some chunks present after
308        an interrupted ingest) has fewer than expected and is re-ingested rather
309        than left broken. find() is metadata-only (no vector search) and only
310        fetches a single entry; the presence check is a count() so no payloads
311        are transferred and large documents never require a single oversized
312        read (Chroma Cloud caps the per-get ``limit`` value).
313        """
314        if doc.source is None:
315            return False
316        stale_filters = {
317            "source_path": doc.source,
318            "content_hash": doc.content_hash,
319            **(scope.to_payload_filters() if scope is not None else {}),
320        }
321        existing = await self._store.find(stale_filters, limit=1)
322        if not existing:
323            return False
324        expected = existing[0].chunk_metadata.get("doc_chunk_count")
325        if expected is None:
326            # Not written by this runtime (every ingest stamps the count):
327            # completeness can't be verified, so re-ingest.
328            return False
329        return await self._store.count(stale_filters) >= expected
330
331    @staticmethod
332    def _stamp_staleness_metadata(doc: Document, chunks: list[Chunk]) -> None:
333        """Inject staleness-detection metadata into every chunk so future
334        `find` calls can identify whether this document has changed."""
335        for chunk in chunks:
336            if doc.source is not None:
337                chunk.metadata.setdefault("source_path", doc.source)
338            if doc.content_hash is not None:
339                chunk.metadata.setdefault("content_hash", doc.content_hash)
340
341    def _split_oversized(
342        self, chunks: list[Chunk], stats: IngestionStats
343    ) -> tuple[list[Chunk], list[EmbeddingFailure]]:
344        """Partition chunks into embeddable ones and per-chunk failures.
345
346        Returns ``(ok_chunks, failures)``; each oversize chunk becomes a
347        single-chunk ``EmbeddingFailure`` and is recorded in ``stats``.
348        """
349        if self._max_tokens is None or self._tokenizer is None:
350            return chunks, []
351        ok_chunks: list[Chunk] = []
352        failures: list[EmbeddingFailure] = []
353        for chunk in chunks:
354            tokens = self._tokenizer.count(chunk.content)
355            if tokens > self._max_tokens:
356                err = ValueError(
357                    f"chunk {chunk.id} has {tokens} tokens "
358                    f"(>{self._max_tokens}); dropped before embedding"
359                )
360                stats.batches_failed += 1
361                failure = EmbeddingFailure(chunks=[chunk], errors=[err])
362                stats.batch_failures.append(failure)
363                failures.append(failure)
364            else:
365                ok_chunks.append(chunk)
366        return ok_chunks, failures
367
368    async def _embed_and_store(
369        self,
370        doc: Document,
371        chunks: list[Chunk],
372        stats: IngestionStats,
373        doc_errors: list[Exception],
374        scope: StoreScope | None,
375    ) -> AsyncGenerator[
376        BatchIngested | EmbeddingFailure | DocumentFailed | DocumentSkipped, None
377    ]:
378        # batch_index is per-document: it counts batches (successful and
379        # failed) within this document and resets for the next one.
380        batch_index = 0
381        delete_done = False
382        async for batch in self._embedder.astream_batches(
383            chunks, batch_size=self._batch_size
384        ):
385            if isinstance(batch, EmbeddingResult):
386                # Check model BEFORE delete_where / write — a mismatch here
387                # must not corrupt the store by clearing prior chunks first.
388                self._check_model(batch.metrics.model)
389                if not delete_done:
390                    await self._store.delete_where(
391                        {
392                            "document_id": str(doc.id),
393                            **(scope.to_payload_filters() if scope is not None else {}),
394                        }
395                    )
396                    delete_done = True
397                for embedded in batch.chunks:
398                    self._capture_model(embedded)
399                    entry = StoreEntry.from_chunk(embedded, scope=scope)
400                    await self._store.write(entry)
401                stats.chunks_embedded += len(batch.chunks)
402                stats.total_metrics = stats.total_metrics + batch.metrics
403                yield BatchIngested(
404                    document_id=doc.id,
405                    embedded_chunks=batch.chunks,
406                    batch_index=batch_index,
407                    metrics=batch.metrics,
408                )
409            else:
410                doc_errors.extend(batch.errors)
411                stats.batches_failed += 1
412                stats.batch_failures.append(batch)
413                yield batch
414            batch_index += 1
415
416    def _capture_model(self, embedded: EmbeddedChunk) -> None:
417        """Record the embedding model from the first successful chunk so later
418        retrieve() calls can enforce model consistency."""
419        if self._captured_model is None and embedded.embedding_model:
420            self._captured_model = embedded.embedding_model
421            logger.info(
422                "RetrievalRuntime captured embedding model %r "
423                "from first successful batch; subsequent retrieve() "
424                "calls will enforce this model.",
425                self._captured_model,
426            )
427
428    async def _ensure_captured_model_seeded(self) -> None:
429        """Lazily seed ``_captured_model`` from an existing store entry so the
430        guard survives across process restarts. ``StoreEntry.embedding_model``
431        is recorded on every persisted entry, so a single ``find`` call is
432        enough — no schema change required. Runs at most once per runtime;
433        a miss against an empty store sets ``_seed_attempted`` so we don't
434        re-query on every doc."""
435        if self._captured_model is not None or self._seed_attempted:
436            return
437        self._seed_attempted = True
438        existing = await self._store.find({}, limit=1)
439        if existing and existing[0].embedding_model:
440            self._captured_model = existing[0].embedding_model
441            logger.info(
442                "RetrievalRuntime seeded captured embedding model %r from "
443                "an existing store entry; mismatched embedders will raise.",
444                self._captured_model,
445            )
446
447    def _check_model(self, embed_model: str | None) -> None:
448        """Raise if ``embed_model`` disagrees with the captured model."""
449        if (
450            self._captured_model is not None
451            and embed_model
452            and embed_model != self._captured_model
453        ):
454            raise EmbeddingModelMismatchError(
455                f"Embedder produced vectors with model {embed_model!r} but "
456                f"store was built with {self._captured_model!r}. Similarity "
457                "scores across models are meaningless; rebuild the store "
458                "with the correct embedder or switch embedders."
459            )
460
461    @staticmethod
462    def _record_document_failed(
463        doc: Document, doc_errors: list[Exception], stats: IngestionStats
464    ) -> DocumentFailed:
465        failed = DocumentFailed(
466            document_id=doc.id,
467            source=doc.source,
468            errors=doc_errors,
469        )
470        stats.documents_failed += 1
471        stats.failed_documents.append(failed)
472        return failed
473
474    async def delete_document(self, document_id: UUID) -> None:
475        """Remove all chunks for a document from the store.
476
477        Convenience wrapper around ``store.delete_where({"document_id": ...})``
478        so callers don't need to know the metadata key.
479        """
480        await self._store.delete_where({"document_id": str(document_id)})
481
482    async def retrieve(
483        self,
484        query: str,
485        top_k: int = 5,
486        metadata_filters: dict[str, Any] | None = None,
487        scope: StoreScope | None = None,
488    ) -> RetrievalResult:
489        """Embed ``query`` and return the top ``top_k`` matches from the store.
490
491        Args:
492            query: The text to embed and search with.
493            top_k: Maximum number of results.
494            metadata_filters: Additional equality filters on chunk metadata.
495            scope: Restricts the search to entries written with the same
496                scope. Leave ``None`` to search across all scopes.
497
498        Raises:
499            EmbeddingModelMismatchError: When the embedder reports a model
500                different from the one captured on first ingest.
501        """
502        await self._ensure_captured_model_seeded()
503        text_result = await self._embedder.aembed([query])
504        self._check_model(text_result.metrics.model)
505
506        store_query = StoreQuery(
507            text=query,
508            scope=scope,
509            embedding=text_result.vectors[0],
510            top_k=top_k,
511            metadata_filters=metadata_filters,
512        )
513        store_hits = await self._store.read(store_query)
514        chunks = [
515            RetrievedChunk(
516                chunk=_entry_to_chunk(hit.entry),
517                score=hit.score,
518                rank=hit.rank,
519                source_retriever=hit.source_retriever,
520                rerank_score=hit.rerank_score,
521            )
522            for hit in store_hits
523        ]
524        result = RetrievalResult(query=query, chunks=chunks)
525        if self._on_retrieve is not None:
526            self._on_retrieve(query, result)
527        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: ...