agent-platform/tests/test_model_service.py

from pathlib import Path

from tests.conftest import (
    build_fastapi_test_client,
    build_postgres_database_url,
    build_postgres_engine,
    prepare_known_service_import,
)


def test_model_service_post_contract_supports_models_and_providers(
    tmp_path: Path,
    monkeypatch,
) -> None:
    prepare_known_service_import("model-gateway-service")

    from app.bootstrap.app import create_app
    from app.db.models import Base
    from core_db import create_session_factory

    database_url = build_postgres_database_url(tmp_path, "models")
    monkeypatch.setenv("AGENT_PLATFORM_DATABASE_URL", database_url)
    engine = build_postgres_engine(database_url)
    Base.metadata.create_all(engine)

    app = create_app()
    app.state.session_factory = create_session_factory(engine)
    client = build_fastapi_test_client(app)

    provider_response = client.post(
        "/models/providers/create",
        json={
            "name": "Local OpenAI Compatible",
            "providerType": "openai_compatible",
            "baseUrl": "http://127.0.0.1:11434/v1",
            "apiKey": "local-secret",
            "models": [
                {
                    "modelId": "llama3.1",
                    "displayName": "Llama 3.1",
                    "modelType": "chat",
                }
            ],
            "defaultModel": "llama3.1",
        },
    )

    assert provider_response.status_code == 200
    provider_payload = provider_response.json()["data"]
    assert provider_payload["apiKeyRef"] == "loc***masked"
    assert provider_payload["models"][0]["modelId"] == "llama3.1"

    providers_response = client.post(
        "/models/providers/list",
        json={"page": 1, "pageSize": 20},
    )
    assert providers_response.status_code == 200
    assert providers_response.json()["data"]["total"] == 1

    discover_response = client.post(
        "/models/providers/discover",
        json={"providerId": provider_payload["id"]},
    )
    assert discover_response.status_code == 200
    assert discover_response.json()["data"]["models"][0]["modelId"] == "llama3.1"

    model_response = client.post(
        "/models/create",
        json={
            "name": "Local Chat",
            "providerId": provider_payload["id"],
            "providerType": "openai_compatible",
            "modelName": "llama3.1",
            "capabilities": ["chat"],
            "timeoutSeconds": 30,
        },
    )

    assert model_response.status_code == 200
    model_payload = model_response.json()["data"]
    assert model_payload["modelName"] == "llama3.1"
    assert model_payload["providerId"] == provider_payload["id"]
    assert model_payload["providerBaseUrl"] == "http://127.0.0.1:11434/v1"
    assert model_payload["hasProviderApiKey"] is True
    assert "code" not in model_payload

    models_response = client.post(
        "/models/list",
        json={"page": 1, "pageSize": 20, "keyword": "local"},
    )
    assert models_response.status_code == 200
    assert models_response.json()["data"]["total"] == 1

    update_response = client.post(
        "/models/update",
        json={
            "modelId": model_payload["id"],
            "name": "Local Chat Updated",
            "defaultTemperature": 0.2,
        },
    )
    assert update_response.status_code == 200
    assert update_response.json()["data"]["defaultTemperature"] == 0.2

    delete_response = client.post(
        "/models/delete",
        json={"modelId": model_payload["id"]},
    )
    assert delete_response.status_code == 200
    assert delete_response.json()["data"]["deleted"] is True


def test_model_provider_client_supports_anthropic_messages(monkeypatch) -> None:
    prepare_known_service_import("model-gateway-service")

    import app.infrastructure.provider as provider_module
    from app.bootstrap.settings import ModelGatewayServiceSettings
    from app.infrastructure.provider import ModelProviderClient
    from core_domain import ChatCompletionRequestContract

    captured: dict[str, object] = {}

    class FakeResponse:
        text = "{}"

        def raise_for_status(self) -> None:
            return None

        def json(self) -> dict[str, object]:
            return {
                "model": "claude-3-5-sonnet-20241022",
                "content": [{"type": "text", "text": "ready"}],
                "stop_reason": "end_turn",
                "usage": {"input_tokens": 12, "output_tokens": 3},
            }

    class FakeClient:
        def __init__(self, *, timeout: float) -> None:
            captured["timeout"] = timeout

        def __enter__(self) -> "FakeClient":
            return self

        def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
            return None

        def post(
            self,
            url: str,
            *,
            json: dict[str, object],
            headers: dict[str, str]) -> FakeResponse:
            captured["url"] = url
            captured["json"] = json
            captured["headers"] = headers
            return FakeResponse()

    monkeypatch.setattr(provider_module.httpx, "Client", FakeClient)

    client = ModelProviderClient(settings=ModelGatewayServiceSettings())
    response = client.create_chat_completion(
        ChatCompletionRequestContract(
            model="claude-3-5-sonnet-20241022",
            messages=[
                {"role": "system", "content": "Be concise."},
                {"role": "user", "content": "Ping"},
            ],
            max_tokens=128),
        provider_type="anthropic",
        provider_base_url="https://api.anthropic.com",
        provider_api_key="sk-test",
        timeout_seconds=15)

    assert captured["url"] == "https://api.anthropic.com/v1/messages"
    assert captured["headers"] == {
        "content-type": "application/json",
        "x-api-key": "sk-test",
        "anthropic-version": "2023-06-01",
    }
    assert captured["json"] == {
        "model": "claude-3-5-sonnet-20241022",
        "max_tokens": 128,
        "messages": [{"role": "user", "content": "Ping"}],
        "system": "Be concise.",
    }
    assert response.content == "ready"
    assert response.finish_reason == "end_turn"


def test_model_service_backfills_legacy_model_connections_as_providers(
    tmp_path: Path,
    monkeypatch,
) -> None:
    prepare_known_service_import("model-gateway-service")

    from app.bootstrap.app import create_app
    from app.db.models import Base
    from core_db import create_session_factory

    database_url = build_postgres_database_url(tmp_path, "models-backfill")
    monkeypatch.setenv("AGENT_PLATFORM_DATABASE_URL", database_url)
    engine = build_postgres_engine(database_url)
    Base.metadata.create_all(engine)

    app = create_app()
    app.state.session_factory = create_session_factory(engine)
    client = build_fastapi_test_client(app)

    model_response = client.post(
        "/models/create",
        json={
            "name": "Legacy Anthropic",
            "providerType": "anthropic",
            "providerBaseUrl": "https://api.anthropic.com",
            "providerApiKey": "sk-legacy",
            "modelName": "claude-3-5-sonnet-20241022",
            "capabilities": ["chat"],
        },
    )
    assert model_response.status_code == 200
    assert model_response.json()["data"]["providerId"] is None

    providers_response = client.post(
        "/models/providers/list",
        json={"page": 1, "pageSize": 20},
    )

    assert providers_response.status_code == 200
    providers_payload = providers_response.json()["data"]
    assert providers_payload["total"] == 1
    provider_payload = providers_payload["items"][0]
    assert provider_payload["providerType"] == "anthropic"
    assert provider_payload["baseUrl"] == "https://api.anthropic.com"
    assert provider_payload["models"][0]["modelId"] == "claude-3-5-sonnet-20241022"

    models_response = client.post(
        "/models/list",
        json={"page": 1, "pageSize": 20},
    )
    assert models_response.status_code == 200
    assert models_response.json()["data"]["items"][0]["providerId"] == provider_payload["id"]