wuyang
161 lines
4.9 KiB
Python
161 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
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
|
|
|
|
|
|
async def test_complete_aggregates_streamed_tool_call(settings: Settings) -> None:
|
|
chunks = [
|
|
{
|
|
"id": "completion-1",
|
|
"object": "chat.completion.chunk",
|
|
"created": 123,
|
|
"model": "test-model",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"role": "assistant",
|
|
"tool_calls": [
|
|
{
|
|
"index": 0,
|
|
"id": "call-1",
|
|
"type": "function",
|
|
"function": {"name": "ping", "arguments": '{"val'},
|
|
}
|
|
],
|
|
},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"tool_calls": [
|
|
{
|
|
"index": 0,
|
|
"function": {"arguments": 'ue":"ok"}'},
|
|
}
|
|
]
|
|
},
|
|
"finish_reason": None,
|
|
}
|
|
]
|
|
},
|
|
{
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {},
|
|
"finish_reason": "tool_calls",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7},
|
|
},
|
|
]
|
|
body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
request_payload = json.loads(request.content)
|
|
assert request_payload["stream"] is True
|
|
assert request_payload["max_tokens"] == 4096
|
|
return httpx.Response(
|
|
200,
|
|
request=request,
|
|
headers={"Content-Type": "text/event-stream"},
|
|
content=body,
|
|
)
|
|
|
|
result = await _complete_with_handler(settings, handler)
|
|
|
|
choice = result["choices"][0]
|
|
assert choice["finish_reason"] == "tool_calls"
|
|
assert choice["message"]["tool_calls"] == [
|
|
{
|
|
"id": "call-1",
|
|
"type": "function",
|
|
"function": {"name": "ping", "arguments": '{"value":"ok"}'},
|
|
}
|
|
]
|
|
assert result["usage"]["total_tokens"] == 7
|