litellm/tests/unit/proxy/proxy_server/test_routes_queue.py

"""Behavior pins for ``proxy_server.py`` queue routes.

Pins (PR2):
    - POST /queue/chat/completions
"""

from __future__ import annotations

from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from litellm import Router
from litellm.proxy import proxy_server

from .conftest import normalize  # type: ignore[import-not-found]

HAPPY_RESPONSE = {
    "id": "chatcmpl-queue",
    "object": "chat.completion",
    "created": 0,
    "model": "gpt-4",
    "choices": [
        {
            "index": 0,
            "finish_reason": "stop",
            "message": {"role": "assistant", "content": "queued reply"},
        }
    ],
    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
    "priority": 0,
}


@pytest.fixture
def patched_queue(monkeypatch):
    router = MagicMock()
    router.split_request_priority = MagicMock(return_value=(0, {}))
    router.schedule_acompletion = AsyncMock(return_value=dict(HAPPY_RESPONSE))
    monkeypatch.setattr(proxy_server, "llm_router", router)
    monkeypatch.setattr(
        proxy_server,
        "proxy_logging_obj",
        MagicMock(post_call_failure_hook=AsyncMock()),
    )
    return router


@pytest.fixture
def queue_no_router(monkeypatch):
    monkeypatch.setattr(proxy_server, "llm_router", None)
    monkeypatch.setattr(
        proxy_server,
        "proxy_logging_obj",
        MagicMock(post_call_failure_hook=AsyncMock()),
    )
    yield


def test_queue_chat_completions_happy(client, auth_as, patched_queue):
    """Pins ``POST /queue/chat/completions`` (happy)."""
    payload = {
        "model": "gpt-4",
        "messages": [{"role": "user", "content": "hi"}],
        "priority": 0,
    }
    with auth_as():
        response = client.post("/queue/chat/completions", json=payload)
    assert response.status_code == 200
    assert normalize(response.json()) == {
        "id": "<VOLATILE>",
        "object": "chat.completion",
        "created": "<VOLATILE>",
        "model": "gpt-4",
        "choices": [
            {
                "index": 0,
                "finish_reason": "stop",
                "message": {"role": "assistant", "content": "queued reply"},
            }
        ],
        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
        "priority": 0,
    }


def test_queue_chat_completions_no_router_error(client, auth_as, queue_no_router):
    """Pins ``POST /queue/chat/completions`` (error: no llm_router)."""
    payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
    with auth_as():
        response = client.post("/queue/chat/completions", json=payload)
    assert response.status_code == 500
    assert len(response.content) > 0


def _queue_router(
    monkeypatch: pytest.MonkeyPatch, default_priority: int | None, mock_response: str = "queued"
) -> Router:
    router: Final = Router(
        model_list=[
            {
                "model_name": "gpt-4",
                "litellm_params": {"model": "openai/gpt-4", "api_key": "sk-queue-test", "mock_response": mock_response},
            }
        ],
        default_priority=default_priority,
    )
    monkeypatch.setattr(proxy_server, "llm_router", router)
    monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock()))
    return router


@pytest.mark.parametrize(
    ("default_priority", "priority_fields", "expected_priority"),
    [
        (3, {"priority": "high", "drop_params": True}, 3),
        (3, {"priority": "1"}, 1),
        (None, {"priority": 1.0}, 1),
        (3, {"priority": 1}, 1),
        (3, {}, 3),
    ],
    ids=("dropped-word-takes-default", "digit-string", "integral-float", "explicit", "absent-takes-default"),
)
def test_queue_chat_completions_reports_the_priority_the_router_queued_at(
    client,
    auth_as,
    monkeypatch: pytest.MonkeyPatch,
    default_priority: int | None,
    priority_fields: dict[str, object],
    expected_priority: int,
) -> None:
    router: Final = _queue_router(monkeypatch, default_priority)
    payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
        with auth_as():
            response: Final = client.post("/queue/chat/completions", json=payload)
    assert response.status_code == 200, response.text
    assert [call.kwargs["request"].priority for call in add_request.await_args_list] == [expected_priority]
    assert response.headers.get("x-litellm-priority") == str(expected_priority)


@pytest.mark.parametrize(
    ("default_priority", "priority_fields", "message"),
    [
        (None, {}, "priority is required"),
        (3, {"priority": "high"}, "priority must be an integer"),
        (None, {"priority": "high", "drop_params": True}, "priority must be an integer, got str 'high'"),
    ],
    ids=("no-priority-and-no-default", "malformed", "dropped-word-and-no-default"),
)
def test_queue_chat_completions_rejects_a_request_it_cannot_queue_as_an_invalid_request(
    client,
    auth_as,
    monkeypatch: pytest.MonkeyPatch,
    default_priority: int | None,
    priority_fields: dict[str, object],
    message: str,
) -> None:
    router: Final = _queue_router(monkeypatch, default_priority)
    payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
        with auth_as():
            response: Final = client.post("/queue/chat/completions", json=payload)
    assert response.status_code == 400, response.text
    error: Final = response.json()["error"]
    assert (error["type"], error["param"]) == ("invalid_request_error", "priority")
    assert message in error["message"]
    assert add_request.await_count == 0


def test_streamed_queue_responses_report_the_priority_the_router_queued_at(
    client, auth_as, monkeypatch: pytest.MonkeyPatch
) -> None:
    _queue_router(monkeypatch, default_priority=3)
    payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "priority": "1", "stream": True}
    with auth_as():
        response: Final = client.post("/queue/chat/completions", json=payload)
    assert response.status_code == 200, response.text
    assert response.headers.get("x-litellm-priority") == "1"


def test_a_provider_bad_request_on_the_queue_route_keeps_its_own_message(
    client, auth_as, monkeypatch: pytest.MonkeyPatch
) -> None:
    _queue_router(monkeypatch, default_priority=3, mock_response="litellm.ContextWindowExceededError")
    payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
    with auth_as():
        response: Final = client.post("/queue/chat/completions", json=payload)
    assert response.status_code == 400, response.text
    error: Final = response.json()["error"]
    assert error["type"] == "invalid_request_error"
    assert "context window" in error["message"].lower()
    assert "Authentication Error" not in error["message"]