litellm/tests/router_unit_tests/test_router_helper_utils.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

2824 lines
97 KiB
Python

import os
import traceback
from dotenv import load_dotenv
from fastapi import Request
from datetime import datetime
from litellm import Router
import pytest
import litellm
from unittest.mock import patch, MagicMock, AsyncMock
from create_mock_standard_logging_payload import create_standard_logging_payload
from litellm.types.utils import StandardLoggingPayload
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
@pytest.fixture
def model_list():
return [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": os.getenv("OPENAI_API_KEY"),
"tpm": 1000, # Add TPM limit so async method doesn't return early
"rpm": 100, # Add RPM limit so async method doesn't return early
},
"model_info": {
"access_groups": ["group1", "group2"],
},
},
{
"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": "*",
"litellm_params": {
"model": "openai/*",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "claude-*",
"litellm_params": {
"model": "anthropic/*",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
},
]
def test_validate_fallbacks(model_list):
router = Router(model_list=model_list, fallbacks=[{"gpt-5.5": "gpt-5-mini"}])
router.validate_fallbacks(fallback_param=[{"gpt-5.5": "gpt-5-mini"}])
def test_routing_strategy_init(model_list):
"""Test if all routing strategies are initialized correctly"""
from litellm.types.router import RoutingStrategy
router = Router(model_list=model_list)
for strategy in RoutingStrategy:
router.routing_strategy_init(
routing_strategy=strategy, routing_strategy_args={}
)
def test_routing_strategy_init_invalid_strategy(model_list):
"""Test that invalid routing_strategy raises ValueError with helpful message.
See: https://github.com/BerriAI/litellm/issues/11330
Invalid strategies like 'simple' (without '-shuffle') should fail fast
with a clear error, not silently cause 'No deployments available' errors.
"""
router = Router(model_list=model_list)
# Test common mistake: "simple" instead of "simple-shuffle"
with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info:
router.routing_strategy_init(
routing_strategy="simple", routing_strategy_args={}
)
# Verify error message is helpful
error_msg = str(exc_info.value)
assert "Invalid routing_strategy" in error_msg
assert "simple" in error_msg
assert "simple-shuffle" in error_msg # Suggests the correct option
# Verify error message tells user WHERE to fix it
assert "config.yaml" in error_msg
assert "router_settings.routing_strategy" in error_msg
assert "Router SDK" in error_msg
# Test completely invalid strategy
with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info:
router.routing_strategy_init(
routing_strategy="not-a-real-strategy", routing_strategy_args={}
)
assert "Invalid routing_strategy" in str(exc_info.value)
def test_routing_strategy_init_valid_string_strategies(model_list):
"""Test that all valid string routing strategies work without error.
Valid strategies are derived from RoutingStrategy enum values plus 'simple-shuffle'.
"""
from litellm.types.router import RoutingStrategy
router = Router(model_list=model_list)
# All strategies from enum + simple-shuffle (default, not in enum)
valid_strategies = ["simple-shuffle"] + [s.value for s in RoutingStrategy]
for strategy in valid_strategies:
# Should not raise
router.routing_strategy_init(
routing_strategy=strategy, routing_strategy_args={}
)
def test_print_deployment(model_list):
"""Test if the api key is masked correctly"""
router = Router(model_list=model_list)
deployment = {
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": os.getenv("OPENAI_API_KEY"),
},
}
printed_deployment = router.print_deployment(deployment)
assert 10 * "*" in printed_deployment["litellm_params"]["api_key"]
def test_print_deployment_with_redact_enabled(model_list):
"""Test if sensitive credentials are masked when redact_user_api_key_info is enabled"""
import litellm
router = Router(model_list=model_list)
deployment = {
"model_name": "bedrock-claude",
"litellm_params": {
"model": "bedrock/anthropic.claude-v2",
"aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
"aws_region_name": "us-west-2",
},
}
original_setting = litellm.redact_user_api_key_info
try:
litellm.redact_user_api_key_info = True
printed_deployment = router.print_deployment(deployment)
assert "*" in printed_deployment["litellm_params"]["aws_access_key_id"]
assert "*" in printed_deployment["litellm_params"]["aws_secret_access_key"]
assert "us-west-2" == printed_deployment["litellm_params"]["aws_region_name"]
finally:
litellm.redact_user_api_key_info = original_setting
def test_completion(model_list):
"""Test if the completion function is working correctly"""
router = Router(model_list=model_list)
response = router._completion(
model="gpt-5-mini",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="I'm fine, thank you!",
)
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=6, delay=1)
@pytest.mark.asyncio
async def test_image_generation(model_list, sync_mode):
"""Test if the underlying '_image_generation' function is working correctly"""
from litellm.types.utils import ImageResponse
router = Router(model_list=model_list)
if sync_mode:
response = router._image_generation(
model="gpt-image-1",
prompt="A cute baby sea otter",
)
else:
response = await router._aimage_generation(
model="gpt-image-1",
prompt="A cute baby sea otter",
)
ImageResponse.model_validate(response)
@pytest.mark.asyncio
async def test_router_acompletion_util(model_list):
"""Test if the underlying '_acompletion' function is working correctly"""
router = Router(model_list=model_list)
response = await router._acompletion(
model="gpt-5-mini",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="I'm fine, thank you!",
)
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
@pytest.mark.asyncio
async def test_router_abatch_completion_one_model_multiple_requests_util(model_list):
"""Test if the 'abatch_completion_one_model_multiple_requests' function is working correctly"""
router = Router(model_list=model_list)
response = await router.abatch_completion_one_model_multiple_requests(
model="gpt-5-mini",
messages=[
[{"role": "user", "content": "Hello, how are you?"}],
[{"role": "user", "content": "Hello, how are you?"}],
],
mock_response="I'm fine, thank you!",
)
print(response)
assert response[0]["choices"][0]["message"]["content"] == "I'm fine, thank you!"
assert response[1]["choices"][0]["message"]["content"] == "I'm fine, thank you!"
@pytest.mark.asyncio
async def test_router_schedule_acompletion(model_list):
"""Test if the 'schedule_acompletion' function is working correctly"""
router = Router(model_list=model_list)
response = await router.schedule_acompletion(
model="gpt-5-mini",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="I'm fine, thank you!",
priority=1,
)
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
@pytest.mark.asyncio
async def test_router_schedule_atext_completion(model_list):
"""Test if the 'schedule_atext_completion' function is working correctly"""
from litellm.types.utils import TextCompletionResponse
router = Router(model_list=model_list)
with patch.object(
router, "_atext_completion", AsyncMock()
) as mock_atext_completion:
mock_atext_completion.return_value = TextCompletionResponse()
response = await router.atext_completion(
model="gpt-5-mini",
prompt="Hello, how are you?",
priority=1,
)
mock_atext_completion.assert_awaited_once()
assert "priority" not in mock_atext_completion.call_args.kwargs
@pytest.mark.asyncio
async def test_router_schedule_factory(model_list):
"""Test if the 'schedule_atext_completion' function is working correctly"""
from litellm.types.utils import TextCompletionResponse
router = Router(model_list=model_list)
with patch.object(
router, "_atext_completion", AsyncMock()
) as mock_atext_completion:
mock_atext_completion.return_value = TextCompletionResponse()
response = await router._schedule_factory(
model="gpt-5-mini",
args=(
"gpt-5-mini",
"Hello, how are you?",
),
priority=1,
kwargs={},
original_function=router.atext_completion,
)
mock_atext_completion.assert_awaited_once()
assert "priority" not in mock_atext_completion.call_args.kwargs
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_router_function_with_fallbacks(model_list, sync_mode):
"""Test if the router 'async_function_with_fallbacks' + 'function_with_fallbacks' are working correctly"""
router = Router(model_list=model_list)
data = {
"model": "gpt-5-mini",
"messages": [{"role": "user", "content": "Hello, how are you?"}],
"mock_response": "I'm fine, thank you!",
"num_retries": 0,
}
if sync_mode:
response = router.function_with_fallbacks(
original_function=router._completion,
**data,
)
else:
response = await router.async_function_with_fallbacks(
original_function=router._acompletion,
**data,
)
assert response.choices[0].message.content == "I'm fine, thank you!"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_router_function_with_retries(model_list, sync_mode):
"""Test if the router 'async_function_with_retries' + 'function_with_retries' are working correctly"""
router = Router(model_list=model_list)
data = {
"model": "gpt-5-mini",
"messages": [{"role": "user", "content": "Hello, how are you?"}],
"mock_response": "I'm fine, thank you!",
"num_retries": 0,
}
response = await router.async_function_with_retries(
original_function=router._acompletion,
**data,
)
assert response.choices[0].message.content == "I'm fine, thank you!"
@pytest.mark.asyncio
async def test_router_make_call(model_list):
"""Test if the router 'make_call' function is working correctly"""
## ACOMPLETION
router = Router(model_list=model_list)
response = await router.make_call(
original_function=router._acompletion,
model="gpt-5-mini",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="I'm fine, thank you!",
)
assert response.choices[0].message.content == "I'm fine, thank you!"
## ATEXT_COMPLETION
response = await router.make_call(
original_function=router._atext_completion,
model="gpt-5-mini",
prompt="Hello, how are you?",
mock_response="I'm fine, thank you!",
)
assert response.choices[0].text == "I'm fine, thank you!"
## AEMBEDDING
response = await router.make_call(
original_function=router._aembedding,
model="gpt-5-mini",
input="Hello, how are you?",
mock_response=[0.1, 0.2, 0.3],
)
assert response.data[0].embedding == [0.1, 0.2, 0.3]
## AIMAGE_GENERATION
response = await router.make_call(
original_function=router._aimage_generation,
model="gpt-image-1",
prompt="A cute baby sea otter",
mock_response="https://example.com/image.png",
)
assert response.data[0].url == "https://example.com/image.png"
def test_update_kwargs_with_deployment(model_list):
"""Test if the '_update_kwargs_with_deployment' function is working correctly"""
router = Router(model_list=model_list)
kwargs: dict = {"metadata": {}}
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
router._update_kwargs_with_deployment(
deployment=deployment,
kwargs=kwargs,
)
set_fields = ["deployment", "api_base", "model_info"]
assert all(field in kwargs["metadata"] for field in set_fields)
def test_update_kwargs_with_default_litellm_params(model_list):
"""Test if the '_update_kwargs_with_default_litellm_params' function is working correctly"""
router = Router(
model_list=model_list,
default_litellm_params={"api_key": "test", "metadata": {"key": "value"}},
)
kwargs: dict = {"metadata": {"key2": "value2"}}
router._update_kwargs_with_default_litellm_params(kwargs=kwargs)
assert kwargs["api_key"] == "test"
assert kwargs["metadata"]["key"] == "value"
assert kwargs["metadata"]["key2"] == "value2"
def test_get_timeout(model_list):
"""Test if the '_get_timeout' function is working correctly"""
router = Router(model_list=model_list)
timeout = router._get_timeout(kwargs={}, data={"timeout": 100})
assert timeout == 100
@pytest.mark.parametrize(
"fallback_kwarg, expected_error",
[
("mock_testing_fallbacks", litellm.InternalServerError),
("mock_testing_context_fallbacks", litellm.ContextWindowExceededError),
("mock_testing_content_policy_fallbacks", litellm.ContentPolicyViolationError),
],
)
def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_error):
"""Test if the '_handle_mock_testing_fallbacks' function is working correctly"""
router = Router(model_list=model_list)
data = {
fallback_kwarg: True,
}
with pytest.raises(expected_error):
router._handle_mock_testing_fallbacks(
kwargs=data,
)
def test_handle_mock_testing_rate_limit_error(model_list):
"""Test if the '_handle_mock_testing_rate_limit_error' function is working correctly"""
router = Router(model_list=model_list)
data = {
"mock_testing_rate_limit_error": True,
}
with pytest.raises(litellm.RateLimitError):
router._handle_mock_testing_rate_limit_error(
kwargs=data,
)
def test_get_fallback_model_group_from_fallbacks(model_list):
"""Test if the '_get_fallback_model_group_from_fallbacks' function is working correctly"""
router = Router(model_list=model_list)
fallback_model_group_name = router._get_fallback_model_group_from_fallbacks(
model_group="gpt-5.5",
fallbacks=[{"gpt-5.5": "gpt-5-mini"}],
)
assert fallback_model_group_name == "gpt-5-mini"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_deployment_callback_on_success(sync_mode):
"""Test if the '_deployment_callback_on_success' function is working correctly"""
import time
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": os.getenv("OPENAI_API_KEY"),
"rpm": 100,
},
"model_info": {"id": "100"},
}
]
router = Router(model_list=model_list)
# Get the actual deployment ID that was generated
gpt_deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
deployment_id = gpt_deployment["model_info"]["id"]
standard_logging_payload = create_standard_logging_payload()
standard_logging_payload["total_tokens"] = 100
standard_logging_payload["model_id"] = "100"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-5-mini",
},
"model_info": {"id": deployment_id},
},
"standard_logging_object": standard_logging_payload,
}
response = litellm.ModelResponse(
model="gpt-5-mini",
usage={"total_tokens": 100},
)
if sync_mode:
tpm_key = router.sync_deployment_callback_on_success(
kwargs=kwargs,
completion_response=response,
start_time=time.time(),
end_time=time.time(),
)
else:
tpm_key = await router.deployment_callback_on_success(
kwargs=kwargs,
completion_response=response,
start_time=time.time(),
end_time=time.time(),
)
assert tpm_key is not None
@pytest.mark.asyncio
async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
"""
An IO-limited deployment (itpm/otpm, no tpm/rpm) must still record TPM usage
in the router's routing counter so TPM-aware routing strategies see its real
load in mixed model groups; its itpm/otpm enforcement runs separately.
"""
import time
model_list = [
{
"model_name": "opus",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "sk-fake",
"itpm": 1000,
},
"model_info": {"id": "io-100"},
}
]
router = Router(model_list=model_list)
standard_logging_payload = create_standard_logging_payload()
standard_logging_payload["total_tokens"] = 100
standard_logging_payload["model_id"] = "io-100"
kwargs = {
"litellm_params": {
"metadata": {
"deployment": "openai/gpt-4o-mini",
"model_group": "opus",
},
"model_info": {"id": "io-100"},
},
"standard_logging_object": standard_logging_payload,
}
response = litellm.ModelResponse(model="openai/gpt-4o-mini", usage={"total_tokens": 100})
tpm_key = await router.deployment_callback_on_success(
kwargs=kwargs,
completion_response=response,
start_time=time.time(),
end_time=time.time(),
)
# The IO deployment is no longer skipped: its TPM routing counter is tracked.
assert tpm_key is not None
assert await router.cache.async_get_cache(key=tpm_key) == 100
@pytest.mark.asyncio
async def test_deployment_callback_on_failure(model_list):
"""Test if the '_deployment_callback_on_failure' function is working correctly"""
import time
router = Router(model_list=model_list)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-5-mini",
},
"model_info": {"id": 100},
},
}
result = router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
start_time=time.time(),
end_time=time.time(),
)
assert isinstance(result, bool)
assert result is False
model_response = router.completion(
model="gpt-5-mini",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="I'm fine, thank you!",
)
result = await router.async_deployment_callback_on_failure(
kwargs=kwargs,
completion_response=model_response,
start_time=time.time(),
end_time=time.time(),
)
def test_deployment_callback_respects_cooldown_time(model_list):
"""Ensure per-model cooldown_time is honored even when exception headers are present."""
import httpx
import time
from unittest.mock import patch
router = Router(model_list=model_list)
class FakeException(Exception):
def __init__(self):
self.status_code = 429
self.headers = httpx.Headers({"x-test": "1"})
kwargs = {
"exception": FakeException(),
"litellm_params": {
"metadata": {"model_group": "gpt-5-mini"},
"model_info": {"id": 100},
"cooldown_time": 0,
},
}
with patch("litellm.router._set_cooldown_deployments") as mock_set:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
start_time=time.time(),
end_time=time.time(),
)
mock_set.assert_called_once()
assert mock_set.call_args.kwargs["time_to_cooldown"] == 0
def test_log_retry(model_list):
"""Test if the '_log_retry' function is working correctly"""
import time
router = Router(model_list=model_list)
new_kwargs = router.log_retry(
kwargs={"metadata": {}},
e=Exception(),
)
assert "metadata" in new_kwargs
assert "previous_models" in new_kwargs["metadata"]
def test_update_usage(model_list):
"""Test if the '_update_usage' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
deployment_id = deployment["model_info"]["id"]
request_count = router._update_usage(
deployment_id=deployment_id, parent_otel_span=None
)
assert request_count == 1
request_count = router._update_usage(
deployment_id=deployment_id, parent_otel_span=None
)
assert request_count == 2
@pytest.mark.parametrize(
"finish_reason, expected_fallback", [("content_filter", True), ("stop", False)]
)
@pytest.mark.parametrize("fallback_type", ["model-specific", "default"])
def test_should_raise_content_policy_error(
model_list, finish_reason, expected_fallback, fallback_type
):
"""Test if the '_should_raise_content_policy_error' function is working correctly"""
router = Router(
model_list=model_list,
default_fallbacks=["gpt-5.5"] if fallback_type == "default" else None,
)
assert (
router._should_raise_content_policy_error(
model="gpt-5-mini",
response=litellm.ModelResponse(
model="gpt-5-mini",
choices=[
{
"finish_reason": finish_reason,
"message": {"content": "I'm fine, thank you!"},
}
],
usage={"total_tokens": 100},
),
kwargs={
"content_policy_fallbacks": (
[{"gpt-5-mini": "gpt-5.5"}]
if fallback_type == "model-specific"
else None
)
},
)
is expected_fallback
)
def test_get_healthy_deployments(model_list):
"""Test if the '_get_healthy_deployments' function is working correctly"""
router = Router(model_list=model_list)
deployments = router._get_healthy_deployments(
model="gpt-5-mini", parent_otel_span=None
)
assert len(deployments) > 0
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_routing_strategy_pre_call_checks(model_list, sync_mode):
"""Test if the '_routing_strategy_pre_call_checks' function is working correctly"""
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
callback = CustomLogger()
litellm.callbacks = [callback]
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
litellm_logging_obj = Logging(
model="gpt-5-mini",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
litellm_call_id="1234",
start_time=datetime.now(),
function_id="1234",
)
if sync_mode:
router.routing_strategy_pre_call_checks(deployment)
else:
## NO EXCEPTION
await router.async_routing_strategy_pre_call_checks(
deployment, litellm_logging_obj
)
## WITH EXCEPTION - rate limit error
with patch.object(
callback,
"async_pre_call_check",
AsyncMock(
side_effect=litellm.RateLimitError(
message="Rate limit error",
llm_provider="openai",
model="gpt-5-mini",
)
),
):
with pytest.raises(litellm.RateLimitError):
await router.async_routing_strategy_pre_call_checks(
deployment, litellm_logging_obj
)
## WITH EXCEPTION - generic error
with patch.object(
callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))
):
with pytest.raises(Exception, match="Error"):
await router.async_routing_strategy_pre_call_checks(
deployment, litellm_logging_obj
)
@pytest.mark.parametrize(
"set_supported_environments, supported_environments, is_supported",
[(True, ["staging"], True), (False, None, True), (True, ["development"], False)],
)
def test_create_deployment(
model_list, set_supported_environments, supported_environments, is_supported
):
"""Test if the '_create_deployment' function is working correctly"""
router = Router(model_list=model_list)
if set_supported_environments:
os.environ["LITELLM_ENVIRONMENT"] = "staging"
deployment = router._create_deployment(
deployment_info={},
_model_name="gpt-5-mini",
_litellm_params={
"model": "gpt-5-mini",
"api_key": "test",
"custom_llm_provider": "openai",
},
_model_info={
"id": 100,
"supported_environments": supported_environments,
},
)
if is_supported:
assert deployment is not None
else:
assert deployment is None
@pytest.mark.parametrize(
"set_supported_environments, supported_environments, is_supported",
[(True, ["staging"], True), (False, None, True), (True, ["development"], False)],
)
def test_deployment_is_active_for_environment(
model_list, set_supported_environments, supported_environments, is_supported
):
"""Test if the '_deployment_is_active_for_environment' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
if set_supported_environments:
os.environ["LITELLM_ENVIRONMENT"] = "staging"
deployment["model_info"]["supported_environments"] = supported_environments
if is_supported:
assert (
router.deployment_is_active_for_environment(deployment=deployment) is True
)
else:
assert (
router.deployment_is_active_for_environment(deployment=deployment) is False
)
def test_set_model_list(model_list):
"""Test if the '_set_model_list' function is working correctly"""
router = Router(model_list=model_list)
router.set_model_list(model_list=model_list)
assert len(router.model_list) == len(model_list)
def test_add_deployment(model_list):
"""Test if the '_add_deployment' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
deployment["model_info"]["id"] = "100"
## Test 1: call user facing function
router.add_deployment(deployment=deployment)
## Test 2: call internal function
router._add_deployment(deployment=deployment)
assert len(router.model_list) == len(model_list) + 1
def test_upsert_deployment(model_list):
"""Test if the 'upsert_deployment' function is working correctly"""
router = Router(model_list=model_list)
print("model list", len(router.model_list))
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
deployment.litellm_params.model = "gpt-5.5"
router.upsert_deployment(deployment=deployment)
assert len(router.model_list) == len(model_list)
def test_delete_deployment(model_list):
"""Test if the 'delete_deployment' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
router.delete_deployment(id=deployment["model_info"]["id"])
assert len(router.model_list) == len(model_list) - 1
def test_get_model_info(model_list):
"""Test if the 'get_model_info' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
model_info = router.get_model_info(id=deployment["model_info"]["id"])
assert model_info is not None
def test_get_model_group(model_list):
"""Test if the 'get_model_group' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
model_group = router.get_model_group(id=deployment["model_info"]["id"])
assert model_group is not None
assert model_group[0]["model_name"] == "gpt-5-mini"
@pytest.mark.parametrize("user_facing_model_group_name", ["gpt-5-mini", "gpt-5.5"])
def test_set_model_group_info(model_list, user_facing_model_group_name):
"""Test if the 'set_model_group_info' function is working correctly"""
router = Router(model_list=model_list)
resp = router._set_model_group_info(
model_group="gpt-5-mini",
user_facing_model_group_name=user_facing_model_group_name,
)
assert resp is not None
assert resp.model_group == user_facing_model_group_name
@pytest.mark.asyncio
async def test_set_response_headers(model_list):
"""Test if the 'set_response_headers' function is working correctly"""
router = Router(model_list=model_list)
resp = await router.set_response_headers(response=None, model_group=None)
assert resp is None
@pytest.mark.asyncio
async def test_set_response_headers_subtracts_in_flight_delta(model_list):
"""
LIT-2719: router-derived `x-ratelimit-remaining-*` headers must be
post-decrement (match OpenAI/Anthropic vendor semantics) so the proxy's
HTTP response headers and the prometheus gauges that read them stay
comparable across providers.
Router's TPM/RPM counter is incremented post-response by
`deployment_callback_on_success`, so `get_remaining_model_group_usage`
sees pre-decrement values. `set_response_headers` must replay the
in-flight increment before writing the headers.
"""
from pydantic import BaseModel
class _Usage(BaseModel):
total_tokens: int = 42
class _Resp(BaseModel):
usage: _Usage = _Usage()
_hidden_params: dict = {}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-limit-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
"x-ratelimit-limit-requests": 100,
}
)
resp = _Resp()
resp._hidden_params = {}
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 958
assert headers["x-ratelimit-remaining-requests"] == 99
# Limit headers pass through unmodified.
assert headers["x-ratelimit-limit-tokens"] == 1000
assert headers["x-ratelimit-limit-requests"] == 100
@pytest.mark.asyncio
async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list):
"""
The in-flight replay applies only to the post-incremented TPM/RPM counters
(`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are
incremented at reservation time (pre-call), so the input/output token
headers already reflect this request and must pass through untouched.
"""
from pydantic import BaseModel
class _Usage(BaseModel):
total_tokens: int = 30
prompt_tokens: int = 20
completion_tokens: int = 10
class _Resp(BaseModel):
usage: _Usage = _Usage()
_hidden_params: dict = {}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
"x-ratelimit-remaining-input-tokens": 1000,
"x-ratelimit-remaining-output-tokens": 500,
}
)
resp = _Resp()
resp._hidden_params = {}
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
# TPM/RPM headers replay the in-flight increment...
assert headers["x-ratelimit-remaining-tokens"] == 970
assert headers["x-ratelimit-remaining-requests"] == 99
# ...but the reservation-based input/output headers pass through unchanged.
assert headers["x-ratelimit-remaining-input-tokens"] == 1000
assert headers["x-ratelimit-remaining-output-tokens"] == 500
@pytest.mark.asyncio
async def test_get_model_group_io_token_usage_sums_across_deployments():
"""
get_model_group_io_token_usage must sum ITPM/OTPM across every deployment
in the model group (not just the first), reading the same per-deployment
cache keys the pre-call reservation writes to.
"""
from litellm.types.router import RouterCacheEnum
from litellm.utils import get_utc_datetime
router = Router(
model_list=[
{
"model_name": "opus",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"itpm": 1000,
"otpm": 500,
},
"model_info": {"id": "io-usage-dep-1"},
},
{
"model_name": "opus",
"litellm_params": {
"model": "openai/gpt-4o",
"itpm": 1000,
"otpm": 500,
},
"model_info": {"id": "io-usage-dep-2"},
},
]
)
minute = get_utc_datetime().strftime("%H-%M")
keys_and_values = [
(
RouterCacheEnum.ITPM.value.format(
id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute
),
30,
),
(
RouterCacheEnum.OTPM.value.format(
id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute
),
10,
),
(
RouterCacheEnum.ITPM.value.format(
id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute
),
70,
),
(
RouterCacheEnum.OTPM.value.format(
id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute
),
20,
),
]
for key, value in keys_and_values:
await router.cache.async_increment_cache(key=key, value=value, ttl=60)
current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus")
assert current_itpm == 100
assert current_otpm == 30
@pytest.mark.asyncio
async def test_get_model_group_io_token_usage_no_deployments_returns_none():
router = Router(model_list=[])
current_itpm, current_otpm = await router.get_model_group_io_token_usage(
"nonexistent-group"
)
assert current_itpm is None
assert current_otpm is None
@pytest.mark.asyncio
async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(model_list):
"""
A model group with both itpm/otpm and tpm/rpm limits must expose the
standard remaining-tokens/requests headers alongside the input/output token
headers, so clients and prometheus gauges relying on either still get data.
"""
from unittest.mock import Mock
from litellm.types.router import ModelGroupInfo
router = Router(model_list=model_list)
router._cached_get_model_group_info = Mock(
return_value=ModelGroupInfo(
model_group="gpt-3.5-turbo",
providers=["openai"],
itpm=2000,
otpm=1000,
tpm=5000,
rpm=50,
)
)
router.get_model_group_io_token_usage = AsyncMock(return_value=(100, 40))
router.get_model_group_usage = AsyncMock(return_value=(500, 5))
headers = await router.get_remaining_model_group_usage("gpt-3.5-turbo")
assert headers["x-ratelimit-remaining-input-tokens"] == 1900
assert headers["x-ratelimit-remaining-output-tokens"] == 960
assert headers["x-ratelimit-remaining-tokens"] == 4500
assert headers["x-ratelimit-remaining-requests"] == 45
@pytest.mark.asyncio
async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(model_list):
"""
A provider that natively returns `x-ratelimit-remaining-input-tokens` must
not suppress the router's own remaining-tokens/requests headers for a
non-IO model group.
"""
from pydantic import BaseModel
class _Usage(BaseModel):
total_tokens: int = 42
class _Resp(BaseModel):
usage: _Usage = _Usage()
_hidden_params: dict = {}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
}
)
resp = _Resp()
resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-input-tokens": 5}}
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 958
assert headers["x-ratelimit-remaining-requests"] == 99
# the provider's native header is left untouched
assert headers["x-ratelimit-remaining-input-tokens"] == 5
@pytest.mark.asyncio
async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(model_list):
from pydantic import BaseModel
class _Usage(BaseModel):
total_tokens: int = 42
class _Resp(BaseModel):
usage: _Usage = _Usage()
_hidden_params: dict = {}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
"x-ratelimit-remaining-input-tokens": 900,
"x-ratelimit-remaining-output-tokens": 450,
}
)
resp = _Resp()
resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-tokens": 5}}
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 5
assert headers["x-ratelimit-remaining-requests"] == 99
assert headers["x-ratelimit-remaining-input-tokens"] == 900
assert headers["x-ratelimit-remaining-output-tokens"] == 450
@pytest.mark.asyncio
async def test_set_response_headers_handles_missing_usage(model_list):
"""
Streaming chunks and some response shapes may lack a `usage` attribute or
populated `total_tokens`. The in-flight subtraction must default to 0
tokens (still subtract 1 from requests) and never raise.
"""
from pydantic import BaseModel
class _Resp(BaseModel):
_hidden_params: dict = {}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
}
)
resp = _Resp()
resp._hidden_params = {}
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 1000
assert headers["x-ratelimit-remaining-requests"] == 99
@pytest.mark.asyncio
async def test_set_response_headers_dict_anthropic_messages_response(model_list):
"""Anthropic /v1/messages returns a dict; IO rate-limit headers must attach."""
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-limit-input-tokens": 25,
"x-ratelimit-remaining-input-tokens": 20,
"x-ratelimit-limit-output-tokens": 100,
"x-ratelimit-remaining-output-tokens": 95,
}
)
resp = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "hi"}],
"usage": {"input_tokens": 5, "output_tokens": 1},
}
await router.set_response_headers(response=resp, model_group="io-itpm-strict")
assert "_hidden_params" in resp
headers = resp["_hidden_params"]["additional_headers"]
assert headers["x-litellm-model-group"] == "io-itpm-strict"
assert headers["x-ratelimit-limit-input-tokens"] == 25
assert headers["x-ratelimit-remaining-input-tokens"] == 20
assert headers["x-ratelimit-remaining-output-tokens"] == 95
@pytest.mark.asyncio
async def test_set_response_headers_wraps_bare_async_generator(model_list):
"""
Streaming responses that never go through Router.make_call's usual
object-based wrappers (e.g. the Anthropic /v1/messages -> Responses API
bridge, which yields a raw async generator with no `_hidden_params` slot)
must still get IO rate-limit headers attached via a thin wrapper.
"""
async def _raw_generator():
yield {"type": "message_start"}
yield {"type": "message_stop"}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-limit-input-tokens": 25,
"x-ratelimit-remaining-input-tokens": 20,
}
)
wrapped = await router.set_response_headers(response=_raw_generator(), model_group="io-itpm-strict")
assert hasattr(wrapped, "_hidden_params")
headers = wrapped._hidden_params["additional_headers"]
assert headers["x-litellm-model-group"] == "io-itpm-strict"
assert headers["x-ratelimit-limit-input-tokens"] == 25
assert headers["x-ratelimit-remaining-input-tokens"] == 20
from collections.abc import AsyncIterator
assert isinstance(wrapped, AsyncIterator)
chunks = [chunk async for chunk in wrapped]
assert chunks == [{"type": "message_start"}, {"type": "message_stop"}]
def test_get_all_deployments(model_list):
"""Test if the 'get_all_deployments' function is working correctly"""
router = Router(model_list=model_list)
deployments = router._get_all_deployments(
model_name="gpt-5-mini", model_alias="gpt-5-mini"
)
assert len(deployments) > 0
def test_get_model_access_groups(model_list):
"""Test if the 'get_model_access_groups' function is working correctly"""
router = Router(model_list=model_list)
access_groups = router.get_model_access_groups()
assert len(access_groups) == 2
def test_update_settings(model_list):
"""Test if the 'update_settings' function is working correctly"""
router = Router(model_list=model_list)
pre_update_allowed_fails = router.allowed_fails
router.update_settings(**{"allowed_fails": 20})
assert router.allowed_fails != pre_update_allowed_fails
assert router.allowed_fails == 20
def test_common_checks_available_deployment(model_list):
"""Test if the 'common_checks_available_deployment' function is working correctly"""
router = Router(model_list=model_list)
_, available_deployments = router._common_checks_available_deployment(
model="gpt-5-mini",
messages=[{"role": "user", "content": "hi"}],
input="hi",
specific_deployment=False,
)
assert len(available_deployments) > 0
def test_filter_cooldown_deployments(model_list):
"""Test if the 'filter_cooldown_deployments' function is working correctly"""
router = Router(model_list=model_list)
deployments = router._filter_cooldown_deployments(
healthy_deployments=router._get_all_deployments(model_name="gpt-5-mini"), # type: ignore
cooldown_deployments=[],
)
assert len(deployments) == len(router._get_all_deployments(model_name="gpt-5-mini"))
def test_track_deployment_metrics(model_list):
"""Test if the 'track_deployment_metrics' function is working correctly"""
from litellm.types.utils import ModelResponse
router = Router(model_list=model_list)
router._track_deployment_metrics(
deployment=router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
),
response=ModelResponse(
model="gpt-5-mini",
usage={"total_tokens": 100},
),
parent_otel_span=None,
)
@pytest.mark.parametrize(
"exception_type, exception_name, num_retries",
[
(litellm.exceptions.BadRequestError, "BadRequestError", 3),
(litellm.exceptions.AuthenticationError, "AuthenticationError", 4),
(litellm.exceptions.RateLimitError, "RateLimitError", 6),
(
litellm.exceptions.ContentPolicyViolationError,
"ContentPolicyViolationError",
7,
),
],
)
def test_get_num_retries_from_retry_policy(
model_list, exception_type, exception_name, num_retries
):
"""Test if the 'get_num_retries_from_retry_policy' function is working correctly"""
from litellm.router import RetryPolicy
data = {exception_name + "Retries": num_retries}
print("data", data)
router = Router(
model_list=model_list,
retry_policy=RetryPolicy(**data),
)
print("exception_type", exception_type)
calc_num_retries = router.get_num_retries_from_retry_policy(
exception=exception_type(
message="test", llm_provider="openai", model="gpt-5-mini"
)
)
assert calc_num_retries == num_retries
@pytest.mark.parametrize(
"exception_type, exception_name, allowed_fails",
[
(litellm.exceptions.BadRequestError, "BadRequestError", 3),
(litellm.exceptions.AuthenticationError, "AuthenticationError", 4),
(litellm.exceptions.RateLimitError, "RateLimitError", 6),
(
litellm.exceptions.ContentPolicyViolationError,
"ContentPolicyViolationError",
7,
),
],
)
def test_get_allowed_fails_from_policy(
model_list, exception_type, exception_name, allowed_fails
):
"""Test if the 'get_allowed_fails_from_policy' function is working correctly"""
from litellm.types.router import AllowedFailsPolicy
data = {exception_name + "AllowedFails": allowed_fails}
router = Router(
model_list=model_list, allowed_fails_policy=AllowedFailsPolicy(**data)
)
calc_allowed_fails = router.get_allowed_fails_from_policy(
exception=exception_type(
message="test", llm_provider="openai", model="gpt-5-mini"
)
)
assert calc_allowed_fails == allowed_fails
def test_initialize_alerting(model_list):
"""Test if the 'initialize_alerting' function is working correctly"""
from litellm.types.router import AlertingConfig
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
router = Router(
model_list=model_list, alerting_config=AlertingConfig(webhook_url="test")
)
router._initialize_alerting()
callback_added = False
for callback in litellm.callbacks:
if isinstance(callback, SlackAlerting):
callback_added = True
assert callback_added is True
def test_flush_cache(model_list):
"""Test if the 'flush_cache' function is working correctly"""
router = Router(model_list=model_list)
router.cache.set_cache("test", "test")
assert router.cache.get_cache("test") == "test"
router.flush_cache()
assert router.cache.get_cache("test") is None
def test_discard(model_list):
"""
Test that discard properly removes a Router from the callback lists
"""
litellm.callbacks = []
litellm.success_callback = []
litellm._async_success_callback = []
litellm.failure_callback = []
litellm._async_failure_callback = []
litellm.input_callback = []
litellm.service_callback = []
router = Router(model_list=model_list)
router.discard()
# Verify all callback lists are empty
assert len(litellm.callbacks) == 0
assert len(litellm.success_callback) == 0
assert len(litellm.failure_callback) == 0
assert len(litellm._async_success_callback) == 0
assert len(litellm._async_failure_callback) == 0
assert len(litellm.input_callback) == 0
assert len(litellm.service_callback) == 0
def test_initialize_assistants_endpoint(model_list):
"""Test if the 'initialize_assistants_endpoint' function is working correctly"""
router = Router(model_list=model_list)
router.initialize_assistants_endpoint()
assert router.acreate_assistants is not None
assert router.adelete_assistant is not None
assert router.aget_assistants is not None
assert router.acreate_thread is not None
assert router.aget_thread is not None
assert router.arun_thread is not None
assert router.aget_messages is not None
assert router.a_add_message is not None
def test_pass_through_assistants_endpoint_factory(model_list):
"""Test if the 'pass_through_assistants_endpoint_factory' function is working correctly"""
router = Router(model_list=model_list)
router._pass_through_assistants_endpoint_factory(
original_function=litellm.acreate_assistants,
custom_llm_provider="openai",
client=None,
**{},
)
def test_factory_function(model_list):
"""Test if the 'factory_function' function is working correctly"""
router = Router(model_list=model_list)
router.factory_function(litellm.acreate_assistants)
def test_get_model_from_alias(model_list):
"""Test if the 'get_model_from_alias' function is working correctly"""
router = Router(
model_list=model_list,
model_group_alias={"gpt-5.5": "gpt-5-mini"},
)
model = router._get_model_from_alias(model="gpt-5.5")
assert model == "gpt-5-mini"
def test_get_deployment_by_litellm_model(model_list):
"""Test if the 'get_deployment_by_litellm_model' function is working correctly"""
router = Router(model_list=model_list)
deployment = router._get_deployment_by_litellm_model(model="gpt-5-mini")
assert deployment is not None
def test_get_pattern(model_list):
router = Router(model_list=model_list)
pattern = router.pattern_router.get_pattern(model="claude-3")
assert pattern is not None
def test_deployments_by_pattern(model_list):
router = Router(model_list=model_list)
deployments = router.pattern_router.get_deployments_by_pattern(model="claude-3")
assert deployments is not None
# def test_pattern_match_deployments(model_list):
# from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
# import re
# patter_router = PatternMatchRouter()
# request = "fo::hi::static::hello"
# model_name = "fo::*:static::*"
# model_name_regex = patter_router._pattern_to_regex(model_name)
# # Match against the request
# match = re.match(model_name_regex, request)
# print(f"match: {match}")
# print(f"match.end: {match.end()}")
# if match is None:
# raise ValueError("Match not found")
# updated_model = patter_router.set_deployment_model_name(
# matched_pattern=match, litellm_deployment_litellm_model="openai/*"
# )
# assert updated_model == "openai/fo::hi:static::hello"
@pytest.mark.parametrize(
"user_request_model, model_name, litellm_model, expected_model",
[
("llmengine/foo", "llmengine/*", "openai/foo", "openai/foo"),
("llmengine/foo", "llmengine/*", "openai/*", "openai/foo"),
(
"fo::hi::static::hello",
"fo::*::static::*",
"openai/fo::*:static::*",
"openai/fo::hi:static::hello",
),
(
"fo::hi::static::hello",
"fo::*::static::*",
"openai/gpt-5-mini",
"openai/gpt-5-mini",
),
(
"bedrock/meta.llama3-70b",
"*meta.llama3*",
"bedrock/meta.llama3-*",
"bedrock/meta.llama3-70b",
),
(
"meta.llama3-70b",
"*meta.llama3*",
"bedrock/meta.llama3-*",
"meta.llama3-70b",
),
],
)
def test_pattern_match_deployment_set_model_name(
user_request_model, model_name, litellm_model, expected_model
):
from re import Match
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
pattern_router = PatternMatchRouter()
import re
# Convert model_name into a proper regex
model_name_regex = pattern_router._pattern_to_regex(model_name)
# Match against the request
match = re.match(model_name_regex, user_request_model)
if match is None:
raise ValueError("Match not found")
# Call the set_deployment_model_name function
updated_model = pattern_router.set_deployment_model_name(match, litellm_model)
print(updated_model) # Expected output: "openai/fo::hi:static::hello"
assert updated_model == expected_model
updated_models = pattern_router._return_pattern_matched_deployments(
match,
deployments=[
{
"model_name": model_name,
"litellm_params": {"model": litellm_model},
}
],
)
for model in updated_models:
assert model["litellm_params"]["model"] == expected_model
@pytest.mark.asyncio
async def test_pass_through_moderation_endpoint_factory(model_list):
router = Router(model_list=model_list)
response = await router._pass_through_moderation_endpoint_factory(
original_function=litellm.amoderation,
input="this is valid good text",
model=None,
)
assert response is not None
@pytest.mark.parametrize(
"has_default_fallbacks, expected_result",
[(True, True), (False, False)],
)
def test_has_default_fallbacks(model_list, has_default_fallbacks, expected_result):
router = Router(
model_list=model_list,
default_fallbacks=(
["my-default-fallback-model"] if has_default_fallbacks else None
),
)
assert router._has_default_fallbacks() is expected_result
def test_add_optional_pre_call_checks(model_list):
router = Router(model_list=model_list)
router.add_optional_pre_call_checks(["prompt_caching"])
assert len(litellm.callbacks) > 0
@pytest.mark.asyncio
async def test_async_callback_filter_deployments(model_list):
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
router = Router(model_list=model_list)
healthy_deployments = router.get_model_list(model_name="gpt-5-mini")
new_healthy_deployments = await router.async_callback_filter_deployments(
model="gpt-5-mini",
healthy_deployments=healthy_deployments,
messages=[],
parent_otel_span=None,
)
assert len(new_healthy_deployments) == len(healthy_deployments)
def test_cached_get_model_group_info(model_list):
"""Test if the '_cached_get_model_group_info' function is working correctly with LRU cache"""
router = Router(model_list=model_list)
# First call - should hit the actual function
result1 = router._cached_get_model_group_info("gpt-5-mini")
# Second call with same argument - should hit the cache
result2 = router._cached_get_model_group_info("gpt-5-mini")
# Verify results are the same
assert result1 == result2
# Verify the cache info shows hits
cache_info = router._cached_get_model_group_info.cache_info()
assert cache_info.hits > 0 # Should have at least one cache hit
def test_init_responses_api_endpoints(model_list):
"""Test if the '_init_responses_api_endpoints' function is working correctly"""
from typing import Callable
router = Router(model_list=model_list)
assert router.aget_responses is not None
assert isinstance(router.aget_responses, Callable)
assert router.adelete_responses is not None
assert isinstance(router.adelete_responses, Callable)
@pytest.mark.parametrize(
"mock_testing_fallbacks, mock_testing_context_fallbacks, mock_testing_content_policy_fallbacks, expected_fallbacks, expected_context, expected_content_policy",
[
# Test string to bool conversion
("true", "false", "True", True, False, True),
("TRUE", "FALSE", "False", True, False, False),
("false", "true", "false", False, True, False),
# Test actual boolean values (should pass through unchanged)
(True, False, True, True, False, True),
(False, True, False, False, True, False),
# Test None values
(None, None, None, None, None, None),
# Test mixed types
("true", False, None, True, False, None),
],
)
def test_mock_router_testing_params_str_to_bool_conversion(
mock_testing_fallbacks,
mock_testing_context_fallbacks,
mock_testing_content_policy_fallbacks,
expected_fallbacks,
expected_context,
expected_content_policy,
):
"""Test if MockRouterTestingParams.from_kwargs correctly converts string values to booleans using str_to_bool"""
from litellm.types.router import MockRouterTestingParams
kwargs = {
"mock_testing_fallbacks": mock_testing_fallbacks,
"mock_testing_context_fallbacks": mock_testing_context_fallbacks,
"mock_testing_content_policy_fallbacks": mock_testing_content_policy_fallbacks,
"other_param": "should_remain", # This should not be affected
}
# Make a copy to verify kwargs are properly popped
original_kwargs = kwargs.copy()
mock_params = MockRouterTestingParams.from_kwargs(kwargs)
# Verify the converted values
assert mock_params.mock_testing_fallbacks == expected_fallbacks
assert mock_params.mock_testing_context_fallbacks == expected_context
assert mock_params.mock_testing_content_policy_fallbacks == expected_content_policy
# Verify that the mock testing params were popped from kwargs
assert "mock_testing_fallbacks" not in kwargs
assert "mock_testing_context_fallbacks" not in kwargs
assert "mock_testing_content_policy_fallbacks" not in kwargs
# Verify other params remain unchanged
assert kwargs["other_param"] == "should_remain"
def test_is_auto_router_deployment(model_list):
"""Test if the '_is_auto_router_deployment' function correctly identifies auto-router deployments"""
router = Router(model_list=model_list)
# Test case 1: Model starts with "auto_router/" - should return True
litellm_params_auto = LiteLLM_Params(model="auto_router/my-auto-router")
assert router._is_auto_router_deployment(litellm_params_auto) is True
# Test case 2: Model doesn't start with "auto_router/" - should return False
litellm_params_regular = LiteLLM_Params(model="gpt-5-mini")
assert router._is_auto_router_deployment(litellm_params_regular) is False
# Test case 3: Model is empty string - should return False
litellm_params_empty = LiteLLM_Params(model="")
assert router._is_auto_router_deployment(litellm_params_empty) is False
# Test case 4: Model contains "auto_router/" but doesn't start with it - should return False
litellm_params_contains = LiteLLM_Params(model="prefix_auto_router/something")
assert router._is_auto_router_deployment(litellm_params_contains) is False
@patch("litellm.router_strategy.auto_router.auto_router.AutoRouter")
def test_init_auto_router_deployment_success(mock_auto_router, model_list):
"""Test if the 'init_auto_router_deployment' function successfully initializes auto-router when all params provided"""
router = Router(model_list=model_list)
# Create a mock AutoRouter instance
mock_auto_router_instance = MagicMock()
mock_auto_router.return_value = mock_auto_router_instance
# Test case: All required parameters provided
litellm_params = LiteLLM_Params(
model="auto_router/test",
auto_router_config_path="/path/to/config",
auto_router_default_model="gpt-5-mini",
auto_router_embedding_model="text-embedding-3-small",
)
deployment = Deployment(
model_name="test-auto-router",
litellm_params=litellm_params,
model_info={"id": "test-id"},
)
# Should not raise any exception
router.init_auto_router_deployment(deployment)
# Verify AutoRouter was called with correct parameters
mock_auto_router.assert_called_once_with(
model_name="test-auto-router",
auto_router_config_path="/path/to/config",
auto_router_config=None,
default_model="gpt-5-mini",
embedding_model="text-embedding-3-small",
litellm_router_instance=router,
max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS,
)
# Verify the auto-router was added to the router's auto_routers dict
assert "test-auto-router" in router.auto_routers
assert router.auto_routers["test-auto-router"][0].strategy == mock_auto_router_instance
@patch("litellm.router_strategy.auto_router.auto_router.AutoRouter")
def test_init_auto_router_deployment_duplicate_model_name(mock_auto_router, model_list):
"""Test if the 'init_auto_router_deployment' function raises ValueError when model_name already exists"""
router = Router(model_list=model_list)
# Create a mock AutoRouter instance
mock_auto_router_instance = MagicMock()
mock_auto_router.return_value = mock_auto_router_instance
# Add an existing auto-router
from litellm.types.router import TaggedPreRoutingStrategy
router.auto_routers["test-auto-router"] = [
TaggedPreRoutingStrategy(tags=(), strategy=mock_auto_router_instance)
]
# Try to add another auto-router with the same name
litellm_params = LiteLLM_Params(
model="auto_router/test",
auto_router_config_path="/path/to/config",
auto_router_default_model="gpt-5-mini",
auto_router_embedding_model="text-embedding-3-small",
)
deployment = Deployment(
model_name="test-auto-router",
litellm_params=litellm_params,
model_info={"id": "test-id"},
)
with pytest.raises(
ValueError, match=r"Auto-router deployment test-auto-router with tags .* already exists"
):
router.init_auto_router_deployment(deployment)
def testgenerate_model_id_with_deployment_model_name(model_list):
"""Test that generate_model_id works correctly with deployment model_name and handles None values properly"""
router = Router(model_list=model_list)
# Test case 1: Normal case with valid model_group and litellm_params
model_group = "gpt-4.1"
litellm_params = {
"model": "gpt-4.1",
"api_key": "test_key",
"api_base": "https://api.openai.com/v1",
}
try:
result = router.generate_model_id(
model_group=model_group, litellm_params=litellm_params
)
assert isinstance(result, str)
assert len(result) > 0
print(f"✓ Success with valid model_group: {result}")
except Exception as e:
pytest.fail(f"Failed with valid model_group: {e}")
# Test case 2: Edge case with None model_group (this should fail as expected - our fix prevents this from happening)
with pytest.raises(TypeError) as exc_info:
router.generate_model_id(model_group=None, litellm_params=litellm_params)
# After optimization, error message changed but still fails appropriately on None
error_str = str(exc_info.value)
assert (
"unsupported operand type(s) for +=" in error_str
or "expected str instance, NoneType found" in error_str
)
# Test case 3: Edge case with None key in litellm_params
litellm_params_with_none_key = {
"model": "gpt-4.1",
"api_key": "test_key",
None: "should_be_skipped", # This should be handled gracefully
}
try:
result = router.generate_model_id(
model_group=model_group, litellm_params=litellm_params_with_none_key
)
assert isinstance(result, str)
assert len(result) > 0
print(f"✓ Success with None key in litellm_params: {result}")
except Exception as e:
pytest.fail(f"Failed with None key in litellm_params: {e}")
# Test case 4: Edge case with empty litellm_params
try:
result = router.generate_model_id(model_group=model_group, litellm_params={})
assert isinstance(result, str)
assert len(result) > 0
print(f"✓ Success with empty litellm_params: {result}")
except Exception as e:
pytest.fail(f"Failed with empty litellm_params: {e}")
# Test case 5: Verify that the same inputs produce the same result (deterministic)
result1 = router.generate_model_id(
model_group=model_group, litellm_params=litellm_params
)
result2 = router.generate_model_id(
model_group=model_group, litellm_params=litellm_params
)
assert result1 == result2, "Model ID generation should be deterministic"
print("✓ All generate_model_id tests passed!")
def test_handle_clientside_credential_with_deployment_model_name(model_list):
"""Test that _handle_clientside_credential uses deployment model_name correctly"""
router = Router(model_list=model_list)
# Mock deployment with model_name
deployment = {
"model_name": "gpt-4.1",
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
}
# Mock kwargs with empty metadata (simulating the original issue)
kwargs = {
"metadata": {}, # Empty metadata, no model_group
"litellm_params": {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
},
}
# Mock dynamic_litellm_params that would be returned by get_dynamic_litellm_params
dynamic_litellm_params = {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
}
# Test that the method doesn't fail when metadata is empty
try:
# This would normally call generate_model_id internally
# We're testing that the fix prevents the TypeError
model_group = deployment["model_name"] # This is what our fix does
assert model_group == "gpt-4.1"
# Verify that generate_model_id works with this model_group
result = router.generate_model_id(
model_group=model_group, litellm_params=dynamic_litellm_params
)
assert isinstance(result, str)
assert len(result) > 0
print(f"✓ Success with deployment model_name: {result}")
except Exception as e:
pytest.fail(f"Failed with deployment model_name: {e}")
print("✓ _handle_clientside_credential test passed!")
def test_sync_generic_api_call_preserves_requested_model_group_in_logs():
router = Router(
model_list=[
{
"model_name": "claude-sonnet-4-6",
"litellm_params": {
"model": "bedrock/global.anthropic.claude-sonnet-4-6",
"aws_access_key_id": "test-access-key",
"aws_secret_access_key": "test-secret-key",
"aws_region_name": "us-west-2",
},
}
]
)
try:
captured_kwargs = {}
def mock_original_function(**kwargs):
captured_kwargs.update(kwargs)
return {"status": "ok"}
response = router._generic_api_call_with_fallbacks(
model="claude-sonnet-4-6",
original_function=mock_original_function,
)
assert response == {"status": "ok"}
assert captured_kwargs["model"] == "bedrock/global.anthropic.claude-sonnet-4-6"
assert captured_kwargs["litellm_metadata"]["model_group"] == "claude-sonnet-4-6"
assert (
captured_kwargs["litellm_metadata"]["deployment"]
== "bedrock/global.anthropic.claude-sonnet-4-6"
)
finally:
router.discard()
def test_sync_generic_api_call_uses_request_kwargs_for_deployment_selection():
router = Router(
model_list=[
{
"model_name": "regional-model",
"litellm_params": {
"model": "anthropic/us-model",
"api_key": "test-api-key",
"region_name": "us",
},
},
{
"model_name": "regional-model",
"litellm_params": {
"model": "anthropic/eu-model",
"api_key": "test-api-key",
"region_name": "eu",
},
},
],
enable_pre_call_checks=True,
)
try:
captured_kwargs = {}
def mock_original_function(**kwargs):
captured_kwargs.update(kwargs)
return {"status": "ok"}
response = router._generic_api_call_with_fallbacks(
model="regional-model",
original_function=mock_original_function,
messages=[{"role": "user", "content": "Hello from Europe"}],
allowed_model_region="eu",
)
assert response == {"status": "ok"}
assert captured_kwargs["model"] == "anthropic/eu-model"
finally:
router.discard()
@pytest.mark.parametrize(
"function_name, expected_metadata_key",
[
("acompletion", "metadata"),
("_ageneric_api_call_with_fallbacks", "litellm_metadata"),
("batch", "litellm_metadata"),
("completion", "metadata"),
("acreate_file", "litellm_metadata"),
("aget_file", "litellm_metadata"),
],
)
def test_handle_clientside_credential_metadata_loading(
model_list, function_name, expected_metadata_key
):
"""Test that _handle_clientside_credential correctly loads metadata based on function name"""
router = Router(model_list=model_list)
# Mock deployment
deployment = {
"model_name": "gpt-4.1",
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
"model_info": {"id": "original-id-123"},
}
# Mock kwargs with clientside credentials and metadata
kwargs = {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
expected_metadata_key: {"model_group": "gpt-4.1", "custom_field": "test_value"},
}
# Call the function
result_deployment = router._handle_clientside_credential(
deployment=deployment, kwargs=kwargs, function_name=function_name
)
# Verify the result is a Deployment object
assert isinstance(result_deployment, Deployment)
# Verify the deployment has the correct model_name (should be the model_group from metadata)
assert result_deployment.model_name == "gpt-4.1"
# Verify the litellm_params contain the clientside credentials
assert result_deployment.litellm_params.api_key == "client_side_key"
assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1"
# Verify the model_info has been updated with a new ID
assert result_deployment.model_info.id != "original-id-123"
assert result_deployment.model_info.original_model_id == "original-id-123"
# Verify the deployment was added to the router
assert len(router.model_list) == len(model_list) + 1
# Test that the function correctly uses the right metadata key
# For acompletion, it should use "metadata"
# For _ageneric_api_call_with_fallbacks/batch, it should use "litellm_metadata"
if function_name == "acompletion":
assert "metadata" in kwargs
assert "litellm_metadata" not in kwargs
elif function_name in [
"_ageneric_api_call_with_fallbacks",
"batch",
"acreate_file",
"aget_file",
]:
assert "litellm_metadata" in kwargs
# Note: acompletion would not have litellm_metadata, but other functions might have both
print(
f"✓ Success with function_name '{function_name}' using '{expected_metadata_key}' metadata key"
)
@pytest.mark.parametrize(
"function_name, metadata_key",
[
("acompletion", "metadata"),
("_ageneric_api_call_with_fallbacks", "litellm_metadata"),
],
)
def test_handle_clientside_credential_metadata_variable_name(
model_list, function_name, metadata_key
):
"""Test that _handle_clientside_credential uses the correct metadata variable name based on function name"""
from litellm.router_utils.batch_utils import _get_router_metadata_variable_name
router = Router(model_list=model_list)
# Verify the metadata variable name is correct for each function
expected_metadata_key = _get_router_metadata_variable_name(
function_name=function_name
)
assert expected_metadata_key == metadata_key
# Mock deployment
deployment = {
"model_name": "gpt-4.1",
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
"model_info": {"id": "original-id-456"},
}
# Mock kwargs with clientside credentials and the correct metadata key
kwargs = {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
metadata_key: {"model_group": "gpt-4.1", "test_field": "test_value"},
}
# Call the function
result_deployment = router._handle_clientside_credential(
deployment=deployment, kwargs=kwargs, function_name=function_name
)
# Verify the function correctly extracted model_group from the right metadata key
assert result_deployment.model_name == "gpt-4.1"
# Verify the deployment was created with the correct metadata
assert result_deployment.litellm_params.api_key == "client_side_key"
assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1"
print(
f"✓ Success with function_name '{function_name}' correctly using '{metadata_key}' for metadata"
)
def test_handle_clientside_credential_no_metadata(model_list):
"""Test that _handle_clientside_credential handles cases where no metadata is provided"""
router = Router(model_list=model_list)
# Mock deployment
deployment = {
"model_name": "gpt-4.1",
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
"model_info": {"id": "original-id-789"},
}
# Mock kwargs with clientside credentials but NO metadata
kwargs = {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
# No metadata key at all
}
# This should fail because there's no model_group in metadata
# The function expects to find model_group in the metadata
try:
result_deployment = router._handle_clientside_credential(
deployment=deployment, kwargs=kwargs, function_name="acompletion"
)
# If we get here, the function should have used deployment.model_name as fallback
assert result_deployment.model_name == "gpt-4.1"
print("✓ Success with no metadata - used deployment.model_name as fallback")
except Exception as e:
# This is expected behavior - the function needs model_group to generate model_id
print(f"✓ Correctly handled no metadata case: {e}")
# Test with empty metadata
kwargs_with_empty_metadata = {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
"metadata": {}, # Empty metadata
}
try:
result_deployment = router._handle_clientside_credential(
deployment=deployment,
kwargs=kwargs_with_empty_metadata,
function_name="acompletion",
)
# Should fail because empty metadata has no model_group
pytest.fail("Expected failure with empty metadata")
except Exception as e:
print(f"✓ Correctly handled empty metadata case: {e}")
def test_handle_clientside_credential_with_responses_function(model_list):
"""Test that _handle_clientside_credential works correctly with responses function name"""
router = Router(model_list=model_list)
# Mock deployment
deployment = {
"model_name": "gpt-4.1",
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
"model_info": {"id": "original-id-responses"},
}
# Mock kwargs with clientside credentials and litellm_metadata (for responses function)
kwargs = {
"api_key": "client_side_key",
"api_base": "https://api.openai.com/v1",
"litellm_metadata": {
"model_group": "gpt-4.1",
"responses_field": "responses_value",
},
}
# Call the function with _ageneric_api_call_with_fallbacks function name (which handles responses)
result_deployment = router._handle_clientside_credential(
deployment=deployment,
kwargs=kwargs,
function_name="_ageneric_api_call_with_fallbacks",
)
# Verify the result
assert isinstance(result_deployment, Deployment)
assert result_deployment.model_name == "gpt-4.1"
assert result_deployment.litellm_params.api_key == "client_side_key"
assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1"
assert result_deployment.model_info.id != "original-id-responses"
assert result_deployment.model_info.original_model_id == "original-id-responses"
# Verify the deployment was added to the router
assert len(router.model_list) == len(model_list) + 1
print(
"✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata"
)
def test_get_metadata_variable_name_from_kwargs(model_list):
"""
Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content.
"""
router = Router(model_list=model_list)
# Test case 1: kwargs contains litellm_metadata - should return "litellm_metadata"
kwargs_with_litellm_metadata = {
"litellm_metadata": {"user": "test"},
"metadata": {"other": "data"},
}
result = router._get_metadata_variable_name_from_kwargs(
kwargs_with_litellm_metadata
)
assert result == "litellm_metadata"
# Test case 2: kwargs only contains metadata - should return "metadata"
kwargs_with_metadata_only = {"metadata": {"user": "test"}}
result = router._get_metadata_variable_name_from_kwargs(kwargs_with_metadata_only)
assert result == "metadata"
# Test case 3: kwargs contains neither - should return "metadata" (default)
kwargs_empty = {}
result = router._get_metadata_variable_name_from_kwargs(kwargs_empty)
assert result == "metadata"
# Test case 4: kwargs contains other keys but no metadata keys - should return "metadata"
kwargs_other = {
"model": "gpt-5.5",
"messages": [{"role": "user", "content": "hello"}],
}
result = router._get_metadata_variable_name_from_kwargs(kwargs_other)
assert result == "metadata"
@pytest.fixture
def search_tools():
"""Fixture for search tools configuration"""
return [
{
"search_tool_name": "test-search-tool",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "test-api-key",
"api_base": "https://api.perplexity.ai",
},
},
{
"search_tool_name": "test-search-tool",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "test-api-key-2",
"api_base": "https://api.perplexity.ai",
},
},
]
@pytest.mark.asyncio
async def test_asearch_with_fallbacks(search_tools):
"""
Test _asearch_with_fallbacks method of Router.
Tests that the _asearch_with_fallbacks method correctly:
- Accepts search parameters
- Calls async_function_with_fallbacks with correct configuration
- Returns SearchResponse
"""
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
router = Router(search_tools=search_tools)
# Create a mock search response
mock_response = SearchResponse(
object="search",
results=[
SearchResult(
title="Test Result",
url="https://example.com",
snippet="Test snippet content",
)
],
)
# Mock the async_function_with_fallbacks to return our mock response
with patch.object(
router, "async_function_with_fallbacks", new_callable=AsyncMock
) as mock_fallbacks:
mock_fallbacks.return_value = mock_response
# Mock original function
async def mock_asearch(**kwargs):
return mock_response
# Call _asearch_with_fallbacks
response = await router._asearch_with_fallbacks(
original_function=mock_asearch,
search_tool_name="test-search-tool",
query="test query",
max_results=5,
)
# Verify async_function_with_fallbacks was called
assert mock_fallbacks.called
# Verify the response
assert isinstance(response, SearchResponse)
assert response.object == "search"
assert len(response.results) == 1
assert response.results[0].title == "Test Result"
@pytest.mark.asyncio
async def test_asearch_with_fallbacks_helper(search_tools):
"""
Test _asearch_with_fallbacks_helper method of Router.
Tests that the _asearch_with_fallbacks_helper method correctly:
- Selects a search tool from available options
- Calls the original search function with correct provider parameters
- Returns SearchResponse
"""
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
router = Router(search_tools=search_tools)
# Create a mock search response
mock_response = SearchResponse(
object="search",
results=[
SearchResult(
title="Helper Test Result",
url="https://example.com/helper",
snippet="Helper test snippet",
)
],
)
# Mock the original generic function
async def mock_original_function(**kwargs):
# Verify correct parameters are passed
assert "search_provider" in kwargs
assert kwargs["search_provider"] == "perplexity"
assert "api_key" in kwargs
assert kwargs["query"] == "helper test query"
return mock_response
# Call _asearch_with_fallbacks_helper
response = await router._asearch_with_fallbacks_helper(
model="test-search-tool",
original_generic_function=mock_original_function,
query="helper test query",
max_results=3,
)
# Verify the response
assert isinstance(response, SearchResponse)
assert response.object == "search"
assert len(response.results) == 1
assert response.results[0].title == "Helper Test Result"
assert response.results[0].url == "https://example.com/helper"
@pytest.mark.asyncio
async def test_asearch_with_fallbacks_helper_missing_search_tool():
"""
Test _asearch_with_fallbacks_helper raises error when search tool not found.
Tests that the helper method raises a ValueError when the requested
search tool name doesn't exist in the router's search_tools configuration.
"""
# Create router with no search tools
router = Router(model_list=[])
async def mock_original_function(**kwargs):
return None
# Should raise ValueError for missing search tool
with pytest.raises(ValueError, match="Search tool 'nonexistent-tool' not found"):
await router._asearch_with_fallbacks_helper(
model="nonexistent-tool",
original_generic_function=mock_original_function,
query="test query",
)
@pytest.mark.asyncio
async def test_asearch_with_fallbacks_helper_missing_search_provider():
"""
Test _asearch_with_fallbacks_helper raises error when search_provider not configured.
Tests that the helper method raises a ValueError when a search tool
is found but doesn't have search_provider in its litellm_params.
"""
# Create router with misconfigured search tool (missing search_provider)
search_tools_bad = [
{
"search_tool_name": "bad-tool",
"litellm_params": {
"api_key": "test-key"
# Missing search_provider
},
}
]
router = Router(search_tools=search_tools_bad)
async def mock_original_function(**kwargs):
return None
# Should raise ValueError for missing search_provider
with pytest.raises(ValueError, match="search_provider not found in litellm_params"):
await router._asearch_with_fallbacks_helper(
model="bad-tool",
original_generic_function=mock_original_function,
query="test query",
)
def test_get_first_default_fallback():
"""Test _get_first_default_fallback method"""
# Test with default fallback ("*")
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {"model": "gpt-5-mini", "api_key": "fake-key"},
}
]
router = Router(model_list=model_list, fallbacks=[{"*": ["gpt-5-mini"]}])
result = router._get_first_default_fallback()
assert result == "gpt-5-mini"
# Test with no fallbacks
router_no_fallbacks = Router(model_list=model_list)
result = router_no_fallbacks._get_first_default_fallback()
assert result is None
# Test with fallbacks but no default
router_no_default = Router(
model_list=model_list, fallbacks=[{"gpt-5.5": ["gpt-5-mini"]}]
)
result = router_no_default._get_first_default_fallback()
assert result is None
# Test with empty default list
router_empty_list = Router(model_list=model_list, fallbacks=[{"*": []}])
result = router_empty_list._get_first_default_fallback()
assert result is None
def test_resolve_model_name_from_model_id():
"""Test resolve_model_name_from_model_id function with various scenarios"""
# Test case 1: model_id is None
router = Router(model_list=[])
result = router.resolve_model_name_from_model_id(None)
assert result is None
# Test case 2: model_id directly matches a model_name
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("gpt-5-mini")
assert result == "gpt-5-mini"
# Test case 3: model_id matches litellm_params.model exactly
model_list = [
{
"model_name": "vertex-ai-sora-2",
"litellm_params": {
"model": "vertex_ai/veo-2.0-generate-001",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("vertex_ai/veo-2.0-generate-001")
assert result == "vertex-ai-sora-2"
# Test case 4: model_id matches when actual_model ends with /model_id
model_list = [
{
"model_name": "vertex-ai-sora-2",
"litellm_params": {
"model": "vertex_ai/veo-2.0-generate-001",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("veo-2.0-generate-001")
assert result == "vertex-ai-sora-2"
# Test case 5: model_id matches when actual_model ends with :model_id
# Note: We use a valid model format for router initialization, but test the function
# with a model_id that would match the pattern vertex_ai:model_id
# Since the router validates models on init, we'll test this by manually setting up
# the model_list after initialization or using a valid format
model_list = [
{
"model_name": "vertex-ai-sora-2",
"litellm_params": {
"model": "vertex_ai/veo-2.0-generate-001",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
# Test that the function can handle model_id that would match if the format was vertex_ai:model_id
# We'll test with a model_id that matches the end of the actual_model
result = router.resolve_model_name_from_model_id("veo-2.0-generate-001")
assert result == "vertex-ai-sora-2"
# Test case 6: model_id doesn't match anything
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("non-existent-model")
assert result is None
# Test case 7: Empty model_list
router = Router(model_list=[])
result = router.resolve_model_name_from_model_id("some-model")
assert result is None
# Test case 8: Multiple models, find the correct one
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": "test-key",
},
},
{
"model_name": "vertex-ai-sora-2",
"litellm_params": {
"model": "vertex_ai/veo-2.0-generate-001",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("veo-2.0-generate-001")
assert result == "vertex-ai-sora-2"
# Test case 9: model_id matches deployment ID (has_model_id check)
# This tests the has_model_id path in Strategy 1
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": "test-key",
},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("gpt-5-mini")
assert result == "gpt-5-mini"
# Test case 10: model_id is a deployment ID (hash) that differs from the
# public model_name. Regression for #32580: managed batch/file IDs embed the
# deployment model_id, and it must resolve back to the public model_name so
# team model-access checks compare against the model group, not the hash.
model_list = [
{
"model_name": "bedrock-batch-model",
"litellm_params": {
"model": "bedrock/global.anthropic.claude-haiku-4-5-20251001-v1:0",
},
"model_info": {"id": "8d0eaa7e6c6f54a425dfd0062cb6b0dc"},
},
]
router = Router(model_list=model_list)
result = router.resolve_model_name_from_model_id("8d0eaa7e6c6f54a425dfd0062cb6b0dc")
assert result == "bedrock-batch-model"
def test_get_valid_args():
"""Test get_valid_args static method returns valid Router.__init__ arguments"""
# Call the static method
valid_args = Router.get_valid_args()
# Verify it returns a list
assert isinstance(valid_args, list)
assert len(valid_args) > 0
# Verify it contains expected Router.__init__ arguments
expected_args = [
"model_list",
"routing_strategy",
"cache_responses",
"num_retries",
"timeout",
"fallbacks",
]
for arg in expected_args:
assert arg in valid_args, f"Expected argument '{arg}' not found in valid_args"
# Verify "self" is not in the list (since it's removed)
assert "self" not in valid_args
# Verify it contains keyword-only arguments too
# These are common Router.__init__ parameters
assert "assistants_config" in valid_args or "search_tools" in valid_args
def test_get_router_model_info_with_deployment_object():
"""Test get_router_model_info accepts Deployment object directly and reuses LiteLLM_Params"""
router = Router(
model_list=[
{
"model_name": "gpt-5.5",
"litellm_params": {"model": "gpt-5.5", "api_key": "test-key"},
"model_info": {"id": "test-id"},
}
]
)
# Get the Deployment object (not dict)
deployment = router.get_deployment(model_id="test-id")
assert deployment is not None
assert isinstance(deployment, Deployment)
assert isinstance(deployment.litellm_params, LiteLLM_Params)
# Pass Deployment directly (not .model_dump()) - this exercises the isinstance check
# that reuses the existing LiteLLM_Params instead of reconstructing it
model_info = router.get_router_model_info(
deployment=deployment,
received_model_name="gpt-5.5",
)
# Verify we got valid model info back
assert model_info is not None
assert isinstance(model_info, dict)
def test_deployment_has_budget_limits():
router = Router(model_list=[])
with_budget = Deployment(
model_name="budgeted-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
max_budget=0.001,
budget_duration="1d",
),
model_info=ModelInfo(id="budget-deployment-id"),
)
without_budget = Deployment(
model_name="unbudgeted-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
model_info=ModelInfo(id="no-budget-deployment-id"),
)
assert router._deployment_has_budget_limits(deployment=with_budget) is True
assert router._deployment_has_budget_limits(deployment=without_budget) is False
def test_sync_deployment_budget_config(monkeypatch):
import asyncio
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
router = Router(model_list=[], optional_pre_call_checks=[])
deployment = Deployment(
model_name="dynamic-budget-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
api_key="fake-key",
max_budget=0.000000000001,
budget_duration="1d",
),
model_info=ModelInfo(id="runtime-budget-deployment"),
)
router._sync_deployment_budget_config(deployment=deployment)
budget_limiter = router._get_router_deployment_budget_limiter()
assert budget_limiter is not None
config = budget_limiter._get_budget_config_for_deployment(
"runtime-budget-deployment"
)
assert config is not None
assert config.max_budget == 0.000000000001
def test_sync_deployment_budget_config_clears_removed_limits(monkeypatch):
import asyncio
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
router = Router(model_list=[], optional_pre_call_checks=[])
model_id = "runtime-budget-deployment"
budgeted = Deployment(
model_name="dynamic-budget-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
api_key="fake-key",
max_budget=0.000000000001,
budget_duration="1d",
),
model_info=ModelInfo(id=model_id),
)
unbudgeted = Deployment(
model_name="dynamic-budget-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
api_key="fake-key",
),
model_info=ModelInfo(id=model_id),
)
router._sync_deployment_budget_config(deployment=budgeted)
budget_limiter = router._get_router_deployment_budget_limiter()
assert budget_limiter is not None
assert budget_limiter._get_budget_config_for_deployment(model_id) is not None
router._sync_deployment_budget_config(deployment=unbudgeted)
assert budget_limiter._get_budget_config_for_deployment(model_id) is None
def test_upsert_deployment_clears_stale_budget_config(monkeypatch):
import asyncio
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
router = Router(model_list=[], optional_pre_call_checks=[])
model_id = "upsert-budget-deployment"
budgeted = Deployment(
model_name="dynamic-budget-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
api_key="fake-key",
max_budget=0.000000000001,
budget_duration="1d",
),
model_info=ModelInfo(id=model_id),
)
unbudgeted = Deployment(
model_name="dynamic-budget-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
api_key="fake-key",
),
model_info=ModelInfo(id=model_id),
)
router.upsert_deployment(deployment=budgeted)
budget_limiter = router._get_router_deployment_budget_limiter()
assert budget_limiter is not None
assert budget_limiter._get_budget_config_for_deployment(model_id) is not None
router.upsert_deployment(deployment=unbudgeted)
assert budget_limiter._get_budget_config_for_deployment(model_id) is None