orchestrata_ultimate.semantic_graph.embedders

Reference › orchestrata_ultimate › semantic_graph › embedders

The embedding backends for the semantic graph: a small protocol and two implementations behind one async interface, so ingestion does not need to know which is in use. Embedder is a runtime-checkable protocol with a dimension property and embed(texts, prompt_type), which returns one vector per text for a document or query prompt.

OpenRouterEmbedder is the cloud backend over OpenRouter's embeddings endpoint (default model mistralai/codestral-embed-2505, dimension 1536, 60-second timeout, batches of up to 96 texts). embed splits inputs with _split_batches, bounded by item count and estimated tokens, then _request_embeddings sends each batch through the recorded-cassette contract; only when that allows a live call does _request_embeddings_live reach the provider. _post_with_retry retries transient failures only.

SentenceTransformerEmbedder is the local backend (default Qwen/Qwen3-Embedding-0.6B, dimension 1024). _load_model loads the model lazily, _prompt_kwargs_for supplies the query or document prompt, and embed runs the model in a worker thread so the event loop is not blocked. Which backend a process gets is decided in embedder_factory. The module does not touch the graph.

At a glance

Class diagram

Dependency diagram

API

Classes

Embedder (protocol)

Bases: Protocol · Decorators: runtime_checkable · Line 39

Protocol for text embedding backends.

OpenRouterEmbedder

Line 65

Cloud embedding backend using the OpenRouter embeddings endpoint.

SentenceTransformerEmbedder

Line 301

Local embedding backend using sentence-transformers.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/semantic_graph/embedders.py

