vecshift 0.1.0__py3-none-any.whl

vecshift/embeddings/providers.py

@@ -0,0 +1,244 @@
+"""Embedding providers that satisfy :class:`vecshift.core.contracts.EmbeddingProvider`."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import math
+import os
+import random
+import re
+import time
+from collections.abc import Sequence
+from typing import Any, Literal
+
+import httpx
+
+from vecshift.core.fingerprint import EmbeddingFingerprint
+from vecshift.embeddings.spec import ModelSpec
+
+EmbedMode = Literal["document", "query"]
+
+RETRY_STATUSES = {408, 409, 429, 500, 502, 503, 504}
+MAX_ATTEMPTS = 6
+CONCURRENCY = 4
+
+_STOPWORDS = (
+    "a an and are as at be by for from has have in is it its of on or that the this to was "
+    "were will with which you your can not but if then than so such these those there their "
+    "they them we our may also into when where how what"
+)
+
+
+class EmbeddingError(Exception):
+    """A provider call failed for good. The message never contains credentials."""
+
+
+def api_key(spec: ModelSpec, error: type[Exception]) -> str | None:
+    """The API key for ``spec``, after checking it can be sent safely."""
+    key = os.environ.get(spec.key_env) if spec.key_env else None
+    if spec.provider == "openai" and not key:
+        raise error(
+            f"{spec.name}: set {spec.key_env} to your OpenAI API key "
+            "(or key_env=NAME in the spec to read it from another variable)"
+        )
+    if key and spec.sends_plaintext:
+        raise error(
+            f"{spec.name}: refusing to send an API key over unencrypted http:// to a remote "
+            "host. Use an https:// URL."
+        )
+    return key
+
+
+class _Base:
+    def __init__(self, spec: ModelSpec) -> None:
+        self.spec = spec
+        self.tokens = 0
+        """Tokens billed so far, as reported by the provider."""
+        self.tokens_estimated = False
+        """Whether any token count had to be estimated."""
+        self.request_seconds: list[float] = []
+        self._dims: int | None = spec.dimensions
+
+    @property
+    def dimensions(self) -> int | None:
+        return self._dims
+
+    @property
+    def fingerprint(self) -> EmbeddingFingerprint:
+        if not self._dims:
+            raise EmbeddingError(f"{self.spec.name}: dimensions are unknown until the first call")
+        return EmbeddingFingerprint(
+            provider=self.spec.provider,
+            model=self.spec.model,
+            dimensions=self._dims,
+            task=None,
+            prefix=(self.spec.query_prefix + "|" + self.spec.doc_prefix) or None,
+            normalized=True,
+        )
+
+    def _prefixed(self, texts: Sequence[str], mode: EmbedMode) -> list[str]:
+        prefix = self.spec.query_prefix if mode == "query" else self.spec.doc_prefix
+        return [prefix + t for t in texts] if prefix else list(texts)
+
+    async def aclose(self) -> None:
+        return None
+
+
+class HashingEmbedder(_Base):
+    """A lexical baseline: hashed bag of words, L2-normalized.
+
+    Free, offline, and deterministic. It captures word overlap only, so it's a floor
+    that real models should beat, and it makes a first run cost nothing.
+    """
+
+    _WORD = re.compile(r"\w+", re.UNICODE)
+    _STOP = frozenset(_STOPWORDS.split())
+
+    def __init__(self, spec: ModelSpec) -> None:
+        super().__init__(spec)
+        if not spec.dimensions:
+            raise EmbeddingError(f"{spec.name}: hash models need a size, e.g. hash/512")
+        self._dims = spec.dimensions
+
+    def _vector(self, text: str) -> list[float]:
+        dims = self._dims or 1
+        counts: dict[int, float] = {}
+        words = self._WORD.findall(text.lower())
+        self.tokens += len(words)
+        words = [w for w in words if w not in self._STOP] or words
+        for word in words:
+            h = hashlib.blake2b(word.encode(), digest_size=8).digest()
+            index = int.from_bytes(h[:4], "little") % dims
+            sign = 1.0 if h[4] & 1 else -1.0
+            counts[index] = counts.get(index, 0.0) + sign
+        vec = [0.0] * dims
+        for i, c in counts.items():
+            vec[i] = math.copysign(1 + math.log(abs(c)), c) if c else 0.0
+        norm = math.sqrt(sum(v * v for v in vec)) or 1.0
+        return [v / norm for v in vec]
+
+    async def embed(self, texts: Sequence[str], mode: EmbedMode = "document") -> list[list[float]]:
+        start = time.perf_counter()
+        out = [self._vector(t) for t in self._prefixed(texts, mode)]
+        self.request_seconds.append(time.perf_counter() - start)
+        return out
+
+
+class OpenAICompatEmbedder(_Base):
+    """Any server that speaks OpenAI's ``POST /embeddings``: OpenAI itself, Ollama, vLLM,
+    TEI, LM Studio, LiteLLM, and most hosted inference platforms."""
+
+    def __init__(
+        self,
+        spec: ModelSpec,
+        client: httpx.AsyncClient | None = None,
+        sleep: Any = asyncio.sleep,
+    ) -> None:
+        super().__init__(spec)
+        if not spec.url:
+            raise EmbeddingError(f"{spec.name}: no URL configured")
+        key = api_key(spec, EmbeddingError)
+        headers = {"Authorization": f"Bearer {key}"} if key else {}
+        # TLS certificates are always verified, and redirects are never followed, so a
+        # key can't be carried to a host other than the one in the spec.
+        self._client = client or httpx.AsyncClient(
+            timeout=httpx.Timeout(120, connect=15),
+            headers=headers,
+            verify=True,
+            follow_redirects=False,
+        )
+        if client is not None and key:
+            self._client.headers.update(headers)
+        self._owns_client = client is None
+        self._sleep = sleep
+        self._limit = asyncio.Semaphore(CONCURRENCY)
+
+    async def aclose(self) -> None:
+        if self._owns_client:
+            await self._client.aclose()
+
+    def _describe(self, response: httpx.Response) -> str:
+        try:
+            body: Any = response.json()
+            message = body.get("error", {})
+            message = message.get("message") if isinstance(message, dict) else message
+            message = message or body.get("detail") or response.text
+        except ValueError:
+            message = response.text
+        message = " ".join(str(message).split())[:300]
+        return f"{self.spec.name}: HTTP {response.status_code}: {message}"
+
+    async def _post(self, batch: list[str]) -> list[list[float]]:
+        payload: dict[str, Any] = {"model": self.spec.model, "input": batch}
+        if self.spec.dimensions:
+            payload["dimensions"] = self.spec.dimensions
+        delay = 1.0
+        for attempt in range(1, MAX_ATTEMPTS + 1):
+            start = time.perf_counter()
+            try:
+                response = await self._client.post(f"{self.spec.url}/embeddings", json=payload)
+            except httpx.TransportError as exc:
+                if attempt == MAX_ATTEMPTS:
+                    raise EmbeddingError(
+                        f"{self.spec.name}: couldn't reach {self.spec.url} ({type(exc).__name__})"
+                    ) from exc
+                await self._sleep(delay + random.uniform(0, delay / 2))
+                delay = min(delay * 2, 30)
+                continue
+            elapsed = time.perf_counter() - start
+            if response.status_code in RETRY_STATUSES and attempt < MAX_ATTEMPTS:
+                retry_after = response.headers.get("retry-after", "")
+                wait = float(retry_after) if retry_after.replace(".", "", 1).isdigit() else delay
+                await self._sleep(min(wait, 60) + random.uniform(0, 0.5))
+                delay = min(delay * 2, 30)
+                continue
+            if response.status_code >= 400:
+                raise EmbeddingError(self._describe(response))
+            self.request_seconds.append(elapsed)
+            return self._parse(response, batch)
+        raise EmbeddingError(
+            f"{self.spec.name}: gave up after {MAX_ATTEMPTS} attempts"
+        )  # pragma: no cover
+
+    def _parse(self, response: httpx.Response, batch: list[str]) -> list[list[float]]:
+        try:
+            body = response.json()
+            items = sorted(body["data"], key=lambda d: d.get("index", 0))
+            vectors = [list(map(float, item["embedding"])) for item in items]
+        except (ValueError, KeyError, TypeError) as exc:
+            raise EmbeddingError(f"{self.spec.name}: unexpected response shape") from exc
+        if len(vectors) != len(batch):
+            raise EmbeddingError(
+                f"{self.spec.name}: sent {len(batch)} texts, got {len(vectors)} vectors"
+            )
+        usage = body.get("usage") or {}
+        tokens = usage.get("prompt_tokens") or usage.get("total_tokens")
+        if isinstance(tokens, int):
+            self.tokens += tokens
+        else:
+            self.tokens += sum(max(1, len(t) // 4) for t in batch)
+            self.tokens_estimated = True
+        if vectors:
+            self._dims = len(vectors[0])
+        return vectors
+
+    async def _batch(self, batch: list[str]) -> list[list[float]]:
+        async with self._limit:
+            return await self._post(batch)
+
+    async def embed(self, texts: Sequence[str], mode: EmbedMode = "document") -> list[list[float]]:
+        prepared = self._prefixed(texts, mode)
+        size = self.spec.batch_size
+        batches = [prepared[i : i + size] for i in range(0, len(prepared), size)]
+        results = await asyncio.gather(*(self._batch(b) for b in batches))
+        return [vec for batch in results for vec in batch]
+
+
+Embedder = HashingEmbedder | OpenAICompatEmbedder
+
+
+def create_embedder(spec: ModelSpec, client: httpx.AsyncClient | None = None) -> Embedder:
+    if spec.provider == "hash":
+        return HashingEmbedder(spec)
+    return OpenAICompatEmbedder(spec, client=client)

vecshift/embeddings/spec.py

@@ -0,0 +1,240 @@
+"""Model specs: one string that says which model to call and how.
+
+    provider/model[,key=value,...]
+
+Examples:
+
+    openai/text-embedding-3-small
+    openai/text-embedding-3-large,dims=256
+    ollama/nomic-embed-text
+    compat/BAAI/bge-m3,url=http://localhost:8080/v1,price=0
+    hash/512
+
+Options: ``url``, ``dims``, ``price`` (USD per million tokens), ``key_env`` (the
+environment variable holding the API key), ``query_prefix``, ``doc_prefix``, and
+``batch`` (texts per request).
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from urllib.parse import urlparse
+
+# Hosts that keep traffic on this machine. ("0.0.0.0" is a destination here, not a bind.)
+LOCAL_HOSTS = {"localhost", "127.0.0.1", "::1", "0.0.0.0", "host.docker.internal"}  # noqa: S104
+
+
+class SpecError(ValueError):
+    """A model spec couldn't be understood."""
+
+
+@dataclass(frozen=True, slots=True)
+class Preset:
+    url: str | None
+    key_env: str | None
+    note: str
+
+
+PRESETS = {
+    "openai": Preset("https://api.openai.com/v1", "OPENAI_API_KEY", "OpenAI"),
+    "ollama": Preset("http://localhost:11434/v1", None, "Ollama on this machine"),
+    # No default key: a key is only sent to a compat server when key_env= names one, so
+    # pointing a spec at a new URL never hands it a key meant for another server.
+    "compat": Preset(None, None, "any OpenAI-compatible server"),
+    "hash": Preset(None, None, "built-in hashing baseline, free and offline"),
+}
+
+# USD per million input tokens for OpenAI's embedding models. Prices change; pass
+# price= in the spec to override.
+KNOWN_PRICES = {
+    "text-embedding-3-small": 0.02,
+    "text-embedding-3-large": 0.13,
+    "text-embedding-ada-002": 0.10,
+}
+
+# Native output size of well-known models, so plans know the target column size without
+# calling the API. Matched on the model's base name, without any path prefix or tag.
+KNOWN_DIMENSIONS = {
+    "text-embedding-3-small": 1536,
+    "text-embedding-3-large": 3072,
+    "text-embedding-ada-002": 1536,
+    "nomic-embed-text": 768,
+    "mxbai-embed-large": 1024,
+    "bge-m3": 1024,
+    "all-minilm": 384,
+}
+
+# Models known to accept a smaller ``dimensions`` request (Matryoshka truncation).
+SUPPORTS_DIMENSIONS = {"text-embedding-3-small", "text-embedding-3-large"}
+
+_SEARCH_INSTRUCTION = "Represent this sentence for searching relevant passages: "
+
+# (pattern, query prefix, document prefix) for models trained with prompts. Using the
+# wrong prefix quietly costs these models a lot of quality.
+KNOWN_PREFIXES: tuple[tuple[re.Pattern[str], str, str], ...] = (
+    (re.compile(r"nomic-embed-text"), "search_query: ", "search_document: "),
+    (re.compile(r"(^|[/_-])e5-"), "query: ", "passage: "),
+    (re.compile(r"bge-(small|base|large)-(en|zh)"), _SEARCH_INSTRUCTION, ""),
+    (re.compile(r"mxbai-embed-large"), _SEARCH_INSTRUCTION, ""),
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    raw: str
+    provider: str
+    model: str
+    url: str | None = None
+    dimensions: int | None = None
+    price: float | None = None
+    """USD per million tokens, or ``None`` if unknown."""
+    key_env: str | None = None
+    query_prefix: str = ""
+    doc_prefix: str = ""
+    batch_size: int = 64
+
+    @property
+    def name(self) -> str:
+        """Short display name, such as ``openai/text-embedding-3-large (256d)``."""
+        base = f"{self.provider}/{self.model}"
+        explicit = self.dimensions and self.provider != "hash"
+        return f"{base} ({self.dimensions}d)" if explicit else base
+
+    @property
+    def is_local(self) -> bool:
+        """Whether text stays on this machine."""
+        if self.provider == "hash":
+            return True
+        host = urlparse(self.url or "").hostname or ""
+        return host in LOCAL_HOSTS
+
+    @property
+    def safe_raw(self) -> str:
+        """The spec as written, minus any URL query string, which can carry secrets."""
+        if not self.url or "?" not in self.url:
+            return self.raw
+        return self.raw.replace(self.url, self.url.split("?", 1)[0] + "?…")
+
+    @property
+    def sends_plaintext(self) -> bool:
+        """Whether requests would cross the network unencrypted."""
+        return (self.url or "").startswith("http://") and not self.is_local
+
+    @property
+    def base_model(self) -> str:
+        """``BAAI/bge-m3`` → ``bge-m3``; ``nomic-embed-text:latest`` → ``nomic-embed-text``."""
+        return self.model.rsplit("/", 1)[-1].split(":", 1)[0].lower()
+
+    @property
+    def native_dimensions(self) -> int | None:
+        return KNOWN_DIMENSIONS.get(self.base_model)
+
+    @property
+    def output_dimensions(self) -> int | None:
+        """What this spec will produce, if known without calling the model."""
+        return self.dimensions or self.native_dimensions
+
+    @property
+    def identity(self) -> str:
+        """Everything that changes the vectors this spec produces, for caching."""
+        return "|".join(
+            str(x)
+            for x in (
+                self.provider,
+                self.model,
+                self.url,
+                self.dimensions,
+                self.query_prefix,
+                self.doc_prefix,
+            )
+        )
+
+
+def _number(key: str, value: str, kind: type[int] | type[float]) -> float:
+    try:
+        number = kind(value)
+    except ValueError:
+        raise SpecError(f"{key}= must be a number, got {value!r}") from None
+    if number < 0 or (kind is int and number == 0):
+        raise SpecError(f"{key}= must be positive, got {value!r}")
+    return number
+
+
+def parse_spec(raw: str) -> ModelSpec:
+    """Parse ``provider/model[,key=value,...]``."""
+    head, *options = [part.strip() for part in raw.split(",")]
+    provider, sep, model = head.partition("/")
+    provider = provider.lower()
+    if not sep or not model:
+        raise SpecError(
+            f"Expected provider/model, got {raw!r}. For example: openai/text-embedding-3-small"
+        )
+    if provider not in PRESETS:
+        raise SpecError(f"Unknown provider {provider!r}. Use one of: {', '.join(sorted(PRESETS))}.")
+    preset = PRESETS[provider]
+    values: dict[str, str] = {}
+    for option in options:
+        key, eq, value = option.partition("=")
+        key = key.strip().lower()
+        if not eq or key not in {
+            "url",
+            "dims",
+            "price",
+            "key_env",
+            "query_prefix",
+            "doc_prefix",
+            "batch",
+        }:
+            raise SpecError(f"Unknown option {option!r} in {raw!r}")
+        values[key] = value
+
+    url = values.get("url", preset.url)
+    if provider == "compat" and not url:
+        raise SpecError(
+            f"compat models need a url=, for example {raw},url=http://localhost:8080/v1"
+        )
+    if url:
+        url = url.rstrip("/")
+        parsed = urlparse(url)
+        if parsed.scheme not in {"http", "https"}:
+            raise SpecError(f"url= must start with http:// or https://, got {url!r}")
+        if parsed.username or parsed.password:
+            raise SpecError(
+                "Don't put credentials in url=: they would end up in reports. Put the key in "
+                "an environment variable and name it with key_env=."
+            )
+
+    dimensions = int(_number("dims", values["dims"], int)) if "dims" in values else None
+    if provider == "hash":
+        dimensions = int(_number("hash dimensions", model, int))
+
+    price: float | None
+    if "price" in values:
+        price = float(_number("price", values["price"], float))
+    elif provider == "openai":
+        price = KNOWN_PRICES.get(model)
+    else:
+        price = 0.0 if provider in {"hash", "ollama"} else None
+
+    query_prefix, doc_prefix = "", ""
+    for pattern, qp, dp in KNOWN_PREFIXES:
+        if pattern.search(model.lower()):
+            query_prefix, doc_prefix = qp, dp
+            break
+    query_prefix = values.get("query_prefix", query_prefix)
+    doc_prefix = values.get("doc_prefix", doc_prefix)
+
+    batch = int(_number("batch", values["batch"], int)) if "batch" in values else 64
+    return ModelSpec(
+        raw=raw,
+        provider=provider,
+        model=model,
+        url=url,
+        dimensions=dimensions,
+        price=price,
+        key_env=values.get("key_env", preset.key_env),
+        query_prefix=query_prefix,
+        doc_prefix=doc_prefix,
+        batch_size=batch,
+    )

vecshift/eval/__init__.py

@@ -0,0 +1,20 @@
+"""Compare a migration's old and new vectors before cutover: quality and latency."""
+
+from vecshift.eval.metrics import Latency, Scores
+from vecshift.eval.queries import EvalQuery, proxy_queries, script
+from vecshift.eval.runner import EvalReport, Gates, SideReport, SweepPoint, Verdict, decide, run
+
+__all__ = [
+    "EvalQuery",
+    "EvalReport",
+    "Gates",
+    "Latency",
+    "Scores",
+    "SideReport",
+    "SweepPoint",
+    "Verdict",
+    "decide",
+    "proxy_queries",
+    "run",
+    "script",
+]