mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 07:12:32 +00:00
Yassin Kortam
a6494e6fe3
- Introduce `_CallbackCapabilities` dataclass and `ProxyLogging._callback_capabilities()` static method that inspects `litellm.callbacks` once and caches capability flags keyed on (list length, member ids); invalidates automatically when the callback list mutates without per-request iteration overhead - Replace O(n) `litellm.callbacks` walks in `async_pre_call_hook`, `during_call_hook`, `async_post_call_streaming_iterator_hook`, `async_post_call_streaming_hook`, and `post_call_response_headers_hook` with fast-path exits when no relevant callbacks are registered - Add `needs_iterator_wrap()` and `needs_per_chunk_streaming_hook()` instance methods to decouple iterator-level wrapping from per-chunk hook execution; avoids `get_response_string` materialization per chunk when no guardrail or chunk-hook callback is active - Introduce `_fast_serialize_simple_model_response_stream()` using `orjson` for common single-choice text streaming chunks, bypassing the full Pydantic serializer; falls back to `model_dump_json` for tool calls, logprobs, usage, and provider-specific fields - Add early-return in `_restamp_streaming_chunk_model` when downstream model already matches the requested model, avoiding unnecessary string comparisons on every chunk - Fix stale zero-cost cache bug in `_is_model_cost_zero`: move the per-router `_zero_cost_cache` dict onto the `Router` instance and clear it in `_invalidate_model_group_info_cache` so in-place pricing updates via `upsert_deployment` immediately resume budget enforcement - Add `scripts/benchmark_chat_completions_perf.py`: standalone async benchmarking tool with a mock OpenAI provider, LiteLLM proxy process management, non-streaming RPS, streaming TTFT, and full-stream latency measurements with repeat/median run support - Add comprehensive unit tests covering capability detection, cache invalidation, fast-path correctness, zero-cost cache regression, and the no-callback streaming fast path Co-authored-by: Yassin Kortam <yassinkortam@g.ucla.edu>
487 lines
15 KiB
Python
487 lines
15 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import AsyncGenerator
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
import yaml
|
|
from fastapi.testclient import TestClient
|
|
|
|
sys.path.insert(0, os.path.abspath("../../.."))
|
|
|
|
import litellm
|
|
|
|
pytestmark = pytest.mark.flaky(condition=False)
|
|
|
|
|
|
def _initialize_proxy_with_config(config: dict, tmp_path) -> TestClient:
|
|
"""
|
|
Initialize the proxy server with a temporary config file and return a TestClient.
|
|
|
|
IMPORTANT: proxy_server.initialize() mutates module-level globals. We must call
|
|
cleanup_router_config_variables() before initializing to prevent cross-test bleed.
|
|
"""
|
|
from litellm.proxy.proxy_server import (
|
|
app,
|
|
cleanup_router_config_variables,
|
|
initialize,
|
|
)
|
|
|
|
cleanup_router_config_variables()
|
|
|
|
config_fp = tmp_path / "proxy_config.yaml"
|
|
config_fp.write_text(yaml.safe_dump(config))
|
|
|
|
asyncio.run(initialize(config=str(config_fp), debug=True))
|
|
return TestClient(app)
|
|
|
|
|
|
def _make_minimal_chat_completion_response(model: str) -> litellm.ModelResponse:
|
|
response = litellm.ModelResponse()
|
|
response.model = model
|
|
response.choices[0].message.content = "hello" # type: ignore[union-attr]
|
|
response.choices[0].finish_reason = "stop" # type: ignore[union-attr]
|
|
return response
|
|
|
|
|
|
def _make_model_response_stream_chunk(model: str) -> litellm.ModelResponseStream:
|
|
"""
|
|
Create a minimal OpenAI-compatible chat.completion.chunk object.
|
|
"""
|
|
chunk_dict = {
|
|
"id": "chatcmpl-test",
|
|
"object": "chat.completion.chunk",
|
|
"created": 0,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {"role": "assistant", "content": "hi"},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
return litellm.ModelResponseStream(**chunk_dict)
|
|
|
|
|
|
def _decode_sse_chunk(chunk) -> str:
|
|
return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
|
|
|
|
|
|
def test_restamp_streaming_chunk_skips_matching_model():
|
|
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
|
|
|
chunk = _make_model_response_stream_chunk("client-model")
|
|
|
|
result, model_mismatch_logged = _restamp_streaming_chunk_model(
|
|
chunk=chunk,
|
|
requested_model_from_client="client-model",
|
|
request_data={"litellm_call_id": "test-call-id"},
|
|
model_mismatch_logged=False,
|
|
)
|
|
|
|
assert result is chunk
|
|
assert result.model == "client-model"
|
|
assert model_mismatch_logged is False
|
|
|
|
|
|
def test_fast_serialize_simple_streaming_chunk_matches_model_dump_json():
|
|
from litellm.proxy.proxy_server import _serialize_streaming_chunk
|
|
|
|
chunk = _make_model_response_stream_chunk("client-model")
|
|
|
|
assert json.loads(_serialize_streaming_chunk(chunk)) == json.loads(
|
|
chunk.model_dump_json(exclude_none=True, exclude_unset=True)
|
|
)
|
|
|
|
|
|
def test_fast_serialize_returns_none_when_model_field_is_missing():
|
|
"""
|
|
The fast path must mirror ``model_dump_json(exclude_none=True)``: when
|
|
``chunk.model`` is ``None`` the slow path omits the field entirely.
|
|
Emitting ``"model": null`` would diverge and trip strict OpenAI-
|
|
compatible clients that reject ``null`` for optional string fields.
|
|
Falling back to ``None`` lets the canonical serializer handle the edge.
|
|
"""
|
|
from litellm.proxy.proxy_server import (
|
|
_fast_serialize_simple_model_response_stream,
|
|
_serialize_streaming_chunk,
|
|
)
|
|
|
|
chunk = _make_model_response_stream_chunk("client-model")
|
|
chunk.model = None # type: ignore[assignment]
|
|
|
|
assert _fast_serialize_simple_model_response_stream(chunk) is None
|
|
|
|
# Going through the public ``_serialize_streaming_chunk`` should still
|
|
# produce a serialized result via the slow-path fallback, and it must
|
|
# not contain ``"model": null``.
|
|
serialized = _serialize_streaming_chunk(chunk)
|
|
payload_str = (
|
|
serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized
|
|
)
|
|
assert '"model": null' not in payload_str
|
|
assert '"model":null' not in payload_str
|
|
assert json.loads(payload_str) == json.loads(
|
|
chunk.model_dump_json(exclude_none=True, exclude_unset=True)
|
|
)
|
|
|
|
|
|
def test_proxy_chat_completion_does_not_return_provider_prefixed_model(
|
|
tmp_path, monkeypatch
|
|
):
|
|
"""
|
|
Regression test:
|
|
|
|
- Client asks for `model="vllm-model"` (no provider prefix)
|
|
- Internal provider path uses `hosted_vllm/...`
|
|
- Proxy should not leak `hosted_vllm/` in the client-facing `model` field.
|
|
"""
|
|
client_model = "vllm-model"
|
|
internal_model = f"hosted_vllm/{client_model}"
|
|
|
|
client = _initialize_proxy_with_config(
|
|
config={
|
|
"general_settings": {"master_key": "sk-1234"},
|
|
"model_list": [
|
|
{
|
|
"model_name": client_model,
|
|
"litellm_params": {"model": internal_model},
|
|
}
|
|
],
|
|
},
|
|
tmp_path=tmp_path,
|
|
)
|
|
|
|
# Patch router call to avoid making any real network request.
|
|
from litellm.proxy import proxy_server
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server.llm_router, # type: ignore[arg-type]
|
|
"acompletion",
|
|
AsyncMock(
|
|
return_value=_make_minimal_chat_completion_response(model=internal_model)
|
|
),
|
|
)
|
|
|
|
# Also no-op proxy logging hooks to keep this test focused and deterministic.
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None)
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"update_request_status",
|
|
AsyncMock(return_value=None),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"post_call_success_hook",
|
|
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
|
|
)
|
|
|
|
resp = client.post(
|
|
"/v1/chat/completions",
|
|
headers={"Authorization": "Bearer sk-1234"},
|
|
json={"model": client_model, "messages": [{"role": "user", "content": "hi"}]},
|
|
)
|
|
|
|
assert resp.status_code == 200, resp.text
|
|
body = resp.json()
|
|
assert body["model"] == client_model
|
|
assert not body["model"].startswith("hosted_vllm/")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_streaming_chunks_do_not_return_provider_prefixed_model(
|
|
monkeypatch,
|
|
):
|
|
"""
|
|
Regression test for streaming:
|
|
|
|
Even if a streaming chunk contains `model="hosted_vllm/<...>"`, the proxy SSE layer
|
|
should not leak the provider prefix to the client.
|
|
"""
|
|
client_model = "vllm-model"
|
|
internal_model = f"hosted_vllm/{client_model}"
|
|
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
# Patch proxy_logging_obj hooks so async_data_generator yields exactly our chunk.
|
|
async def _iterator_hook(
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: AsyncGenerator,
|
|
request_data: dict,
|
|
):
|
|
yield _make_model_response_stream_chunk(model=internal_model)
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_iterator_hook",
|
|
_iterator_hook,
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_hook",
|
|
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"has_streaming_callbacks",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_iterator_wrap",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_per_chunk_streaming_hook",
|
|
MagicMock(return_value=True),
|
|
)
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="sk-1234")
|
|
|
|
gen = proxy_server.async_data_generator(
|
|
response=MagicMock(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data={"model": client_model},
|
|
)
|
|
|
|
chunks = []
|
|
async for item in gen:
|
|
chunks.append(item)
|
|
|
|
# First chunk is expected to be JSON, last chunk is [DONE]
|
|
assert len(chunks) >= 2
|
|
first = _decode_sse_chunk(chunks[0])
|
|
assert first.startswith("data: ")
|
|
|
|
payload = json.loads(first[len("data: ") :].strip())
|
|
assert payload["model"] == client_model
|
|
assert not payload["model"].startswith("hosted_vllm/")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_streaming_chunks_use_client_requested_model_before_alias_mapping(
|
|
monkeypatch,
|
|
):
|
|
"""
|
|
Regression test for alias mapping on streaming:
|
|
|
|
- `common_processing_pre_call_logic` can rewrite `request_data["model"]` via model_alias_map / key-specific aliases.
|
|
- Non-streaming responses are restamped using the original client-requested model (captured before the rewrite).
|
|
- Streaming chunks must do the same to avoid mismatched `model` values between streaming and non-streaming.
|
|
"""
|
|
client_model_alias = "alias-model"
|
|
canonical_model = "vllm-model"
|
|
internal_model = f"hosted_vllm/{canonical_model}"
|
|
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
async def _iterator_hook(
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: AsyncGenerator,
|
|
request_data: dict,
|
|
):
|
|
yield _make_model_response_stream_chunk(model=internal_model)
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_iterator_hook",
|
|
_iterator_hook,
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_hook",
|
|
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"has_streaming_callbacks",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_iterator_wrap",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_per_chunk_streaming_hook",
|
|
MagicMock(return_value=True),
|
|
)
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="sk-1234")
|
|
|
|
gen = proxy_server.async_data_generator(
|
|
response=MagicMock(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data={
|
|
"model": canonical_model,
|
|
"_litellm_client_requested_model": client_model_alias,
|
|
},
|
|
)
|
|
|
|
chunks = []
|
|
async for item in gen:
|
|
chunks.append(item)
|
|
|
|
assert len(chunks) >= 2
|
|
first = _decode_sse_chunk(chunks[0])
|
|
assert first.startswith("data: ")
|
|
|
|
payload = json.loads(first[len("data: ") :].strip())
|
|
assert payload["model"] == client_model_alias
|
|
assert not payload["model"].startswith("hosted_vllm/")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_streaming_azure_model_router_preserves_actual_model(monkeypatch):
|
|
"""
|
|
Regression test for Azure Model Router streaming:
|
|
|
|
When the client requests azure_ai/model_router, the streaming chunks should
|
|
preserve the actual model used (e.g., azure_ai/gpt-5-nano-2025-08-07) from
|
|
the downstream response, NOT override to the router model.
|
|
"""
|
|
router_model = "azure_ai/model_router"
|
|
actual_model_used = "azure_ai/gpt-5-nano-2025-08-07"
|
|
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
async def _iterator_hook(
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: AsyncGenerator,
|
|
request_data: dict,
|
|
):
|
|
yield _make_model_response_stream_chunk(model=actual_model_used)
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_iterator_hook",
|
|
_iterator_hook,
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_hook",
|
|
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"has_streaming_callbacks",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_iterator_wrap",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_per_chunk_streaming_hook",
|
|
MagicMock(return_value=True),
|
|
)
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="sk-1234")
|
|
|
|
gen = proxy_server.async_data_generator(
|
|
response=MagicMock(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data={
|
|
"model": router_model,
|
|
"_litellm_client_requested_model": router_model,
|
|
},
|
|
)
|
|
|
|
chunks = []
|
|
async for item in gen:
|
|
chunks.append(item)
|
|
|
|
assert len(chunks) >= 2
|
|
first = _decode_sse_chunk(chunks[0])
|
|
assert first.startswith("data: ")
|
|
|
|
payload = json.loads(first[len("data: ") :].strip())
|
|
# Azure Model Router: preserve actual model used, not the router model
|
|
assert payload["model"] == actual_model_used
|
|
assert payload["model"] != router_model
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_streaming_fastest_response_preserves_winning_model(monkeypatch):
|
|
"""
|
|
Regression test for fastest_response streaming:
|
|
|
|
When the client sends a comma-separated model list with fastest_response=True,
|
|
the streaming chunks should preserve the winning model's name from the
|
|
downstream response, NOT override to the comma-separated list.
|
|
"""
|
|
comma_separated_models = "openai/gpt-4o,gemini/gemini-2.5-flash"
|
|
winning_model = "gemini-2.5-flash"
|
|
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
async def _iterator_hook(
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: AsyncGenerator,
|
|
request_data: dict,
|
|
):
|
|
yield _make_model_response_stream_chunk(model=winning_model)
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_iterator_hook",
|
|
_iterator_hook,
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"async_post_call_streaming_hook",
|
|
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"has_streaming_callbacks",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_iterator_wrap",
|
|
MagicMock(return_value=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server.proxy_logging_obj,
|
|
"needs_per_chunk_streaming_hook",
|
|
MagicMock(return_value=True),
|
|
)
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="sk-1234")
|
|
|
|
gen = proxy_server.async_data_generator(
|
|
response=MagicMock(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data={
|
|
"model": comma_separated_models,
|
|
"_litellm_client_requested_model": comma_separated_models,
|
|
"fastest_response": True,
|
|
},
|
|
)
|
|
|
|
chunks = []
|
|
async for item in gen:
|
|
chunks.append(item)
|
|
|
|
assert len(chunks) >= 2
|
|
first = _decode_sse_chunk(chunks[0])
|
|
assert first.startswith("data: ")
|
|
|
|
payload = json.loads(first[len("data: ") :].strip())
|
|
assert payload["model"] == winning_model
|
|
assert payload["model"] != comma_separated_models
|