"""The ``VectorStore`` sink family base.
A ``VectorStore`` is the pipeline's storage backend, not a ``PipelineStep``
itself: a caller opens it once (``async with MilvusVectorStore(url=...) as
vector_store:``) and hands an upstream ``DocumentStep`` to
:meth:`~fastrag.vector_store.base.VectorStore.save`, which returns the
terminal ``PipelineStep[Document]`` that indexes when evaluated. Constructing
that step performs no I/O.
Batching mirrors ``fastrag.ingestion.steps.embedding.base.Embedder``: each
document is queued into a private micro-batcher that coalesces pending
documents into one existence check plus (at most) one
:meth:`~fastrag.vector_store.base.VectorStore.upsert` call once
``policy.batch_size`` documents have queued, or after a short window lets
slower siblings catch up. The store is the source of truth for "already
indexed": :meth:`~fastrag.vector_store.base.VectorStore.existing_ids` is
asked once per batch, present ids are skipped and reported ``cached=True``,
the rest are upserted and reported ``cached=False``. Under a policy with
``recreate_cache``, nothing counts as already indexed: every document is
upserted, overwriting what the store holds for its id, as a rebuild expects.
:meth:`~fastrag.vector_store.base.VectorStore.sync` wraps that same step and,
once the upstream is exhausted, deletes every stored document the run did not
yield, so the store ends up holding exactly the run's documents. It deletes
only after a clean run: a failed item's documents are missing from the run
without being gone from the source.
``VectorStore`` and its concrete stores live at the top level, in
:mod:`fastrag.vector_store`, rather than under :mod:`fastrag.ingestion`: a
vector store is queried at read time as well as written at ingest time, so it
is a storage backend the whole SDK depends on.
"""
from __future__ import annotations
import asyncio
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager
from dataclasses import replace
from typing import Any, Self
from fastrag.exceptions import ConfigurationError, VectorStoreError
from fastrag.ingestion.cache.base import _public_item
from fastrag.ingestion.events import (
DocumentEmbedded,
DocumentSaved,
DocumentsPruned,
ItemFailed,
PipelineEvent,
)
from fastrag.ingestion.policy import Policy
from fastrag.ingestion.steps.base import (
DocumentStep,
ExecutionContext,
PipelineStep,
RoutedStep,
_with_retries,
run_step,
stream_batches,
)
from fastrag.models import Document
#: How long a batch waits for slower siblings before flushing short. Without
#: it, a run with fewer than ``batch_size`` documents would hang (default
#: ``concurrency=8 < batch_size=16``).
_BATCH_WINDOW_SECONDS = 0.05
#: A queued document, the future resolving whether the store already held
#: it, and whether to overwrite it if so.
_PendingSave = tuple[Document, "asyncio.Future[bool]", bool]
#: Ids per :meth:`VectorStore.delete` call when a sync prunes stale documents.
_DELETE_BATCH_SIZE = 1000
[docs]
class VectorStore(AbstractAsyncContextManager["VectorStore"], ABC):
"""Sink backend for persisting processed and embedded documents.
Opened once per run (``async with MilvusVectorStore(url=...) as
vector_store:``) and handed the final upstream step via :meth:`save`,
which returns a lazy terminal :class:`~fastrag.ingestion.steps.base.PipelineStep`
— nothing is indexed until that step is evaluated.
"""
def __init__(self, *, policy: Policy | None = None) -> None:
"""Store the default policy; no I/O happens until ``__aenter__``.
Args:
policy: Execution policy for the terminal step :meth:`save`
returns — governs its concurrency, retries, and the size of
each :meth:`upsert` batch.
"""
self.policy = policy or Policy()
async def __aenter__(self) -> Self:
"""Open backend resources. Overridden by concrete stores that need to."""
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: Any,
) -> None:
"""Release backend resources. Overridden by concrete stores that need to."""
return None
[docs]
@abstractmethod
async def ensure_ready(self, dimension: int) -> None:
"""Create the backing collection for ``dimension``-wide vectors if absent.
Idempotent and cheap to call again on a collection that already
exists; called once, lazily, before the first upsert, with the vector
width taken from the first document flowing through the terminal step.
Args:
dimension: Width of the vectors this run's documents carry.
"""
...
[docs]
@abstractmethod
async def upsert(self, documents: Sequence[Document]) -> None:
"""Insert or replace ``documents`` by :attr:`~fastrag.models.Document.id`, in one call.
Idempotent: rerunning with the same ids converges rather than
duplicating rows. Called with at most ``self.policy.batch_size``
documents at a time — the batcher in :meth:`save`'s returned step
enforces that ceiling, not the caller — and possibly fewer, since
documents already reported by :meth:`existing_ids` are excluded.
Args:
documents: The batch to upsert. Every document carries a
non-``None`` :attr:`~fastrag.models.Document.embedding`
— the terminal step never lets one through otherwise.
"""
...
[docs]
@abstractmethod
async def existing_ids(self, ids: Sequence[str]) -> set[str]:
"""Return the subset of ``ids`` already present in the store.
Called once per pending batch, never once per document, so the round
trips it adds scale with the number of batches. An id returned here
is treated as already indexed (skipped, reported ``cached=True``); an
absent id is upserted and reported ``cached=False``. Skipping a
present document is an optimisation, not a correctness requirement —
a store unable to answer cheaply may upsert everything, but must then
report which ids were genuinely present rather than guessing.
Args:
ids: Candidate document ids to check, at most
``self.policy.batch_size`` at a time.
"""
...
[docs]
@abstractmethod
async def search(
self, embedding: Sequence[float], *, k: int, filter: str | None = None
) -> list[Document]:
"""Return at most ``k`` documents nearest ``embedding``, best first.
The read-path counterpart to :meth:`upsert`, queried by
:class:`~fastrag.retrieval.vector.VectorRetriever` and, through it,
:class:`~fastrag.retrieval.reranking_retriever.RerankingRetriever`.
Never opens the store — the caller (a retriever) holds an
already-open one.
Stamps two reserved metadata keys on every returned document,
alongside its stored ``content``/``metadata`` (including
``source_uri``):
- ``score``: the backend's own relevance score, on the backend's own
scale — never comparable across backends or metrics.
- ``rank``: 0-based position in the returned (best-first) list.
The list is best-first regardless of the configured metric's raw sort
direction, so position 0 is always the best match.
Args:
embedding: Query vector, the same width the store's collection
was created with.
k: Maximum number of documents to return. Fewer come back if the
collection holds fewer than ``k`` matching documents — this
is not an error.
filter: Backend-native filter expression, applied server-side.
``None`` returns the ``k`` nearest documents unfiltered.
"""
...
[docs]
async def all_ids(self) -> set[str]:
"""Return the id of every document in the store.
Needed only by :meth:`sync`, so a store that does not support it may
leave this unimplemented. Returns an empty set when the backing
collection does not exist yet.
"""
raise NotImplementedError(f"{type(self).__name__} does not support sync()")
[docs]
async def delete(self, ids: Sequence[str]) -> None:
"""Remove the documents with ``ids``; ids the store does not hold are ignored.
Needed only by :meth:`sync`, which calls it with at most 1000 ids at
a time.
"""
raise NotImplementedError(f"{type(self).__name__} does not support sync()")
[docs]
def sync(
self, upstream: PipelineStep[Document], *, skip_existing: bool = False
) -> PipelineStep[Document]:
"""Like :meth:`save`, then delete every stored document ``upstream`` did not yield.
After a clean run the store holds exactly the documents the run
produced: documents whose source disappeared or changed (a changed
chunk gets a new id) are removed. Deletion happens once the upstream
is exhausted, so a run that raises or is stopped early deletes
nothing. A run with any ``ItemFailed`` deletes nothing either: a
failed item's documents are missing from the run, not from the
source. Either way a
:class:`~fastrag.ingestion.events.DocumentsPruned` event reports the
outcome.
The run must cover everything the store should keep — use one store
(collection/index) per synced pipeline.
Args:
upstream: Final upstream step yielding embedded documents.
skip_existing: Passed through to :meth:`save`.
"""
return _SyncStep(self)(self.save(upstream, skip_existing=skip_existing))
[docs]
def save(
self, upstream: PipelineStep[Document], *, skip_existing: bool = False
) -> PipelineStep[Document]:
"""Return a lazy terminal step that indexes ``upstream``'s documents.
Constructing the returned step performs no I/O — evaluating it
(``async for document in vector_store.save(upstream):`` or
``.run()``) is what upserts. Each document is queued into a batcher
that checks :meth:`existing_ids` and flushes to :meth:`upsert` once
``self.policy.batch_size`` documents have queued, or after a short
window; every document is yielded back once its batch is resolved,
with :attr:`~fastrag.models.Document.cached` set to whether the
store already held it, and
:class:`~fastrag.ingestion.events.DocumentSaved` emitted for it.
Args:
upstream: Final upstream step yielding embedded documents.
skip_existing: For an identity-preserving Embedder, check destination
ids before embedding. Indexed items bypass vector retrieval and
produce cached DocumentEmbedded and DocumentSaved events. Cached events
and results may contain only identity/metadata with empty content
and no vectors. Custom content filters disable header projection.
Use a separate destination for each embedding configuration.
Either way, ``recreate_cache`` (on the store's policy or the
pipeline's) skips the existence check and upserts every document.
"""
if skip_existing:
from fastrag.ingestion.steps.embedding.base import Embedder
if not isinstance(upstream, Embedder) or upstream._upstream is None:
raise ConfigurationError("skip_existing requires a bound Embedder")
return _IndexStep(self, upstream)(upstream)
return _SaveStep(self)(upstream)
def _headers_safe(root: PipelineStep[Any]) -> bool:
"""Content-aware custom routing requires complete cached records."""
pending = [root]
seen: set[int] = set()
while pending:
step = pending.pop()
if id(step) in seen:
continue
seen.add(id(step))
if step.upstream is not None:
pending.append(step.upstream)
if isinstance(step, RoutedStep):
if step.routing_requires_content:
return False
for predicate, child in step.routes:
if not predicate.header_safe:
return False
pending.append(child)
if step._default_step is not None:
pending.append(step._default_step)
return True
class _EmbeddingWindow(DocumentStep):
"""Run a missing-document window through normal embedding lifecycle events."""
def __init__(self, embedder: Any, items: list[Document]) -> None:
super().__init__(cache=False, policy=embedder.policy)
self.name = embedder.name
self.event_name = embedder.event_name
self._embedder = embedder
self._items = items
def config(self) -> Mapping[str, Any]:
return {}
async def transform(
self, item: Any, context: ExecutionContext
) -> AsyncIterator[Document]:
raise AssertionError("embedding window evaluates the existing embedder")
yield # pragma: no cover
async def evaluate(self, context: ExecutionContext) -> AsyncIterator[Document]:
async def inputs() -> AsyncIterator[Document]:
for item in self._items:
yield item
async for item in self._embedder.evaluate_items(inputs(), context):
yield item
class _IndexStep(DocumentStep):
"""Index-only orchestration: check text ids before retrieving vectors."""
event_name = "Indexing"
def __init__(self, store: VectorStore, embedder: Any) -> None:
super().__init__(cache=False, policy=store.policy)
self.name = f"{type(store).__name__}.save"
self._store = store
self._embedder = embedder
def config(self) -> Mapping[str, Any]:
"""Indexing is never represented by a processing-cache bucket."""
return {}
async def transform(
self, item: Any, context: ExecutionContext
) -> AsyncIterator[Document]:
"""The index-only path is driven in bounded windows by evaluate."""
raise AssertionError("_IndexStep is evaluated as a stream")
yield # pragma: no cover
async def evaluate(self, context: ExecutionContext) -> AsyncIterator[Document]:
"""Skip present documents before the embedder's cache or provider call."""
assert self._upstream is not None
# Preserve the embedder in the declared graph for observers, but pull
# text directly so destination hits still bypass embedding entirely.
text_upstream = self._upstream.upstream
assert text_upstream is not None
policy = self.policy or context.policy
overwrite = self.effective_policy(context).recreate_cache
ready = False
index_context = replace(
context, cache_headers_only=_headers_safe(text_upstream)
)
async for batch in stream_batches(
run_step(text_upstream, index_context), policy.cache_batch_size
):
try:
if not ready:
await _with_retries(
lambda: self._store.ensure_ready(self._embedder.dimension),
policy,
)
ready = True
present = (
set()
if overwrite
else await _with_retries(
lambda: self._store.existing_ids([item.id for item in batch]),
policy,
)
)
except Exception as error:
await self._failed_batch(batch, error, context, policy)
continue
missing: list[Document] = []
for item in batch:
if item.id in present:
saved = replace(_public_item(item), cached=True)
await context.emit(
DocumentEmbedded(
document=saved,
step=self._embedder.name,
event_name=self._embedder.event_name,
)
)
await context.emit(
DocumentSaved(
document=saved, step=self.name, event_name=self.event_name
)
)
yield saved
else:
missing.append(item)
if not missing:
continue
embedded = run_step(_EmbeddingWindow(self._embedder, missing), context)
async for documents in stream_batches(embedded, policy.batch_size):
try:
await _with_retries(lambda: self._store.upsert(documents), policy)
except Exception as error:
await self._failed_batch(documents, error, context, policy)
continue
for document in documents:
saved = replace(document, cached=False)
await context.emit(
DocumentSaved(
document=saved, step=self.name, event_name=self.event_name
)
)
yield saved
async def _failed_batch(
self,
documents: Sequence[Document],
error: Exception,
context: ExecutionContext,
policy: Policy,
) -> None:
"""Report every affected item and honor the sink's terminal error policy."""
for document in documents:
await context.emit(
ItemFailed(
uri=document.id,
source_id=document.source_id,
error=error,
step=self.name,
event_name=self.event_name,
)
)
if policy.on_error == "raise":
raise error
class _SaveStep(DocumentStep):
"""The terminal ``PipelineStep[Document]`` :meth:`VectorStore.save` returns.
Never caches (a sink has nothing to serve from a cache) and always runs
under the owning :class:`VectorStore`'s own policy — except for
``recreate_cache``, which a pipeline-wide rebuild also sets for the sink.
Declares ``event_name = "Indexing"``: unlike a parser/chunker/embedder,
the sink family's base class is :class:`VectorStore` — not a
``PipelineStep`` at all — so the stage label is declared here, on the
step :meth:`VectorStore.save` actually returns. Its :attr:`name` still
names the concrete store (``MilvusVectorStore.save``), so several stores
in one graph stay distinguishable while sharing the one stage.
"""
event_name = "Indexing"
def __init__(self, store: VectorStore) -> None:
super().__init__(cache=False, policy=store.policy)
self.name = f"{type(store).__name__}.save"
self._store = store
self._pending: list[_PendingSave] = []
self._batch_lock = asyncio.Lock()
self._ready_lock = asyncio.Lock()
self._ready = False
def config(self) -> Mapping[str, Any]:
"""No cache-identity config: this step never caches (see above)."""
return {}
async def transform(
self, item: Document, context: ExecutionContext
) -> AsyncIterator[Document]:
"""Upsert one document (batched) and emit ``DocumentSaved``.
The store is the source of truth for ``cached``: a document the store
already holds is not re-upserted and is yielded with ``cached=True``;
one it does not yet have is upserted and yielded with
``cached=False`` — regardless of what ``item.cached`` was coming in.
Raises:
fastrag.exceptions.VectorStoreError: ``item.embedding`` is
``None`` — an unembedded document is a step error, handled
like any other by ``DocumentStep``'s evaluator, not silently
skipped here.
"""
if item.embedding is None:
raise VectorStoreError(
f"document {item.id} has no embedding; cannot index it"
)
effective_policy = self.policy or context.policy
await self._ensure_ready(len(item.embedding))
already_indexed = await self._save_one(
item,
effective_policy.batch_size,
overwrite=self.effective_policy(context).recreate_cache,
)
saved = (
item
if item.cached == already_indexed
else replace(item, cached=already_indexed)
)
await context.emit(
DocumentSaved(document=saved, step=self.name, event_name=self.event_name)
)
yield saved
async def _ensure_ready(self, dimension: int) -> None:
"""Call ``VectorStore.ensure_ready`` exactly once per instance."""
if self._ready:
return
async with self._ready_lock:
if not self._ready:
await self._store.ensure_ready(dimension)
self._ready = True
async def _save_one(
self, document: Document, batch_size: int, *, overwrite: bool
) -> bool:
"""Queue ``document`` for upsert, flushing the batch it lands in.
Mirrors ``Embedder._embed_one``: joins the accumulating batch,
flushing immediately once ``batch_size`` documents have queued, or
scheduling :meth:`_delayed_flush` if it is the first arrival.
Args:
document: The document to index.
batch_size: How many queued documents trigger a flush.
overwrite: Upsert ``document`` even if the store already holds
its id.
Returns:
Whether the store already held this document (skipped, not
upserted) — resolved by :meth:`_flush` once its batch's
existence check comes back. Always ``False`` with ``overwrite``.
"""
loop = asyncio.get_running_loop()
future: asyncio.Future[bool] = loop.create_future()
to_flush: list[_PendingSave] | None = None
async with self._batch_lock:
self._pending.append((document, future, overwrite))
if len(self._pending) >= batch_size:
to_flush, self._pending = self._pending, []
elif len(self._pending) == 1:
asyncio.create_task(self._delayed_flush())
if to_flush is not None:
await self._flush(to_flush)
return await future
async def _delayed_flush(self) -> None:
"""Flush whatever is still pending after the batch window elapses."""
await asyncio.sleep(_BATCH_WINDOW_SECONDS)
to_flush: list[_PendingSave] | None = None
async with self._batch_lock:
if self._pending:
to_flush, self._pending = self._pending, []
if to_flush is not None:
await self._flush(to_flush)
async def _flush(self, batch: list[_PendingSave]) -> None:
"""Check existence once for ``batch``, upsert what's missing, resolve each future.
One :meth:`VectorStore.existing_ids` call decides which of the batch
is already indexed; those are excluded from the (at most one)
:meth:`VectorStore.upsert` call that follows — skipped entirely if
every document is already present. A failure from either call fails
every document in the batch; each document's own ``DocumentStep``
retry calls :meth:`transform` again, resubmitting into a fresh batch.
Documents queued with ``overwrite`` skip the existence check and are
always upserted.
"""
checked = [document.id for document, _, overwrite in batch if not overwrite]
try:
present = await self._store.existing_ids(checked) if checked else set()
missing = [
document for document, _, _ in batch if document.id not in present
]
if missing:
await self._store.upsert(missing)
except Exception as error: # propagated to every waiter, not swallowed
for _, future, _ in batch:
if not future.done():
future.set_exception(error)
return
for document, future, _ in batch:
if not future.done():
future.set_result(document.id in present)
class _SyncStep(DocumentStep):
"""The terminal step :meth:`VectorStore.sync` returns: save, then prune.
Bound to the step :meth:`VectorStore.save` returns. Records the id of
every document that step yields, and counts the run's ``ItemFailed``
events by evaluating it under a context whose sink sees each event on its
way to the pipeline's own. Shares the ``Indexing`` stage with the save
step; its :attr:`name` is ``MilvusVectorStore.sync``.
"""
event_name = "Indexing"
def __init__(self, store: VectorStore) -> None:
super().__init__(cache=False, policy=store.policy)
self.name = f"{type(store).__name__}.sync"
self._store = store
def config(self) -> Mapping[str, Any]:
"""Syncing is never represented by a processing-cache bucket."""
return {}
async def transform(
self, item: Any, context: ExecutionContext
) -> AsyncIterator[Document]:
"""The sync step only wraps the save step's stream; see evaluate."""
raise AssertionError("_SyncStep is evaluated as a stream")
yield # pragma: no cover
async def evaluate(self, context: ExecutionContext) -> AsyncIterator[Document]:
"""Yield every saved document, then delete the stored ones not yielded."""
assert self._upstream is not None
sink = context.event_sink
failures = 0
async def watch(event: PipelineEvent) -> None:
nonlocal failures
if isinstance(event, ItemFailed):
failures += 1
if sink is not None:
await sink(event)
kept: set[str] = set()
watched = replace(context, event_sink=watch)
async for document in run_step(self._upstream, watched):
kept.add(document.id)
yield document
policy = self.policy or context.policy
stale = sorted(await _with_retries(self._store.all_ids, policy) - kept)
if not failures:
for start in range(0, len(stale), _DELETE_BATCH_SIZE):
ids = stale[start : start + _DELETE_BATCH_SIZE]
await _with_retries(lambda: self._store.delete(ids), policy)
await context.emit(
DocumentsPruned(
step=self.name,
stale=len(stale),
deleted=0 if failures else len(stale),
failures=failures,
event_name=self.event_name,
)
)