fix: retry transient model failures
This commit is contained in:
@@ -1,11 +1,16 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from agent_platform.config import Settings
|
||||
|
||||
TRANSIENT_COMPLETE_STATUS_CODES = frozenset({408, 425, 429, 500, 502, 503, 504, 524})
|
||||
COMPLETE_MAX_ATTEMPTS = 2
|
||||
MAX_RETRY_DELAY_SECONDS = 5.0
|
||||
|
||||
|
||||
def completions_url(base_url: str) -> str:
|
||||
base = base_url.rstrip("/")
|
||||
@@ -14,6 +19,17 @@ def completions_url(base_url: str) -> str:
|
||||
return f"{base}/v1/chat/completions"
|
||||
|
||||
|
||||
def retry_delay_seconds(response: httpx.Response | None, attempt: int) -> float:
|
||||
if response is not None:
|
||||
retry_after = response.headers.get("Retry-After")
|
||||
if retry_after:
|
||||
try:
|
||||
return min(MAX_RETRY_DELAY_SECONDS, max(0.0, float(retry_after)))
|
||||
except ValueError:
|
||||
pass
|
||||
return min(MAX_RETRY_DELAY_SECONDS, float(2 ** (attempt - 1)))
|
||||
|
||||
|
||||
class ModelProvider:
|
||||
def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None:
|
||||
self.settings = settings
|
||||
@@ -52,16 +68,31 @@ class ModelProvider:
|
||||
payload["tool_choice"] = "auto"
|
||||
if temperature is not None:
|
||||
payload["temperature"] = temperature
|
||||
response = await self.client.post(
|
||||
completions_url(self.settings.model_api_base_url),
|
||||
headers=self.headers,
|
||||
json=payload,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if not data.get("choices"):
|
||||
raise RuntimeError("Model provider returned no choices")
|
||||
return data
|
||||
for attempt in range(1, COMPLETE_MAX_ATTEMPTS + 1):
|
||||
try:
|
||||
response = await self.client.post(
|
||||
completions_url(self.settings.model_api_base_url),
|
||||
headers=self.headers,
|
||||
json=payload,
|
||||
)
|
||||
except httpx.TransportError:
|
||||
if attempt >= COMPLETE_MAX_ATTEMPTS:
|
||||
raise
|
||||
await asyncio.sleep(retry_delay_seconds(None, attempt))
|
||||
continue
|
||||
|
||||
if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS:
|
||||
await response.aread()
|
||||
await asyncio.sleep(retry_delay_seconds(response, attempt))
|
||||
continue
|
||||
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if not data.get("choices"):
|
||||
raise RuntimeError("Model provider returned no choices")
|
||||
return data
|
||||
|
||||
raise RuntimeError("Model provider retry loop exited unexpectedly")
|
||||
|
||||
async def forward(self, payload: dict[str, Any]) -> httpx.Response:
|
||||
request = self.client.build_request(
|
||||
|
||||
Reference in New Issue
Block a user