mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
399a2d918d
The priority tests spied on Scheduler.add_request to check which priority a request was scheduled at and that it was scheduled exactly once. A request a healthy deployment can take no longer joins the queue, so add_request is never called on that path and the spies saw nothing, while the zero-count checks would have kept passing even if a request without a priority was scheduled. wait_for_turn is the router's single entry into the scheduler and receives the same FlowItem, so the assertions stay as they were.
216 lines
7.9 KiB
Python
216 lines
7.9 KiB
Python
"""Behavior pins for ``proxy_server.py`` queue routes.
|
|
|
|
Pins (PR2):
|
|
- POST /queue/chat/completions
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from contextlib import AbstractContextManager
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
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_drops_client_sent_stream_usage_flags(
|
|
client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], patched_queue: MagicMock
|
|
) -> None:
|
|
payload: Final = {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"priority": 0,
|
|
"_litellm_strip_stream_usage": True,
|
|
"_litellm_stream_usage_injected": True,
|
|
}
|
|
with auth_as():
|
|
response: Final = client.post("/queue/chat/completions", json=payload)
|
|
assert response.status_code == 200
|
|
forwarded: Final = patched_queue.schedule_acompletion.call_args.kwargs
|
|
assert "_litellm_strip_stream_usage" not in forwarded
|
|
assert "_litellm_stream_usage_injected" not in forwarded
|
|
assert forwarded["messages"] == [{"role": "user", "content": "hi"}]
|
|
|
|
|
|
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, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
|
|
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 wait_for_turn.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, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
|
|
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 wait_for_turn.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"]
|