zk-data-agent/tests/test_runtime_api.py

from __future__ import annotations

import json

import httpx
from fastapi.testclient import TestClient

from agent_platform.runtime.app import create_app
from agent_platform.store import RuntimeStore


class FakeModelProvider:
    def __init__(self) -> None:
        self.completed = []
        self.forwarded = []

    async def complete(self, **kwargs):
        self.completed.append(kwargs)
        return {
            "choices": [{"message": {"role": "assistant", "content": "work complete"}, "finish_reason": "stop"}],
            "usage": {"total_tokens": 10},
        }

    async def forward(self, payload):
        self.forwarded.append(payload)
        request = httpx.Request("POST", "https://model.example.test/v1/chat/completions")
        if payload.get("stream"):
            return httpx.Response(
                200,
                request=request,
                headers={"content-type": "text/event-stream"},
                content=(
                    'data: {"model":"ChatGPT-5.6:Luna","choices":[{"index":0,"delta":'
                    '{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":'
                    '{"name":"read_file","arguments":"{}"}}]},"finish_reason":null}]}\n\n'
                    "data: [DONE]\n\n"
                ),
            )
        return httpx.Response(
            200,
            request=request,
            headers={"content-type": "application/json"},
            json={
                "model": payload["model"],
                "choices": [{"message": {"role": "assistant", "content": "chat reply"}, "finish_reason": "stop"}],
            },
        )

    async def close(self):
        return None


class FakeTools:
    async def close(self):
        return None

    def specs(self, *, read_only=False, allow_delegate=True):
        return []

    async def execute(self, name, arguments, context):
        raise AssertionError("No tool should be called")


def headers(settings, identity_jwt, chat_id="chat-1"):
    return {
        "Authorization": f"Bearer {settings.internal_provider_key}",
        "X-OpenWebUI-User-Jwt": identity_jwt,
        "X-OpenWebUI-Chat-Id": chat_id,
    }


def test_model_list_requires_internal_key(settings) -> None:
    app = create_app(
        settings,
        store=RuntimeStore(settings.database_url),
        provider=FakeModelProvider(),
        tools=FakeTools(),
    )
    with TestClient(app) as client:
        assert client.get("/v1/models").status_code == 401
        response = client.get(
            "/v1/models",
            headers={"Authorization": f"Bearer {settings.internal_provider_key}"},
        )
        assert response.status_code == 200
        assert len(response.json()["data"]) == 6


def test_chat_maps_provider_model_and_work_upgrade_is_one_way(settings, identity_jwt) -> None:
    provider = FakeModelProvider()
    app = create_app(settings, store=RuntimeStore(settings.database_url), provider=provider, tools=FakeTools())
    with TestClient(app) as client:
        chat = client.post(
            "/v1/chat/completions",
            headers=headers(settings, identity_jwt),
            json={"model": "chat-light", "messages": [{"role": "user", "content": "hello"}]},
        )
        assert chat.status_code == 200
        assert provider.forwarded[0]["model"] == "ChatGPT-5.6:Luna"
        assert chat.json()["model"] == "chat-light"

        work = client.post(
            "/v1/chat/completions",
            headers=headers(settings, identity_jwt),
            json={"model": "work-high", "messages": [{"role": "user", "content": "do it"}]},
        )
        assert work.status_code == 200
        assert provider.completed[-1]["model"] == "ChatGPT-5.6:Sol"

        downgrade = client.post(
            "/v1/chat/completions",
            headers=headers(settings, identity_jwt),
            json={"model": "chat-medium", "messages": [{"role": "user", "content": "back"}]},
        )
        assert downgrade.status_code == 409


def test_chat_stream_hides_provider_model_and_preserves_native_tool_calls(settings, identity_jwt) -> None:
    provider = FakeModelProvider()
    app = create_app(settings, store=RuntimeStore(settings.database_url), provider=provider, tools=FakeTools())
    with TestClient(app) as client:
        response = client.post(
            "/v1/chat/completions",
            headers=headers(settings, identity_jwt, "chat-stream"),
            json={
                "model": "chat-light",
                "stream": True,
                "messages": [{"role": "user", "content": "inspect"}],
            },
        )
    assert response.status_code == 200
    assert '"model": "chat-light"' in response.text
    assert "ChatGPT-5.6:Luna" not in response.text
    assert '"tool_calls"' in response.text
    assert '"name": "read_file"' in response.text
    assert "data: [DONE]" in response.text


def test_work_stream_never_emits_delta_tool_calls(settings, identity_jwt) -> None:
    app = create_app(
        settings,
        store=RuntimeStore(settings.database_url),
        provider=FakeModelProvider(),
        tools=FakeTools(),
    )
    with TestClient(app) as client:
        response = client.post(
            "/v1/chat/completions",
            headers=headers(settings, identity_jwt, "stream-chat"),
            json={
                "model": "work-medium",
                "stream": True,
                "messages": [{"role": "user", "content": "do it"}],
            },
        )
        assert response.status_code == 200
        assert "work complete" in response.text
        assert "tool_calls" in response.text  # rendered Work status details
        for line in response.text.splitlines():
            if line.startswith("data: {"):
                chunk = json.loads(line.removeprefix("data: "))
                assert "tool_calls" not in chunk["choices"][0]["delta"]
        assert '"finish_reason": "stop"' in response.text