soc_ai/agent/_gateway_retry.py
Rama main · réplica de lectura del repositorio original.
"""A retrying httpx transport for the LiteLLM gateway (primary model path).
The Oracle client (:mod:`soc_ai.oracle.client`) already wraps its calls in an
exponential-backoff retry with a retryable/terminal split. The PRIMARY model path
(investigator / synthesizer / hunt / chat) runs inside pydantic-ai, so we can't
wrap each call the same way — but we CAN make the underlying HTTP transport
resilient, which covers every primary-path call transparently.
This transport retries transient gateway failures — HTTP 429/502/503/504 and
connection/read/timeout transport errors — with exponential backoff + jitter,
honoring a ``Retry-After`` header when present. Terminal 4xx (auth / bad request)
and any 2xx/3xx pass straight through. It is the single retry authority for the
primary path, so the ``AsyncOpenAI`` client is built with ``max_retries=0``.
"""
from __future__ import annotations
import asyncio
import logging
import random
from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any
import httpx
_LOGGER = logging.getLogger(__name__)
# Which BACKEND actually served the most recent gateway call on this run.
#
# soc-ai asks LiteLLM for one alias (``analyst_model``) and the gateway may route
# that to a different deployment, or fall back to another one entirely, without
# changing anything the client can see: the response body's ``model`` field and
# pydantic-ai's ``ModelResponse.model_name`` both echo back the ALIAS. So a
# failure recorded against "the model" names what we asked for, not what ran —
# which is exactly why "the lower-tier fallback models fail more" was
# unfalsifiable from the error log.
#
# LiteLLM does report the truth, but only in response HEADERS, which pydantic-ai
# does not surface (``ModelResponse.provider_details`` carries just finish_reason
# and a timestamp). This transport is the one place soc-ai owns that sees them.
#
# The ContextVar holds a MUTABLE dict rather than a value: a ContextVar set
# inside a child task is invisible to its parent, but mutating a shared dict is
# visible everywhere the reference reached. That keeps attribution working
# regardless of how pydantic-ai schedules the request.
_ATTRIBUTION: ContextVar[dict[str, Any] | None] = ContextVar(
"soc_ai_gateway_attribution", default=None
)
# LiteLLM response headers -> the key we record. `api_base` is the load-bearing
# one (it identifies the serving engine); `attempted_fallbacks` answers "did this
# request get demoted to another deployment?" directly.
_ATTRIBUTION_HEADERS = {
"x-litellm-model-api-base": "api_base",
"x-litellm-model-group": "model_group",
"x-litellm-model-id": "deployment_id",
"x-litellm-attempted-fallbacks": "attempted_fallbacks",
"x-litellm-attempted-retries": "attempted_retries",
}
@contextmanager
def capture_backend_attribution() -> Iterator[dict[str, Any]]:
"""Collect which backend served the gateway calls made inside this block.
Yields the sink dict; it is populated as calls complete, so read it AFTER the
block's work (including from an exception handler — the last attempt's
headers are already in there).
"""
sink: dict[str, Any] = {}
token = _ATTRIBUTION.set(sink)
try:
yield sink
finally:
_ATTRIBUTION.reset(token)
def _record_attribution(response: httpx.Response) -> None:
"""Stash the gateway's backend headers into the active sink (no-op if none).
Fail-soft on purpose: attribution is diagnostics, and must never be able to
break a model call that otherwise succeeded.
"""
try:
sink = _ATTRIBUTION.get()
if sink is None:
return
for header, key in _ATTRIBUTION_HEADERS.items():
value = response.headers.get(header)
if value:
sink[key] = value
except Exception: # pragma: no cover - diagnostics must never raise
_LOGGER.debug("failed to record gateway attribution", exc_info=True)
# Transient gateway statuses worth retrying (mirrors the Oracle client's
# "5xx/transport = retryable, 4xx = terminal" split, plus 429 rate-limit).
_RETRYABLE_STATUS = frozenset({429, 502, 503, 504})
_RETRYABLE_EXC = (
httpx.ConnectError,
httpx.ConnectTimeout,
httpx.ReadError,
httpx.ReadTimeout,
httpx.WriteError,
httpx.PoolTimeout,
httpx.RemoteProtocolError,
)
def _backoff_s(attempt: int, *, base: float, cap: float) -> float:
"""Exponential backoff with full jitter: random in [0, min(cap, base*2**attempt)]."""
ceiling = min(cap, base * (2.0**attempt))
return random.random() * ceiling # noqa: S311 - jitter, not security-sensitive
class RetryingAsyncTransport(httpx.AsyncBaseTransport):
"""Wraps an ``httpx.AsyncHTTPTransport``, retrying transient gateway failures.
``max_retries`` is the number of RETRIES after the first attempt (so a value
of 5 means up to 6 total attempts). Backoff grows 0.5, 1, 2, 4, capped at
``max_backoff_s`` and jittered; a ``Retry-After`` header (429) overrides it.
"""
def __init__(
self,
*,
max_retries: int,
verify: bool = True,
base_delay_s: float = 0.5,
max_backoff_s: float = 8.0,
inner: httpx.AsyncBaseTransport | None = None,
) -> None:
self._inner = inner if inner is not None else httpx.AsyncHTTPTransport(verify=verify)
self._max_retries = max(0, int(max_retries))
self._base = base_delay_s
self._cap = max_backoff_s
async def aclose(self) -> None:
await self._inner.aclose()
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
attempt = 0
while True:
try:
response = await self._inner.handle_async_request(request)
except _RETRYABLE_EXC as exc:
if attempt >= self._max_retries:
raise
delay = _backoff_s(attempt, base=self._base, cap=self._cap)
_LOGGER.warning(
"gateway transport error (%s) — retry %d/%d in %.1fs",
type(exc).__name__,
attempt + 1,
self._max_retries,
delay,
)
else:
# Record on EVERY response, including the ones we go on to retry:
# if the final attempt raises a transport error, the last headers
# we saw are still the best attribution available.
_record_attribution(response)
if response.status_code not in _RETRYABLE_STATUS or attempt >= self._max_retries:
return response
# Must release the connection before retrying the request.
retry_after = _parse_retry_after(response.headers.get("retry-after"))
await response.aclose()
delay = (
retry_after
if retry_after is not None
else _backoff_s(attempt, base=self._base, cap=self._cap)
)
_LOGGER.warning(
"gateway returned %d — retry %d/%d in %.1fs",
response.status_code,
attempt + 1,
self._max_retries,
delay,
)
await asyncio.sleep(delay)
attempt += 1
def _parse_retry_after(value: str | None) -> float | None:
"""Parse a numeric ``Retry-After`` (seconds) header; ignore HTTP-date form."""
if not value:
return None
try:
secs = float(value.strip())
except (TypeError, ValueError):
return None
# Clamp to a sane bound so a hostile/huge value can't stall the run.
return max(0.0, min(secs, 30.0))