Skip to content

embedders.md

Latest commit

 

History

History
180 lines (134 loc) · 5.54 KB

embedders.md

File metadata and controls

180 lines (134 loc) · 5.54 KB
Raw

rag-embedders — reference

Embedder plugins for AgentContextOS. Step 1.8 ships three plugin classes covering four model families, plus a shared BatchingEmbedder base that handles batching, retry, and dimension normalization.

Backends

Plugin Backend Default model Default dim Extra
OpenAIEmbedder OpenAI async API text-embedding-3-small 1536 [openai]
CohereEmbedder Cohere async API embed-english-v3.0 1024 [cohere]
SentenceTransformersEmbedder local PyTorch BAAI/bge-large-en-v1.5 1024 [sentence-transformers]

The [sentence-transformers] extra pulls in PyTorch (~600 MB). Convenience factories bge_large_en() and e5_large_v2() build pre-configured instances for the BGE and E5 model families.

Default install is rag-core only — pick the extras you need.

OpenAIEmbedder

from rag_embedders import OpenAIEmbedder

embedder = OpenAIEmbedder(
    model="text-embedding-3-small",   # 1536 native
    target_dimension=512,             # forwarded server-side
    api_key=None,                     # falls back to OPENAI_API_KEY
)
embeddings = await embedder.bulk_embed(ctx, texts, chunk_ids)

For text-embedding-3-* models, target_dimension is forwarded to the provider via the dimensions request parameter — server-side truncation + renormalization saves bandwidth and latency. For older models (text-embedding-ada-002) the base class handles truncation client-side.

Unknown model ids require an explicit target_dimension so the SPI can advertise an output dimension before the first call.

CohereEmbedder

from rag_embedders import CohereEmbedder

embedder = CohereEmbedder(
    model="embed-english-v3.0",
    input_type="search_document",   # ingest-time default
    api_key=None,                   # falls back to COHERE_API_KEY
)

Cohere v3 requires an input_type — search_document at ingest, search_query at retrieval (build a separate instance for the query side in Phase 2).

The wrapper handles both SDK v4 (embeddings is a plain list of lists) and v5 (embeddings.float is the float vector accessor) so upgrading the SDK doesn't require code changes.

SentenceTransformersEmbedder

from rag_embedders import SentenceTransformersEmbedder, bge_large_en, e5_large_v2

# Direct construction (any HuggingFace sentence-transformers model)
embedder = SentenceTransformersEmbedder(
    model="any/sentence-transformer-model",
    device="cuda",            # autodetect when None
    target_dimension=None,
    normalize_embeddings=True,
    passage_prefix="",        # e.g. "passage: " for E5
)

# Convenience factories
bge = bge_large_en(device="cuda")
e5 = e5_large_v2(device="cuda")

E5 models receive the passage: instruction prefix at ingest time automatically (the e5_large_v2() factory wires it; the _PASSAGE_PREFIXES map covers e5-large-v2, e5-base-v2, e5-small-v2, and multilingual-e5-large).

Encoding runs in asyncio.to_thread so the synchronous sentence-transformers API doesn't block the event loop.

BatchingEmbedder base

Subclass this to write a custom embedder:

from rag_embedders import BatchingEmbedder

class MyEmbedder(BatchingEmbedder):
    @property
    def model(self) -> str:
        return "my-model"

    @property
    def dimension(self) -> int:
        return 768

    async def _embed_one_batch(self, texts: list[str]) -> list[list[float]]:
        return await my_provider.embed(texts)

    async def health(self) -> bool:
        return True

The base class gives you:

  • Sub-batch splitting at max_batch_size.
  • Retry with exponential back-off via the shared RetryPolicy.
  • Optional dimension normalization (truncate + L2 renorm).
  • Input/output order preserved across sub-batches.
  • Wrapping non-EmbeddingError exceptions into EmbeddingError.

RetryPolicy

from rag_embedders import RetryPolicy
from openai import RateLimitError, APITimeoutError

policy = RetryPolicy(
    max_attempts=5,
    base_delay_s=0.5,
    max_delay_s=30.0,
    jitter=0.2,
    retry_on=(RateLimitError, APITimeoutError),
)
embedder = OpenAIEmbedder(retry=policy)

retry_on is what makes the difference between "retry transient" and "raise immediately". The defaults for each plugin map provider-native rate-limit / timeout / connection exceptions.

normalize_dimension

from rag_embedders import normalize_dimension

vec_512 = normalize_dimension(vec_1536, 512)  # truncate + L2-renorm

Pure helper exposed for callers that want post-hoc dim reduction outside the embedder hot path.

CLI

ragctl embed doc.md                                          # noop default
ragctl embed doc.md --backend openai --model text-embedding-3-small
ragctl embed doc.md --backend cohere
ragctl embed doc.md --backend bge --target-dim 256
ragctl embed doc.md --backend e5

Runs the full parse → chunk → enrich → PII → embed pipeline and prints first-N embeddings with dim + first 4 coordinates.

Errors

All embedder failures wrap to rag_core.errors.EmbeddingError:

  • Provider rate-limit / connection failures (after retry exhaustion)
  • Returned vector count mismatches input count
  • Returned vector dimension mismatches advertised dimension
  • Missing extras (clear install hint in the message)

Extension points

Write a custom embedder by subclassing BatchingEmbedder (see above). Plug it into the ingest pipeline at the same point a built-in embedder would — Embedder is the SPI contract; the pipeline doesn't care which class satisfies it.