""" Embedding providers for the retrieval layer. Three providers, all behind the same ``Embedder`` interface: 1. ``RemoteEmbedder`` -- OpenAI-compatible remote endpoint. This is the **production** path: it hits an HTTP endpoint and returns real dense vectors. Configured via ``EMBEDDING_API_BASE`` / ``EMBEDDING_API_KEY`` / ``EMBEDDING_MODEL``. 2. ``LocalEmbedder`` -- real dense embeddings from a small local sentence-transformers model. This is the **default local** path when no API key is present: genuine semantic vectors (synonyms/paraphrases map together), no API cost, modest RAM. ``paraphrase-MiniLM-L3-v2`` by default (~90 MB). Slower than hashing (~1-2 h for the full 3.15M corpus) but far better quality. 3. ``HashingEmbedder`` -- deterministic, dependency-free fallback. No model weights, no GPU, no API key, no internet. Projects a bag of character n-grams into a fixed-dimension vector via the hashing trick. Not semantic (synonyms are not mapped together), but stable and fast -- handy for CI or a quick offline smoke test. ``get_embedder()`` resolves one of the above: an explicit instance/name wins, then the API key if set, otherwise the local semantic model. """ from __future__ import annotations import hashlib import os from abc import ABC, abstractmethod from typing import Sequence import numpy as np # The vector dimension the hashing fallback projects into. The remote provider # returns its own native dimension; the index is built to whatever the embedder # reports via ``dimension``. HASHING_DIM = 256 class Embedder(ABC): """Abstract embedding provider.""" #: Dimensionality of the vectors this provider emits. DIM: int @abstractmethod def embed(self, texts: Sequence[str]) -> np.ndarray: """Embed a batch of texts, returning an (N, DIM) float32 array.""" raise NotImplementedError def embed_one(self, text: str) -> np.ndarray: """Embed a single text, returning a (DIM,) float32 array.""" return self.embed([text])[0] class HashingEmbedder(Embedder): """ Deterministic, dependency-free embedding via the hashing trick. We tokenize into character 3-grams, hash each one into a bucket, and write +1/-1 (sign hashed from the gram) into a fixed-size vector, then L2-normalize. This is *not* semantically meaningful (synonyms are not mapped together), but it is stable, fast, and shares enough lexical overlap with the query for near-neighbors to cluster around the right topics. Perfect for exercising the ANN pipeline without an API key. """ DIM = HASHING_DIM _NGRAM = 3 _SEED = 0xC0FFEE def __init__(self, dim: int = HASHING_DIM, ngram: int = 3): self.DIM = dim self._ngram = ngram @staticmethod def _grams(text: str, ngram: int = 3) -> list[str]: t = text.lower() return [t[i:i + ngram] for i in range(max(0, len(t) - ngram + 1))] def embed(self, texts: Sequence[str]) -> np.ndarray: out = np.zeros((len(texts), self.DIM), dtype=np.float32) ngram = self._ngram for i, text in enumerate(texts): for gram in self._grams(text or "", ngram): h = int.from_bytes( hashlib.sha1(gram.encode("utf-8")).digest()[:8], "little", ) bucket = h % self.DIM sign = 1.0 if (h >> 63) & 1 else -1.0 out[i, bucket] += sign norm = np.linalg.norm(out[i]) if norm > 0: out[i] /= norm return out class LocalEmbedder(Embedder): """ Real dense embeddings from a local sentence-transformers model. Uses a small CPU model (``paraphrase-MiniLM-L3-v2`` by default, ~90 MB) so it runs on a laptop without an API key and without the Strix Halo OOM history we had loading big models. This is the *production local* path: unlike ``HashingEmbedder`` it produces genuinely semantic vectors (synonyms and paraphrases map close together). Trade-off is speed: embedding the full 3.15M-row corpus takes ~1-2 hours on this box. For a quick demo, index a single ``--category`` subset instead. Any model name accepted by ``SentenceTransformer`` works; set ``EMBEDDING_LOCAL_MODEL`` to override (e.g. ``all-MiniLM-L6-v2`` for better quality at ~3x the time). """ DIM = 384 # paraphrase-MiniLM-L3-v2's dimensionality _DEFAULT_MODEL = "paraphrase-MiniLM-L3-v2" def __init__(self, model: str | None = None): from sentence_transformers import SentenceTransformer self._model_name = model or os.environ.get( "EMBEDDING_LOCAL_MODEL", self._DEFAULT_MODEL ) self._model = SentenceTransformer(self._model_name) # Cache the real dim so the index/registry match the model. DIM = getattr(self._model, "get_embedding_dimension", None) or getattr( self._model, "get_sentence_embedding_dimension", None ) self.DIM = int(DIM()) self._normalize = True def embed(self, texts: Sequence[str]) -> np.ndarray: return self._model.encode( list(texts), normalize_embeddings=self._normalize, show_progress_bar=False, ) class RemoteEmbedder(Embedder): """ OpenAI-compatible remote embedding provider. Configured entirely through the environment: EMBEDDING_API_BASE base URL, e.g. https://api.openai.com/v1 EMBEDDING_API_KEY API key (required) EMBEDDING_MODEL model name, default text-embedding-3-small Any endpoint with the OpenAI embeddings shape works (OpenAI, Together, Groq, a local vLLM/Ollama server, etc.). """ _DEFAULT_MODEL = "text-embedding-3-small" _DEFAULT_DIM = 1536 # text-embedding-3-small's default output dim def __init__( self, api_base: str | None = None, api_key: str | None = None, model: str | None = None, dim: int | None = None, ): import requests # imported lazily so the fallback needs no network stack self._requests = requests self._api_base = (api_base or os.environ.get("EMBEDDING_API_BASE", "")).rstrip("/") self._api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "") self._model = model or os.environ.get("EMBEDDING_MODEL", self._DEFAULT_MODEL) self._dim = dim or int( os.environ.get("EMBEDDING_DIM", self._DEFAULT_DIM) ) if not self._api_base: raise ValueError( "EMBEDDING_API_BASE is not set; configure it (or use the local " "HashingEmbedder fallback)." ) if not self._api_key: raise ValueError( "EMBEDDING_API_KEY is not set; set it, or use the local " "HashingEmbedder fallback." ) @property def DIM(self) -> int: return self._dim def embed(self, texts: Sequence[str]) -> np.ndarray: resp = self._requests.post( f"{self._api_base}/embeddings", headers={"Authorization": f"Bearer {self._api_key}"}, json={"model": self._model, "input": list(texts), "dimensions": self._dim}, timeout=60, ) resp.raise_for_status() payload = resp.json() items = payload["data"] # Some backends don't preserve input order; sort by the returned index. items = sorted(items, key=lambda d: d.get("index", 0)) return np.array([d["embedding"] for d in items], dtype=np.float32) def get_embedder(embedder: str | Embedder | None = None) -> Embedder: """ Resolve an embedder from a name, an instance, or the environment. Resolution order: 1. ``embedder="remote"`` or an ``Embedder`` instance -> use it directly. 2. ``embedder="local"`` -> local sentence-transformers model (semantic). 3. ``embedder="hashing"`` -> deterministic local fallback (no deps). 4. If ``EMBEDDING_API_BASE`` and ``EMBEDDING_API_KEY`` are both set -> remote provider (the production path). 5. Otherwise -> the local semantic model (with a printed note). This means the CLI works out of the box with no configuration: it uses real semantic embeddings from a small local model when no API key is present. """ name = (embedder or "").lower() if isinstance(embedder, str) else None if isinstance(embedder, Embedder): return embedder if name == "hashing": return HashingEmbedder() if name == "remote": return RemoteEmbedder() if name == "local": return LocalEmbedder() api_base = os.environ.get("EMBEDDING_API_BASE") api_key = os.environ.get("EMBEDDING_API_KEY") if api_base and api_key: print(f"[embed] using remote provider: {api_base} (model={os.environ.get('EMBEDDING_MODEL')})") return RemoteEmbedder() print( "[embed] no EMBEDDING_API_BASE/KEY configured -> using local " "sentence-transformers model (semantic). Set EMBEDDING_LOCAL_MODEL to " "override, or embedder='hashing' for a dependency-free fallback." ) return LocalEmbedder()