embeddings

zocrates.rag.embeddings

Embedding functions for the vector store.

Supports two modes

• "local" → Jina v5 (jina-embeddings-v5-text-nano), runs entirely offline.
• "cloud" → OpenRouter-hosted embedding models (e.g. voyage-4, nemotron-3-embed-1b:free), requires an OpenRouter API key.

EMBEDDING_DIMENSION module-attribute

EMBEDDING_DIMENSION = int(os.getenv('EMBEDDING_DIMENSION', '768'))

EMBEDDING_MODE module-attribute

EMBEDDING_MODE = os.getenv('EMBEDDING_MODE', 'local')

KNOWN_EMBEDDING_DIMENSIONS module-attribute

KNOWN_EMBEDDING_DIMENSIONS = {'baai/bge-m3': 1024, 'nvidia/nemotron-3-embed-1b:free': 1024, 'nvidia/nemotron-3-embed-1b': 1024, 'text-embedding-3-small': 1536, 'text-embedding-3-large': 3072, 'text-embedding-ada-002': 1536, 'jinaai/jina-embeddings-v5-text-nano': 768, 'text-embedding-004': 768, 'qwen/qwen-2.5-coder-32b-instruct': 1024}

LOCAL_EMBEDDING_MODEL module-attribute

LOCAL_EMBEDDING_MODEL = os.getenv('LOCAL_EMBEDDING_MODEL', 'jinaai/jina-embeddings-v5-text-nano')

MAX_EMBEDDING_DIMENSION module-attribute

MAX_EMBEDDING_DIMENSION = 4096

MIN_EMBEDDING_DIMENSION module-attribute

MIN_EMBEDDING_DIMENSION = 64

OpenRouterEmbeddings module-attribute

OpenRouterEmbeddings = OpenRouterEmbedder

_SHARED_LOCAL_DEVICE module-attribute

_SHARED_LOCAL_DEVICE = None

_SHARED_LOCAL_DIMENSION module-attribute

_SHARED_LOCAL_DIMENSION = 768

_SHARED_LOCAL_LOCK module-attribute

_SHARED_LOCAL_LOCK = threading.Lock()

_SHARED_LOCAL_MODEL module-attribute

_SHARED_LOCAL_MODEL = None

logger module-attribute

logger = logging.getLogger(__name__)

EmbedFn

Bases: Protocol

Callable protocol accepted by VectorStore.

LocalEmbedder

LocalEmbedder(model_name=LOCAL_EMBEDDING_MODEL, dimension=EMBEDDING_DIMENSION)

Wraps Hugging Face jina-embeddings-v5-text-nano model with opportunistic GPU acceleration and graceful CPU fallback.

Initialize the local embedding model wrapper.

Parameters:

Name	Type	Description	Default
model_name	str	Hugging Face model identifier string.	LOCAL_EMBEDDING_MODEL
dimension	int	Target embedding output dimension.	EMBEDDING_DIMENSION

Source code in zocrates/rag/embeddings.py

def __init__(
    self,
    model_name: str = LOCAL_EMBEDDING_MODEL,
    dimension: int = EMBEDDING_DIMENSION,
) -> None:
    """Initialize the local embedding model wrapper.

    Args:
        model_name: Hugging Face model identifier string.
        dimension: Target embedding output dimension.
    """
    self.requested_model_name = "jinaai/jina-embeddings-v5-text-nano"
    self._model_name = "jinaai/jina-embeddings-v5-text-nano"
    self._dimension = 768
    self._model = None
    self._device: Any = None
    self.is_fallback = False
    self.validation_error: str | None = None
    self._load_lock = threading.Lock()

dimension property

dimension

Output dimension of the local embedding model.

OpenRouterEmbedder

OpenRouterEmbedder(model_name=None, dimension=None, api_key=None)

Wraps OpenRouter embedding models via litellm.embedding with dynamic dimension auto-detection.

Initialize the OpenRouter cloud embedding model wrapper.

Parameters:

Name	Type	Description	Default
model_name	str | None	Optional OpenRouter embedding model identifier.	None
dimension	int | None	Optional explicit embedding dimension override.	None
api_key	str | None	Optional OpenRouter API key override.	None

Source code in zocrates/rag/embeddings.py

