Spaces:
Sleeping
Sleeping
File size: 3,275 Bytes
31c360c | 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 | from __future__ import annotations
from time import sleep
from typing import Iterable
from ..config import config
MAX_EMBEDDING_ATTEMPTS = 3
class EmbeddingError(RuntimeError):
pass
class OpenRouterEmbeddingClient:
def __init__(self) -> None:
self._client = None
@property
def client(self):
if self._client is None:
if not config.processing.EMBEDDING_API_KEY:
raise EmbeddingError("OPEN_ROUTER_API_KEY is not configured for embeddings")
from openai import OpenAI
self._client = OpenAI(
api_key=config.processing.EMBEDDING_API_KEY,
base_url=config.processing.EMBEDDING_BASE_URL,
)
return self._client
def embed_query(self, text: str) -> list[float]:
return self.embed_documents([text])[0]
def embed_documents(self, texts: Iterable[str]) -> list[list[float]]:
clean_texts = [(text or " ").strip() or " " for text in texts]
if not clean_texts:
return []
last_error: Exception | None = None
for attempt in range(1, MAX_EMBEDDING_ATTEMPTS + 1):
try:
response = self.client.embeddings.create(
model=config.processing.EMBEDDING_MODEL,
input=clean_texts,
)
embeddings = self._extract_embeddings(response)
self._validate_embeddings(embeddings, expected_count=len(clean_texts))
return embeddings
except Exception as exc:
last_error = exc
if attempt == MAX_EMBEDDING_ATTEMPTS or not self._is_retryable(exc):
break
sleep(min(2 ** (attempt - 1), 8))
raise EmbeddingError(f"Failed to generate OpenRouter embeddings: {last_error}") from last_error
@staticmethod
def _extract_embeddings(response) -> list[list[float]]:
data = list(getattr(response, "data", []) or [])
data.sort(key=lambda item: getattr(item, "index", 0))
return [list(getattr(item, "embedding", []) or []) for item in data]
def _validate_embeddings(self, embeddings: list[list[float]], expected_count: int) -> None:
if len(embeddings) != expected_count:
raise EmbeddingError(
f"Embedding response count mismatch: expected {expected_count}, got {len(embeddings)}"
)
for idx, embedding in enumerate(embeddings):
if len(embedding) != config.processing.EMBEDDING_DIMENSIONS:
raise EmbeddingError(
f"Embedding {idx} has dimension {len(embedding)}; "
f"expected {config.processing.EMBEDDING_DIMENSIONS}"
)
@staticmethod
def _is_retryable(error: Exception) -> bool:
status_code = getattr(error, "status_code", None)
if status_code in {408, 409, 425, 429, 500, 502, 503, 504}:
return True
text = str(error).lower()
return any(
signal in text
for signal in [
"rate limit",
"timeout",
"temporarily unavailable",
"connection",
"server error",
]
)
|