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", "mode": "turbo", }, }, { "search_tool_name": "test-search-tool", "litellm_params": { "search_provider": "perplexity", "api_key": "test-api-key-2", "api_base": "https://api.perplexity.ai", "mode": "turbo", }, }, ] @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["mode"] == "turbo" 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