"""The ``Embedder`` family base.
``docs/specification/subclasses.md``: ``Embedder`` is the ``DocumentStep``
every concrete embedder (``Model2VecEmbedder``, ``OpenAIEmbedder``) extends.
Its one job beyond ``DocumentStep`` itself is attaching a dense vector to each
incoming document while coalescing the underlying model calls into batches of
up to ``policy.batch_size`` texts, rather than issuing one call per document.
``DocumentStep`` (task 003) drives caching, retries and per-item concurrency
around a single-item ``transform`` hook — there is no batch-shaped hook on the
base to plug into, and adding one would mean changing a foundation three
already-finished tasks (parsing, chunking) build on. So batching is a detail
private to this family: :meth:`Embedder.transform` still satisfies the
single-item hook (each cache decision, retry and ``ItemFailed`` stays exactly
what ``DocumentStep`` already gives every other step), but internally it hands
its text to a small per-instance batcher that coalesces concurrently pending
calls into one :meth:`Embedder.embed` call once ``policy.batch_size`` texts
have queued, or after a short window lets slower siblings catch up. Nothing
here re-implements caching or retries — a batch that fails fails every text
in it, and each of those documents' own ``DocumentStep``-level retry (if any)
resubmits into a fresh batch.
"""
from __future__ import annotations
import asyncio
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Sequence
from dataclasses import replace
from typing import Literal
from fastrag.exceptions import EmbeddingError
from fastrag.ingestion._hashing import config_hash
from fastrag.ingestion.cache import CacheStore
from fastrag.ingestion.events import DocumentEmbedded
from fastrag.ingestion.policy import Policy
from fastrag.ingestion.steps.base import DocumentStep, ExecutionContext
from fastrag.models import ContentFormat, Document
#: How long a batch waits for concurrently-arriving siblings before it is
#: flushed with fewer than ``policy.batch_size`` texts. Short relative to any
#: real model call (network round trip or an encode pass), long enough for
#: the handful of upstream items ``policy.concurrency`` admits at once to all
#: reach the batcher.
_BATCH_WINDOW_SECONDS = 0.05
_PendingEmbed = tuple[str, "asyncio.Future[list[float]]"]
[docs]
class Embedder(DocumentStep, ABC):
r"""Attaches a dense vector to each incoming document.
Binding an upstream (``embedder(chunked_documents)``, via the inherited
``__call__``) and
evaluating it (``async for document in embedder`` /
:meth:`~fastrag.ingestion.steps.base.PipelineStep.run`) never embeds
eagerly — only consuming the resulting stream does.
Caching (per document id and this embedder's own configuration) and
retries are entirely :class:`~fastrag.ingestion.steps.base.DocumentStep`'s
job, applied around :meth:`transform`; nothing here re-implements either.
Only the underlying model call is batched, up to ``policy.batch_size``
texts per call — see the module docstring.
Writing one is a class plus three members::
from collections.abc import Sequence
from fastrag.ingestion.steps.embedding import Embedder
class EchoEmbedder(Embedder):
def __init__(self, *, width: int = 8, cache=False, policy=None):
super().__init__(cache=cache, policy=policy)
self.width = width
def config(self):
return {"width": self.width}
@property
def model_id(self) -> str:
return f"echo@{self.width}"
@property
def dimension(self) -> int:
return self.width
async def embed(self, texts: Sequence[str]) -> list[list[float]]:
return [[float(len(text))] * self.width for text in texts]
Every event an embedder emits is labelled ``"Embedding"``, inherited by
every concrete embedder (``Model2VecEmbedder``, ``OpenAIEmbedder``, ...)
unless it declares an ``event_name`` of its own.
"""
event_name = "Embedding"
def __init__(
self,
*,
cache: CacheStore | Literal[False] | None = None,
policy: Policy | None = None,
) -> None:
"""Initialize an ``Embedder``.
Args:
cache: Cache backend, or ``False``/``None`` — see
:class:`~fastrag.ingestion.steps.base.PipelineStep`.
policy: Execution policy overriding the pipeline default.
"""
super().__init__(cache=cache, policy=policy)
self._pending: list[_PendingEmbed] = []
self._batch_lock = asyncio.Lock()
@property
def config_hash(self) -> str:
"""Invalidate embeddings when the shared input serialization changes."""
from fastrag.ingestion.steps.chunking.structure import EMBEDDING_TEXT_VERSION
return config_hash(
type(self).__name__,
{
"provider": self.config(),
"embedding_text_version": EMBEDDING_TEXT_VERSION,
},
)
@property
@abstractmethod
def model_id(self) -> str:
"""Stable identifier of the model configuration.
Changing it — a different model, or a pinned dimension — must change
a subclass's :meth:`~fastrag.ingestion.steps.base.PipelineStep.config`
(and therefore its ``config_hash``), so vectors of different models or
widths never share a document-cache entry.
"""
@property
@abstractmethod
def dimension(self) -> int:
"""The width of every returned vector, known at construction.
Resolved eagerly — never behind a network call — so a caller can name
a vector-store collection from it before anything is embedded.
"""
[docs]
@abstractmethod
async def embed(self, texts: Sequence[str]) -> list[list[float]]:
"""Embed a batch of texts, in order.
Called with at most ``policy.batch_size`` texts at a time — the
batcher in this base class is what enforces that ceiling, not the
caller. Must return exactly one vector per text, in the same order.
"""
async def _on_produced(
self, item: Document, produced: list[Document], context: ExecutionContext
) -> None:
"""Emit one :class:`~fastrag.ingestion.events.DocumentEmbedded` per embedded document.
Called once per upstream document regardless of whether ``produced``
came from the cache or a fresh :meth:`transform` call, so
``document.cached`` is trustworthy either way — see
``DocumentStep._on_produced``.
"""
for document in produced:
await context.emit(
DocumentEmbedded(
document=document,
step=self.name,
event_name=self.event_name,
)
)
async def _embed_one(self, text: str, batch_size: int) -> list[float]:
"""Queue ``text`` for embedding, flushing the batch it lands in.
A call joins whatever batch is currently accumulating. It flushes
immediately once ``batch_size`` texts have queued; otherwise it is
the first arrival responsible for scheduling :func:`_delayed_flush`,
which flushes whatever is still pending after
:data:`_BATCH_WINDOW_SECONDS` so a batch smaller than ``batch_size``
is never left waiting forever for siblings that will never arrive.
"""
loop = asyncio.get_running_loop()
future: asyncio.Future[list[float]] = loop.create_future()
to_flush: list[_PendingEmbed] | None = None
async with self._batch_lock:
self._pending.append((text, future))
if len(self._pending) >= batch_size:
to_flush, self._pending = self._pending, []
elif len(self._pending) == 1:
asyncio.create_task(self._delayed_flush(batch_size))
if to_flush is not None:
await self._flush(to_flush)
return await future
async def _delayed_flush(self, batch_size: int) -> None:
"""Flush whatever is still pending after the batch window elapses."""
await asyncio.sleep(_BATCH_WINDOW_SECONDS)
to_flush: list[_PendingEmbed] | 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)
# `batch_size` is unused once the window fires — flushing at this
# point always takes whatever is pending, full or not — but is kept
# as a parameter so a future revision can re-check the ceiling
# without changing the call site.
del batch_size
async def _flush(self, batch: list[_PendingEmbed]) -> None:
"""Call :meth:`embed` once for ``batch`` and resolve each future.
Runs with the batch lock already released, so other calls keep
accumulating the next batch while this one's model call is in
flight — batches overlap rather than serializing behind one another.
A failure fails every text in this batch; each document's own
``DocumentStep``-level retry (if any) calls :meth:`transform` again,
which resubmits into a fresh batch.
"""
texts = [text for text, _ in batch]
try:
vectors = await self.embed(texts)
except Exception as error: # propagated to every waiter, not swallowed
for _, future in batch:
if not future.done():
future.set_exception(error)
return
if len(vectors) != len(batch):
mismatch = EmbeddingError(
f"{type(self).__name__}.embed returned {len(vectors)} vectors "
f"for {len(batch)} texts"
)
for _, future in batch:
if not future.done():
future.set_exception(mismatch)
return
for (_, future), vector in zip(batch, vectors, strict=True):
if not future.done():
future.set_result(vector)