scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78
litellm/tests/test_litellm/proxy/test_response_model_sanitization.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00

483 lines
15 KiB
Python
Raw Permalink Blame History

import asyncio
import json
from typing import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
import yaml
from fastapi.testclient import TestClient
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 a new issue View git blame Copy permalink