Spaces:
Sleeping
Sleeping
File size: 4,049 Bytes
526950e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """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)
|