mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
* test: move the unit half of 126 mixed legacy files into tests/unit * test: restore litellm globals that moved tests set * test: finalize migration test cleanup * test: restore original bodies of moved legacy tests The move into tests/unit had rewritten 612 test bodies, and some of the rewrites dropped assertions. Each moved test now carries its original body from the legacy file, with only the imports, helpers, fake provider credentials and monkeypatched env it needs to run under tests/unit test_timeout_streaming goes back to tests/local_testing because it needs the fake OpenAI endpoint server. The image payload fixture moves with its only user, and two tests that leaked global state (a registered model cost entry and queued logging tasks) are now isolated * test: drop module imports shadowed by restored local imports * test: assert on LiteLLM output in no-assertion moved tests and isolate leaks Twenty no-assertion candidates get one assertion on the value LiteLLM returns, with the original lines unchanged. Four tests go back to their legacy files because they only check types or imports, write into the working directory, or cannot assert without a body change Two moved tests leaked globals into later tests in the same worker, so monkeypatch fixtures now restore the retry-after header parser and the end user cost tracking flags * test: drain queued logging tasks before the Phoenix span test The moved Phoenix test counted spans from logging tasks that earlier tests had queued, so the drain fixture moves to tests/unit/conftest.py and both it and the Datadog batch test use it. test_factory_function goes back to its legacy file because its returned wrapper calls the real Assistants API and cannot be asserted on without a body change --------- Co-authored-by: yuneng <yuneng@berri.ai>
1224 lines
37 KiB
Python
1224 lines
37 KiB
Python
import asyncio
|
|
import os
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
import litellm
|
|
import pytest
|
|
from litellm import Router
|
|
|
|
|
|
@pytest.fixture
|
|
def model_list():
|
|
return [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-5.5",
|
|
"litellm_params": {
|
|
"model": "gpt-5.5",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-image-1",
|
|
"litellm_params": {
|
|
"model": "gpt-image-1",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "cohere-rerank",
|
|
"litellm_params": {
|
|
"model": "cohere/rerank-english-v3.0",
|
|
"api_key": os.getenv("COHERE_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "claude-sonnet-4-5-20250929",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"mock_response": "hi this is macintosh.",
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aspeech_fallbacks_on_deployment_failure():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "tts-main",
|
|
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
|
|
},
|
|
{
|
|
"model_name": "tts-backup",
|
|
"litellm_params": {"model": "openai/tts-1-hd", "api_key": "fake-key"},
|
|
},
|
|
],
|
|
fallbacks=[{"tts-main": ["tts-backup"]}],
|
|
num_retries=0,
|
|
)
|
|
|
|
called_models = []
|
|
|
|
async def mock_aspeech(*args, **kwargs):
|
|
called_models.append(kwargs["model"])
|
|
if kwargs["model"] == "openai/tts-1":
|
|
raise litellm.InternalServerError(
|
|
message="deployment down",
|
|
llm_provider="openai",
|
|
model="tts-1",
|
|
)
|
|
return MagicMock()
|
|
|
|
with patch("litellm.aspeech", side_effect=mock_aspeech):
|
|
response = await router.aspeech(
|
|
model="tts-main",
|
|
input="the quick brown fox jumped over the lazy dogs",
|
|
voice="alloy",
|
|
)
|
|
|
|
assert response is not None
|
|
assert called_models == ["openai/tts-1", "openai/tts-1-hd"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aspeech_success_returns_response():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "tts",
|
|
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
|
|
},
|
|
]
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
|
|
response = await router.aspeech(
|
|
model="tts",
|
|
input="the quick brown fox jumped over the lazy dogs",
|
|
voice="alloy",
|
|
)
|
|
|
|
assert response is mock_response
|
|
mock_aspeech.assert_called_once()
|
|
assert mock_aspeech.call_args.kwargs["model"] == "openai/tts-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aspeech_sets_deployment_metadata():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "tts",
|
|
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
|
|
},
|
|
]
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
|
|
response = await router._aspeech(
|
|
model="tts",
|
|
input="the quick brown fox jumped over the lazy dogs",
|
|
voice="alloy",
|
|
)
|
|
|
|
assert response is mock_response
|
|
metadata = mock_aspeech.call_args.kwargs["metadata"]
|
|
assert metadata["deployment"] == "openai/tts-1"
|
|
assert metadata["deployment_model_name"] == "tts"
|
|
assert metadata["model_info"]["id"] is not None
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_moderation_endpoint_with_api_base():
|
|
"""
|
|
Test that the moderation endpoint respects api_base configuration
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
custom_api_base = "https://us.api.openai.com/v1"
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "openai/omni-moderation-latest",
|
|
"litellm_params": {
|
|
"model": "openai/omni-moderation-latest",
|
|
"api_base": custom_api_base,
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
)
|
|
|
|
# Mock the OpenAI client to verify api_base is passed
|
|
with patch(
|
|
"litellm.main.openai_chat_completions.get_openai_client"
|
|
) as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_response = MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
"id": "modr-123",
|
|
"model": "omni-moderation-latest",
|
|
"results": [
|
|
{
|
|
"flagged": False,
|
|
"categories": {},
|
|
"category_scores": {},
|
|
"category_applied_input_types": {},
|
|
}
|
|
],
|
|
}
|
|
mock_client.moderations.create = AsyncMock(return_value=mock_response)
|
|
mock_get_client.return_value = mock_client
|
|
|
|
response = await router.amoderation(
|
|
model="openai/omni-moderation-latest", input="hello this is a test"
|
|
)
|
|
|
|
# Verify that get_openai_client was called with the custom api_base
|
|
mock_get_client.assert_called()
|
|
call_kwargs = mock_get_client.call_args.kwargs
|
|
assert (
|
|
call_kwargs.get("api_base") == custom_api_base
|
|
), f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}"
|
|
|
|
print(f"✓ Moderation endpoint correctly uses api_base: {custom_api_base}")
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_aaaaatext_completion_endpoint(model_list, sync_mode):
|
|
router = Router(model_list=model_list)
|
|
|
|
if sync_mode:
|
|
response = router.text_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
else:
|
|
response = await router.atext_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
|
|
response_2 = await router._atext_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
assert response_2.choices[0].text == "I'm fine, thank you!"
|
|
|
|
assert response.choices[0].text == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_with_empty_choices(model_list):
|
|
"""
|
|
https://github.com/BerriAI/litellm/issues/8306
|
|
"""
|
|
router = Router(model_list=model_list)
|
|
mock_response = litellm.ModelResponse(
|
|
choices=[],
|
|
usage=litellm.Usage(
|
|
prompt_tokens=10,
|
|
completion_tokens=10,
|
|
total_tokens=20,
|
|
),
|
|
model="gpt-5-mini",
|
|
object="chat.completion",
|
|
created=1723081200,
|
|
).model_dump()
|
|
response = await router.acompletion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response=mock_response,
|
|
)
|
|
assert response is not None
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
def test_generic_api_call_with_fallbacks_basic(sync_mode):
|
|
"""
|
|
Test both the sync and async versions of generic_api_call_with_fallbacks with a basic successful call
|
|
"""
|
|
# Create a mock function that will be passed to generic_api_call_with_fallbacks
|
|
if sync_mode:
|
|
from unittest.mock import Mock
|
|
|
|
mock_function = Mock()
|
|
mock_function.__name__ = "test_function"
|
|
else:
|
|
mock_function = AsyncMock()
|
|
mock_function.__name__ = "test_function"
|
|
|
|
# Create a mock response
|
|
mock_response = {
|
|
"id": "resp_123456",
|
|
"role": "assistant",
|
|
"content": "This is a test response",
|
|
"model": "test-model",
|
|
"usage": {"input_tokens": 10, "output_tokens": 20},
|
|
}
|
|
mock_function.return_value = mock_response
|
|
|
|
# Create a router with a test model
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model-alias",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
# Call the appropriate generic_api_call_with_fallbacks method
|
|
if sync_mode:
|
|
response = router._generic_api_call_with_fallbacks(
|
|
model="test-model-alias",
|
|
original_function=mock_function,
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
max_tokens=100,
|
|
)
|
|
else:
|
|
response = asyncio.run(
|
|
router._ageneric_api_call_with_fallbacks(
|
|
model="test-model-alias",
|
|
original_function=mock_function,
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
max_tokens=100,
|
|
)
|
|
)
|
|
|
|
# Verify the mock function was called
|
|
mock_function.assert_called_once()
|
|
|
|
# Verify the response
|
|
assert response == mock_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aadapter_completion():
|
|
"""
|
|
Test the aadapter_completion method which uses async_function_with_fallbacks
|
|
"""
|
|
|
|
mock_response = {
|
|
"id": "adapter_resp_123",
|
|
"object": "adapter.completion",
|
|
"created": 1677858242,
|
|
"model": "test-model-with-adapter",
|
|
"choices": [
|
|
{
|
|
"text": "This is a test adapter response",
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
}
|
|
|
|
with patch.object(Router, "_aadapter_completion", new_callable=AsyncMock) as mock_method:
|
|
mock_method.return_value = mock_response
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-adapter-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router.async_function_with_fallbacks = AsyncMock(return_value=mock_response)
|
|
|
|
response = await router.aadapter_completion(
|
|
adapter_id="test-adapter-id",
|
|
model="test-adapter-model",
|
|
prompt="This is a test prompt",
|
|
max_tokens=100,
|
|
)
|
|
|
|
assert response == mock_response
|
|
|
|
router.async_function_with_fallbacks.assert_called_once()
|
|
call_kwargs = router.async_function_with_fallbacks.call_args.kwargs
|
|
assert call_kwargs["adapter_id"] == "test-adapter-id"
|
|
assert call_kwargs["model"] == "test-adapter-model"
|
|
assert call_kwargs["prompt"] == "This is a test prompt"
|
|
assert call_kwargs["max_tokens"] == 100
|
|
assert call_kwargs["original_function"] == router._aadapter_completion
|
|
assert "metadata" in call_kwargs
|
|
assert call_kwargs["metadata"]["model_group"] == "test-adapter-model"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test__aadapter_completion():
|
|
"""
|
|
Test the _aadapter_completion method directly
|
|
"""
|
|
|
|
mock_response = {
|
|
"id": "adapter_resp_123",
|
|
"object": "adapter.completion",
|
|
"created": 1677858242,
|
|
"model": "test-model-with-adapter",
|
|
"choices": [
|
|
{
|
|
"text": "This is a test adapter response",
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
}
|
|
|
|
with patch("litellm.aadapter_completion", new_callable=AsyncMock) as mock_adapter_completion:
|
|
mock_adapter_completion.return_value = mock_response
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-adapter-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router.async_get_available_deployment = AsyncMock(
|
|
return_value={
|
|
"model_name": "test-adapter-model",
|
|
"litellm_params": {
|
|
"model": "test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
"model_info": {
|
|
"id": "test-unique-id",
|
|
},
|
|
}
|
|
)
|
|
|
|
router.async_routing_strategy_pre_call_checks = AsyncMock()
|
|
|
|
response = await router._aadapter_completion(
|
|
adapter_id="test-adapter-id",
|
|
model="test-adapter-model",
|
|
prompt="This is a test prompt",
|
|
max_tokens=100,
|
|
)
|
|
|
|
assert response == mock_response
|
|
|
|
mock_adapter_completion.assert_called_once()
|
|
call_kwargs = mock_adapter_completion.call_args.kwargs
|
|
assert call_kwargs["adapter_id"] == "test-adapter-id"
|
|
assert call_kwargs["model"] == "test-model"
|
|
assert call_kwargs["prompt"] == "This is a test prompt"
|
|
assert call_kwargs["max_tokens"] == 100
|
|
assert call_kwargs["api_key"] == "fake-api-key"
|
|
assert call_kwargs["caching"] == router.cache_responses
|
|
|
|
assert router.success_calls["test-model"] == 1
|
|
assert router.total_calls["test-model"] == 1
|
|
|
|
router.async_routing_strategy_pre_call_checks.assert_called_once()
|
|
|
|
|
|
def test_initialize_router_endpoints():
|
|
"""
|
|
Test that initialize_router_endpoints correctly sets up all router endpoints
|
|
"""
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router.initialize_router_endpoints()
|
|
|
|
assert hasattr(router, "amoderation")
|
|
assert hasattr(router, "aanthropic_messages")
|
|
assert hasattr(router, "aresponses")
|
|
assert hasattr(router, "responses")
|
|
assert hasattr(router, "aget_responses")
|
|
assert hasattr(router, "adelete_responses")
|
|
|
|
assert callable(router.amoderation)
|
|
assert callable(router.aanthropic_messages)
|
|
assert callable(router.aresponses)
|
|
assert callable(router.responses)
|
|
assert callable(router.aget_responses)
|
|
assert callable(router.adelete_responses)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_responses_api_endpoints():
|
|
"""
|
|
A simpler test for _init_responses_api_endpoints that focuses on the basic functionality
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
# Create a router with a basic model
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
# Just mock the _ageneric_api_call_with_fallbacks method
|
|
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
|
|
|
# Add a mock implementation of _get_model_id_from_response_id to the Router instance
|
|
ResponsesAPIRequestUtils.get_model_id_from_response_id = MagicMock(
|
|
return_value=None
|
|
)
|
|
|
|
# Call without a response_id (no model extraction should happen)
|
|
await router._init_responses_api_endpoints(
|
|
original_function=AsyncMock(), thread_id="thread_xyz"
|
|
)
|
|
|
|
# Verify _ageneric_api_call_with_fallbacks was called but model wasn't changed
|
|
first_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
|
assert "model" not in first_call_kwargs
|
|
assert first_call_kwargs["thread_id"] == "thread_xyz"
|
|
|
|
# Reset the mock
|
|
router._ageneric_api_call_with_fallbacks.reset_mock()
|
|
|
|
# Change the return value for the second call
|
|
ResponsesAPIRequestUtils.get_model_id_from_response_id.return_value = (
|
|
"claude-3-sonnet"
|
|
)
|
|
|
|
# Call with a response_id
|
|
await router._init_responses_api_endpoints(
|
|
original_function=AsyncMock(), response_id="resp_claude_123"
|
|
)
|
|
|
|
# Verify model was updated in the kwargs
|
|
second_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
|
assert second_call_kwargs["model"] == "claude-3-sonnet"
|
|
assert second_call_kwargs["response_id"] == "resp_claude_123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_vector_store_api_endpoints():
|
|
"""
|
|
Test that _init_vector_store_api_endpoints correctly passes custom_llm_provider to kwargs
|
|
"""
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
mock_original_function = AsyncMock(return_value={"status": "success"})
|
|
|
|
result = await router._init_vector_store_api_endpoints(
|
|
original_function=mock_original_function, vector_store_id="test-store"
|
|
)
|
|
|
|
mock_original_function.assert_called_once_with(vector_store_id="test-store")
|
|
assert result == {"status": "success"}
|
|
|
|
mock_original_function.reset_mock()
|
|
|
|
await router._init_vector_store_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
vector_store_id="test-store",
|
|
)
|
|
|
|
mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai")
|
|
|
|
|
|
def test_apply_default_settings():
|
|
"""
|
|
Test the apply_default_settings method.
|
|
|
|
This test verifies that apply_default_settings correctly initializes
|
|
default pre-call checks and doesn't modify existing router state.
|
|
"""
|
|
|
|
router = Router()
|
|
initial_optional_callbacks = router.optional_callbacks
|
|
|
|
result = router.apply_default_settings()
|
|
|
|
assert result is None
|
|
|
|
assert router.optional_callbacks == initial_optional_callbacks
|
|
|
|
router_with_callbacks = Router()
|
|
mock_callback = MagicMock()
|
|
router_with_callbacks.optional_callbacks = [mock_callback]
|
|
|
|
result = router_with_callbacks.apply_default_settings()
|
|
|
|
assert result is None
|
|
|
|
assert mock_callback in router_with_callbacks.optional_callbacks
|
|
|
|
with patch.object(Router, "apply_default_settings") as mock_apply:
|
|
Router()
|
|
mock_apply.assert_called_once()
|
|
|
|
router_test = Router()
|
|
with patch.object(router_test, "add_optional_pre_call_checks") as mock_add_checks:
|
|
router_test.apply_default_settings()
|
|
|
|
mock_add_checks.assert_called_once_with([])
|
|
|
|
|
|
def test_initialize_core_endpoints():
|
|
"""
|
|
Test that _initialize_core_endpoints correctly sets up all core router endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_core_endpoints()
|
|
|
|
core_endpoints = [
|
|
"amoderation",
|
|
"aanthropic_messages",
|
|
"agenerate_content",
|
|
"aadapter_generate_content",
|
|
"aresponses",
|
|
"afile_delete",
|
|
"afile_content",
|
|
"responses",
|
|
"aget_responses",
|
|
"acancel_responses",
|
|
"adelete_responses",
|
|
"alist_input_items",
|
|
"_arealtime",
|
|
"acreate_fine_tuning_job",
|
|
"acancel_fine_tuning_job",
|
|
"alist_fine_tuning_jobs",
|
|
"aretrieve_fine_tuning_job",
|
|
"afile_list",
|
|
"aimage_edit",
|
|
"allm_passthrough_route",
|
|
]
|
|
|
|
for endpoint in core_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_specialized_endpoints():
|
|
"""
|
|
Test that _initialize_specialized_endpoints correctly sets up specialized endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_specialized_endpoints()
|
|
|
|
specialized_endpoints = [
|
|
"avector_store_search",
|
|
"avector_store_create",
|
|
"vector_store_search",
|
|
"vector_store_create",
|
|
"agenerate_content",
|
|
"generate_content",
|
|
"agenerate_content_stream",
|
|
"generate_content_stream",
|
|
"aocr",
|
|
"ocr",
|
|
"asearch",
|
|
"search",
|
|
"avideo_generation",
|
|
"video_generation",
|
|
"avideo_list",
|
|
"video_list",
|
|
"avideo_status",
|
|
"video_status",
|
|
"avideo_content",
|
|
"video_content",
|
|
"avideo_remix",
|
|
"video_remix",
|
|
"acreate_container",
|
|
"create_container",
|
|
"alist_containers",
|
|
"list_containers",
|
|
"aretrieve_container",
|
|
"retrieve_container",
|
|
"adelete_container",
|
|
"delete_container",
|
|
"acreate_skill",
|
|
"alist_skills",
|
|
"aget_skill",
|
|
"adelete_skill",
|
|
]
|
|
|
|
for endpoint in specialized_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_vector_store_endpoints():
|
|
"""
|
|
Test that _initialize_vector_store_endpoints correctly sets up vector store endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_vector_store_endpoints()
|
|
|
|
vector_store_endpoints = [
|
|
"avector_store_search",
|
|
"avector_store_create",
|
|
"vector_store_search",
|
|
"vector_store_create",
|
|
]
|
|
|
|
for endpoint in vector_store_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_vector_store_file_endpoints():
|
|
"""
|
|
Test that _initialize_vector_store_file_endpoints correctly sets up vector store file endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_vector_store_file_endpoints()
|
|
|
|
vector_store_file_endpoints = [
|
|
"avector_store_file_create",
|
|
"vector_store_file_create",
|
|
"avector_store_file_list",
|
|
"vector_store_file_list",
|
|
"avector_store_file_retrieve",
|
|
"vector_store_file_retrieve",
|
|
"avector_store_file_content",
|
|
"vector_store_file_content",
|
|
"avector_store_file_update",
|
|
"vector_store_file_update",
|
|
"avector_store_file_delete",
|
|
"vector_store_file_delete",
|
|
]
|
|
|
|
for endpoint in vector_store_file_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_google_genai_endpoints():
|
|
"""
|
|
Test that _initialize_google_genai_endpoints correctly sets up Google GenAI endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_google_genai_endpoints()
|
|
|
|
google_genai_endpoints = [
|
|
"agenerate_content",
|
|
"generate_content",
|
|
"agenerate_content_stream",
|
|
"generate_content_stream",
|
|
]
|
|
|
|
for endpoint in google_genai_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_ocr_search_endpoints():
|
|
"""
|
|
Test that _initialize_ocr_search_endpoints correctly sets up OCR and search endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_ocr_search_endpoints()
|
|
|
|
ocr_search_endpoints = [
|
|
"aocr",
|
|
"ocr",
|
|
"asearch",
|
|
"search",
|
|
]
|
|
|
|
for endpoint in ocr_search_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_video_endpoints():
|
|
"""
|
|
Test that _initialize_video_endpoints correctly sets up video endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_video_endpoints()
|
|
|
|
video_endpoints = [
|
|
"avideo_generation",
|
|
"video_generation",
|
|
"avideo_list",
|
|
"video_list",
|
|
"avideo_status",
|
|
"video_status",
|
|
"avideo_content",
|
|
"video_content",
|
|
"avideo_remix",
|
|
"video_remix",
|
|
]
|
|
|
|
for endpoint in video_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_container_endpoints():
|
|
"""
|
|
Test that _initialize_container_endpoints correctly sets up container endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_container_endpoints()
|
|
|
|
container_endpoints = [
|
|
"acreate_container",
|
|
"create_container",
|
|
"alist_containers",
|
|
"list_containers",
|
|
"aretrieve_container",
|
|
"retrieve_container",
|
|
"adelete_container",
|
|
"delete_container",
|
|
]
|
|
|
|
for endpoint in container_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_skills_endpoints():
|
|
"""
|
|
Test that _initialize_skills_endpoints correctly sets up skills endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_skills_endpoints()
|
|
|
|
skills_endpoints = [
|
|
"acreate_skill",
|
|
"alist_skills",
|
|
"aget_skill",
|
|
"adelete_skill",
|
|
]
|
|
|
|
for endpoint in skills_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints():
|
|
"""
|
|
Test that _init_containers_api_endpoints calls the original function
|
|
directly when there is no managed container ID (no embedded model_id).
|
|
"""
|
|
router = Router(model_list=[])
|
|
|
|
mock_response = {"id": "cntr_test", "name": "Test Container"}
|
|
mock_original_function = AsyncMock(return_value=mock_response)
|
|
|
|
result = await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
)
|
|
|
|
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container")
|
|
assert result == mock_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks():
|
|
"""
|
|
Managed ``cntr_`` IDs embed ``model_id``; router should decode and use
|
|
``_ageneric_api_call_with_fallbacks`` so deployment credentials apply.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "azure-router-model",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-5.5",
|
|
"api_key": "fake-key",
|
|
"api_base": "https://westus.api.cognitive.microsoft.com",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
|
|
|
managed_id = ResponsesAPIRequestUtils.build_container_id(
|
|
custom_llm_provider="azure",
|
|
model_id="azure-router-model",
|
|
container_id="cfile_upstream_abc",
|
|
)
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=AsyncMock(),
|
|
custom_llm_provider="openai",
|
|
container_id=managed_id,
|
|
file_id="cfile_xyz",
|
|
)
|
|
|
|
router._ageneric_api_call_with_fallbacks.assert_called_once()
|
|
call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
|
assert call_kw["model"] == "azure-router-model"
|
|
assert call_kw["container_id"] == "cfile_upstream_abc"
|
|
assert call_kw["file_id"] == "cfile_xyz"
|
|
assert call_kw["custom_llm_provider"] == "azure"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps():
|
|
"""
|
|
Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a
|
|
streaming response had no router metadata). The router must still unwrap the
|
|
managed ID before calling the upstream provider — otherwise the raw
|
|
``cntr_...`` token leaks downstream and the provider rejects it.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
router = Router(model_list=[])
|
|
mock_original_function = AsyncMock(return_value={"ok": True})
|
|
|
|
managed_id = ResponsesAPIRequestUtils.build_container_id(
|
|
custom_llm_provider="openai",
|
|
model_id=None,
|
|
container_id="cfile_upstream_abc",
|
|
)
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
container_id=managed_id,
|
|
file_id="cfile_xyz",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["container_id"] == "cfile_upstream_abc"
|
|
assert call_kw["file_id"] == "cfile_xyz"
|
|
assert call_kw["custom_llm_provider"] == "openai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider():
|
|
"""
|
|
A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with
|
|
an empty ``model_id`` (streaming events without router ``model_info.id``).
|
|
The router must still apply the decoded provider so the request routes to
|
|
the correct upstream — not stay on the default ``openai``.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
router = Router(model_list=[])
|
|
mock_original_function = AsyncMock(return_value={"ok": True})
|
|
|
|
managed_id = ResponsesAPIRequestUtils.build_container_id(
|
|
custom_llm_provider="azure",
|
|
model_id=None,
|
|
container_id="cfile_upstream_abc",
|
|
)
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
container_id=managed_id,
|
|
file_id="cfile_xyz",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["container_id"] == "cfile_upstream_abc"
|
|
assert call_kw["file_id"] == "cfile_xyz"
|
|
assert call_kw["custom_llm_provider"] == "azure"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch):
|
|
"""
|
|
``POST /v1/containers`` carries no container ID, so a ``model`` in the request
|
|
body is the only way to pick a deployment. The upstream call must receive that
|
|
deployment's ``api_key``/``api_base`` instead of falling back to the global
|
|
``OPENAI_API_KEY`` (which may be unset on the proxy).
|
|
"""
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5.4",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-5.4",
|
|
"api_key": "sk-model-list-key",
|
|
"api_base": "https://custom.openai.example/v1",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"})
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
model="gpt-5.4",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["api_key"] == "sk-model-list-key"
|
|
assert call_kw["api_base"] == "https://custom.openai.example/v1"
|
|
assert call_kw["model"] == "openai/gpt-5.4"
|
|
assert call_kw["name"] == "Test Container"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_create_without_model_calls_directly():
|
|
"""
|
|
Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list
|
|
must keep calling the handler directly with global provider credentials.
|
|
"""
|
|
router = Router(model_list=[])
|
|
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
|
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
model=None,
|
|
)
|
|
|
|
router._ageneric_api_call_with_fallbacks.assert_not_called()
|
|
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch):
|
|
"""
|
|
A ``model`` that names no configured deployment must not turn into a 400. The call
|
|
falls through to the handler with the caller's model and no injected deployment
|
|
credentials, matching the behaviour before model-based routing existed.
|
|
"""
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5.4",
|
|
"litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"},
|
|
}
|
|
]
|
|
)
|
|
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
model="does-not-exist",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["model"] == "does-not-exist"
|
|
assert call_kw["name"] == "Test Container"
|
|
assert "api_key" not in call_kw
|
|
assert "api_base" not in call_kw
|
|
|
|
|
|
def test_router_model_group_encrypted_content_affinity_callback_registration():
|
|
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
|
DeploymentAffinityCheck,
|
|
)
|
|
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
|
EncryptedContentAffinityCheck,
|
|
)
|
|
|
|
model_group = "openai.gpt-5.1-codex"
|
|
model_group_affinity_config = {
|
|
model_group: ["encrypted_content_affinity"],
|
|
}
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model_group,
|
|
"litellm_params": {
|
|
"model": "openai/gpt-5.1-codex",
|
|
"api_key": "mock-api-key",
|
|
},
|
|
}
|
|
],
|
|
model_group_affinity_config=model_group_affinity_config,
|
|
num_retries=0,
|
|
)
|
|
|
|
try:
|
|
callbacks = router.optional_callbacks or []
|
|
encrypted_content_callbacks = [
|
|
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
|
]
|
|
deployment_callback = next(
|
|
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
|
|
)
|
|
assert len(encrypted_content_callbacks) == 1
|
|
assert encrypted_content_callbacks[0].enable_global_affinity is False
|
|
assert (
|
|
encrypted_content_callbacks[0].model_group_affinity_config
|
|
== model_group_affinity_config
|
|
)
|
|
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
|
|
deployment_callback
|
|
)
|
|
|
|
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
|
|
|
|
callbacks = router.optional_callbacks or []
|
|
encrypted_content_callbacks = [
|
|
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
|
]
|
|
assert len(encrypted_content_callbacks) == 1
|
|
assert encrypted_content_callbacks[0].enable_global_affinity is True
|
|
assert encrypted_content_callbacks[0].router is router
|
|
finally:
|
|
router.discard()
|