litellm/tests/unit/systemone/test_main.py

from __future__ import annotations

import asyncio
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final

import pytest
import respx
from respx.models import Call

import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse

_STATE: Final = "Customer asked for a refund"
_SYSTEMONE_QUESTIONS: Final[Mapping[str, object]] = MappingProxyType(
    {"is_refund": {"type": "noul", "instructions": "Is this a refund request?"}}
)
_OPENAI_QUESTIONS: Final[tuple[Mapping[str, object], ...]] = (
    {"type": "predicate", "name": "is_refund", "instructions": "Is this a refund request?"},
)
_SYSTEMONE_RESPONSE: Final[Mapping[str, object]] = {
    "model": "jev-1.13",
    "answers": {"is_refund": {"type": "noul", "noul": 0.93}},
    "usage": {"input_tokens": 12, "output_tokens": 1},
}
_TYPESAFE_URL: Final = "https://api.typesafe.ai/v1/systemone"
_CLOUDFLARE_API_BASE: Final = "https://api.cloudflare.com/client/v4/accounts/acct/ai/run"
_CLOUDFLARE_CLEF_URL: Final = f"{_CLOUDFLARE_API_BASE}/@cf/cloudflare/clef"
_PNG_DATA_URL: Final = "data:image/png;base64,AA=="


@pytest.fixture(autouse=True)
def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    litellm.in_memory_llm_clients_cache.flush_cache()


class _CallTypeRecorder(CustomLogger):
    def __init__(self) -> None:
        super().__init__()
        self.call_types: list[str] = []  # mutable-ok: test recorder collects logger callbacks

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
        self.call_types.append(kwargs["call_type"])


@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone"))
async def test_systemone_sends_the_systemone_wire_format_and_returns_a_systemone_response(
    use_async: bool, respx_mock: respx.MockRouter
) -> None:
    route: Final = respx_mock.post(_TYPESAFE_URL).respond(json=_SYSTEMONE_RESPONSE)
    kwargs: Final[Mapping[str, object]] = MappingProxyType(
        {"model": "typesafe/jev-1.13", "state": _STATE, "questions": _SYSTEMONE_QUESTIONS, "api_key": "caller-key"}
    )

    response: Final = await litellm.asystemone(**kwargs) if use_async else litellm.systemone(**kwargs)

    assert route.called
    call: Final[Call] = respx_mock.calls.last
    assert json.loads(call.request.content) == {
        "model": "jev-1.13",
        "state": _STATE,
        "questions": {"is_refund": {"type": "noul", "instructions": "Is this a refund request?"}},
    }
    assert isinstance(response, DecisionsResponse)
    assert response.answers["is_refund"].model_dump(mode="json") == {"type": "noul", "noul": 0.93}


@pytest.mark.asyncio
async def test_systemone_rejects_openai_input_and_points_at_decisions(respx_mock: respx.MockRouter) -> None:
    route: Final = respx_mock.post(_TYPESAFE_URL).respond(json=_SYSTEMONE_RESPONSE)

    with pytest.raises(litellm.BadRequestError, match=r"litellm\.decisions") as caught:
        await litellm.asystemone(
            model="typesafe/jev-1.13", input=_STATE, questions=_SYSTEMONE_QUESTIONS, api_key="caller-key"
        )

    assert caught.value.status_code == 400
    assert not route.called


@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", (True, False), ids=("adecisions", "decisions"))
async def test_decisions_with_state_still_sends_the_systemone_format_and_warns(
    use_async: bool, respx_mock: respx.MockRouter
) -> None:
    route: Final = respx_mock.post(_TYPESAFE_URL).respond(json=_SYSTEMONE_RESPONSE)
    kwargs: Final[Mapping[str, object]] = MappingProxyType(
        {"model": "typesafe/jev-1.13", "state": _STATE, "questions": _SYSTEMONE_QUESTIONS, "api_key": "caller-key"}
    )

    with pytest.warns(UserWarning, match=r"litellm\.systemone"):
        response: Final = await litellm.adecisions(**kwargs) if use_async else litellm.decisions(**kwargs)

    assert route.called
    call: Final[Call] = respx_mock.calls.last
    assert json.loads(call.request.content) == {
        "model": "jev-1.13",
        "state": _STATE,
        "questions": {"is_refund": {"type": "noul", "instructions": "Is this a refund request?"}},
    }
    assert isinstance(response, DecisionsResponse)
    assert response.answers["is_refund"].model_dump(mode="json") == {"type": "noul", "noul": 0.93}


@pytest.mark.asyncio
async def test_decisions_with_state_still_sends_its_images(respx_mock: respx.MockRouter) -> None:
    route: Final = respx_mock.post(_CLOUDFLARE_CLEF_URL).respond(json=_SYSTEMONE_RESPONSE)

    with pytest.warns(UserWarning, match=r"litellm\.systemone"):
        await litellm.adecisions(
            model="cloudflare/clef",
            state=_STATE,
            questions=_SYSTEMONE_QUESTIONS,
            images=[_PNG_DATA_URL],
            api_key="caller-key",
            api_base=_CLOUDFLARE_API_BASE,
        )

    sent: Final = json.loads(respx_mock.calls.last.request.content)
    assert sent["images"] == [_PNG_DATA_URL]
    assert route.called


@pytest.mark.asyncio
async def test_systemone_rejects_openai_question_lists_before_http(respx_mock: respx.MockRouter) -> None:
    route: Final = respx_mock.post(_TYPESAFE_URL).respond(json=_SYSTEMONE_RESPONSE)

    with pytest.raises(litellm.BadRequestError, match="Invalid Decisions request"):
        await litellm.asystemone(
            model="typesafe/jev-1.13", state=_STATE, questions=_OPENAI_QUESTIONS, api_key="caller-key"
        )

    assert not route.called


@pytest.mark.asyncio
async def test_router_asystemone_logs_the_systemone_call_type(
    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
    respx_mock.post(_TYPESAFE_URL).respond(json=_SYSTEMONE_RESPONSE)
    recorder: Final = _CallTypeRecorder()
    monkeypatch.setattr(litellm, "callbacks", [recorder])
    router: Final = litellm.Router(
        model_list=[{"model_name": "jev", "litellm_params": {"model": "typesafe/jev-1.13", "api_key": "caller-key"}}]
    )

    response: Final = await router.asystemone(model="jev", state=_STATE, questions=_SYSTEMONE_QUESTIONS)
    await asyncio.sleep(0)
    GLOBAL_LOGGING_WORKER.start()
    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)

    assert isinstance(response, DecisionsResponse)
    assert recorder.call_types == ["asystemone"]
    assert not isinstance(response, OpenAIDecisionResponse)