sync/embeddings_client.py

"""Client for the university's OpenAI-compatible Qwen3-Embedding-4B endpoint.

The endpoint runs a 4B model, so a single request carrying an entire weekly
delta (or a first-run full anthology) easily exceeds the read timeout. We
therefore chunk the input and retry the transient timeouts the openai SDK
does not retry on its own (it retries 429/5xx but treats read timeouts as
non-retryable connection errors).
"""
import logging
import time

import numpy as np
from openai import APIConnectionError, OpenAI, OpenAIError
from tqdm import tqdm

logger = logging.getLogger(__name__)

# Max texts per embedding request. Small enough that each call to the 4B
# model finishes well inside the read timeout, even on a slow week.
_BATCH_SIZE = 32
# Number of attempts per chunk before giving up (and propagating before any
# writes in run_sync, per its "fail before writes" contract).
_MAX_RETRIES = 3
_RETRY_BACKOFF_SECONDS = 2.0


class EmbeddingsClient:
    def __init__(self, base_url: str, api_key: str, model: str = "qwen3-embedding-4b",
                 timeout: float = 60.0):
        # The hosted qwen3-embedding-4b does NOT support matryoshka representation,
        # so we must not send `dimensions` — sending it 400s. The model returns its
        # fixed native dimensionality; the index dim is derived from the returned
        # vectors in run_sync rather than assumed here.
        self.model = model
        self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout)

    def embed_batch(self, texts: list[str]) -> list[np.ndarray]:
        if not texts:
            return []
        total_chunks = (len(texts) + _BATCH_SIZE - 1) // _BATCH_SIZE
        vectors: list[np.ndarray] = []
        starts = range(0, len(texts), _BATCH_SIZE)
        for start in tqdm(starts, total=total_chunks, desc="embedding", unit="batch"):
            chunk = texts[start:start + _BATCH_SIZE]
            vectors.extend(self._embed_chunk(chunk))
        return vectors

    def _embed_chunk(self, chunk: list[str]) -> list[np.ndarray]:
        last_error: Exception | None = None
        for attempt in range(1, _MAX_RETRIES + 1):
            try:
                response = self._client.embeddings.create(
                    model=self.model, input=chunk
                )
            except APIConnectionError as e:
                # Read timeouts (APITimeoutError is a subclass) and other
                # transient connection errors: the SDK won't retry, so we do.
                last_error = e
                if attempt == _MAX_RETRIES:
                    break
                logger.warning(
                    "embedding request failed (attempt %d/%d), retrying in %.1fs: %s",
                    attempt, _MAX_RETRIES, _RETRY_BACKOFF_SECONDS, e,
                )
                time.sleep(_RETRY_BACKOFF_SECONDS)
                continue
            except OpenAIError as e:
                raise RuntimeError(f"Embedding API request failed: {e}") from e
            data = sorted(response.data, key=lambda d: d.index)
            return [np.array(d.embedding, dtype=np.float32) for d in data]
        raise RuntimeError(
            f"Embedding API request failed after {_MAX_RETRIES} attempts: {last_error}"
        ) from last_error