Chris CoutinhoandClaude Opus 4.8
6c99906ed4
Round-1 review on PR #893: - _drop_reason now descends through nested ExceptionGroups to the first leaf (was single-level), so a doubly-wrapped cause isn't mislabelled "other"; added a nested-group test. Commented why both the httpx and openai isinstance branches exist (raw Nextcloud-API errors vs SDK-wrapped variants). - Documented that generate() intentionally shares the broadened transient retry (RAG sampling path), with the worst-case latency note. - Added a docstring note to process_document on how the provider-level retry (5x) layers over the outer loop (3x in-process / 1x procrastinate). - Added test_embed_batch_retries_on_connection_error for the batch path. - Renamed test_retry_reraises_non_rate_limit_immediately -> test_retry_reraises_when_predicate_returns_false (it tests the predicate, not a specific status). SonarCloud: - S5708 (BLOCKER) on the helper's dynamic `except exception_type`: the type is constrained to BaseException/tuple by the signature; suppressed with a justified NOSONAR. - S7503 (async without await) in the embed-retry test: use AsyncMock side_effect instead of a hand-rolled async function. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
322 lines
11 KiB
Python
322 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
|
|
|
|
# Transient retry intentionally covers generation too (RAG sampling path):
|
|
# a pod rollover breaks generation as readily as embedding. Worst case adds
|
|
# ~30s (5 attempts, 2s→60s backoff) to an interactive call hitting a
|
|
# sustained connection issue, which is preferable to a hard failure.
|
|
@_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()
|