Spaces:
Sleeping
Sleeping
| """Embedder protocol + OpenAIEmbedder production impl + MockEmbedder for tests. | |
| OpenAIEmbedder reads OPENAI_API_KEY from env. Retries with exponential backoff | |
| on transient errors. Raises MissingApiKeyError, EmbeddingApiError as typed | |
| exceptions. | |
| MockEmbedder produces deterministic hash-derived pseudo-vectors of the correct | |
| dimension. Used in unit tests so no network calls are made. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import os | |
| import struct | |
| import time | |
| from typing import Protocol | |
| class MissingApiKeyError(RuntimeError): | |
| """OPENAI_API_KEY not in env (or other usable source).""" | |
| class EmbeddingApiError(RuntimeError): | |
| """Embedding API returned a persistent error after retries.""" | |
| class Embedder(Protocol): | |
| model_name: str | |
| dimensions: int | |
| def embed_one(self, text: str) -> list[float]: ... | |
| def embed_batch(self, texts: list[str]) -> list[list[float]]: ... | |
| class MockEmbedder: | |
| """Deterministic test embedder. Returns hash-derived pseudo-vectors. | |
| Same input always returns the same vector. Different inputs return | |
| different vectors. Used in unit tests. | |
| """ | |
| def __init__(self, *, model_name: str = "mock", dimensions: int = 3072) -> None: | |
| self.model_name = model_name | |
| self.dimensions = dimensions | |
| def embed_one(self, text: str) -> list[float]: | |
| seed = hashlib.sha256(text.encode("utf-8")).digest() | |
| out: list[float] = [] | |
| i = 0 | |
| while len(out) < self.dimensions: | |
| block = hashlib.sha256(seed + i.to_bytes(4, "big")).digest() | |
| for j in range(0, 32, 4): | |
| if len(out) >= self.dimensions: | |
| break | |
| u = struct.unpack(">I", block[j:j+4])[0] | |
| out.append((u / 2**32) * 2 - 1) | |
| i += 1 | |
| return out | |
| def embed_batch(self, texts: list[str]) -> list[list[float]]: | |
| return [self.embed_one(t) for t in texts] | |
| def _make_client(): | |
| """Factory for the OpenAI client. Patched in tests.""" | |
| import openai | |
| return openai.OpenAI() | |
| def _sleep(seconds: float) -> None: | |
| """Sleep helper. Patched in tests to avoid waiting.""" | |
| time.sleep(seconds) | |
| class OpenAIEmbedder: | |
| """Production embedder backed by OpenAI text-embedding-3-* models. | |
| Reads OPENAI_API_KEY from env. Retries on transient errors with | |
| exponential backoff (3 attempts, 1s/2s/4s). | |
| """ | |
| _MODEL_DIMS = { | |
| "text-embedding-3-large": 3072, | |
| "text-embedding-3-small": 1536, | |
| } | |
| _MAX_RETRIES = 3 | |
| def __init__(self, *, model_name: str = "text-embedding-3-large") -> None: | |
| if not os.environ.get("OPENAI_API_KEY"): | |
| raise MissingApiKeyError( | |
| "OPENAI_API_KEY not set. Either export it in your shell, " | |
| "or run ./bootstrap.{ps1,sh} with --accept-key-file-risk to opt into " | |
| "per-machine key file storage." | |
| ) | |
| self.model_name = model_name | |
| self.dimensions = self._MODEL_DIMS.get(model_name, 3072) | |
| self._client = _make_client() | |
| def _create_with_retry(self, inputs: list[str]) -> list[list[float]]: | |
| import openai | |
| last_exc: Exception | None = None | |
| for attempt in range(self._MAX_RETRIES): | |
| try: | |
| resp = self._client.embeddings.create(model=self.model_name, input=inputs) | |
| return [d.embedding for d in resp.data] | |
| except (openai.RateLimitError, openai.APIError) as e: | |
| last_exc = e | |
| if attempt < self._MAX_RETRIES - 1: | |
| _sleep(2 ** attempt) | |
| continue | |
| break | |
| raise EmbeddingApiError( | |
| f"embedding API failed after {self._MAX_RETRIES} attempts: {last_exc}" | |
| ) from last_exc | |
| def embed_one(self, text: str) -> list[float]: | |
| return self._create_with_retry([text])[0] | |
| def embed_batch(self, texts: list[str]) -> list[list[float]]: | |
| if not texts: | |
| return [] | |
| return self._create_with_retry(texts) | |