fix: retry transient model failures
This commit is contained in:
@@ -0,0 +1,80 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from agent_platform.config import Settings
|
||||
from agent_platform.runtime.provider import ModelProvider
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def disable_retry_sleep(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def no_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("agent_platform.runtime.provider.asyncio.sleep", no_sleep)
|
||||
|
||||
|
||||
async def _complete_with_handler(
|
||||
settings: Settings,
|
||||
handler: Callable[[httpx.Request], httpx.Response | Awaitable[httpx.Response]],
|
||||
) -> dict:
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
provider = ModelProvider(settings, client)
|
||||
return await provider.complete(model="test-model", messages=[{"role": "user", "content": "hello"}])
|
||||
|
||||
|
||||
async def test_complete_retries_transient_gateway_response(settings: Settings) -> None:
|
||||
calls = 0
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
return httpx.Response(524, request=request, headers={"Retry-After": "0"})
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
json={"choices": [{"message": {"role": "assistant", "content": "ok"}}]},
|
||||
)
|
||||
|
||||
result = await _complete_with_handler(settings, handler)
|
||||
|
||||
assert calls == 2
|
||||
assert result["choices"][0]["message"]["content"] == "ok"
|
||||
|
||||
|
||||
async def test_complete_retries_transport_error(settings: Settings) -> None:
|
||||
calls = 0
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
raise httpx.ConnectError("temporary disconnect", request=request)
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
json={"choices": [{"message": {"role": "assistant", "content": "recovered"}}]},
|
||||
)
|
||||
|
||||
result = await _complete_with_handler(settings, handler)
|
||||
|
||||
assert calls == 2
|
||||
assert result["choices"][0]["message"]["content"] == "recovered"
|
||||
|
||||
|
||||
async def test_complete_does_not_retry_client_error(settings: Settings) -> None:
|
||||
calls = 0
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return httpx.Response(400, request=request, json={"error": "bad request"})
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
await _complete_with_handler(settings, handler)
|
||||
|
||||
assert calls == 1
|
||||
Reference in New Issue
Block a user