def __init__(
    self,
    model_name: str | None = None,
    dimension: int | None = None,
    api_key: str | None = None,
) -> None:
    """Initialize the OpenRouter cloud embedding model wrapper.

    Args:
        model_name: Optional OpenRouter embedding model identifier.
        dimension: Optional explicit embedding dimension override.
        api_key: Optional OpenRouter API key override.
    """
    from zocrates.config import load_zocrates_config
    cfg = load_zocrates_config()
    self._model_name = model_name or cfg.get("embedding_model", "openrouter/nvidia/nemotron-3-embed-1b:free")
    if not self._model_name.startswith("openrouter/") and not self._model_name.startswith("openai/"):
        self._model_name = f"openrouter/{self._model_name}"
    self._dimension = dimension
    self._api_key = api_key

    if self._dimension is None:
        # Tier 1: Explicit config in zocrates_config.json
        cfg_dim = cfg.get("embedding_dimension")
        if cfg_dim and isinstance(cfg_dim, int):
            self._dimension = cfg_dim

        # Tier 2: Check if ChromaDB already has an existing collection on disk for this model
        if self._dimension is None:
            client = None
            created_client = False
            try:
                import re
                from zocrates.config import get_chroma_dir, resolve_zotero_data_dir

                try:
                    from zocrates.rag.service_manager import get_syncer
                    syncer = get_syncer()
                    if syncer and getattr(syncer, "store", None) and getattr(syncer.store, "_client", None):
                        client = syncer.store._client
                except Exception:
                    pass

                if client is None:
                    zdir = resolve_zotero_data_dir()
                    if zdir:
                        cdir = get_chroma_dir(zdir)
                        if cdir.exists():
                            import chromadb
                            client = chromadb.PersistentClient(path=str(cdir))
                            created_client = True

                if client is not None:
                    sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", self._model_name.strip())
                    sanitized = re.sub(r"_+", "_", sanitized).strip("_")
                    pattern = re.compile(rf"^zotero_library_{re.escape(sanitized)}_(\d+)d$")

                    # In Chroma v1.x, list_collections() returns Collection objects
                    for col_entry in client.list_collections():
                        col_name = (
                            col_entry
                            if isinstance(col_entry, str)
                            else str(getattr(col_entry, "name", col_entry))
                        )
                        m = pattern.match(col_name)
                        if m:
                            self._dimension = int(m.group(1))
                            break

                    # Fallback to collection metadata matching
                    if self._dimension is None:
                        for col_entry in client.list_collections():
                            col_name = (
                                col_entry
                                if isinstance(col_entry, str)
                                else str(getattr(col_entry, "name", col_entry))
                            )
                            try:
                                if hasattr(col_entry, "metadata") and col_entry.metadata:
                                    meta = col_entry.metadata
                                else:
                                    col = client.get_collection(name=col_name)
                                    meta = col.metadata or {}
                                if (
                                    meta.get("model_name") == self._model_name
                                    and meta.get("dimension")
                                ):
                                    self._dimension = int(meta["dimension"])
                                    break
                            except Exception:
                                pass
            except Exception:
                pass
            finally:
                if created_client and client is not None and hasattr(client, "close"):
                    try:
                        client.close()
                    except Exception:
                        pass

        # Tier 3: Static known dimensions registry
        if self._dimension is None:
            norm_name = self._model_name.removeprefix("openrouter/").removeprefix("openai/")
            for k, d in KNOWN_EMBEDDING_DIMENSIONS.items():
                if norm_name == k or norm_name.endswith(k):
                    self._dimension = d
                    break

        # Tier 4: LiteLLM model_cost registry
        if self._dimension is None:
            try:
                import litellm

                norm_name = self._model_name.removeprefix("openrouter/").removeprefix("openai/")
                info = (
                    litellm.model_cost.get(norm_name)
                    or litellm.model_cost.get(self._model_name)
                    or litellm.model_cost.get(f"openai/{norm_name}")
                )
                if info and info.get("output_vector_size"):
                    self._dimension = int(info["output_vector_size"])
            except Exception:
                pass

    if self._dimension is not None and self._dimension > MAX_EMBEDDING_DIMENSION:
        logger.warning(
            "Configured embedding dimension %d exceeds ceiling %d. Clamping to %d.",
            self._dimension,
            MAX_EMBEDDING_DIMENSION,
            MAX_EMBEDDING_DIMENSION,
        )
        self._dimension = MAX_EMBEDDING_DIMENSION

dimension property

dimension

Output dimension of the OpenRouter embedding model.

_detect_best_device

_detect_best_device()

Return the optimal torch device to use for local embedding inference.

Opportunistically probes CUDA and Apple Silicon MPS with a lightweight tensor sanity probe. Gracefully falls back to CPU if probes fail or no GPU accelerator is available.

Returns:

Type	Description
DeviceLike	A torch.device instance for "cuda", "mps", or "cpu".

Source code in zocrates/rag/hardware.py

