From 36f0a0fe6d8b6a2a6dc9b1f2eac983fb2938477d Mon Sep 17 00:00:00 2001 From: JingHao-Leon Date: Sat, 3 Oct 2026 05:16:12 +0800 Subject: [PATCH] fix(router): retry /v1/messages stream drops that happen before first content MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On the aanthropic_messages route, a transport drop before the first content chunk (e.g. aiohttp 'Connection closed' / TransferEncodingError from a Databricks AI Gateway) is raised from the stream iterator — outside the retried call boundary — so num_retries / retry_policy never applied and the request surfaced as a 500 after a single upstream attempt (issue #44238; 9 dropped streams in one week across three Databricks-served models in the report). When nothing has reached the client and no fallbacks are configured, retry the same group once from the stream recovery path; after any frame reached the client (or with fallbacks configured, whose recovery owns the failure) the existing behavior is unchanged. The retry is structurally bounded: a retry that drops again comes back from a raw stream outside the recovery gate. --- litellm/router.py | 29 + tests/unit/test_router/test_router.py | 1449 +++++++++++-------------- 2 files changed, 673 insertions(+), 805 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 26e1ff16152..619dcb66190 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5551,10 +5551,12 @@ class Router: model, initial_kwargs ) buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + chunks_sent_to_client = 0 # any frame that reached the client blocks a same-group retry try: async for chunk in source_iterator: if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): yield chunk + chunks_sent_to_client += 1 continue if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): has_generated_content = True @@ -5609,9 +5611,36 @@ class Router: yield buffered_chunk buffered_lifecycle_chunks = () yield chunk + chunks_sent_to_client += 1 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 + # A transport drop before ANY frame reached the client never + # reaches the retry machinery — the error is raised from the + # stream iterator, outside the retried call boundary — and + # surfaced as a 500 after a single upstream attempt (issue + # #44238). With no fallbacks configured there is nothing to + # hand the failure to, so retry the same group once; once any + # frame reached the client (or fallbacks are configured, whose + # recovery owns the failure) keep the existing behavior. + if ( + chunks_sent_to_client == 0 + and not fallbacks_disabled_for_request(initial_kwargs) + and not initial_kwargs.get("fallbacks", self.fallbacks) + ): + verbose_router_logger.info( + "Anthropic messages stream dropped before first content; retrying the same group once" + ) + retry_kwargs: Final = { + **initial_kwargs, + "original_function": self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt, + } + retry_stream = await self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt( + **retry_kwargs + ) + async for item in retry_stream: + yield item + return async for item in self._aanthropic_messages_recover_stream_error( stream_error, has_generated_content, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 34eb200e7cd..d171e6f4013 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -56,7 +56,14 @@ from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deploymen from litellm.router_utils.fallback_event_handlers import DISABLE_FALLBACKS_METADATA_KEY from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.llms.openai import ChatCompletionRequest -from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy +from litellm.types.router import ( + Deployment, + DeploymentTypedDict, + LiteLLM_Params, + ModelInfo, + PreRoutingHookResponse, + RetryPolicy, +) def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -156,31 +163,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 @@ -211,13 +205,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( @@ -258,10 +248,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, @@ -281,9 +268,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 @@ -330,18 +315,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, @@ -352,9 +329,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}, @@ -540,7 +515,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=[ { @@ -551,9 +528,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) @@ -888,8 +863,6 @@ async def test_arouter_async_get_healthy_deployments(): assert result[0]["litellm_params"]["model"] == "gpt-3.5-turbo" - - def test_arouter_test_team_model(): """ Test that router.test_team_model returns the correct model @@ -1001,9 +974,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", @@ -1263,6 +1234,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1274,8 +1246,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 @@ -1356,9 +1327,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=[ { @@ -1419,7 +1388,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!"}], @@ -1503,9 +1472,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( @@ -1513,9 +1480,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( @@ -1531,30 +1498,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 @@ -1563,9 +1518,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(): @@ -1615,9 +1568,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: @@ -1691,7 +1642,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/*", @@ -1716,7 +1667,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" @@ -1793,15 +1746,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, @@ -1836,7 +1783,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, @@ -1866,15 +1813,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, @@ -1903,16 +1844,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, @@ -1980,9 +1915,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 @@ -2088,14 +2023,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"] @@ -2190,9 +2121,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 @@ -2220,9 +2149,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 @@ -2339,9 +2266,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", @@ -2387,9 +2312,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", @@ -2649,6 +2572,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) @@ -2996,9 +2920,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" @@ -3070,9 +2992,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" @@ -3099,9 +3019,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 @@ -3175,9 +3093,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 @@ -3751,11 +3667,7 @@ def test_completion_streaming_iterator_routes_mid_stream_fallback_through_common class FallbackStream: def __init__(self): self._chunks = iter( - [ - litellm.ModelResponseStream( - choices=[{"index": 0, "delta": {"content": "from the fallback"}}] - ) - ] + [litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "from the fallback"}}])] ) def __iter__(self): @@ -4185,11 +4097,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -4289,9 +4197,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( @@ -4537,9 +4443,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 @@ -4954,9 +4860,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, @@ -4965,9 +4869,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, @@ -5268,9 +5170,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"}}, @@ -5298,14 +5198,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"}}, @@ -5332,14 +5228,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"}}, @@ -5366,9 +5258,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( @@ -5464,14 +5354,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"}}, @@ -5499,9 +5385,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 = [] @@ -5536,9 +5420,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"}}, @@ -5574,9 +5456,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", @@ -5604,9 +5484,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"}}, @@ -5631,9 +5509,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"}}, @@ -5675,9 +5551,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", @@ -5746,9 +5620,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): @@ -5868,7 +5740,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) @@ -5883,9 +5755,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 @@ -5975,9 +5845,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: { @@ -5985,15 +5853,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 @@ -6004,26 +5868,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 = { @@ -6042,9 +5898,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 @@ -6064,9 +5918,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 @@ -6099,9 +5951,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 @@ -6110,12 +5960,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 @@ -6138,9 +5984,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 @@ -6176,15 +6020,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 @@ -6231,18 +6071,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 @@ -6252,29 +6088,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) @@ -6291,7 +6115,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", ), @@ -6346,10 +6174,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 = { @@ -6361,10 +6188,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 = { @@ -6376,9 +6202,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 = { @@ -6389,9 +6215,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(): @@ -6528,14 +6354,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 @@ -6714,15 +6536,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" @@ -6749,9 +6566,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" @@ -6836,9 +6651,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" @@ -6875,9 +6688,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" @@ -6921,9 +6732,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(): @@ -6993,15 +6802,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" @@ -7042,15 +6847,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" @@ -7077,16 +6878,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(): @@ -7126,9 +6919,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" @@ -7138,12 +6929,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_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -7175,9 +6961,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" @@ -7189,21 +6973,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(): @@ -7276,22 +7050,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(): @@ -7480,9 +7245,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", @@ -7561,12 +7324,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(): @@ -7591,9 +7350,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 @@ -7622,9 +7379,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 @@ -7651,9 +7406,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 @@ -7735,9 +7488,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 @@ -7768,9 +7519,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"}] @@ -7799,9 +7548,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) @@ -7903,12 +7650,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"] @@ -7940,12 +7683,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"] @@ -8058,6 +7797,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) @@ -8084,6 +7824,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) @@ -8300,23 +8041,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 @@ -8361,12 +8096,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( @@ -8491,9 +8222,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 @@ -8538,9 +8267,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( @@ -8615,9 +8342,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( @@ -8684,9 +8409,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( @@ -8741,9 +8464,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 ( @@ -8765,10 +8486,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 @@ -8794,10 +8513,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 @@ -8805,10 +8522,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): @@ -8856,10 +8570,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} ) @@ -8926,7 +8637,6 @@ async def test_health_probe_preserves_normal_caller_policy( assert await router.cooldown_cache.async_get_active_cooldowns(["dep-0", "dep-1"], parent_otel_span=None) == [] - @pytest.mark.asyncio async def test_async_get_fully_unhealthy_model_names_keeps_name_when_partial(): router = _router_with_two_deployments([False, False]) @@ -8987,9 +8697,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 @@ -8998,9 +8706,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 @@ -9017,9 +8723,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 @@ -9107,9 +8811,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" @@ -9118,9 +8820,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(): @@ -9162,9 +8862,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(): @@ -9175,12 +8873,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(): @@ -9214,16 +8907,9 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() + assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=missing_blocked) - ) - is False - ) - assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) - ) + litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) is True ) @@ -9263,9 +8949,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 @@ -9283,22 +8967,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 @@ -9314,22 +8992,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) @@ -9648,9 +9317,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 @@ -9659,12 +9326,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): @@ -9679,12 +9341,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) @@ -9963,13 +9620,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 @@ -10013,12 +9670,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(): @@ -10241,13 +9894,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(): @@ -10289,13 +9937,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(): @@ -10354,13 +9997,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 def test_get_configured_display_name_treats_malformed_values_as_absent(): @@ -10593,13 +10231,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", @@ -10628,13 +10269,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=[ { @@ -10668,7 +10309,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 @@ -10684,13 +10327,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=[ { @@ -10699,9 +10342,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() @@ -10715,7 +10356,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=[ { @@ -10750,9 +10393,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 {}), } @@ -11057,9 +10698,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), ) @@ -11168,9 +10807,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"] @@ -11357,22 +10994,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 @@ -12132,18 +11763,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.""" @@ -12267,9 +11894,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) @@ -12543,9 +12168,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" @@ -12611,9 +12234,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 []: @@ -12755,9 +12376,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" @@ -12992,6 +12611,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( @@ -13050,7 +12670,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 @@ -13274,7 +12893,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 @@ -13301,17 +12919,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"] @@ -13319,17 +12934,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 @@ -13356,6 +12968,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=[ @@ -13446,7 +13059,6 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o assert result.supported_reasoning_efforts is None - @pytest.mark.parametrize( "model,provider,expected", [ @@ -13466,11 +13078,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") @@ -13482,16 +13098,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") @@ -13499,24 +13120,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") @@ -13756,6 +13383,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). @@ -13868,10 +13496,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: @@ -13916,9 +13541,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")] @@ -13937,12 +13560,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 @@ -13954,9 +13579,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] @@ -13965,6 +13588,74 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ await wrapped.__anext__() +@pytest.mark.asyncio +async def test_pre_first_chunk_drop_retries_same_group_when_no_fallback(): + """Issue #44238: with no fallback groups configured, a pre-first-chunk + transport drop must be retried on the same group instead of surfacing as + an APIConnectionError after a single attempt.""" + import litellm.router + + router = _anthropic_messages_make_router(fallbacks=[]) + source = _AnthropicMessagesRaisingByteStream( + chunks=[], + error=litellm.exceptions.APIConnectionError( + message="Connection closed", llm_provider="anthropic", model="primary" + ), + ) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("recovered")]) + attempts = [] + + async def _attempt(**kwargs): + attempts.append(kwargs) + return fallback_stream + + with patch.object(router, "_ageneric_api_call_with_fallbacks_anthropic_messages_attempt", side_effect=_attempt): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "original_generic_function": None}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_content_chunk("recovered")] + assert len(attempts) == 1, "the same-group retry must run exactly once" + + +@pytest.mark.asyncio +async def test_pre_first_chunk_same_group_retry_is_bounded(): + """A retried stream that drops again must not trigger another same-group + retry: the controls marker bounds the recovery path.""" + router = _anthropic_messages_make_router(fallbacks=[]) + source = _AnthropicMessagesRaisingByteStream( + chunks=[], + error=litellm.exceptions.APIConnectionError( + message="Connection closed", llm_provider="anthropic", model="primary" + ), + ) + still_dropping = _AnthropicMessagesRaisingByteStream( + chunks=[], + error=litellm.exceptions.APIConnectionError( + message="Connection closed", llm_provider="anthropic", model="primary" + ), + ) + attempts = [] + + async def _attempt(**kwargs): + attempts.append(kwargs) + return still_dropping + + with patch.object(router, "_ageneric_api_call_with_fallbacks_anthropic_messages_attempt", side_effect=_attempt): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "original_generic_function": None}, + ) + with pytest.raises(litellm.exceptions.APIConnectionError): + [chunk async for chunk in wrapped] + + # the second drop came from the raw retried stream, outside the recovery + # gate — it surfaces to the caller instead of triggering another retry + assert len(attempts) == 1 + + @pytest.mark.asyncio async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback(): """Greptile review round: transport-level buffering can coalesce a real @@ -14007,7 +13698,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() @@ -14052,7 +13745,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() @@ -14117,7 +13812,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) @@ -14157,8 +13854,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}, @@ -14171,7 +13875,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): @@ -14244,7 +13950,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) @@ -14289,7 +13997,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() @@ -16405,9 +16115,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() @@ -16487,7 +16197,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.""" @@ -16792,9 +16506,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" @@ -16827,9 +16539,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" @@ -18073,7 +17783,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 @@ -18308,18 +18020,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) @@ -18329,44 +18048,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 @@ -18375,36 +18121,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 @@ -18414,11 +18192,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"}) ) @@ -18429,11 +18219,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")) @@ -18444,35 +18243,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")) @@ -18480,22 +18300,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) @@ -18506,13 +18336,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"): @@ -18528,7 +18361,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: @@ -18573,10 +18408,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) @@ -18639,9 +18476,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", @@ -18808,7 +18647,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( @@ -18868,9 +18709,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: