File size: 4,578 Bytes
44c2f50 4e6fd2a 44c2f50 4e6fd2a 44c2f50 4e6fd2a 44c2f50 4e6fd2a 44c2f50 | 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 | """Thin wrapper around OpenAI's structured-outputs API with retries + metrics.
We use `client.beta.chat.completions.parse` which:
- Takes a Pydantic model as `response_format`
- Handles JSON Schema translation + strict-mode transformation automatically
- Returns a validated Pydantic instance (`.choices[0].message.parsed`)
"""
from __future__ import annotations
from typing import Any
from openai import APIError, OpenAI, RateLimitError
from pydantic import BaseModel
from tenacity import (
retry,
retry_if_exception_type,
stop_after_attempt,
wait_exponential,
)
from src.utils.config import get_settings
from src.utils.cost_tracker import ExtractionMetrics, Timer
from src.utils.logging import logger
# --- Retry policy -----------------------------------------------------------
# Retry on rate limits + transient API errors. Do NOT retry on validation or
# 4xx auth errors — those won't fix themselves.
_RETRYABLE = (RateLimitError, APIError)
class OpenAIExtractionClient:
"""Wraps the OpenAI client. One instance per app (created at startup)."""
def __init__(self, api_key: str | None = None, model: str | None = None):
s = get_settings()
self._client = OpenAI(
api_key=api_key or s.openai_api_key,
timeout=s.openai_request_timeout,
)
self._default_model = model or s.openai_model
self._max_retries = s.openai_max_retries
# ------------------------------------------------------------------
def parse_structured(
self,
*,
response_format: type[BaseModel],
messages: list[dict[str, Any]],
model: str | None = None,
temperature: float | None = None,
reasoning_effort: str | None = None,
) -> tuple[BaseModel, ExtractionMetrics]:
"""Call OpenAI with structured outputs; return (parsed_object, metrics).
- `response_format` is a Pydantic model class (envelope schema).
- `messages` is standard OpenAI messages list, may include vision content.
- `temperature` is only forwarded when explicitly set. gpt-5-family
models reject any override — they run at a fixed sampling profile —
so leaving this None is the right default. For gpt-4o and earlier
you can pass 0.0 explicitly for deterministic extraction.
- `reasoning_effort` (gpt-5 only) — one of "minimal", "low", "medium",
"high". Structured extraction is a well-formed task that rarely
benefits from long chain-of-thought, so "minimal" typically cuts
both cost and latency by ~10-20x with negligible quality loss.
"""
active_model = model or self._default_model
# Only include optional params when the caller opted in. Passing None to
# OpenAI would 400 on some models; omitting the key lets the server
# pick its own default.
extra: dict[str, Any] = {}
if temperature is not None:
extra["temperature"] = temperature
if reasoning_effort is not None:
extra["reasoning_effort"] = reasoning_effort
@retry(
stop=stop_after_attempt(self._max_retries),
wait=wait_exponential(multiplier=1, min=1, max=10),
retry=retry_if_exception_type(_RETRYABLE),
reraise=True,
)
def _call():
return self._client.beta.chat.completions.parse(
model=active_model,
messages=messages,
response_format=response_format,
**extra,
)
with Timer() as t:
response = _call()
parsed = response.choices[0].message.parsed
if parsed is None:
# Model refused or returned unparseable — surface the refusal.
refusal = getattr(response.choices[0].message, "refusal", None)
raise RuntimeError(
f"Model returned no parsed output (model={active_model}). "
f"Refusal: {refusal!r}"
)
usage = response.usage
metrics = ExtractionMetrics(
input_tokens=usage.prompt_tokens if usage else 0,
output_tokens=usage.completion_tokens if usage else 0,
latency_ms=t.elapsed_ms,
model=active_model,
)
logger.info(
f"OpenAI call OK: model={active_model} "
f"tokens={metrics.input_tokens}+{metrics.output_tokens} "
f"latency={metrics.latency_ms:.0f}ms cost=${metrics.cost_usd:.5f}"
)
return parsed, metrics
|