"""Embedding backends for the semantic graph.

Provides a small protocol plus two concrete implementations:

* :class:`OpenRouterEmbedder` - cloud embeddings via OpenRouter (Codestral).
* :class:`SentenceTransformerEmbedder` - local GPU embeddings via
  ``sentence-transformers`` (Qwen3-Embedding, etc.).

Both expose the same async interface so ingestion code does not need to know
which backend is in use.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Mapping
from typing import Any, Protocol, runtime_checkable

import anyio
import anyio.to_thread
import httpx

from orchestrata_ultimate.semantic_graph.exceptions import VectorSearchError

logger = logging.getLogger(__name__)

#: Attempts per embeddings request, including the first. A corpus ingest issues
#: thousands of requests, so a single un-retried transient ends the whole job:
#: request 330 of a 2026-08-09 run failed on an empty-message
#: httpx.RequestError (a read timeout) after 329 consecutive 200s, discarding
#: everything embedded up to that point.
RETRY_ATTEMPTS = 4
#: Seconds before the first retry; doubles after each attempt.
RETRY_BACKOFF_SECONDS = 1.0


@runtime_checkable
class Embedder(Protocol):
    """Protocol for text embedding backends."""

    @property
    def dimension(self) -> int:
        """Return the vector dimension produced by this embedder."""
        ...

    async def embed(
        self, texts: list[str], prompt_type: str = "document"
    ) -> list[list[float]]:
        """Return one embedding vector for each input text.

        Args:
            texts: Input strings to embed. Empty strings are valid and must
                return a vector of ``self.dimension`` floats.
            prompt_type: Semantic role of the texts. Typical values are
                ``"document"`` and ``"query"``. Backends that do not use
                prompt prefixes ignore this value.

        Returns:
            A list of embedding vectors parallel to ``texts``.
        """
        ...


class OpenRouterEmbedder:
    """Cloud embedding backend using the OpenRouter embeddings endpoint.

    Defaults to ``mistralai/codestral-embed-2505`` (1536 dimensions). The
    caller must supply an API key; the class never reads environment variables
    directly.
    """

    def __init__(
        self,
        api_key: str,
        model: str = "mistralai/codestral-embed-2505",
        dimension: int = 1536,
        base_url: str = "https://openrouter.ai/api/v1",
        timeout: float = 60.0,
        batch_size: int = 96,
        max_tokens_per_batch: int = 100_000,
        env: Mapping[str, str] | None = None,
        retry_attempts: int = RETRY_ATTEMPTS,
        retry_backoff_seconds: float = RETRY_BACKOFF_SECONDS,
    ) -> None:
        """Initialize the embedder.

        Args:
            api_key: OpenRouter API key.
            model: Model identifier on OpenRouter.
            dimension: Expected dimension of returned vectors.
            base_url: OpenRouter API base URL.
            timeout: HTTP request timeout in seconds.
            batch_size: Maximum texts per embedding request.
            max_tokens_per_batch: Approximate token ceiling per request. Batches
                are split further when the estimated total exceeds this value.
            retry_attempts: Attempts per request, including the first. Only
                transient failures consume one; a refusal fails at once.
            retry_backoff_seconds: Delay before the first retry, doubling
                after each. Tests pass a small value rather than reassigning a
                class attribute.
            env: Environment the cassette contract is read from. Defaults to
                the process environment. Pass an explicit mapping WITHOUT
                ``EMBEDDING_CASSETTE_DIR`` to force the real HTTP path — which
                is what the error-path tests need, since their subject is this
                client's own failure handling against a local dead port and
                never a provider.
        """
        if not api_key:
            raise VectorSearchError("OpenRouter API key is required")
        self._api_key = api_key
        self._model = model
        self._dimension = dimension
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._batch_size = max(1, batch_size)
        self._max_tokens_per_batch = max(1, max_tokens_per_batch)
        self._env = env
        self._retry_attempts = max(1, retry_attempts)
        self._retry_backoff_seconds = max(0.0, retry_backoff_seconds)

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

    async def embed(
        self, texts: list[str], prompt_type: str = "document"
    ) -> list[list[float]]:
        """Embed texts via OpenRouter in bounded batches."""
        del prompt_type  # Codestral does not use prompt prefixes.
        if not texts:
            return []

        results: list[list[float] | None] = [None] * len(texts)
        semaphore = asyncio.Semaphore(4)

        batches = self._split_batches(texts)

        async def _embed_batch(offset: int, batch: list[str]) -> None:
            async with semaphore:
                embeddings = await self._request_embeddings(batch)
                for idx, embedding in enumerate(embeddings):
                    results[offset + idx] = embedding

        offset = 0
        tasks = []
        for batch in batches:
            tasks.append(_embed_batch(offset, batch))
            offset += len(batch)
        await asyncio.gather(*tasks)
        return [embedding for embedding in results if embedding is not None]

    def _split_batches(self, texts: list[str]) -> list[list[str]]:
        """Split texts into batches bounded by count and estimated tokens.

        The initial split is by ``batch_size``. Each chunk is then subdivided if
        the estimated token count exceeds ``max_tokens_per_batch``. Token counts
        are estimated as ``max(1, len(text) // 4)`` — conservative for code and
        English prose and good enough to avoid provider batch limits.
        """
        batches: list[list[str]] = []
        for i in range(0, len(texts), self._batch_size):
            chunk = texts[i : i + self._batch_size]
            current: list[str] = []
            current_tokens = 0
            for text in chunk:
                tokens = max(1, len(text) // 4)
                if current and current_tokens + tokens > self._max_tokens_per_batch:
                    batches.append(current)
                    current = [text]
                    current_tokens = tokens
                else:
                    current.append(text)
                    current_tokens += tokens
            if current:
                batches.append(current)
        return batches

    async def _request_embeddings(self, texts: list[str]) -> list[list[float]]:
        """Embeddings for *texts*, through the cassette contract.

        This path holds its own httpx client and never reaches
        ``llm/client.py``, so until the cassette grew a non-chat form every
        embedding the suite computed was a real billed OpenRouter request.
        """
        from orchestrata_ultimate.llm._cassette import (  # noqa: PLC0415
            cassette_exchange,
            synthetic_embeddings,
        )

        embeddings: list[list[float]] = await cassette_exchange(
            kind="embeddings",
            payload={"model": self._model, "input": list(texts)},
            synth=lambda: synthetic_embeddings(texts, self._dimension),
            live=lambda: self._request_embeddings_live(texts),
            env=self._env,
        )
        return embeddings

    #: Status codes worth trying again. Everything else — 400, 401, 402, 403 —
    #: is a refusal that a retry reproduces exactly, at the cost of the wait.
    RETRY_STATUS = frozenset({429, 500, 502, 503, 504})

    async def _post_with_retry(
        self,
        url: str,
        headers: dict[str, str],
        payload: dict[str, Any],
    ) -> Any:
        """POST the embeddings request, retrying only transient failures.

        Raises:
            VectorSearchError: On a non-transient refusal, or once the
                transient attempts are exhausted. The message names which of
                the two happened, so a quota problem is never reported as
                flakiness.
        """
        delay = self._retry_backoff_seconds
        for attempt in range(1, self._retry_attempts + 1):
            last = attempt == self._retry_attempts
            try:
                async with httpx.AsyncClient(timeout=self._timeout) as client:
                    response = await client.post(url, headers=headers, json=payload)
                    response.raise_for_status()
                    return response.json()
            except httpx.HTTPStatusError as exc:
                status = exc.response.status_code if exc.response else 0
                body = exc.response.text if exc.response else ""
                if status not in self.RETRY_STATUS or last:
                    logger.error("OpenRouter embeddings request failed: %s", body)
                    raise VectorSearchError(
                        f"OpenRouter embeddings HTTP error {status or 'unknown'}"
                    ) from exc
                logger.warning(
                    "OpenRouter embeddings HTTP %s (attempt %s/%s); retrying in %.1fs",
                    status,
                    attempt,
                    self._retry_attempts,
                    delay,
                )
            except httpx.RequestError as exc:
                if last:
                    logger.error("OpenRouter embeddings request failed: %s", exc)
                    raise VectorSearchError(
                        f"OpenRouter embeddings request failed after "
                        f"{self._retry_attempts} attempts: {exc!r}"
                    ) from exc
                logger.warning(
                    "OpenRouter embeddings transport error %r (attempt %s/%s); "
                    "retrying in %.1fs",
                    exc,
                    attempt,
                    self._retry_attempts,
                    delay,
                )
            await asyncio.sleep(delay)
            delay *= 2
        # Unreachable: the final attempt either returns or raises above.
        raise VectorSearchError("OpenRouter embeddings retry loop exhausted")

    async def _request_embeddings_live(self, texts: list[str]) -> list[list[float]]:
        """The real provider call. Reached only when the cassette contract
        says so — an unset LLM_CASSETTE_DIR, or an explicit re-record."""
        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": self._model,
            "input": texts,
        }
        url = f"{self._base_url}/embeddings"
        data = await self._post_with_retry(url, headers, payload)

        try:
            items = data["data"]
        except KeyError as exc:
            raise VectorSearchError(
                f"OpenRouter embeddings response missing 'data': {data}"
            ) from exc

        # Sort by index to preserve input order; some providers return results
        # out of order under load.
        indexed = sorted(
            ((item.get("index", 0), item["embedding"]) for item in items),
            key=lambda pair: pair[0],
        )
        embeddings = [embedding for _, embedding in indexed]
        if len(embeddings) != len(texts):
            raise VectorSearchError(
                f"OpenRouter returned {len(embeddings)} embeddings for {len(texts)} texts"
            )
        for embedding in embeddings:
            if len(embedding) != self._dimension:
                raise VectorSearchError(
                    f"Expected dimension {self._dimension}, got {len(embedding)}"
                )
        return embeddings


class SentenceTransformerEmbedder:
    """Local embedding backend using ``sentence-transformers``.

    Defaults to ``Qwen/Qwen3-Embedding-0.6B`` on CUDA if available. The model
    is loaded lazily on first use so import-time overhead stays low.

    Query texts may use an optional prompt prefix; documents are embedded
    without a prefix by default. Both prompts are configurable.
    """

    def __init__(
        self,
        model_name: str = "Qwen/Qwen3-Embedding-0.6B",
        dimension: int = 1024,
        device: str | None = None,
        query_prompt: str | None = None,
        document_prompt: str = "",
        trust_remote_code: bool = True,
        batch_size: int = 32,
        normalize_embeddings: bool = True,
    ) -> None:
        """Initialize the embedder.

        Args:
            model_name: Hugging Face model identifier or local path.
            dimension: Expected vector dimension.
            device: Torch device (``"cuda"``, ``"cpu"``, etc.). ``None`` picks
                CUDA when available.
            query_prompt: Prompt prefix for ``prompt_type="query"``. ``None``
                uses ``prompt_name="query"`` if the model provides it.
            document_prompt: Prompt prefix for ``prompt_type="document"``.
            trust_remote_code: Passed to ``SentenceTransformer``.
            batch_size: Encoding batch size.
            normalize_embeddings: L2-normalize output vectors.
        """
        self._model_name = model_name
        self._dimension = dimension
        self._device = device
        self._query_prompt = query_prompt
        self._document_prompt = document_prompt
        self._trust_remote_code = trust_remote_code
        self._batch_size = max(1, batch_size)
        self._normalize_embeddings = normalize_embeddings
        self._model: Any | None = None

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

    async def embed(
        self, texts: list[str], prompt_type: str = "document"
    ) -> list[list[float]]:
        """Embed texts locally, running the model in a worker thread."""
        if not texts:
            return []

        model = self._load_model()
        prompt_kwargs = self._prompt_kwargs_for(prompt_type)

        def _encode() -> Any:
            kwargs: dict[str, Any] = {
                "batch_size": self._batch_size,
                "normalize_embeddings": self._normalize_embeddings,
                "show_progress_bar": False,
            }
            kwargs.update(prompt_kwargs)
            return model.encode(texts, **kwargs)

        embeddings = await anyio.to_thread.run_sync(_encode)
        return [list(vector) for vector in embeddings]

    def _load_model(self) -> Any:
        """Lazy-load the sentence-transformers model."""
        if self._model is None:
            from sentence_transformers import SentenceTransformer

            device = self._device
            if device is None:
                import torch

                device = "cuda" if torch.cuda.is_available() else "cpu"
            self._model = SentenceTransformer(
                self._model_name,
                device=device,
                trust_remote_code=self._trust_remote_code,
            )
        return self._model

    def _prompt_kwargs_for(self, prompt_type: str) -> dict[str, str]:
        """Return sentence-transformers prompt arguments for ``prompt_type``."""
        if prompt_type == "query":
            if self._query_prompt is not None:
                return {"prompt": self._query_prompt}
            model = self._load_model()
            if hasattr(model, "prompts") and "query" in model.prompts:
                return {"prompt_name": "query"}
            return {}
        if prompt_type == "document" and self._document_prompt:
            return {"prompt": self._document_prompt}
        return {}