def _detect_best_device() -> DeviceLike:
    """Return the optimal torch device to use for local embedding inference.

    Opportunistically probes CUDA and Apple Silicon MPS with a lightweight
    tensor sanity probe. Gracefully falls back to CPU if probes fail or no
    GPU accelerator is available.

    Returns:
        A ``torch.device`` instance for ``"cuda"``, ``"mps"``, or ``"cpu"``.
    """
    import torch

    # 1. Probe NVIDIA CUDA
    try:
        cuda_avail = torch.cuda.is_available()
    except Exception as exc:
        logger.warning("torch.cuda.is_available() check failed: %s", exc)
        cuda_avail = False

    if cuda_avail:
        try:
            probe = torch.zeros(1, device="cuda")
            del probe
            logger.info("Opportunistic GPU acceleration enabled for embeddings (cuda).")
            return torch.device("cuda")
        except Exception as exc:
            logger.warning(
                "CUDA detected but tensor sanity probe failed (%s: %s). Falling back.",
                type(exc).__name__,
                exc,
            )

    # 2. Probe Apple Silicon MPS
    try:
        mps_avail = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
    except Exception as exc:
        logger.warning("torch.backends.mps check failed: %s", exc)
        mps_avail = False

    if mps_avail:
        try:
            probe = torch.zeros(1, device="mps")
            del probe
            logger.info("Opportunistic GPU acceleration enabled for embeddings (mps).")
            return torch.device("mps")
        except Exception as exc:
            logger.warning(
                "MPS detected but tensor sanity probe failed (%s: %s). Falling back to CPU.",
                type(exc).__name__,
                exc,
            )

    # 3. Default fallback: CPU
    logger.info("Using CPU hardware execution for local embeddings.")
    return torch.device("cpu")

_safe_to

_safe_to(module, device, dtype)

Recursively moves module parameters/buffers to device/dtype, skipping meta tensors.

Parameters:

Name	Type	Description	Default
module	Any	PyTorch module to move.	required
device	Any	Target torch device or device string.	required
dtype	Any	Target torch data type.	required

Returns:

Type	Description
Any	The moved module.

Raises:

Type	Description
Exception	If moving to a non-CPU device fails due to hardware or driver errors.

Source code in zocrates/rag/embeddings.py

def _safe_to(module: Any, device: Any, dtype: Any) -> Any:
    """Recursively moves module parameters/buffers to device/dtype, skipping meta tensors.

    Args:
        module: PyTorch module to move.
        device: Target torch device or device string.
        dtype: Target torch data type.

    Returns:
        The moved module.

    Raises:
        Exception: If moving to a non-CPU device fails due to hardware or driver errors.
    """
    try:
        return module.to(device=device, dtype=dtype)
    except Exception as exc:
        # If moving to a non-CPU device failed due to hardware/runtime issues (e.g. OOM, driver),
        # raise so that opportunistic GPU callers can gracefully fall back to CPU.
        dev_type = getattr(device, "type", str(device))
        if dev_type != "cpu" and "meta" not in str(exc).lower():
            raise
        for child in module.children():
            _safe_to(child, device, dtype)
        for name, param in getattr(module, "_parameters", {}).items():
            if param is not None and getattr(param, "device", None) and getattr(param.device, "type", "") != "meta":
                module._parameters[name] = param.to(device=device, dtype=dtype)
        for name, buf in getattr(module, "_buffers", {}).items():
            if buf is not None and getattr(buf, "device", None) and getattr(buf.device, "type", "") != "meta":
                module._buffers[name] = buf.to(device=device, dtype=dtype)
        return module

get_embedder

get_embedder(mode=None)

Return the configured embedding function.

Reads the active mode from the EMBEDDING_MODE environment variable, falling back to the persisted config file value, then to "local".

Parameters:

Name	Type	Description	Default
mode	str | None	Override for the embedding mode ("local" or "cloud"). If None, the mode is determined from the environment and the Zocrates config file.	None

Returns:

Type	Description
LocalEmbedder | OpenRouterEmbedder	A LocalEmbedder for CPU-based Jina embeddings, or an
LocalEmbedder | OpenRouterEmbedder	OpenRouterEmbedder for OpenRouter cloud embeddings.

Source code in zocrates/rag/embeddings.py

def get_embedder(
    mode: str | None = None,
) -> LocalEmbedder | OpenRouterEmbedder:
    """Return the configured embedding function.

    Reads the active mode from the ``EMBEDDING_MODE`` environment variable,
    falling back to the persisted config file value, then to ``"local"``.

    Args:
        mode: Override for the embedding mode (``"local"`` or ``"cloud"``).
            If ``None``, the mode is determined from the environment and
            the Zocrates config file.

    Returns:
        A ``LocalEmbedder`` for CPU-based Jina embeddings, or an
        ``OpenRouterEmbedder`` for OpenRouter cloud embeddings.
    """
    from zocrates.config import load_zocrates_config

    cfg = load_zocrates_config()
    active_mode = mode or os.getenv(
        "EMBEDDING_MODE",
        cfg.get("embedding_provider", "local"),
    )

    if active_mode == "cloud":
        logger.info("Using cloud embedding mode (OpenRouter).")
        return OpenRouterEmbedder()

    logger.info("Using local CPU embedding mode with model jinaai/jina-embeddings-v5-text-nano (768d).")
    return LocalEmbedder()