diff --git a/litellm/router.py b/litellm/router.py index 7d6499cf7d2..1831f6f677a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2095,6 +2095,9 @@ class Router: async for item in model_response: yield item except MidStreamFallbackError as e: + if not e.is_pre_first_chunk and e.generated_content: + raise + from litellm.main import stream_chunk_builder complete_response_object = stream_chunk_builder(chunks=model_response.chunks) @@ -2113,24 +2116,7 @@ class Router: "content_policy_fallbacks", self.content_policy_fallbacks ) initial_kwargs["original_function"] = self._acompletion - if e.is_pre_first_chunk or not e.generated_content: - # No content was generated before the error (e.g. a - # rate-limit 429 on the very first chunk). Retry with - # the original messages — adding a continuation prompt - # would waste tokens and confuse the model. - initial_kwargs["messages"] = messages - else: - initial_kwargs["messages"] = messages + [ - { - "role": "system", - "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", - }, - { - "role": "assistant", - "content": e.generated_content, - "prefix": True, - }, - ] + initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( e=e, @@ -2650,6 +2636,9 @@ class Router: for item in model_response: yield item except MidStreamFallbackError as e: + if not e.is_pre_first_chunk and e.generated_content: + raise + from litellm.main import stream_chunk_builder complete_response_object = stream_chunk_builder(chunks=model_response.chunks) @@ -2669,20 +2658,7 @@ class Router: router_self.content_policy_fallbacks, ) initial_kwargs["original_function"] = router_self._completion - if e.is_pre_first_chunk or not e.generated_content: - initial_kwargs["messages"] = messages - else: - initial_kwargs["messages"] = messages + [ - { - "role": "system", - "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", - }, - { - "role": "assistant", - "content": e.generated_content, - "prefix": True, - }, - ] + initial_kwargs["messages"] = messages router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = router_self.function_with_fallbacks( **initial_kwargs, @@ -2890,6 +2866,21 @@ class Router: ) if isinstance(response, CustomStreamWrapper): + if response.completion_stream is None and response.make_call is not None: + try: + await response.fetch_stream() + except Exception as fetch_err: + _headers = getattr(fetch_err, "headers", None) + if isinstance(_headers, dict): + _framing_headers = frozenset( + {"content-length", "transfer-encoding", "content-encoding", "content-type"} + ) + setattr( + fetch_err, + "headers", + {k: v for k, v in _headers.items() if k.lower() not in _framing_headers}, + ) + raise fetch_err return await self._acompletion_streaming_iterator( model_response=response, messages=messages, @@ -6117,7 +6108,7 @@ class Router: """ Common utilities for async_function_with_fallbacks """ - verbose_router_logger.debug("Traceback%s", traceback.format_exc()) + verbose_router_logger.debug("Traceback", exc_info=True) original_exception = e fallback_model_group = None original_model_group: str | None = kwargs.get("model") # type: ignore @@ -6333,15 +6324,17 @@ class Router: except Exception as new_exception: parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) fallback_failure_exception_str = redact_string(str(new_exception)) + cooldown_info = await _async_get_cooldown_deployments_with_debug_info( + litellm_router_instance=self, + parent_otel_span=parent_otel_span, + ) verbose_router_logger.error( - "litellm.router.py::async_function_with_fallbacks() - Error occurred while trying to do fallbacks - {}\n{}\n\nDebug Information:\nCooldown Deployments={}".format( - fallback_failure_exception_str, - redact_string(traceback.format_exc()), - await _async_get_cooldown_deployments_with_debug_info( - litellm_router_instance=self, - parent_otel_span=parent_otel_span, - ), - ) + "litellm.router.py::async_function_with_fallbacks() - " + "Error occurred while trying to do fallbacks - %s\n" + "Debug Information:\nCooldown Deployments=%s", + fallback_failure_exception_str, + cooldown_info, + exc_info=True, ) if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..437bac25e5f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7,9 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm @@ -113,31 +111,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 @@ -168,13 +153,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( @@ -215,10 +196,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, @@ -238,9 +216,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 @@ -287,18 +263,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, @@ -309,9 +277,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}, @@ -696,9 +662,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", @@ -719,9 +683,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=[ { @@ -866,9 +828,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( @@ -876,9 +836,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( @@ -894,30 +854,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 @@ -926,9 +874,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(): @@ -978,9 +924,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: @@ -1106,15 +1050,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, @@ -1177,15 +1115,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1214,15 +1146,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, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: with pytest.raises(Exception) as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", @@ -1291,9 +1217,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( + f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + ) def test_router_get_model_access_groups_team_only_models(): @@ -1314,14 +1240,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"] @@ -1416,9 +1338,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 @@ -1446,9 +1366,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 @@ -1782,10 +1700,12 @@ async def test_acompletion_streaming_iterator(): assert all(chunk in mock_chunks for chunk in collected_chunks) print("✓ Successfully streamed all chunks") - # Test 2: MidStreamFallbackError with fallback - print("\n=== Test 2: MidStreamFallbackError with fallback ===") + # Test 2: MidStreamFallbackError with generated content is re-raised, not silently continued + print("\n=== Test 2: MidStreamFallbackError re-raises when content already generated ===") - # Create error that should trigger after first chunk + # Error with generated content and is_pre_first_chunk=False (the default): + # the router must re-raise instead of attempting a continuation-prompt fallback, + # because partial content has already been sent to the client. error = MidStreamFallbackError( message="Connection lost", model="gpt-4", @@ -1812,62 +1732,26 @@ async def test_acompletion_streaming_iterator(): self.index += 1 return item - mock_error_response = AsyncIteratorWithError( - mock_chunks, 1 - ) # Error after first chunk + mock_error_response = AsyncIteratorWithError(mock_chunks, 1) # Error after first chunk setattr(mock_error_response, "model", "gpt-4") setattr(mock_error_response, "custom_llm_provider", "openai") setattr(mock_error_response, "logging_obj", MagicMock()) - # Mock the fallback response - fallback_chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content=" world"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]), - ] - - mock_fallback_response = AsyncIterator(fallback_chunks) - - # Mock the fallback function - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=mock_fallback_response, - ) as mock_fallback_utils: - collected_chunks = [] - result = await router._acompletion_streaming_iterator( - model_response=mock_error_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) + result = await router._acompletion_streaming_iterator( + model_response=mock_error_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + # Collect streamed chunks — the first chunk succeeds, then the error re-raises + collected_chunks = [] + with pytest.raises(MidStreamFallbackError): async for chunk in result: collected_chunks.append(chunk) - # Verify fallback was called - assert mock_fallback_utils.called - call_args = mock_fallback_utils.call_args - - # Check that generated content was added to messages - fallback_kwargs = call_args.kwargs["kwargs"] - modified_messages = fallback_kwargs["messages"] - - # Should have original message + system message + assistant message with prefix - assert len(modified_messages) == 3 - assert modified_messages[0] == {"role": "user", "content": "Hello"} - assert modified_messages[1]["role"] == "system" - assert "continuation" in modified_messages[1]["content"] - assert modified_messages[2]["role"] == "assistant" - assert modified_messages[2]["content"] == "Hello" - assert modified_messages[2]["prefix"] == True - - # Verify fallback parameters - assert call_args.kwargs["disable_fallbacks"] == False - assert call_args.kwargs["model_group"] == "gpt-4" - - # Should get original chunk + fallback chunks - assert len(collected_chunks) == 3 # 1 original + 2 fallback - print("✓ Fallback system called correctly with proper message modification") + assert len(collected_chunks) == 1, "one chunk yielded before the error" + print("✓ MidStreamFallbackError re-raised correctly when content was already generated") print("\n=== All tests passed! ===") @@ -2208,11 +2092,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -2292,9 +2172,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( @@ -2377,9 +2255,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 @@ -2497,9 +2375,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, @@ -2508,9 +2384,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, @@ -2811,9 +2685,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"}}, @@ -2841,14 +2713,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"}}, @@ -3089,9 +2957,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: { @@ -3099,15 +2965,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 @@ -3118,26 +2980,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 = { @@ -3156,9 +3010,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 @@ -3178,9 +3030,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 @@ -3213,9 +3063,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 @@ -3224,12 +3072,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 @@ -3252,9 +3096,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 @@ -3290,15 +3132,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 @@ -3345,18 +3183,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 @@ -3366,29 +3200,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) @@ -3425,10 +3247,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 = { @@ -3440,10 +3261,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 = { @@ -3455,9 +3275,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 = { @@ -3468,9 +3288,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(): @@ -3522,14 +3342,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 @@ -3621,15 +3437,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" @@ -3656,9 +3467,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" @@ -3698,9 +3507,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" @@ -4195,9 +4002,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", @@ -4276,12 +4081,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(): @@ -4306,9 +4107,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 @@ -4337,9 +4136,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 @@ -4366,9 +4163,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 @@ -4406,9 +4201,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 @@ -4439,9 +4232,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"}] @@ -4470,9 +4261,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) @@ -4574,12 +4363,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"] @@ -4611,12 +4396,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"] @@ -4683,9 +4464,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f StreamingChoices( finish_reason=None, index=0, - delta=Delta( - content="The Roman Empire began when", role="assistant" - ), + delta=Delta(content="The Roman Empire began when", role="assistant"), ) ], usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), @@ -4738,56 +4517,28 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f assert len(collected) == 1 logging_obj.dispatch_success_handlers.assert_not_called() - # Fallback success: the fallback stream owns success accounting via - # _combine_fallback_usage, so this iterator must not dispatch its own. + # Mid-stream errors with generated content are now re-raised immediately; + # no continuation-prompt fallback is attempted. Success handlers must + # still not be dispatched in this path. model_response, logging_obj = _make_interrupted_model_response() - class _FallbackStream: - def __init__(self, items): - self.items = items - self.index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index >= len(self.items): - raise StopAsyncIteration - item = self.items[self.index] - self.index += 1 - return item - - fallback_stream = _FallbackStream( - [ - litellm.ModelResponseStream( - id="chatcmpl-fallback-1", - model="gpt-3.5-turbo", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content=" continued", role="assistant"), - ) - ], - ) - ] - ) with patch.object( router, "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ): + new=AsyncMock(), + ) as mock_fallback: result = await router._acompletion_streaming_iterator( model_response=model_response, messages=messages, initial_kwargs=dict(initial_kwargs), ) collected = [] - async for chunk in result: - collected.append(chunk) + with pytest.raises(MidStreamFallbackError): + async for chunk in result: + collected.append(chunk) - assert len(collected) == 2 + assert len(collected) == 1, "only the partial chunk before the error" + mock_fallback.assert_not_called(), "fallback must not be called when content already generated" logging_obj.dispatch_success_handlers.assert_not_called() @@ -4995,23 +4746,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 @@ -5056,12 +4801,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( @@ -5186,9 +4927,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 @@ -5233,9 +4972,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( @@ -5310,9 +5047,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( @@ -5379,9 +5114,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( @@ -5436,9 +5169,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 ( @@ -5460,10 +5191,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 @@ -5489,10 +5218,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 @@ -5500,10 +5227,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): @@ -5551,10 +5275,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} ) @@ -5625,9 +5346,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 @@ -5636,9 +5355,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 @@ -5655,9 +5372,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 @@ -5699,9 +5414,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" @@ -5710,9 +5423,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_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -5736,9 +5447,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_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): @@ -5750,9 +5459,7 @@ def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): router = _router_with_two_pass_through_deployments([False, False]) assert len(router.get_model_list()) == 2 assert ( - passthrough_endpoint_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) + passthrough_endpoint_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" ) @@ -5788,16 +5495,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 ) @@ -5837,9 +5537,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 @@ -5857,22 +5555,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 @@ -5888,23 +5580,206 @@ 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} + assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 + + +# --------------------------------------------------------------------------- +# Deferred-stream eager-fetch tests +# --------------------------------------------------------------------------- + + +def _make_deferred_stream_wrapper(make_call_fn): + """Return a CustomStreamWrapper with completion_stream=None and the given make_call.""" + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}} + return CustomStreamWrapper( + completion_stream=None, + model="vertex_ai/gemini-2.0-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=make_call_fn, + ) + + +def _make_router_with_vertex_and_fallback(): + return litellm.Router( + model_list=[ + { + "model_name": "my-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + }, + { + "model_name": "my-fallback", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + }, + }, + ], + fallbacks=[{"my-gemini": ["my-fallback"]}], + num_retries=0, + ) + + +@pytest.mark.asyncio +async def test_acompletion_deferred_stream_error_propagates_through_acompletion(): + """Regression: a deferred-stream CustomStreamWrapper whose make_call raises a 429 + must propagate the exception from within _acompletion's except block so that + fail_calls is incremented (i.e., deployment cooldown fires) and the standard + router fallback chain can handle it. + + Before the fix, the HTTP call happened inside __anext__ (outside the except block), + so fail_calls was never incremented. + """ + import litellm as _litellm + + rate_limit_err = _litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + + async def failing_make_call(**kwargs): + raise rate_limit_err + + router = _make_router_with_vertex_and_fallback() + deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) + + with patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=deferred_wrapper, + ): + with pytest.raises(_litellm.RateLimitError): + await router._acompletion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], ) - == 60 + + model_name = router.model_list[0]["litellm_params"]["model"] + assert router.fail_calls[model_name] == 1, ( + "fail_calls must be incremented when the deferred HTTP call fails; " + "without the eager fetch_stream() fix this stays at 0" + ) + + +@pytest.mark.asyncio +async def test_acompletion_deferred_stream_strips_framing_headers_on_error(): + """Content-Length / Transfer-Encoding / Content-Encoding / Content-Type from a + provider error response are stripped before the exception propagates, preventing + HTTP framing mismatches when LiteLLM builds its own error body. + + Non-framing headers (e.g. x-request-id) must be preserved. + """ + import litellm as _litellm + + err = _litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + err.headers = { + "content-length": "42", + "transfer-encoding": "chunked", + "content-encoding": "gzip", + "content-type": "application/json", + "x-request-id": "abc-123", + } + + async def failing_make_call(**kwargs): + raise err + + router = _make_router_with_vertex_and_fallback() + deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) + + with patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=deferred_wrapper, + ): + with pytest.raises(_litellm.RateLimitError) as exc_info: + await router._acompletion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], + ) + + raised = exc_info.value + headers = getattr(raised, "headers", {}) + assert "content-length" not in headers, "content-length must be stripped" + assert "transfer-encoding" not in headers, "transfer-encoding must be stripped" + assert "content-encoding" not in headers, "content-encoding must be stripped" + assert "content-type" not in headers, "content-type must be stripped" + assert headers.get("x-request-id") == "abc-123", "x-request-id must be preserved" + + +@pytest.mark.asyncio +async def test_acompletion_deferred_stream_skipped_when_stream_already_set(): + """When completion_stream is already populated (non-deferred provider), the eager + fetch_stream() call must be skipped entirely; no exception should be raised even + if make_call would fail. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + async def would_fail(**kwargs): + raise RuntimeError("should not be called") + + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}} + + async def noop_aiter(): + return + yield + + already_set_wrapper = CustomStreamWrapper( + completion_stream=noop_aiter(), + model="openai/gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + make_call=would_fail, + ) + + router = litellm.Router( + model_list=[ + { + "model_name": "my-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + }, + } + ], + ) + + with patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=already_set_wrapper, + ): + result = await router._acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], ) + assert result is not None, "should return a streaming wrapper without errors" + class TestAdvisorSubCallCooldown: """Regression for LIT-4565: an advisor orchestration failure must not cool