tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-12 21:04:10 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
f5b11b72a6dcc8f4e7a16f00a61bdd124cc171d2
litellm/tests/test_litellm/proxy/test_response_model_sanitization.py
T
Yassin Kortam a6494e6fe3 perf: eliminate per-request callback scanning on proxy hot path (#27858) 
- 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>
2026-05-14 09:28:31 -07:00

487 lines
15 KiB
Python
Raw Blame History

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
Reference in New Issue View Git Blame Copy Permalink