orchestrata_ultimate.semantic_graph.embedder_factory

Reference › orchestrata_ultimate › semantic_graph › embedder_factory

Selects the embedding backend from the environment, once, where three scripts used to hardcode the same OpenRouter model and a 1536 dimension next to a direct OpenRouterEmbedder import. The variable is SEMANTIC_EMBED_BACKEND. The value local selects Qwen3-Embedding-0.6B on this machine through SentenceTransformerEmbedder, with no key, network or cost, which is what the dev and test profiles use so the suite can rebuild its index for free. The default, openrouter, uses codestral-embed through OpenRouterEmbedder; it stays the default because the production semantic index is 1536-dimensional and switching under it would invalidate every stored vector until a full re-index. Accessors report the choice: embed_backend, embed_model (overridable per deployment), embed_dimension, requires_api_key, chunk_label and vector_index_name, and create_embedder builds the matching embedder. The two backends are not interchangeable against one index, because their vectors differ in dimension and geometry, so whatever wrote the index must be what queries it; this is why the dimension, chunk label and index name come from the same selection. All functions accept an explicit environment mapping. The evaluation and ingestion scripts import it, and the vectors land in the shared memgraph semantic graph. See embedders.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

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

"""Which embedding backend — a choice, where it used to be a constant.

``scripts/semantic_search.py``, ``scripts/ingest_orchestrata_ultimate.py`` and
``scripts/evaluate_semantic_graph_qa.py`` each hardcoded
``EMBED_MODEL = "mistralai/codestral-embed-2505"`` and
``EMBED_DIMENSION = 1536`` next to a direct ``OpenRouterEmbedder`` import.
Three copies of one decision, and the decision was "always pay a provider".

:class:`~orchestrata_ultimate.semantic_graph.embedders.SentenceTransformerEmbedder`
had been sitting beside it the whole time — tested, declared in
``pyproject.toml``, model cached on disk — with no production caller at all.

So the backend is selected here, once, from the environment:

``SEMANTIC_EMBED_BACKEND=local``
    Qwen3-Embedding-0.6B on this machine. No API key, no network, no cost.
    The dev and test profiles use it: an index the suite can rebuild for free
    is the difference between a live-graph test that can be repaired and one
    that waits on somebody topping up a credit balance.

``SEMANTIC_EMBED_BACKEND=openrouter`` (the default)
    codestral-embed via OpenRouter. Left as the default deliberately — prod's
    semantic index is 1536-dimensional and switching the backend under it
    would invalidate every stored vector until a full re-index finished.

The two are NOT interchangeable against one index: their vectors have
different dimensions and different geometry. Whatever wrote the index has to
be what queries it, which is exactly why this returns the dimension too.
"""

from __future__ import annotations

import os
from collections.abc import Mapping

from orchestrata_ultimate.semantic_graph.embedders import (
    Embedder,
    OpenRouterEmbedder,
    SentenceTransformerEmbedder,
)

__all__ = [
    "DEFAULT_BACKEND",
    "LOCAL_CHUNK_LABEL",
    "LOCAL_DEVICE",
    "LOCAL_DIMENSION",
    "LOCAL_MODEL",
    "OPENROUTER_DIMENSION",
    "OPENROUTER_MODEL",
    "chunk_label",
    "create_embedder",
    "embed_backend",
    "embed_dimension",
    "embed_model",
    "vector_index_name",
    "requires_api_key",
]

LOCAL_MODEL = "Qwen/Qwen3-Embedding-0.6B"
LOCAL_CHUNK_LABEL = "QwenChunk"
LOCAL_INDEX_NAME = "qwen_chunk_embedding_idx"
LOCAL_DIMENSION = 1024
#: The local backend runs on the CPU unless told otherwise. Left to choose for
#: itself it grabs the discrete GPU, and a corpus-sized ingest exhausts a 4 GB
#: card partway through — a slower run that finishes beats a faster one that
#: dies at 60%. Override per machine with SEMANTIC_EMBED_DEVICE.
LOCAL_DEVICE = "cpu"
OPENROUTER_MODEL = "mistralai/codestral-embed-2505"
OPENROUTER_CHUNK_LABEL = "CodestralChunk"
OPENROUTER_INDEX_NAME = "codestral_chunk_embedding_idx"
OPENROUTER_DIMENSION = 1536
DEFAULT_BACKEND = "openrouter"

_BACKENDS = ("local", "openrouter")


def _env(env: Mapping[str, str] | None) -> Mapping[str, str]:
    return env if env is not None else os.environ


