Chris CoutinhoandClaude Opus 4.8
258ee96f4c
From card 309 (OHR-Bench smoke-test triage): during a backend-pod rollover the
embedding endpoint was briefly unreachable, and openai.APIConnectionError /
ConnectError propagated unretried (the provider only retried 429). Documents
exhausted the 3 in-process retries and were dropped for that scan cycle.
Broaden the provider-level retry to the transient set -- APIConnectionError,
APITimeoutError, 429, and 5xx -- on the existing exponential backoff (2s->60s,
5 attempts), so a few seconds of retry rides through the rollover. Permanent
4xx (auth, bad request) still re-raise immediately. Generalize the shared
_retry helper (retry_on_rate_limit -> retry_on_transient, predicate renamed to
should_retry, accurate log label) with a back-compat alias; Mistral gets 429+5xx
for parity. The production gateway path inherits this via GatewayProvider, which
delegates to the decorated OpenAIProvider methods.
Add astrolabe_vector_ingest_dropped_total{reason}, incremented when a document
exhausts retries, classified (connection|timeout|rate_limit|server|qdrant|other)
by _drop_reason so the embed-drop rate is alertable per cause. Dropped docs are
NOT marked failed, so the next full scan re-picks them (re-queue via scan loop).
Refs: Deck board 12 card 309 (AC #1 no permanently-dropped docs; embed-drop
metric for AC #5).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
318 lines
11 KiB
Python
318 lines
11 KiB
Python
"""Unified OpenAI provider for embeddings and text generation.
|
|
|
|
Supports:
|
|
- OpenAI's standard API
|
|
- GitHub Models API (models.github.ai)
|
|
- Any OpenAI-compatible API via base_url override
|
|
"""
|
|
|
|
import logging
|
|
|
|
from openai import APIConnectionError, APIError, APIStatusError, AsyncOpenAI
|
|
|
|
from ._retry import retry_on_transient
|
|
from .base import Provider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _is_transient(exc: BaseException) -> bool:
|
|
"""Whether an OpenAI APIError is transient and worth retrying.
|
|
|
|
Covers the failures seen dropping documents during a backend-pod rollover
|
|
(card 309): ``APIConnectionError`` / ``APITimeoutError`` (brief gateway
|
|
unreachability) and 429 / 5xx status errors. Permanent 4xx (auth, bad
|
|
request) are NOT retried — they would fail identically every attempt.
|
|
"""
|
|
if isinstance(exc, APIConnectionError): # incl. APITimeoutError
|
|
return True
|
|
if isinstance(exc, APIStatusError):
|
|
return exc.status_code == 429 or exc.status_code >= 500
|
|
return False
|
|
|
|
|
|
# Catch the APIError base (parent of both APIConnectionError and APIStatusError)
|
|
# and let the predicate decide; non-transient errors re-raise immediately.
|
|
_retry_transient = retry_on_transient(
|
|
APIError,
|
|
should_retry=_is_transient,
|
|
provider_name="OpenAI",
|
|
label="transient error",
|
|
)
|
|
|
|
|
|
# Well-known embedding dimensions for OpenAI models
|
|
OPENAI_EMBEDDING_DIMENSIONS: dict[str, int] = {
|
|
"text-embedding-3-small": 1536,
|
|
"text-embedding-3-large": 3072,
|
|
"text-embedding-ada-002": 1536,
|
|
# GitHub Models API uses openai/ prefix
|
|
"openai/text-embedding-3-small": 1536,
|
|
"openai/text-embedding-3-large": 3072,
|
|
}
|
|
|
|
|
|
class OpenAIProvider(Provider):
|
|
"""
|
|
OpenAI provider supporting both embeddings and text generation.
|
|
|
|
Works with:
|
|
- OpenAI's standard API (api.openai.com)
|
|
- GitHub Models API (models.github.ai)
|
|
- Any OpenAI-compatible API (via base_url)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str,
|
|
base_url: str | None = None,
|
|
embedding_model: str | None = None,
|
|
generation_model: str | None = None,
|
|
timeout: float = 120.0,
|
|
):
|
|
"""
|
|
Initialize OpenAI provider.
|
|
|
|
Args:
|
|
api_key: OpenAI API key (or GITHUB_TOKEN for GitHub Models)
|
|
base_url: Base URL override (e.g., "https://models.github.ai/inference")
|
|
embedding_model: Model for embeddings (e.g., "text-embedding-3-small").
|
|
None disables embeddings.
|
|
generation_model: Model for text generation (e.g., "gpt-4o-mini").
|
|
None disables generation.
|
|
timeout: HTTP timeout in seconds (default: 120)
|
|
"""
|
|
self.embedding_model = embedding_model
|
|
self.generation_model = generation_model
|
|
self._dimension: int | None = None
|
|
|
|
# Initialize async client
|
|
self.client = AsyncOpenAI(
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
timeout=timeout,
|
|
)
|
|
|
|
# Try to get known dimension without API call
|
|
if embedding_model and embedding_model in OPENAI_EMBEDDING_DIMENSIONS:
|
|
self._dimension = OPENAI_EMBEDDING_DIMENSIONS[embedding_model]
|
|
|
|
logger.info(
|
|
"Initialized OpenAI provider: base_url=%s "
|
|
"(embedding_model=%s, generation_model=%s, dimension=%s)",
|
|
base_url or "default",
|
|
embedding_model,
|
|
generation_model,
|
|
self._dimension,
|
|
)
|
|
|
|
@property
|
|
def supports_embeddings(self) -> bool:
|
|
"""Whether this provider supports embedding generation."""
|
|
return self.embedding_model is not None
|
|
|
|
@property
|
|
def supports_generation(self) -> bool:
|
|
"""Whether this provider supports text generation."""
|
|
return self.generation_model is not None
|
|
|
|
@_retry_transient
|
|
async def embed(self, text: str) -> list[float]:
|
|
"""
|
|
Generate embedding vector for text.
|
|
|
|
Args:
|
|
text: Input text to embed
|
|
|
|
Returns:
|
|
Vector embedding as list of floats
|
|
|
|
Raises:
|
|
NotImplementedError: If embeddings not enabled (no embedding_model)
|
|
"""
|
|
if not self.supports_embeddings:
|
|
raise NotImplementedError(
|
|
"Embedding not supported - no embedding_model configured"
|
|
)
|
|
|
|
assert self.embedding_model is not None # Type narrowing
|
|
response = await self.client.embeddings.create(
|
|
input=text,
|
|
model=self.embedding_model,
|
|
)
|
|
|
|
embedding = response.data[0].embedding
|
|
|
|
# Update dimension if not set
|
|
if self._dimension is None:
|
|
self._dimension = len(embedding)
|
|
logger.info(
|
|
"Detected embedding dimension: %d for model %s",
|
|
self._dimension,
|
|
self.embedding_model,
|
|
)
|
|
|
|
return embedding
|
|
|
|
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
|
|
"""
|
|
Generate embeddings for multiple texts using OpenAI's batch API.
|
|
|
|
OpenAI supports up to 2048 inputs per request.
|
|
|
|
Args:
|
|
texts: List of texts to embed
|
|
|
|
Returns:
|
|
List of vector embeddings
|
|
|
|
Raises:
|
|
NotImplementedError: If embeddings not enabled (no embedding_model)
|
|
"""
|
|
if not self.supports_embeddings:
|
|
raise NotImplementedError(
|
|
"Embedding not supported - no embedding_model configured"
|
|
)
|
|
|
|
embeddings, _ = await self.embed_batch_with_usage(texts)
|
|
return embeddings
|
|
|
|
async def embed_with_usage(self, text: str) -> tuple[list[float], int]:
|
|
"""Embed one text, reporting the request's token count."""
|
|
embeddings, tokens = await self.embed_batch_with_usage([text])
|
|
if not embeddings:
|
|
raise RuntimeError(
|
|
"OpenAI embeddings API returned no embedding for model "
|
|
f"{self.embedding_model}"
|
|
)
|
|
return embeddings[0], tokens
|
|
|
|
async def embed_batch_with_usage(
|
|
self, texts: list[str]
|
|
) -> tuple[list[list[float]], int]:
|
|
"""Embed multiple texts, summing the API-reported token usage.
|
|
|
|
Returns ``(embeddings, total_tokens)`` where ``total_tokens`` sums
|
|
``response.usage.total_tokens`` across the sub-requests (the unit the
|
|
provider bills on). Used by the usage-metering hooks (Deck #67). Also
|
|
serves the gateway path via :class:`GatewayProvider`.
|
|
"""
|
|
if not self.supports_embeddings:
|
|
raise NotImplementedError(
|
|
"Embedding not supported - no embedding_model configured"
|
|
)
|
|
|
|
if not texts:
|
|
return [], 0
|
|
|
|
# OpenAI supports batches up to 2048, but use smaller batches for safety
|
|
batch_size = 100
|
|
all_embeddings: list[list[float]] = []
|
|
total_tokens = 0
|
|
|
|
for i in range(0, len(texts), batch_size):
|
|
batch = texts[i : i + batch_size]
|
|
|
|
# Use helper method with retry logic for each batch
|
|
batch_embeddings, batch_tokens = await self._embed_batch_request(batch)
|
|
all_embeddings.extend(batch_embeddings)
|
|
total_tokens += batch_tokens
|
|
|
|
# Update dimension if not set
|
|
if self._dimension is None and batch_embeddings:
|
|
self._dimension = len(batch_embeddings[0])
|
|
logger.info(
|
|
"Detected embedding dimension: %d for model %s",
|
|
self._dimension,
|
|
self.embedding_model,
|
|
)
|
|
|
|
return all_embeddings, total_tokens
|
|
|
|
@_retry_transient
|
|
async def _embed_batch_request(
|
|
self, batch: list[str]
|
|
) -> tuple[list[list[float]], int]:
|
|
"""Make a single batch embedding request with retry logic.
|
|
|
|
Returns ``(embeddings, token_count)``; ``token_count`` comes from the
|
|
response's ``usage.total_tokens`` and falls back to a char-based
|
|
estimate if the API omits usage.
|
|
"""
|
|
assert self.embedding_model is not None # Type narrowing
|
|
response = await self.client.embeddings.create(
|
|
input=batch,
|
|
model=self.embedding_model,
|
|
)
|
|
# Sort by index to maintain order
|
|
sorted_data = sorted(response.data, key=lambda x: x.index)
|
|
embeddings = [item.embedding for item in sorted_data]
|
|
|
|
usage = getattr(response, "usage", None)
|
|
total_tokens = getattr(usage, "total_tokens", None) if usage else None
|
|
# Guard on numeric type (not just ``is not None``): a real response
|
|
# gives an int, but test doubles / partial responses can surface a
|
|
# non-numeric attribute — fall back to the estimate there.
|
|
tokens = (
|
|
round(total_tokens)
|
|
if isinstance(total_tokens, (int, float))
|
|
else self._estimate_tokens(batch)
|
|
)
|
|
return embeddings, tokens
|
|
|
|
def get_dimension(self) -> int:
|
|
"""
|
|
Get embedding dimension.
|
|
|
|
Returns:
|
|
Vector dimension for the configured embedding model
|
|
|
|
Raises:
|
|
NotImplementedError: If embeddings not enabled (no embedding_model)
|
|
RuntimeError: If dimension not detected yet (call embed first)
|
|
"""
|
|
if not self.supports_embeddings:
|
|
raise NotImplementedError(
|
|
"Embedding not supported - no embedding_model configured"
|
|
)
|
|
|
|
if self._dimension is None:
|
|
raise RuntimeError(
|
|
f"Embedding dimension not detected yet for model {self.embedding_model}. "
|
|
"Call embed() first or use a known model."
|
|
)
|
|
return self._dimension
|
|
|
|
@_retry_transient
|
|
async def generate(self, prompt: str, max_tokens: int = 500) -> str:
|
|
"""
|
|
Generate text from a prompt.
|
|
|
|
Args:
|
|
prompt: The prompt to generate from
|
|
max_tokens: Maximum tokens to generate
|
|
|
|
Returns:
|
|
Generated text
|
|
|
|
Raises:
|
|
NotImplementedError: If generation not enabled (no generation_model)
|
|
"""
|
|
if not self.supports_generation:
|
|
raise NotImplementedError(
|
|
"Text generation not supported - no generation_model configured"
|
|
)
|
|
|
|
response = await self.client.chat.completions.create(
|
|
model=self.generation_model,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
max_tokens=max_tokens,
|
|
temperature=0.7,
|
|
)
|
|
|
|
return response.choices[0].message.content or ""
|
|
|
|
async def close(self) -> None:
|
|
"""Close HTTP client."""
|
|
await self.client.close()
|