diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a0973b1f898..56057817847 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7,7 +7,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path import litellm @@ -111,18 +113,31 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] - deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( + litellm.callbacks.index(deployment_callback) + ) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -153,9 +168,13 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + {model_group: ["encrypted_content_affinity"]} + ) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -196,7 +215,10 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] + assert ( + "encrypted_content_affinity_enabled" + not in disabled_request_kwargs["litellm_metadata"] + ) global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -216,7 +238,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert global_request_kwargs["litellm_metadata"][ + "encrypted_content_affinity_enabled" + ] @pytest.mark.asyncio @@ -263,10 +287,18 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) - encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) - assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) - assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callback) < ( + litellm.callbacks.index(deployment_callback) + ) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -277,7 +309,9 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -662,7 +696,9 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: + with patch.object( + litellm, "aretrieve_batch", return_value=AsyncMock() + ) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -683,7 +719,9 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: + with patch.object( + litellm, "afile_content", return_value=AsyncMock() + ) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -828,7 +866,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert result is True, "Should return True when team_id and team_public_model_name match" + assert ( + result is True + ), "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -836,9 +876,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert result is False, ( - "Should return False when team_id matches but model_name doesn't match team_public_model_name" - ) + assert ( + result is False + ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -854,18 +894,30 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert result is True, "Should return True when team deployment has no team_public_model_name to match" + assert ( + result is True + ), "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) - assert result is True, "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id=None + ) + assert ( + result is True + ), "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") - assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id="any-team" + ) + assert ( + result is True + ), "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) + result = router.should_include_deployment( + model_name="different-model", model=deployment_without_team, team_id=None + ) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -874,7 +926,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert result is True, "Should return True when matching model with exact model_name" + assert ( + result is True + ), "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -924,7 +978,9 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + mock_response.text = ( + '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + ) with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1050,9 +1106,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: - with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: + with patch.object( + router, "_get_client", return_value=None + ) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1115,9 +1177,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "_get_client", return_value=mock_semaphore + ) as mock_get_client: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1146,9 +1214,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "_get_client", return_value=None) as mock_get_client: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "_get_client", return_value=None + ) as mock_get_client: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: with pytest.raises(Exception) as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", @@ -1217,9 +1291,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( - f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" - ) + assert ( + captured["model"] == "vertex_ai/gemini-2.5-flash" + ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" def test_router_get_model_access_groups_team_only_models(): @@ -1240,10 +1314,14 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) + access_groups = router.get_model_access_groups( + model_name="gpt-3.5-turbo", team_id=None + ) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") + access_groups = router.get_model_access_groups( + model_name="gpt-3.5-turbo", team_id="team_1" + ) assert list(access_groups.keys()) == ["default-models"] @@ -1338,7 +1416,9 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -1366,7 +1446,9 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -2145,7 +2227,11 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator + base = ( + LiteLLMCompletionStreamingIterator + if bridge + else BaseResponsesAPIStreamingIterator + ) class _Iter(base): def __init__(self): @@ -2225,7 +2311,9 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") + router = _make_router_with_fallback( + "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" + ) src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -2308,9 +2396,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get("metadata", {}), ( - "model_group leaked into 'metadata' instead of 'litellm_metadata'" - ) + assert "model_group" not in fbk.get( + "metadata", {} + ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" @pytest.mark.asyncio @@ -2428,7 +2516,9 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) + fallback_response_object.usage = ResponseAPIUsage( + input_tokens=20, output_tokens=15, total_tokens=35 + ) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -2437,7 +2527,9 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), + return_value=SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) + ), ), patch.object( router, @@ -2738,7 +2830,9 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -2766,10 +2860,14 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3010,7 +3108,9 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): + with patch.object( + litellm, "model_cost", {"test-custom-model": mock_custom_model_info} + ): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -3018,11 +3118,15 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model", model_name="test-model" + ) # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call + mock_get_model_info.assert_any_call( + model="gpt-3.5-turbo" + ) # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -3033,18 +3137,26 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) - assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) + assert ( + result["input_cost_per_token"] == 0.001 + ) # From custom model (overrides base 0.0015) + assert ( + result["output_cost_per_token"] == 0.002 + ) # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert result["mode"] == "chat" # From base model (overrides litellm "completion") + assert ( + result["mode"] == "chat" + ) # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) + assert ( + result["key"] == "gpt-3.5-turbo" + ) # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -3063,7 +3175,9 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model-no-base", model_name="test-model" + ) # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -3083,7 +3197,9 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") + result = router.get_deployment_model_info( + model_id="non-existent-model", model_name="test-model" + ) # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -3116,7 +3232,9 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model-invalid", model_name="test-model" + ) # Should handle exception gracefully and still return merged result assert result is not None @@ -3125,8 +3243,12 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): - result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") + with patch.object( + litellm, "get_model_info", side_effect=Exception("Not found") + ): + result = router.get_deployment_model_info( + model_id="non-existent", model_name="non-existent" + ) # Should return None when no model info is found assert result is None @@ -3149,7 +3271,9 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") + result = router.get_deployment_model_info( + model_id="custom-model-id", model_name="unknown-model" + ) # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -3185,11 +3309,15 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") + result = router.get_deployment_model_info( + model_id="custom-with-base", model_name="unknown-model" + ) # Should return custom_model_info merged with base model info assert result is not None - assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) + assert ( + result["input_cost_per_token"] == 0.01 + ) # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -3236,14 +3364,18 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): + with patch.object( + litellm, "model_cost", {"custom-model-id": mock_custom_model_info} + ): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") + result = router.get_deployment_model_info( + model_id="custom-model-id", model_name="test-model" + ) assert result is not None @@ -3253,17 +3385,29 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) - assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) + assert ( + result["input_cost_per_token"] == 0.01 + ) # From custom model (overrides base 0.03) + assert ( + result["max_tokens"] == 8000 + ) # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) - assert result["litellm_provider"] == "openai" # From base model (not in custom) - assert result["base_only_field"] == "base_value" # From base model (not in custom) + assert ( + result["output_cost_per_token"] == 0.06 + ) # From base model (not in custom) + assert ( + result["litellm_provider"] == "openai" + ) # From base model (not in custom) + assert ( + result["base_only_field"] == "base_value" + ) # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert result["mode"] == "completion" # From litellm model name info (not in custom or base) + assert ( + result["mode"] == "completion" + ) # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -3300,9 +3444,10 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( - f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] + == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" + ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -3314,9 +3459,10 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( - f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] + == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" + ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" # Test Case 3: Bedrock converse endpoint kwargs = { @@ -3328,9 +3474,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( - f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" + ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -3341,9 +3487,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( - f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" + ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -3395,10 +3541,14 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [{"role": "user", "content": "This call should succeed by falling back."}] + messages = [ + {"role": "user", "content": "This call should succeed by falling back."} + ] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion(model="completely-unknown-model", messages=messages) + response = await router.acompletion( + model="completely-unknown-model", messages=messages + ) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -3490,10 +3640,15 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-claude-model" + ) assert credentials is not None - assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert ( + credentials["aws_bedrock_runtime_endpoint"] + == "https://bedrock-runtime.us-east-1.amazonaws.com" + ) assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -3520,7 +3675,9 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + credentials = router.get_deployment_credentials_with_provider( + model_id="vertex-gemini" + ) assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -3560,7 +3717,9 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") + credentials = router.get_deployment_credentials_with_provider( + model_id="azure-gpt-4" + ) assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -4055,7 +4214,9 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ), model="gpt-3.5-turbo", model_id="model-123", model_group="openai-gpt", @@ -4134,8 +4295,12 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" - assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" + assert ( + cached_result is not None + ), "Model ID should be cached for anthropic_messages call type" + assert ( + cached_result["model_id"] == test_model_id + ), f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -4160,7 +4325,9 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -4189,7 +4356,9 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -4216,7 +4385,9 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -4254,7 +4425,9 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -4285,7 +4458,9 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -4314,7 +4489,9 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -4416,8 +4593,12 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name="generic_api_call" + ) assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -4449,8 +4630,12 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name=None + ) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -4799,17 +4984,23 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" + assert ( + len(deployment_ids) == 2 + ), "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="other-team" + ) assert len(deployments) == 0 @@ -4854,8 +5045,12 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") - assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert ( + len(deployments) == 2 + ), "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -4980,7 +5175,9 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] + deployment_groups = [ + d.get("model_info", {}).get("access_groups") for d in deployments + ] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -5025,7 +5222,9 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -5100,7 +5299,9 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -5167,7 +5368,9 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -5222,7 +5425,9 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) + router_in_names._try_early_resolve_deployments_for_model_not_in_names( + model="gpt-5", request_team_id=None + ) is None ) assert ( @@ -5244,8 +5449,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None + pattern_result = ( + pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None + ) ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -5271,8 +5478,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None + default_result = ( + default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None + ) ) assert default_result is not None resolved_model, default_deployment = default_result @@ -5280,7 +5489,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" + assert ( + default_router.default_deployment["litellm_params"]["model"] + == "openai/will-be-overridden" + ) def _router_with_two_deployments(blocked_flags): @@ -5328,7 +5540,10 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} + { + uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} + for uid in unhealthy_ids + } ) @@ -5399,7 +5614,9 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) + healthy, all_dep = await router._async_get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -5408,7 +5625,9 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) + healthy, all_dep = router._get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -5425,7 +5644,9 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) + deployments = await router.async_get_healthy_deployments( + model="gpt-4o", request_kwargs={} + ) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -5467,7 +5688,9 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) + deployment = router.get_available_deployment_for_pass_through( + model="gpt-4o", request_kwargs={} + ) assert deployment["model_info"]["id"] == "pt-1" @@ -5476,7 +5699,9 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) + router.get_available_deployment_for_pass_through( + model="pt-0", request_kwargs={} + ) def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -5500,7 +5725,9 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] + assert [m["model_info"]["id"] for m in router.get_model_list()] == [ + "bedrock-iam-pt" + ] def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): @@ -5512,7 +5739,9 @@ def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): router = _router_with_two_pass_through_deployments([False, False]) assert len(router.get_model_list()) == 2 assert ( - passthrough_endpoint_router.get_credentials(custom_llm_provider="openai", region_name=None) + passthrough_endpoint_router.get_credentials( + custom_llm_provider="openai", region_name=None + ) == "sk-fake-for-tests" ) @@ -5548,9 +5777,16 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() - assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=missing_blocked) + ) + is False + ) + assert ( + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) is True ) @@ -5590,7 +5826,9 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): + def test_request_timeout_stored_independently_when_both_set( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -5608,16 +5846,22 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): + def test_non_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): + def test_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): + def test_explicit_stream_timeout_still_wins_over_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -5633,13 +5877,22 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): + def test_per_deployment_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): + def test_per_request_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 + assert ( + router._get_non_stream_timeout( + kwargs={"timeout": 60}, data={"timeout": 120} + ) + == 60 + ) # ---------------------------------------------------------------------------