def embed_backend(env: Mapping[str, str] | None = None) -> str:
    """The selected backend name.

    Raises:
        ValueError: On an unrecognised name. A typo must not silently fall
            back to the paid default and re-introduce the cost this exists to
            remove.
    """
    backend = _env(env).get("SEMANTIC_EMBED_BACKEND", DEFAULT_BACKEND).strip().lower()
    if backend not in _BACKENDS:
        raise ValueError(
            f"SEMANTIC_EMBED_BACKEND={backend!r} is not one of {_BACKENDS}"
        )
    return backend


def embed_model(env: Mapping[str, str] | None = None) -> str:
    """The model id for the selected backend, overridable per deployment."""
    explicit = _env(env).get("SEMANTIC_EMBED_MODEL")
    if explicit:
        return explicit
    return LOCAL_MODEL if embed_backend(env) == "local" else OPENROUTER_MODEL


def embed_dimension(env: Mapping[str, str] | None = None) -> int:
    """The vector width the selected backend produces.

    Callers need this to build the index and to size a stand-in vector, so a
    mismatch surfaces as a configuration error rather than as an empty search
    result nobody can explain.
    """
    explicit = _env(env).get("SEMANTIC_EMBED_DIMENSION")
    if explicit:
        return int(explicit)
    return LOCAL_DIMENSION if embed_backend(env) == "local" else OPENROUTER_DIMENSION


def requires_api_key(env: Mapping[str, str] | None = None) -> bool:
    """Does the selected backend need a provider key at all?

    The local backend does not, and a script that demanded one anyway would
    refuse to run on a machine that never needed an account.
    """
    return embed_backend(env) == "openrouter"


def create_embedder(
    *,
    api_key: str = "",
    batch_size: int | None = None,
    max_tokens_per_batch: int | None = None,
    env: Mapping[str, str] | None = None,
) -> Embedder:
    """The embedder the environment asked for.

    Args:
        api_key: Provider key. Ignored by the local backend, which needs none.
        batch_size: Texts per batch. Each backend keeps its own default when
            this is None — they are not comparable numbers, one bounding an
            HTTP request and the other a forward pass.
        max_tokens_per_batch: Approximate per-request token ceiling. Applies
            to the remote backend only; the local one is not billed or rate
            limited by token count.
        env: Environment to read the selection from. Defaults to the process
            environment.
    """
    environ = _env(env)
    if embed_backend(env) == "local":
        local_kwargs: dict[str, object] = {
            "model_name": embed_model(env),
            "dimension": embed_dimension(env),
            "device": environ.get("SEMANTIC_EMBED_DEVICE", LOCAL_DEVICE),
        }
        if batch_size is not None:
            local_kwargs["batch_size"] = batch_size
        return SentenceTransformerEmbedder(**local_kwargs)  # type: ignore[arg-type]

    remote_kwargs: dict[str, object] = {
        "api_key": api_key,
        "model": embed_model(env),
        "dimension": embed_dimension(env),
        "base_url": environ.get(
            "OPENROUTER_EMBED_BASE_URL", "https://openrouter.ai/api/v1"
        ),
    }
    if batch_size is not None:
        remote_kwargs["batch_size"] = batch_size
    if max_tokens_per_batch is not None:
        remote_kwargs["max_tokens_per_batch"] = max_tokens_per_batch
    return OpenRouterEmbedder(**remote_kwargs)  # type: ignore[arg-type]


def chunk_label(env: Mapping[str, str] | None = None) -> str:
    """The node label the selected backend's vectors are stored under.

    The label is not decoration: it is what says "these vectors are comparable
    to each other". The semantic graph is shared — the knowledge-wiki writes
    into it too — so storing Qwen vectors under ``CodestralChunk`` would put
    two incompatible vector spaces behind one name, and a later codestral
    query would match them and return confident nonsense.

    ``MEMGRAPH_CODESTRAL_CHUNK_LABEL`` still wins when set, for deployments
    pinned to it before the backend became a choice.
    """
    environ = _env(env)
    legacy = environ.get("MEMGRAPH_CODESTRAL_CHUNK_LABEL")
    if legacy:
        return legacy
    explicit = environ.get("SEMANTIC_CHUNK_LABEL")
    if explicit:
        return explicit
    return (
        LOCAL_CHUNK_LABEL if embed_backend(env) == "local" else OPENROUTER_CHUNK_LABEL
    )


def vector_index_name(env: Mapping[str, str] | None = None) -> str:
    """The vector index for the selected backend.

    Separate per backend for the same reason as the label, plus one of its
    own: a vector index is built at a fixed width, so pointing a 1024-wide
    backend at an index created for 1536 is a configuration error rather than
    a slow path.
    """
    environ = _env(env)
    legacy = environ.get("MEMGRAPH_CODESTRAL_INDEX_NAME")
    if legacy:
        return legacy
    explicit = environ.get("SEMANTIC_VECTOR_INDEX")
    if explicit:
        return explicit
    return LOCAL_INDEX_NAME if embed_backend(env) == "local" else OPENROUTER_INDEX_NAME