ai-resume/corpus/embeddings.py

"""The one place a vector is made.

`docs/phase-4-corpus.md` §4.3: *"Record the embedding model and dimension per
collection. A provider change is a governed migration, not a silent corruption.
Provider decision pending Sam — build provider-agnostic behind one interface."*

So: one interface (`Provider`), two implementations, and the provider name,
model and dimension are written to the collection metadata by the pipeline. A
later run with a different model refuses to mix vectors into the same
collection; it reports a governed migration instead.

**local-hashing** is deterministic, offline, and *not semantic*. It hashes word
tokens into a fixed-width vector. It exists so that the pipeline, its tests and
a from-scratch rebuild never depend on a network or on an unmade decision. It is
the default for `make corpus-test` and for the Qdrant memory measurement.

**omniroute** is the production path: an OpenAI-compatible `/v1/embeddings`
call through OmniRoute, which is the only permitted model gateway
(AGENTS.md §4.5). No provider SDK, no direct vendor call, key from the
environment only.
"""

from __future__ import annotations

import hashlib
import json
import math
import time
import urllib.error
import urllib.request


class EmbeddingError(RuntimeError):
    """Embedding failed. The pipeline aborts rather than writing partial vectors."""


class Provider:
    """Interface every embedding provider implements."""

    name = "provider"
    model = "model"
    dim = 0

    def embed(self, texts: list) -> list:
        raise NotImplementedError

    def stats(self) -> dict:
        """Call counters. The DoD proves incremental mode with these numbers."""
        return {"calls": 0, "texts": 0}

    def reset_stats(self) -> None:
        pass

    def describe(self) -> dict:
        return {"provider": self.name, "model": self.model, "dim": self.dim}


class LocalHashingProvider(Provider):
    """Deterministic feature hashing. Offline, fast, not semantic.

    Every token is hashed with sha256 (never Python's salted `hash()`), so the
    same text produces the same vector on every host and every run. Vectors are
    L2-normalised, which suits cosine distance.
    """

    name = "local-hashing"

    def __init__(self, model: str = "local-hashing-v1", dim: int = 256) -> None:
        if dim <= 0:
            raise EmbeddingError(f"local-hashing needs a positive dimension, got {dim}")
        self.model = model
        self.dim = dim
        self._calls = 0
        self._texts = 0

    def embed(self, texts: list) -> list:
        self._calls += 1
        self._texts += len(texts)
        return [self._vector(text) for text in texts]

    def _vector(self, text: str) -> list:
        vector = [0.0] * self.dim
        for token in _tokens(text):
            digest = hashlib.sha256(token.encode("utf-8")).digest()
            index = int.from_bytes(digest[:8], "big") % self.dim
            sign = 1.0 if digest[8] & 1 else -1.0
            vector[index] += sign
        norm = math.sqrt(sum(value * value for value in vector))
        if norm == 0.0:
            # An empty text still needs a unit vector of the right width.
            vector[0] = 1.0
            return vector
        return [value / norm for value in vector]

    def stats(self) -> dict:
        return {"calls": self._calls, "texts": self._texts}

    def reset_stats(self) -> None:
        self._calls = 0
        self._texts = 0


class OmniRouteProvider(Provider):
    """OpenAI-compatible `/v1/embeddings` through OmniRoute.

    Retries on 429 and 5xx with exponential backoff (the risk table in
    `docs/phase-4-corpus.md` calls this out for bulk ingest). The key is never
    logged, never put in an error message, and never read from this repository.
    """

    name = "omniroute"

    def __init__(
        self,
        api_url: str,
        api_key: str,
        model: str,
        dim: int,
        batch_size: int = 32,
        timeout: float = 30.0,
        attempts: int = 4,
        backoff: float = 1.5,
    ) -> None:
        if not api_key:
            raise EmbeddingError(
                "OmniRoute needs OMNIROUTE_API_KEY in the environment "
                "(~/.config/environment.d/10-secrets.conf)"
            )
        self.api_url = api_url.rstrip("/")
        self._api_key = api_key
        self.model = model
        self.dim = dim
        self.batch_size = max(1, batch_size)
        self.timeout = timeout
        self.attempts = max(1, attempts)
        self.backoff = backoff
        self._calls = 0
        self._texts = 0

    def embed(self, texts: list) -> list:
        out = []
        for start in range(0, len(texts), self.batch_size):
            batch = texts[start : start + self.batch_size]
            out.extend(self._embed_batch(batch))
        return out

    def _embed_batch(self, batch: list) -> list:
        body = json.dumps({"model": self.model, "input": batch}).encode("utf-8")
        last_error = None
        for attempt in range(self.attempts):
            request = urllib.request.Request(
                f"{self.api_url}/v1/embeddings",
                data=body,
                method="POST",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self._api_key}",
                },
            )
            self._calls += 1
            try:
                with urllib.request.urlopen(request, timeout=self.timeout) as response:
                    payload = json.loads(response.read().decode("utf-8"))
                vectors = _vectors_from_response(payload, self.model, self.dim)
                self._texts += len(batch)
                return vectors
            except urllib.error.HTTPError as exc:
                last_error = f"HTTP {exc.code} from {self.api_url}/v1/embeddings"
                exc.close()
                # 4xx other than 429 is our fault: do not hammer the gateway.
                if exc.code != 429 and exc.code < 500:
                    raise EmbeddingError(f"{last_error} (not retryable)") from None
            except (urllib.error.URLError, TimeoutError) as exc:
                last_error = f"cannot reach {self.api_url}: {exc.reason}"
            except EmbeddingError:
                raise
            if attempt + 1 < self.attempts:
                time.sleep(self.backoff ** attempt)
        raise EmbeddingError(f"{last_error} after {self.attempts} attempts")

    def stats(self) -> dict:
        return {"calls": self._calls, "texts": self._texts}

    def reset_stats(self) -> None:
        self._calls = 0
        self._texts = 0


def _vectors_from_response(payload: dict, model: str, dim: int) -> list:
    data = payload.get("data")
    if not isinstance(data, list) or not data:
        raise EmbeddingError("embedding response has no 'data' array")
    ordered = sorted(data, key=lambda item: item.get("index", 0))
    vectors = []
    for item in ordered:
        vector = item.get("embedding")
        if not isinstance(vector, list) or not vector:
            raise EmbeddingError("embedding response item has no 'embedding' array")
        if dim and len(vector) != dim:
            raise EmbeddingError(
                f"provider returned dimension {len(vector)} but the collection is "
                f"configured for {dim}; refusing to mix dimensions"
            )
        vectors.append([float(value) for value in vector])
    returned_model = payload.get("model")
    if returned_model and model and returned_model != model:
        raise EmbeddingError(
            f"provider answered as model {returned_model!r}, not {model!r}; "
            "a model change must be a governed migration"
        )
    return vectors


def build_provider(cfg) -> Provider:
    """Construct the provider named by the configuration."""
    if cfg.provider == "local-hashing":
        return LocalHashingProvider(model=cfg.model, dim=cfg.dim)
    if cfg.provider == "omniroute":
        return OmniRouteProvider(
            api_url=cfg.api_url,
            api_key=cfg.api_key,
            model=cfg.model,
            dim=cfg.dim,
            batch_size=cfg.batch_size,
            timeout=cfg.timeout,
        )
    raise EmbeddingError(f"unknown embedding provider {cfg.provider!r}")


def _tokens(text: str) -> list:
    out = []
    current = []
    for char in text.lower():
        if char.isalnum():
            current.append(char)
        elif current:
            out.append("".join(current))
            current = []
    if current:
        out.append("".join(current))
    return out