zk-data-agent/tests/test_provider.py

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]],
    **kwargs,
) -> 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"}],
            **kwargs,
        )


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


async def test_deepseek_route_enables_thinking_and_uses_separate_credentials(settings: Settings) -> None:
    body = (
        'data: {"choices":[{"index":0,"delta":{"role":"assistant",'
        '"reasoning_content":"private state","tool_calls":[{"index":0,"id":"call-1",'
        '"type":"function","function":{"name":"ping","arguments":"{}"}}]},'
        '"finish_reason":"tool_calls"}]}\n\n'
        "data: [DONE]\n\n"
    )

    def handler(request: httpx.Request) -> httpx.Response:
        assert str(request.url) == "https://api.deepseek.example.test/v1/chat/completions"
        assert request.headers["Authorization"] == "Bearer deepseek-secret"
        payload = json.loads(request.content)
        assert payload["model"] == "test-model"
        assert payload["thinking"] == {"type": "enabled"}
        assert payload["reasoning_effort"] == "max"
        assert payload["max_tokens"] == 16_384
        return httpx.Response(
            200,
            request=request,
            headers={"Content-Type": "text/event-stream"},
            content=body,
        )

    result = await _complete_with_handler(
        settings,
        handler,
        provider="deepseek",
        thinking_enabled=True,
        reasoning_effort="max",
        max_output_tokens=16_384,
    )

    message = result["choices"][0]["message"]
    assert message["reasoning_content"] == "private state"
    assert message["tool_calls"][0]["function"]["name"] == "ping"