diff --git a/litellm/router.py b/litellm/router.py index d6730db8748..3e23a074e7a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5530,6 +5530,7 @@ class Router: self, response: AsyncIterator[bytes], initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + allow_fallback: bool = True, ) -> AsyncIterator[bytes]: """ Wrap an anthropic_messages (/v1/messages) streaming response so a @@ -5668,6 +5669,10 @@ class Router: for buffered_chunk in buffered_lifecycle_chunks: yield buffered_chunk except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate + if not allow_fallback: + if isinstance(stream_error, MidStreamFallbackError) and stream_error.original_exception is not None: + raise stream_error.original_exception from stream_error + raise async for item in self._aanthropic_messages_recover_stream_error( stream_error, has_generated_content, @@ -6005,7 +6010,16 @@ class Router: kwargs=initial_kwargs, include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, ) - async for fallback_item in self._aanthropic_messages_yield_recovered(fallback_response, wrapper): + recovered: Final = ( + await self._aanthropic_messages_streaming_iterator( + response=cast("AsyncIterator[bytes]", fallback_response), # cast-ok: __aiter__ checked below + initial_kwargs=initial_kwargs, + allow_fallback=False, + ) + if hasattr(fallback_response, "__aiter__") + else fallback_response + ) + async for fallback_item in self._aanthropic_messages_yield_recovered(recovered, wrapper): yield fallback_item except Exception as fallback_error: verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index ec9e4861fdd..a2a68146c53 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -79,7 +79,7 @@ from litellm.types.router import ( RoutingContext, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.base_llm.vector_store.transformation import( +from litellm.llms.base_llm.vector_store.transformation import ( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) @@ -194,31 +194,18 @@ 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 @@ -249,13 +236,9 @@ 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( @@ -296,10 +279,7 @@ 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, @@ -319,9 +299,7 @@ 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 @@ -368,18 +346,10 @@ 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, @@ -390,9 +360,7 @@ 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}, @@ -578,7 +546,9 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards from io import BytesIO from unittest.mock import MagicMock, patch - jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + jsonl_content = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + ) router = litellm.Router( model_list=[ { @@ -589,9 +559,7 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards ) with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: - await router.acreate_file( - model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True - ) + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True) forwarded = mock_acreate_file.call_args.kwargs assert forwarded["passthrough"] is True forwarded["file"].seek(0) @@ -1037,9 +1005,7 @@ 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", @@ -1299,6 +1265,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1310,8 +1277,7 @@ async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: moved = sorted( f"{key}={cache_dict[key]}" for key in cache_dict - if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) - and cache_dict[key] + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) and cache_dict[key] ) if moved: return moved @@ -1392,9 +1358,7 @@ 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=[ { @@ -1455,7 +1419,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1539,9 +1503,7 @@ 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( @@ -1549,9 +1511,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( @@ -1567,30 +1529,18 @@ 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 @@ -1599,9 +1549,7 @@ 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(): @@ -1651,9 +1599,7 @@ 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: @@ -1727,7 +1673,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1752,7 +1698,9 @@ def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler" ) - monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}]) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}] + ) monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) return "test-onprem-llm" @@ -1829,15 +1777,9 @@ 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, @@ -1872,7 +1814,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match='No deployment available') as exc_info: + with pytest.raises(Exception, match="No deployment available") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1902,15 +1844,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): max_parallel_requests=1, model_id="deployment-1", model_group="gpt-3.5-turbo" ) - 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, @@ -1939,16 +1875,10 @@ 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 pytest.raises(Exception, match='Mock failure') as exc_info: + 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, match="Mock failure") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -2016,9 +1946,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']}'" + ) @pytest.mark.asyncio @@ -2124,14 +2054,10 @@ 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"] @@ -2226,9 +2152,7 @@ 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 @@ -2256,9 +2180,7 @@ 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 @@ -2552,9 +2474,7 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object( - router, "get_deployment_model_info", side_effect=_model_info_with_str_costs - ): + with patch.object(router, "get_deployment_model_info", side_effect=_model_info_with_str_costs): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -2600,9 +2520,7 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - 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._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -2862,6 +2780,7 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] + async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -3209,9 +3128,7 @@ def test_adopt_fallback_response_headers_replaces_rather_than_merges(): "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, } - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} assert wrapper._hidden_params["model_id"] == "fallback-deployment" @@ -3283,9 +3200,7 @@ def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_repla fallback._response_headers = None fallback._hidden_params = {"model_id": "fallback-deployment"} - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers is None assert wrapper._hidden_params["model_id"] == "fallback-deployment" @@ -3312,9 +3227,7 @@ def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none() hidden_params_before = wrapper._hidden_params fallback = object() - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers is None assert wrapper._hidden_params is hidden_params_before @@ -3388,9 +3301,7 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro **additional_headers, ) - assert not { - key for key in proxy_headers if key.startswith("x-litellm-complexity-router-") - } + assert not {key for key in proxy_headers if key.startswith("x-litellm-complexity-router-")} @pytest.mark.asyncio @@ -4437,11 +4348,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -4541,9 +4448,7 @@ 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( @@ -4789,9 +4694,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 @@ -5206,9 +5111,7 @@ 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, @@ -5217,9 +5120,7 @@ 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, @@ -5520,9 +5421,7 @@ 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"}}, @@ -5550,14 +5449,10 @@ 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"}}, @@ -5584,14 +5479,10 @@ def test_pre_call_checks_uses_precounted_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 1 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -5618,9 +5509,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) counting_threads = [] monkeypatch.setattr( @@ -5716,14 +5605,10 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], 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"}}, @@ -5751,9 +5636,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], 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}) counting_threads = [] @@ -5788,9 +5671,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], 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}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -5826,9 +5707,7 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke 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) count = await router._acount_pre_call_check_tokens( model="m", @@ -5856,9 +5735,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -5883,9 +5760,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -5927,9 +5802,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -5998,9 +5871,7 @@ def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwarg prompt_only_tokens = router._count_pre_call_check_tokens( messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -6120,7 +5991,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -6135,9 +6006,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(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}) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -6228,9 +6097,7 @@ 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: { @@ -6238,15 +6105,11 @@ 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 @@ -6257,26 +6120,18 @@ 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 = { @@ -6295,9 +6150,7 @@ 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 @@ -6317,9 +6170,7 @@ 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 @@ -6352,9 +6203,7 @@ 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 @@ -6363,12 +6212,8 @@ 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 @@ -6391,9 +6236,7 @@ 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 @@ -6429,15 +6272,11 @@ 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 @@ -6484,18 +6323,14 @@ 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 @@ -6505,29 +6340,17 @@ 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) @@ -6544,7 +6367,11 @@ def test_get_deployment_model_info_base_model_merge_priority(): [ ( "gpt", - {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, + { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.services.ai.azure.com", + "api_key": "key", + }, "gpt/openai/deployments/gpt-5.4-mini/chat/completions", "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", ), @@ -6599,10 +6426,9 @@ 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 = { @@ -6614,10 +6440,9 @@ 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 = { @@ -6629,9 +6454,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 = { @@ -6642,9 +6467,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(): @@ -6781,14 +6606,10 @@ 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 @@ -6967,15 +6788,10 @@ 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" @@ -7002,9 +6818,7 @@ 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" @@ -7089,9 +6903,7 @@ 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" @@ -7128,9 +6940,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -7174,9 +6984,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -7246,15 +7054,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -7295,15 +7099,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -7330,16 +7130,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_deployment_usable_by_team_helpers(): @@ -7379,9 +7171,7 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) + picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -7391,12 +7181,7 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) + assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None def test_deployment_usable_by_team_uses_dynamic_model_info_get(): @@ -7436,9 +7221,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -7450,21 +7233,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -7537,22 +7310,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -7741,9 +7505,7 @@ 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", @@ -7822,12 +7584,8 @@ 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(): @@ -7852,9 +7610,7 @@ 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 @@ -7883,9 +7639,7 @@ 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 @@ -7912,9 +7666,7 @@ 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 @@ -7996,9 +7748,7 @@ 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 @@ -8029,9 +7779,7 @@ 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"}] @@ -8060,9 +7808,7 @@ 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) @@ -8164,12 +7910,8 @@ 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"] @@ -8201,12 +7943,8 @@ 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"] @@ -8319,6 +8057,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -8345,6 +8084,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -8561,23 +8301,17 @@ 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 @@ -8622,12 +8356,8 @@ 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( @@ -8752,9 +8482,7 @@ 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 @@ -8799,9 +8527,7 @@ 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( @@ -8876,9 +8602,7 @@ 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( @@ -8945,9 +8669,7 @@ 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( @@ -9002,9 +8724,7 @@ 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 ( @@ -9026,10 +8746,8 @@ 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 @@ -9055,10 +8773,8 @@ 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 @@ -9066,10 +8782,7 @@ 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): @@ -9117,10 +8830,7 @@ 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} ) @@ -9247,9 +8957,7 @@ 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 @@ -9258,9 +8966,7 @@ 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 @@ -9277,9 +8983,7 @@ 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 @@ -9367,9 +9071,7 @@ 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" @@ -9378,9 +9080,7 @@ 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_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): @@ -9422,9 +9122,7 @@ 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_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -9435,12 +9133,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -9527,9 +9220,7 @@ 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 @@ -9547,22 +9238,16 @@ 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 @@ -9578,22 +9263,13 @@ 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 def test_passthrough_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) @@ -9912,9 +9588,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -9923,12 +9597,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -9943,12 +9612,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -10227,13 +9891,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -10277,12 +9941,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -10505,13 +10165,8 @@ def test_get_model_listing_info_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_reads_deployment_model_info(): @@ -10553,13 +10208,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -10618,13 +10268,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None @pytest.mark.parametrize( @@ -10653,7 +10298,10 @@ def test_get_configured_service_tiers_returns_one_value_per_deployment_in_model_ "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"service_tiers": ["ultrafast"]}, }, - {"model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}}, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}, + }, { "model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://b.example"}, @@ -11051,13 +10699,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -11086,13 +10737,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -11126,7 +10777,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -11142,13 +10795,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -11157,9 +10810,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -11173,7 +10824,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -11208,9 +10861,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -11515,9 +11166,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -11626,9 +11275,7 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) + actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) assert actual is expected, params["model"] @@ -11815,22 +11462,16 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") assert len(router.model_list) == 1 @@ -12590,18 +12231,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -12733,9 +12370,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -13009,9 +12644,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -13077,9 +12710,7 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] + affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -13221,9 +12852,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -13465,6 +13094,7 @@ class TestTeamPublicNameReachesPreRoutingStrategies: assert response is not None assert response.model == "gemini-flash" + @pytest.mark.asyncio async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): router = self._router( @@ -13523,7 +13153,6 @@ class TestTeamPublicNameReachesPreRoutingStrategies: with pytest.raises(litellm.BadRequestError, match="multiple teams"): two_teams._team_deployments_across_teams(self.PUBLIC_NAME) - def test_compression_policy_follows_the_same_resolution_for_every_principal(self): from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model @@ -13747,7 +13376,6 @@ class TestAutoRouterCompressionDecoupling: @pytest.mark.usefixtures("local_model_cost_map") - @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -13774,17 +13402,14 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" + assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + f"unexpected error log: {mock_error.call_args_list}" + ) # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -13792,17 +13417,14 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" + assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + "expected the error log for an unmappable azure deployment name" + ) # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -13829,6 +13451,7 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -13938,11 +13561,15 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o def test_model_group_info_fast_mode_uses_exact_provider_catalog( local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool ) -> None: - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": operator_flag}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": operator_flag}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -13954,16 +13581,21 @@ def test_model_group_info_fast_mode_uses_exact_provider_catalog( def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object ) -> None: - entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items() - if key != "supports_fast_mode"} + entry: Final = { + key: value for key, value in litellm.model_cost["claude-opus-5"].items() if key != "supports_fast_mode" + } if flag is not None: entry["supports_fast_mode"] = flag monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": True}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": True}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -13971,24 +13603,30 @@ def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( assert result.supports_fast_mode is False -@pytest.mark.parametrize("other_model,expected", [ - ("anthropic/claude-opus-4-8", True), - ("anthropic/claude-opus-4-7", False), - ("anthropic/off-map-opus", False), - ("vertex_ai/claude-opus-5", False), - ("bedrock/claude-opus-5", False), -]) +@pytest.mark.parametrize( + "other_model,expected", + [ + ("anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus-4-7", False), + ("anthropic/off-map-opus", False), + ("vertex_ai/claude-opus-5", False), + ("bedrock/claude-opus-5", False), + ], +) @pytest.mark.parametrize("reverse", [True, False]) def test_model_group_info_fast_mode_requires_every_deployment( local_model_cost_map: None, other_model: str, expected: bool, reverse: bool ) -> None: - models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ( - "anthropic/claude-opus-5", other_model + models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ("anthropic/claude-opus-5", other_model) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "api_key": "fake-key"}, + } + for model in models + ] ) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "api_key": "fake-key"}, - } for model in models]) result: Final = router.get_model_group_info("fast-group") @@ -14228,6 +13866,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -14342,10 +13981,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -14390,9 +14026,7 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -14411,12 +14045,14 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + assert collected == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + message_stop, + ] @pytest.mark.asyncio @@ -14428,9 +14064,7 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -14481,7 +14115,9 @@ async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwar await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -14526,7 +14162,9 @@ async def test_anthropic_messages_no_fallback_message_start_reaches_client_befor await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() content_released.set() @@ -14591,7 +14229,9 @@ async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecy await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) pending = asyncio.ensure_future(wrapped.__anext__()) await asyncio.sleep(0.2) @@ -14631,8 +14271,15 @@ def _anthropic_messages_two_order_primary_model_list() -> list: id="wildcard-overridden-by-request-none", ), pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"), + pytest.param( + {"fallbacks": None}, + {"model": "primary", "fallbacks": [{"model": "fallback"}]}, + True, + id="request-dict-fallback", + ), + pytest.param( + {"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback" + ), pytest.param( {"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary", "disable_fallbacks": True}, @@ -14645,7 +14292,9 @@ def _anthropic_messages_two_order_primary_model_list() -> list: True, id="content-policy-fallback", ), - pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), + pytest.param( + {"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover" + ), ], ) def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): @@ -14718,7 +14367,9 @@ async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames( await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) pending = asyncio.ensure_future(wrapped.__anext__()) await asyncio.sleep(0.2) @@ -14763,7 +14414,9 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -15460,7 +15113,9 @@ class _AnthropicMessagesScriptedProvider: async def __call__(self, **kwargs): litellm_metadata = kwargs.get("litellm_metadata") or {} - self.calls.append((kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries"))) + self.calls.append( + (kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries")) + ) assert self._streams, "provider called more times than scripted" return self._streams.pop(0)() @@ -15478,7 +15133,9 @@ def _anthropic_messages_transport_drop(original_exception: Exception | None = No def _anthropic_messages_dropped_before_content(): - return _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop()) + return _AnthropicMessagesRaisingByteStream( + [_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop() + ) def _anthropic_messages_bridge_error_chunk() -> bytes: @@ -16083,20 +15740,28 @@ async def test_anthropic_messages_retries_keep_the_budget_the_first_drop_committ }, { "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-b", "api_key": "sk-test", "num_retries": sibling_num_retries}, + "litellm_params": { + "model": "anthropic/glm-b", + "api_key": "sk-test", + "num_retries": sibling_num_retries, + }, }, ], num_retries=0, fallbacks=None, ) router.set_custom_routing_strategy(_AnthropicMessagesAlternatingDeployments(router, "glm")) - provider = _AnthropicMessagesScriptedProvider(*[_anthropic_messages_dropped_before_content] * len(expected_counters)) + provider = _AnthropicMessagesScriptedProvider( + *[_anthropic_messages_dropped_before_content] * len(expected_counters) + ) stream = await _anthropic_messages_stream_through_router(router, provider) with pytest.raises(litellm.APIConnectionError): [chunk async for chunk in stream] - assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[: len(expected_counters)] + assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[ + : len(expected_counters) + ] assert [(attempted, budget) for _, attempted, budget in provider.calls] == expected_counters @@ -17583,9 +17248,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) + assert litellm.Router._declared_param_allowlist( + {"allowed_openai_params": ["reasoning_effort", 3]} + ) == frozenset({"reasoning_effort"}) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -17665,7 +17330,11 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], + [ + {"model_name": "x"}, + {"model_name": "x", "litellm_params": {}}, + {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, + ], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -17970,9 +17639,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -18005,9 +17672,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -19253,7 +18918,9 @@ async def test_router_deployment_slot_rejects_while_held_and_frees_slot_on_exit( async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): with pytest.raises(litellm.RateLimitError) as overflow: - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): + async with router._deployment_slot( + deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None + ): pass assert overflow.value.status_code == 429 assert "slot-deployment" in overflow.value.message @@ -19488,18 +19155,25 @@ class TestMemberAutoRouterInference: self.cache = UserApiKeyCache() self.team = LiteLLM_TeamTable( - team_id="router-team", models=["member-router", "permitted-model"], + team_id="router-team", + models=["member-router", "permitted-model"], members_with_roles=[Member(user_id="router-member", role="user")], ) self.actor = UserAPIKeyAuth( - user_id="router-member", team_id="router-team", user_role=LitellmUserRoles.INTERNAL_USER, - models=["member-router", "permitted-model"], api_key="test-key-hash", config={"timeout": 60}, + user_id="router-member", + team_id="router-team", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["member-router", "permitted-model"], + api_key="test-key-hash", + config={"timeout": 60}, + ) + self.database = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), + litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), + ) ) - self.database = SimpleNamespace(db=SimpleNamespace( - litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), - litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), - litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), - )) monkeypatch.setattr(proxy_server, "user_api_key_cache", self.cache) monkeypatch.setattr(proxy_server, "prisma_client", self.database) @@ -19509,44 +19183,71 @@ class TestMemberAutoRouterInference: return { "model_name": "model_name_router-team_member-router", "litellm_params": { - "model": "auto_router/complexity_router", "complexity_router_default_model": target, + "model": "auto_router/complexity_router", + "complexity_router_default_model": target, "complexity_router_config": { - "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), "adaptive": False, + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), + "adaptive": False, **({"classifier_type": "llm", "classifier_llm_config": {"model": target}} if classifier else {}), }, - "tags": ["member" if member else "admin"], "timeout": 13.0 if member else 29.0, + "tags": ["member" if member else "admin"], + "timeout": 13.0 if member else 29.0, }, "model_info": { - "team_id": "router-team", "team_public_model_name": "member-router", "member_auto_router": member, + "team_id": "router-team", + "team_public_model_name": "member-router", + "member_auto_router": member, }, } @classmethod def _router(cls, *markers: dict[str, object]) -> Router: - return Router(model_list=[ - *(markers or (cls._marker(),)), - {"model_name": "permitted-model", "litellm_params": { - "model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://api.openai.com/v1", - }}, - {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, - ]) + return Router( + model_list=[ + *(markers or (cls._marker(),)), + { + "model_name": "permitted-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://api.openai.com/v1", + }, + }, + {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, + ] + ) def _request( - self, *, actor: UserAPIKeyAuth | None = None, metadata_name: str = "metadata", tag: str = "member", + self, + *, + actor: UserAPIKeyAuth | None = None, + metadata_name: str = "metadata", + tag: str = "member", ) -> dict[str, object]: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data={metadata_name: {"tags": [tag]}, **({"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} - if metadata_name == "litellm_metadata" else {})}, - user_api_key_dict=actor or self.actor, _metadata_variable_name=metadata_name, + data={ + metadata_name: {"tags": [tag]}, + **( + {"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} + if metadata_name == "litellm_metadata" + else {} + ), + }, + user_api_key_dict=actor or self.actor, + _metadata_variable_name=metadata_name, ) async def _route( - self, router: Router, request: dict[str, object] | None = None, model: str = "member-router", + self, + router: Router, + request: dict[str, object] | None = None, + model: str = "member-router", ) -> PreRoutingHookResponse: response: Final = await router.async_pre_routing_hook( - model=model, request_kwargs=request if request is not None else self._request(), + model=model, + request_kwargs=request if request is not None else self._request(), messages=[{"role": "user", "content": "Hello"}], ) assert response is not None @@ -19555,36 +19256,68 @@ class TestMemberAutoRouterInference: @pytest.mark.asyncio @pytest.mark.parametrize("metadata_name", ("metadata", "litellm_metadata")) async def test_cached_roster_revocation_blocks_classifier_and_session_rebinding( - self, metadata_name: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + self, + metadata_name: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_checks import delete_cache_team_object monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") router: Final = self._router(self._marker(classifier=True)) - classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond(200, json={ - "id": "classifier", "object": "chat.completion", "created": 0, "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, - "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - }) - request: Final = {**self._request(metadata_name=metadata_name), "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", - }}} + classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 200, + json={ + "id": "classifier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + request: Final = { + **self._request(metadata_name=metadata_name), + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + } + }, + } first: Final = await self._route(router, request) assert first.model == "permitted-model" and first.routing_decision is not None assert first.routing_decision["cause"] == "llm_classifier" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_teamtable.find_unique.await_count == 1 assert self.database.db.litellm_teammembership.find_unique.await_count == 1 - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={"members_with_roles": []}) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={"members_with_roles": []} + ) await delete_cache_team_object( - team_id=self.team.team_id, team_alias=None, user_api_key_cache=self.cache, proxy_logging_obj=None, + team_id=self.team.team_id, + team_alias=None, + user_api_key_cache=self.cache, + proxy_logging_obj=None, ) with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, request) - rebound: Final = {**request, "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", "x-claude-code-agent-id": "subagent", - }}} + rebound: Final = { + **request, + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + "x-claude-code-agent-id": "subagent", + } + }, + } with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, rebound, model="restricted-model") assert classify.call_count == 2 @@ -19594,11 +19327,23 @@ class TestMemberAutoRouterInference: async def test_member_router_fails_closed(self, state: str, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server - request: Final = {"metadata": {"user_api_key_team_id": "router-team", "user_api_key_auth": { - "team_id": "router-team", "user_role": "proxy_admin", - }}} if state == "forged" else self._request(actor=self.actor.model_copy( - update={"user_id": ""} if state == "empty-user" else {}, - )) + request: Final = ( + { + "metadata": { + "user_api_key_team_id": "router-team", + "user_api_key_auth": { + "team_id": "router-team", + "user_role": "proxy_admin", + }, + } + } + if state == "forged" + else self._request( + actor=self.actor.model_copy( + update={"user_id": ""} if state == "empty-user" else {}, + ) + ) + ) self.database.db.litellm_teamtable.find_unique.return_value = ( None if state == "deleted" else self.team.model_copy(update={"blocked": state == "blocked"}) ) @@ -19609,11 +19354,20 @@ class TestMemberAutoRouterInference: assert error.value.status_code == (503 if state == "unavailable" else 403) @pytest.mark.asyncio - @pytest.mark.parametrize("user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)]) - async def test_service_key_and_admin_preserve_runtime_access(self, user_id: str | None, role: LitellmUserRoles) -> None: - assert (await self._route(self._router(), self._request( - actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), - ))).model == "permitted-model" + @pytest.mark.parametrize( + "user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)] + ) + async def test_service_key_and_admin_preserve_runtime_access( + self, user_id: str | None, role: LitellmUserRoles + ) -> None: + assert ( + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), + ), + ) + ).model == "permitted-model" @pytest.mark.asyncio @pytest.mark.parametrize("ceiling", ("team", "key", "member", "organization", "project")) @@ -19624,35 +19378,56 @@ class TestMemberAutoRouterInference: from litellm.proxy._types import LiteLLM_ProjectTableCachedObj from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if ceiling == "team" else self.team.models, - "organization_id": "router-org" if ceiling == "organization" else None, - }) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if ceiling == "team" else self.team.models, + "organization_id": "router-org" if ceiling == "organization" else None, + } + ) if ceiling == "member": await self.cache.async_set_cache( key=team_membership_reservation_cache_key(user_id="router-member", team_id="router-team"), - value=LiteLLM_TeamMembership(user_id="router-member", team_id="router-team", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"])), + value=LiteLLM_TeamMembership( + user_id="router-member", + team_id="router-team", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"]), + ), model_type=LiteLLM_TeamMembership, ) elif ceiling == "organization": await self.cache.async_set_cache( - key="org_id:router-org", value=LiteLLM_OrganizationTable( - organization_id="router-org", budget_id="org-budget", created_by="admin", updated_by="admin", + key="org_id:router-org", + value=LiteLLM_OrganizationTable( + organization_id="router-org", + budget_id="org-budget", + created_by="admin", + updated_by="admin", models=["restricted-model"], - ), model_type=LiteLLM_OrganizationTable, + ), + model_type=LiteLLM_OrganizationTable, ) elif ceiling == "project": await self.cache.async_set_cache( - key="project_id:router-project", value=LiteLLM_ProjectTableCachedObj( - project_id="router-project", team_id="router-team", models=["restricted-model"], - ), model_type=LiteLLM_ProjectTableCachedObj, + key="project_id:router-project", + value=LiteLLM_ProjectTableCachedObj( + project_id="router-project", + team_id="router-team", + models=["restricted-model"], + ), + model_type=LiteLLM_ProjectTableCachedObj, ) with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(self._router(), self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if ceiling == "key" else self.actor.models, - "project_id": "router-project" if ceiling == "project" else None, - }))) + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if ceiling == "key" else self.actor.models, + "project_id": "router-project" if ceiling == "project" else None, + } + ) + ), + ) @pytest.mark.asyncio @pytest.mark.parametrize("group_owner", ("team", "key")) @@ -19660,22 +19435,32 @@ class TestMemberAutoRouterInference: from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast group: Final = LiteLLM_AccessGroupTable( - access_group_id="router-group", access_group_name="Router targets", access_model_names=["permitted-model"], + access_group_id="router-group", + access_group_name="Router targets", + access_model_names=["permitted-model"], ) self.database.db.litellm_accessgrouptable.find_unique.return_value = group - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if group_owner == "team" else self.team.models, - "access_group_ids": ["router-group"] if group_owner == "team" else [], - }) - request: Final = self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if group_owner == "key" else self.actor.models, - "access_group_ids": ["router-group"] if group_owner == "key" else [], - })) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if group_owner == "team" else self.team.models, + "access_group_ids": ["router-group"] if group_owner == "team" else [], + } + ) + request: Final = self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if group_owner == "key" else self.actor.models, + "access_group_ids": ["router-group"] if group_owner == "key" else [], + } + ) + ) router: Final = self._router() assert (await self._route(router, request)).model == "permitted-model" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 1 - self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy(update={"access_model_names": []}) + self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy( + update={"access_model_names": []} + ) await evict_and_broadcast(cache_keys=("access_group_id:router-group",), user_api_key_cache=self.cache) with pytest.raises(ProxyException, match="is not available for this API key"): await self._route(router, request) @@ -19686,13 +19471,16 @@ class TestMemberAutoRouterInference: router: Final = self._router(self._marker(member=False), self._marker()) request: Final = self._request() selected: Final = router._selected_strategy_marker_deployment( - model="model_name_router-team_member-router", strategy_tags=("member",), request_kwargs=request, + model="model_name_router-team_member-router", + strategy_tags=("member",), + request_kwargs=request, ) assert selected is not None and selected["model_info"]["member_auto_router"] is True assert (await self._route(router, request)).model == "permitted-model" assert request["timeout"] == 13.0 await self.cache.async_set_cache( - key="team_id:router-team", model_type=LiteLLM_TeamTable, + key="team_id:router-team", + model_type=LiteLLM_TeamTable, value=self.team.model_copy(update={"models": ["member-router"]}), ) with pytest.raises(ProxyException, match="is not available for this API key"): @@ -19708,7 +19496,9 @@ class TestMemberAutoRouterInference: router: Final = self._router(self._marker(member=False)) monkeypatch.setitem(sys.modules, "fastapi", None) monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False) - assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model" + assert ( + await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}}) + ).model == "restricted-model" def _access_window_offsets(start_hours: float, end_hours: float, team_ids: list) -> dict: @@ -19753,10 +19543,12 @@ def test_access_windows_hide_reserved_deployment_from_other_teams(): def test_access_windows_raise_when_only_reserved_deployments_remain(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], - )[:1]) + router = Router( + model_list=_reserved_model_list( + windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], + windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], + )[:1] + ) for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-b"}}, {}): with pytest.raises(litellm.BadRequestError, match="reserved for another team"): router._common_checks_available_deployment(model="gpt-4o-ptu", request_kwargs=request_kwargs) @@ -19819,9 +19611,11 @@ def test_access_windows_inactive_window_leaves_deployments_available(): def test_access_windows_apply_when_calling_by_model_id(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - )) + router = Router( + model_list=_reserved_model_list( + windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], + ) + ) with pytest.raises(litellm.BadRequestError, match="reserved for another team"): router._common_checks_available_deployment( model="reserved-deployment", @@ -20046,7 +19840,9 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref @pytest.mark.asyncio -async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> None: +async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> ( + None +): """The context-window chain is keyed the same way the ordinary chain is, so a key spelled like the wildcard deployment ("anthropic/claude-sonnet-4-6") must catch the bare group's context-window error too.""" router = litellm.Router( @@ -20106,9 +19902,7 @@ def test_bare_model_group_served_by_wildcard_deployment_has_provider_prefixed_co [ GuardrailRaisedException(guardrail_name="chunk-scanner", message="blocked"), HTTPException(status_code=403, detail={"error": "blocked", "guardrail_name": "chunk-scanner"}), - ModifyResponseException( - message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner" - ), + ModifyResponseException(message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner"), ], ) async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: Exception) -> None: @@ -20309,6 +20103,7 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(): assert seen == ["router_usage"] assert current_service_target() is None + class _SpanRecordingInMemoryCache(InMemoryCache): """Records the live OTel span each read runs under, so the test sees what a Redis span would nest in.""" @@ -20503,6 +20298,8 @@ async def test_non_chat_surfaces_mark_their_deployment_pick(monkeypatch: pytest. router.completion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) assert events == [_pick("embed", "initial", 1), _pick("gpt-4o", "initial", 1)] + + class TestAutoRouterTraceProvenance: @pytest.fixture def _tracing(self, monkeypatch: pytest.MonkeyPatch) -> "tuple[OpenTelemetryV2, InMemorySpanExporter]": @@ -20514,7 +20311,9 @@ class TestAutoRouterTraceProvenance: config: Final = OpenTelemetryV2Config(exporter="in_memory") exporter: Final = InMemorySpanExporter() - logger: Final = OpenTelemetryV2(config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter)) + logger: Final = OpenTelemetryV2( + config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter) + ) monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) return logger, exporter @@ -20524,7 +20323,8 @@ class TestAutoRouterTraceProvenance: "model_name": "traced-router", "model_info": {"id": f"definition-{tag}", "updated_at": "2026-10-01T12:00:00+00:00"}, "litellm_params": { - "model": "auto_router/complexity_router", "tags": [tag], + "model": "auto_router/complexity_router", + "tags": [tag], "complexity_router_config": { "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), }, @@ -20533,17 +20333,25 @@ class TestAutoRouterTraceProvenance: def test_definition_fingerprint_is_canonical_and_tracks_config_not_credentials(self): original: Final = Deployment.model_validate(self._marker("blue")) - reordered: Final = Deployment.model_validate({ - **self._marker("blue"), - "litellm_params": { - "api_key": "credential-only-change", "tags": ["blue"], "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer")}, - }, - }) + reordered: Final = Deployment.model_validate( + { + **self._marker("blue"), + "litellm_params": { + "api_key": "credential-only-change", + "tags": ["blue"], + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer") + }, + }, + } + ) fingerprint: Final = Router._routing_definition_fingerprint(original) assert fingerprint is not None assert fingerprint == Router._routing_definition_fingerprint(reordered) - assert fingerprint != Router._routing_definition_fingerprint(Deployment.model_validate(self._marker("blue", "other"))) + assert fingerprint != Router._routing_definition_fingerprint( + Deployment.model_validate(self._marker("blue", "other")) + ) @pytest.mark.asyncio async def test_direct_classifier_and_plugin_prelude_never_inherit_parent_decisions( @@ -20561,7 +20369,9 @@ class TestAutoRouterTraceProvenance: async def classify(self, context: RoutingContext) -> str: await litellm.acompletion( - model="openai/gpt-5.6-luna", api_key="test", mock_response="SIMPLE", + model="openai/gpt-5.6-luna", + api_key="test", + mock_response="SIMPLE", messages=[{"role": "user", "content": "choose a tier"}], metadata={**context.metadata, "internal_call_origin": "autorouter_classifier"}, ) @@ -20575,16 +20385,25 @@ class TestAutoRouterTraceProvenance: monkeypatch.setattr(litellm, "_async_success_callback", [logger, CallbackComplete()]) classifier: Final = ObservedClassifier() - router: Final = Router(model_list=[ - {**self._marker("blue"), "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": { - "classifier_type": "custom", "classifier_plugin": classifier, - "tiers": {"SIMPLE": "answer"}, "session_affinity": False, "deployment_affinity": False, + router: Final = Router( + model_list=[ + { + **self._marker("blue"), + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "custom", + "classifier_plugin": classifier, + "tiers": {"SIMPLE": "answer"}, + "session_affinity": False, + "deployment_affinity": False, + }, + }, }, - }}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], plugins=[classifier]) + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + plugins=[classifier], + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale, "user_api_key_team_id": "team"}} with logger.tracer.start_as_current_span("request") as root: @@ -20594,7 +20413,8 @@ class TestAutoRouterTraceProvenance: ) await completed.wait() classifier_span: Final = next( - span for span in exporter.get_finished_spans() + span + for span in exporter.get_finished_spans() if (span.attributes or {}).get("litellm.request.purpose") == "autorouter_classifier" ) assert not any(key.startswith("litellm.routing.") for key in classifier_span.attributes) @@ -20608,36 +20428,66 @@ class TestAutoRouterTraceProvenance: @pytest.mark.asyncio @pytest.mark.parametrize("build_fails", [False, True]) async def test_semantic_internal_failure_keeps_phase_identity_without_parent_retry_attribution( - self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, build_fails: bool, + self, + _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + build_fails: bool, ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) query_statuses: Final = iter((503, 200)) def embedding_response(request: httpx.Request) -> httpx.Response: - status: Final = 400 if build_fails else ( - next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200 + status: Final = ( + 400 + if build_fails + else (next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200) + ) + return httpx.Response( + status, + headers={"x-should-retry": "false"}, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, ) - return httpx.Response(status, headers={"x-should-retry": "false"}, json={ - "object": "list", "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock(side_effect=embedding_response) - router: Final = Router(model_list=[ - {"model_name": "semantic", "model_info": {"id": "semantic-definition"}, "litellm_params": { - "model": "auto_router/semantic-router", "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - }}, - *({"model_name": "embed", "model_info": {"id": f"embed-{index}"}, "litellm_params": { - "model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", - "api_key": "test", "max_retries": 0, - }} for index in range(2)), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], num_retries=1, retry_after=0) + upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock( + side_effect=embedding_response + ) + router: Final = Router( + model_list=[ + { + "model_name": "semantic", + "model_info": {"id": "semantic-definition"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + }, + }, + *( + { + "model_name": "embed", + "model_info": {"id": f"embed-{index}"}, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + "max_retries": 0, + }, + } + for index in range(2) + ), + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + num_retries=1, + retry_after=0, + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale}} with logger.tracer.start_as_current_span("request") as root: @@ -20669,20 +20519,22 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]" ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), - { - "model_name": "traced-router", - "model_info": {"id": "definition-semantic"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "tags": ["blue"], - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", + router: Final = Router( + model_list=[ + self._marker("blue"), + { + "model_name": "traced-router", + "model_info": {"id": "definition-semantic"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "tags": ["blue"], + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, }, - }, - ]) + ] + ) with logger.start_phase_span("route traced-router"): response: Final = await router.async_pre_routing_hook( model="traced-router", request_kwargs={"metadata": {"tags": ["blue"]}} @@ -20702,26 +20554,42 @@ class TestAutoRouterTraceProvenance: ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - respx_mock.post("https://embedding.example/v1/embeddings").respond(200, json={ - "object": "list", - "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - router: Final = Router(model_list=[ - { - "model_name": "internal-semantic-router", - "model_info": {"id": "semantic-definition", "team_id": "trace-team", "team_public_model_name": "public-router"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - }, + respx_mock.post("https://embedding.example/v1/embeddings").respond( + 200, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, }, - {"model_name": "embed", "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", "api_key": "test"}}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ]) + ) + router: Final = Router( + model_list=[ + { + "model_name": "internal-semantic-router", + "model_info": { + "id": "semantic-definition", + "team_id": "trace-team", + "team_public_model_name": "public-router", + }, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + }, + }, + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ] + ) request: Final = {"metadata": {"user_api_key_team_id": "trace-team"}} with logger.tracer.start_as_current_span("request"): await router.async_get_available_deployment( @@ -20733,8 +20601,13 @@ class TestAutoRouterTraceProvenance: assert decision["router_config_id"] == router.model_list[0]["model_info"]["id"] assert decision["cause"] == "semantic_match" spans: Final = {span.name: span for span in exporter.get_finished_spans()} - assert spans["route public-router"].attributes["litellm.routing.router_model_name"] == decision["router_model_name"] - assert spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + assert ( + spans["route public-router"].attributes["litellm.routing.router_model_name"] + == decision["router_model_name"] + ) + assert ( + spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + ) retry: Final = next(event for event in spans["request"].events if event.name == "litellm.routing.retry") assert retry.attributes["litellm.routing.router_model_name"] == decision["router_model_name"] @@ -20744,18 +20617,37 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", passthrough: bool ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), self._marker("green"), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test", "use_in_pass_through": True, "tags": ["blue", "green"]}}, - ], enable_tag_filtering=True) + router: Final = Router( + model_list=[ + self._marker("blue"), + self._marker("green"), + { + "model_name": "answer", + "litellm_params": { + "model": "openai/gpt-5.6-luna", + "api_key": "test", + "use_in_pass_through": True, + "tags": ["blue", "green"], + }, + }, + ], + enable_tag_filtering=True, + ) request: Final = {"metadata": {"tags": ["green"]}, "litellm_metadata": {"tags": ["green"]}} - pick: Final = router.async_get_available_deployment_for_pass_through if passthrough else router.async_get_available_deployment + pick: Final = ( + router.async_get_available_deployment_for_pass_through + if passthrough + else router.async_get_available_deployment + ) with logger.tracer.start_as_current_span("request"): await pick(model="traced-router", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) decision: Final = request["litellm_metadata"]["routing_decision"] assert decision["router_config_id"] == "definition-green" assert decision["router_config_updated_at"] == "2026-10-01T12:00:00+00:00" - assert decision["router_config_fingerprint"] == router.complexity_routers["traced-router"][1].definition_fingerprint + assert ( + decision["router_config_fingerprint"] + == router.complexity_routers["traced-router"][1].definition_fingerprint + ) assert "routing_decision" not in request["metadata"] router.log_retry({**request, "model": "answer"}, ValueError("private-error-text")) await pick(model="answer", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) @@ -20776,6 +20668,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -20828,6 +20721,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -20848,6 +20742,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -20870,6 +20765,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @patch("azure.identity.get_bearer_token_provider") @patch("azure.identity.ClientSecretCredential") @@ -20951,6 +20847,7 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab # finally verify if the mocked token was used by Azure SDK mocked_func_generating_token.assert_called() + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_audio_speech_router(): @@ -20994,6 +20891,7 @@ async def test_audio_speech_router(): client_passed_in_request = mock_aspeech.call_args.kwargs["client"] assert client_passed_in_request == expected_openai_client + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_update_kwargs_before_fallbacks_unit_test(): router = Router( @@ -21019,6 +20917,7 @@ def test_update_kwargs_before_fallbacks_unit_test(): assert kwargs["litellm_trace_id"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "call_type", @@ -21078,6 +20977,7 @@ async def test_update_kwargs_before_fallbacks(call_type): print(mock_client.call_args.kwargs) assert mock_client.call_args.kwargs["litellm_trace_id"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_model_info_wildcard_routes(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -21097,6 +20997,7 @@ def test_router_get_model_info_wildcard_routes(): assert model_info["tpm"] is not None assert model_info["rpm"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) @@ -21127,6 +21028,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): assert tpm is not None, "tpm is None" assert rpm is not None, "rpm is None" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_call_router_callbacks_on_success(): @@ -21160,6 +21062,7 @@ async def test_call_router_callbacks_on_success(): assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") assert increment["increment_value"] == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.serial @pytest.mark.asyncio @@ -21188,6 +21091,7 @@ async def test_call_router_callbacks_on_failure(): assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") + @pytest.mark.asyncio async def test_rate_limit_error_callback(): """ @@ -21233,13 +21137,8 @@ async def test_rate_limit_error_callback(): except Exception: pass - with patch.object( - customHandler, "log_model_group_rate_limit_error", new=AsyncMock() - ) as mock_client: - - print( - f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}" - ) + with patch.object(customHandler, "log_model_group_rate_limit_error", new=AsyncMock()) as mock_client: + print(f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}") try: _ = await router.acompletion( @@ -21287,6 +21186,7 @@ async def test_router_model_group_headers(): assert "x-ratelimit-remaining-requests" in resp._hidden_params["additional_headers"] assert "x-ratelimit-remaining-tokens" in resp._hidden_params["additional_headers"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_remaining_model_group_usage(): @@ -21318,6 +21218,7 @@ async def test_get_remaining_model_group_usage(): assert "x-ratelimit-remaining-requests" in remaining_usage assert "x-ratelimit-remaining-tokens" in remaining_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "potential_access_group, expected_result", @@ -21336,6 +21237,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result) access_groups = router.is_model_access_group_for_wildcard_route(model_access_group=potential_access_group) assert access_groups == expected_result + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_redis_cache(): router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}]) @@ -21346,6 +21248,7 @@ def test_router_redis_cache(): assert router.cache.redis_cache == redis_cache + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_handle_clientside_credential(): """A caller-supplied credential must stay scoped to the current call: it must @@ -21374,6 +21277,7 @@ def test_router_handle_clientside_credential(): assert len(router.get_model_list()) == 1 assert router.get_deployment(model_id=new_deployment.model_info.id) is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_clientside_credential_not_reused_by_other_callers(respx_mock, monkeypatch: pytest.MonkeyPatch): """End-to-end regression test for LIT-7811. @@ -21431,6 +21335,7 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router( @@ -21447,6 +21352,7 @@ def test_router_get_async_openai_model_client(): model_client = router._get_async_openai_model_client(deployment=MagicMock(), kwargs={}) assert model_client is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials(): router = Router( @@ -21462,6 +21368,7 @@ def test_router_get_deployment_credentials(): assert credentials is not None assert credentials["api_key"] == "123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider(): """ @@ -21506,6 +21413,7 @@ def test_router_get_deployment_credentials_with_provider(): credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider_wildcard(): """ @@ -21553,6 +21461,7 @@ def test_router_get_deployment_credentials_with_provider_wildcard(): credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_model_info(): router = Router( @@ -21567,12 +21476,14 @@ def test_router_get_deployment_model_info(): model_info = router.get_deployment_model_info(model_id="1", model_name="gemini/gemini-1.5-flash") assert model_info is not None + @pytest.fixture() def _vcr_outcome_gate_router_unit(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown_router_unit(): """ @@ -21594,6 +21505,7 @@ def setup_and_teardown_router_unit(): loop.close() asyncio.set_event_loop(None) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_acompletion_deployment_not_mutated(): @@ -21643,6 +21555,7 @@ async def test_acompletion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_completion_deployment_not_mutated(): """ @@ -21691,6 +21604,7 @@ def test_completion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_default_deployment_isolation(): """ @@ -21763,10 +21677,12 @@ def test_default_deployment_isolation(): assert deployment2["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore assert router.default_deployment["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore + class NoItemsAliasDict(dict): def items(self): raise AssertionError("Unexpected full alias iteration via items()") + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup(): router = Router( @@ -21783,6 +21699,7 @@ def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup model_alias_list = router.get_model_list_from_model_alias(model_name="gpt-5-mini") assert model_alias_list == [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name(): router = Router( @@ -21805,6 +21722,7 @@ def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name result = router.map_team_model(team_model_name="team-model", team_id="team-1") assert result == "team-model", f"Expected public name 'team-model', got {result}" + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestPreCallChecksOptimization: """ @@ -21912,9 +21830,11 @@ class TestPreCallChecksOptimization: assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + if __name__ == "__main__": pytest.main([__file__, "-v"]) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_is_prompt_management_model_optimization(): """ @@ -21969,6 +21889,7 @@ def test_is_prompt_management_model_optimization(): finally: litellm._known_custom_logger_compatible_callbacks = original_callbacks + @pytest.fixture def router(): """Create a router with a mock deployment""" @@ -21984,6 +21905,7 @@ def router(): ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch(router): @@ -22007,6 +21929,7 @@ async def test_router_acancel_batch(router): assert response.id == "batch_123" assert response.status == "cancelled" + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_resolves_credential_name(): @@ -22047,6 +21970,7 @@ async def test_router_acancel_batch_resolves_credential_name(): finally: litellm.credential_list = [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_removes_unresolved_credential_name(): @@ -22079,6 +22003,7 @@ async def test_router_acancel_batch_removes_unresolved_credential_name(): call_kwargs = mock_cancel.call_args.kwargs assert "litellm_credential_name" not in call_kwargs + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingHeaders: """Test that embedding methods properly propagate headers from router configuration.""" @@ -22393,6 +22318,7 @@ class TestRouterEmbeddingHeaders: assert "litellm_trace_id" in completion_kwargs assert "litellm_trace_id" in embedding_kwargs + if __name__ == "__main__": # Run a simple test test = TestRouterEmbeddingHeaders() @@ -22409,6 +22335,7 @@ OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: return respx_mock.post(url).mock( return_value=httpx.Response( @@ -22422,11 +22349,13 @@ def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route ) ) + def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: request = route.calls[index].request body = json.loads(request.read()) return request.headers["authorization"], body["model"], body["input"] + def _alias_router() -> Router: return Router( model_list=[ @@ -22440,6 +22369,7 @@ def _alias_router() -> Router: ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" @@ -22926,10 +22856,12 @@ class TestRouterEmbeddingIntegration: assert "headers" in call_kwargs assert call_kwargs["headers"]["X-Custom-Azure-Header"] == "azure-value" + if __name__ == "__main__": # Run tests pytest.main([__file__, "-v"]) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterIndexManagement: """Test cases for router index management functions""" @@ -23337,7 +23269,6 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -23391,7 +23322,6 @@ def test_router_context_window_check_pre_call_check(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -23445,7 +23375,6 @@ def test_router_context_window_check_pre_call_check_out_group(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -23471,7 +23400,13 @@ def test_router_context_window_check_pre_call_check_out_group(): }, ] - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0, context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}]) # type: ignore + router = Router( + model_list=model_list, + set_verbose=True, + enable_pre_call_checks=True, + num_retries=0, + context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}], + ) # type: ignore response = router.completion( model="gpt-3.5-turbo-small", @@ -23567,13 +23502,9 @@ def test_router_region_pre_call_check(allowed_model_region): if allowed_model_region is None: assert len(_healthy_deployments) == 2 else: - assert len(_healthy_deployments) == 1, "{} models selected as healthy".format( - len(_healthy_deployments) - ) - assert ( - _healthy_deployments[0]["model_info"]["id"] == "1" - ), "Incorrect model id picked. Got id={}, expected id=1".format( - _healthy_deployments[0]["model_info"]["id"] + assert len(_healthy_deployments) == 1, "{} models selected as healthy".format(len(_healthy_deployments)) + assert _healthy_deployments[0]["model_info"]["id"] == "1", ( + "Incorrect model id picked. Got id={}, expected id=1".format(_healthy_deployments[0]["model_info"]["id"]) ) @@ -23746,9 +23677,7 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider): else: pytest.fail("Expected to raise an error. Got={}".format(updated_list)) except Exception as e: - if ( - llm_provider == "azure" and base_model is not None - ) or llm_provider == "openai": + if (llm_provider == "azure" and base_model is not None) or llm_provider == "openai": pass except Exception as e: pytest.fail(f"Got unexpected exception on router! - {str(e)}") @@ -23825,9 +23754,7 @@ def test_router_dynamic_cooldown_correct_retry_after_time(): ): new_retry_after_mock_client = MagicMock(return_value=-1) - litellm.utils._get_retry_after_from_exception_header = ( - new_retry_after_mock_client - ) + litellm.utils._get_retry_after_from_exception_header = new_retry_after_mock_client try: router.embedding( @@ -23913,9 +23840,7 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): ) if sync_mode: - cooldown_deployments = get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None - ) + cooldown_deployments = get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) else: cooldown_deployments = await async_get_cooldown_deployments( litellm_router_instance=router, parent_otel_span=None @@ -24301,12 +24226,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception as e: print(f"Assertion Error: {traceback.format_exc()}") @@ -24332,9 +24253,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.iscoroutine(kwargs["original_response"]) or inspect.isasyncgen(kwargs["original_response"]) ) @@ -24346,12 +24265,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -24368,23 +24283,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -24405,22 +24315,17 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponse) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) ) + assert isinstance(kwargs["api_key"], (str, type(None))) + assert isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) assert isinstance(kwargs["additional_args"], (dict, type(None))) assert isinstance(kwargs["log_event_type"], str) assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) @@ -24439,23 +24344,18 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -24488,29 +24388,19 @@ class CompletionCustomHandler( ## END TIME assert isinstance(end_time, datetime) ## RESPONSE OBJECT - assert isinstance( - response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse) - ) + assert isinstance(response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)) ## KWARGS assert isinstance(kwargs["model"], str) # checking we use base_model for azure cost calculation - base_model = litellm.utils.get_base_model_from_metadata( - model_call_details=kwargs - ) + base_model = litellm.utils.get_base_model_from_metadata(model_call_details=kwargs) - if ( - kwargs["model"] == "chatgpt-v-3" - and base_model is not None - and kwargs["stream"] != True - ): + if kwargs["model"] == "chatgpt-v-3" and base_model is not None and kwargs["stream"] != True: # when base_model is set for azure, we should use pricing for the base_model # this checks response_cost == litellm.cost_per_token(model=base_model) assert isinstance(kwargs["response_cost"], float) response_cost = kwargs["response_cost"] - print( - f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}" - ) + print(f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}") prompt_tokens = response_obj.usage.prompt_tokens completion_tokens = response_obj.usage.completion_tokens # ensure the pricing is based on the base_model here @@ -24521,9 +24411,9 @@ class CompletionCustomHandler( ) expected_price = prompt_price + completion_price print(f"expected price: {expected_price}") - assert ( - response_cost == expected_price - ), f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + assert response_cost == expected_price, ( + f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + ) assert isinstance(kwargs["messages"], list) assert isinstance(kwargs["optional_params"], dict) @@ -24534,9 +24424,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -24549,12 +24437,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -24581,9 +24465,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -24593,3 +24475,44 @@ class CompletionCustomHandler( except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) + + +@pytest.mark.asyncio +async def test_anthropic_messages_double_failure_holds_back_fallback_lifecycle_frames(): + router = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "anthropic/" + name, "api_key": "sk-test"}} + for name in ["primary", "fallback"] + ], + num_retries=0, + fallbacks=[{"primary": ["fallback"]}], + ) + calls = [] + fallback_error = litellm.APIError( + status_code=503, message="fallback also overloaded", llm_provider="anthropic", model="fallback" + ) + + async def provider(**kwargs): + calls.append(kwargs["model"]) + if kwargs["model"] == "anthropic/primary": + return _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) + return _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], fallback_error) + + stream = await router._aanthropic_messages_with_streaming_fallbacks( + original_function=provider, + model="primary", + stream=True, + messages=[{"role": "user", "content": "hi"}], + max_tokens=10, + ) + chunks = [] + + async def collect_chunks(): + async for chunk in stream: + chunks.append(chunk) + + with pytest.raises(litellm.APIError, match="fallback also overloaded"): + await collect_chunks() + + assert chunks == [] + assert calls == ["anthropic/primary", "anthropic/fallback"]