embeddings.py

"""Pluggable text embeddings.

Three backends, selected by ``EMBED_BACKEND``:

* ``local``  — sentence-transformers (free, no API key); the default for deploy.
* ``gemini`` — Gemini embedding API (tiny install, needs ``GEMINI_API_KEY``).
* ``hash``   — deterministic, dependency-free hashing embedder. Low quality but
  lets the whole pipeline run offline for smoke tests.

All backends return L2-normalized ``float32`` vectors, so cosine similarity is a
plain dot product.
"""

from __future__ import annotations

import hashlib
import re
from functools import lru_cache
from typing import List

import numpy as np

import config

_TOKEN_RE = re.compile(r"[a-z0-9]+")


def _normalize(mat: np.ndarray) -> np.ndarray:
    mat = np.asarray(mat, dtype=np.float32)
    if mat.ndim == 1:
        mat = mat[None, :]
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return mat / norms


# --------------------------------------------------------------------------- #
# local (sentence-transformers)
# --------------------------------------------------------------------------- #
@lru_cache(maxsize=1)
def _local_model():
    from sentence_transformers import SentenceTransformer

    return SentenceTransformer(config.EMBED_MODEL_LOCAL)


def _embed_local(texts: List[str]) -> np.ndarray:
    model = _local_model()
    vecs = model.encode(texts, normalize_embeddings=True, show_progress_bar=False)
    return np.asarray(vecs, dtype=np.float32)


# --------------------------------------------------------------------------- #
# gemini
# --------------------------------------------------------------------------- #
def _embed_gemini(texts: List[str]) -> np.ndarray:
    from google import genai

    client = genai.Client(api_key=config.GEMINI_API_KEY)
    out = []
    # The API accepts batches; keep them modest to stay within limits.
    for i in range(0, len(texts), 64):
        batch = texts[i : i + 64]
        resp = client.models.embed_content(model=config.EMBED_MODEL_GEMINI, contents=batch)
        out.extend([e.values for e in resp.embeddings])
    return _normalize(np.asarray(out, dtype=np.float32))


# --------------------------------------------------------------------------- #
# hash (offline fallback)
# --------------------------------------------------------------------------- #
def _embed_hash(texts: List[str]) -> np.ndarray:
    dim = config.EMBED_DIM_HASH
    out = np.zeros((len(texts), dim), dtype=np.float32)
    for row, text in enumerate(texts):
        for tok in _TOKEN_RE.findall(text.lower()):
            h = int(hashlib.md5(tok.encode("utf-8")).hexdigest(), 16)
            out[row, h % dim] += 1.0
            # a second hashed slot reduces collisions
            out[row, (h // dim) % dim] += 0.5
    return _normalize(out)


_BACKENDS = {"local": _embed_local, "gemini": _embed_gemini, "hash": _embed_hash}


def _resolve_backend() -> str:
    backend = config.EMBED_BACKEND
    if backend == "local":
        try:
            import sentence_transformers  # noqa: F401
        except Exception:
            print("[embeddings] sentence-transformers not installed; falling back to 'hash'.")
            return "hash"
    if backend == "gemini" and not config.GEMINI_API_KEY:
        print("[embeddings] GEMINI_API_KEY missing; falling back to 'hash'.")
        return "hash"
    return backend


def active_backend() -> str:
    return _resolve_backend()


def embed_texts(texts: List[str]) -> np.ndarray:
    if not texts:
        return np.zeros((0, config.EMBED_DIM_HASH), dtype=np.float32)
    return _BACKENDS[_resolve_backend()](list(texts))


def embed_query(text: str) -> np.ndarray:
    return embed_texts([text])[0]