fix: retry transient model failures

This commit is contained in:
wuyang
2026-07-26 14:46:43 +08:00
parent 1216e96038
commit 2c17a74203
2 changed files with 121 additions and 10 deletions
+41 -10
View File
@@ -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(