diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index eaa2711bff6..7da843b4d5e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1859,18 +1859,15 @@ class CustomStreamWrapper: if self.logging_obj._is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) + _PROVIDERS_WITHOUT_STREAM_FINISH_REASON = frozenset({"baseten", "vllm"}) def _has_provider_finish_reason(self) -> bool: - """True when the provider (or an intermittent mid-stream reason) supplied a terminal finish_reason.""" return self.received_finish_reason is not None or self.intermittent_finish_reason is not None + def _should_require_provider_finish_reason(self) -> bool: + return self.custom_llm_provider not in self._PROVIDERS_WITHOUT_STREAM_FINISH_REASON + def _raise_incomplete_stream_without_finish_reason(self) -> "NoReturn": - """ - OpenAI/Azure-compatible chat streams that end without any provider finish_reason - must not be labeled as a successful completion (finish_reason="stop"). Raise a - MidStreamFallbackError so callers can retry/fallback, while retaining partial - content via generated_content / failure usage recovery. - """ message = ( "Stream ended without a finish_reason from the provider. " "Partial content was received but the response was not successfully completed." @@ -2092,7 +2089,7 @@ class CustomStreamWrapper: self._restore_consumer_correlation_context() raise # Re-raise StopIteration else: - if not self._has_provider_finish_reason(): + if self._should_require_provider_finish_reason() and not self._has_provider_finish_reason(): self._raise_incomplete_stream_without_finish_reason() self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() @@ -2356,7 +2353,7 @@ class CustomStreamWrapper: self._restore_consumer_correlation_context() raise StopAsyncIteration # Re-raise StopIteration else: - if not self._has_provider_finish_reason(): + if self._should_require_provider_finish_reason() and not self._has_provider_finish_reason(): self._raise_incomplete_stream_without_finish_reason() self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() diff --git a/litellm/utils.py b/litellm/utils.py index 52c1859b525..e49916d025e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6975,6 +6975,7 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int | completion_obj = Delta(role="assistant", content=mock_response[i : i + 3]) if n is None: model_response.choices[0].delta = completion_obj + model_response.choices[0].finish_reason = None else: _all_choices = [] for j in range(n): @@ -6985,6 +6986,19 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int | _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if n is None: + model_response.choices[0].delta = Delta(role="assistant", content=None) + model_response.choices[0].finish_reason = "stop" + else: + model_response.choices = [ + litellm.utils.StreamingChoices( + index=j, + delta=litellm.utils.Delta(role="assistant", content=None), + finish_reason="stop", + ) + for j in range(n) + ] + yield model_response async def async_mock_completion_streaming_obj( @@ -7002,6 +7016,7 @@ async def async_mock_completion_streaming_obj( completion_obj = Delta(role="assistant", content=mock_response[i : i + 3]) if n is None: model_response.choices[0].delta = completion_obj + model_response.choices[0].finish_reason = None else: _all_choices = [] for j in range(n): @@ -7012,6 +7027,19 @@ async def async_mock_completion_streaming_obj( _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if n is None: + model_response.choices[0].delta = Delta(role="assistant", content=None) + model_response.choices[0].finish_reason = "stop" + else: + model_response.choices = [ + litellm.utils.StreamingChoices( + index=j, + delta=litellm.utils.Delta(role="assistant", content=None), + finish_reason="stop", + ) + for j in range(n) + ] + yield model_response ########## Reading Config File ############################ diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 8e3ffb048a4..e85fb8da86a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -173,9 +173,7 @@ def test_is_chunk_non_empty_with_annotations( "index": 0, "delta": { "content": None, - "annotations": [ - {"type": "url_citation", "url": "https://www.google.com"} - ], + "annotations": [{"type": "url_citation", "url": "https://www.google.com"}], }, "logprobs": None, "finish_reason": None, @@ -249,9 +247,7 @@ def test_optional_combine_thinking_block_in_choices( # Process first chunk first_response = ModelResponseStream(**first_chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - first_response - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(first_response) print("first_response", json.dumps(first_response, indent=4, default=str)) assert first_response.choices[0].delta.content == "Let me think about this" # assert the response does not have attribute reasoning_content @@ -261,18 +257,14 @@ def test_optional_combine_thinking_block_in_choices( # Process middle chunk middle_response = ModelResponseStream(**middle_chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - middle_response - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(middle_response) print("middle_response", json.dumps(middle_response, indent=4, default=str)) assert middle_response.choices[0].delta.content == " step by step" assert not hasattr(middle_response.choices[0].delta, "reasoning_content") # Process final chunk final_response = ModelResponseStream(**final_chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - final_response - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(final_response) print("final_response", json.dumps(final_response, indent=4, default=str)) assert final_response.choices[0].delta.content == "The answer is 42" assert initialized_custom_stream_wrapper.sent_last_thinking_block is True @@ -383,19 +375,15 @@ def test_multi_chunk_reasoning_and_content( # Process each chunk and verify results for i, (chunk, expected_content) in enumerate(zip(chunks, expected_contents)): response = ModelResponseStream(**chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - response - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(response) # Check content - assert ( - response.choices[0].delta.content == expected_content - ), f"Chunk {i+1}: content mismatch" + assert response.choices[0].delta.content == expected_content, f"Chunk {i + 1}: content mismatch" # Check reasoning_content was removed - assert not hasattr( - response.choices[0].delta, "reasoning_content" - ), f"Chunk {i+1}: reasoning_content should be removed" + assert not hasattr(response.choices[0].delta, "reasoning_content"), ( + f"Chunk {i + 1}: reasoning_content should be removed" + ) # Verify final state assert initialized_custom_stream_wrapper.sent_first_thinking_block is True @@ -406,18 +394,9 @@ def test_strip_sse_data_from_chunk(): """Test the static method that strips 'data: ' prefix from SSE chunks""" # Test with string inputs assert CustomStreamWrapper._strip_sse_data_from_chunk("data: content") == "content" - assert ( - CustomStreamWrapper._strip_sse_data_from_chunk("data: spaced content") - == " spaced content" - ) - assert ( - CustomStreamWrapper._strip_sse_data_from_chunk("regular content") - == "regular content" - ) - assert ( - CustomStreamWrapper._strip_sse_data_from_chunk("regular content with data:") - == "regular content with data:" - ) + assert CustomStreamWrapper._strip_sse_data_from_chunk("data: spaced content") == " spaced content" + assert CustomStreamWrapper._strip_sse_data_from_chunk("regular content") == "regular content" + assert CustomStreamWrapper._strip_sse_data_from_chunk("regular content with data:") == "regular content with data:" # Test with None input assert CustomStreamWrapper._strip_sse_data_from_chunk(None) is None @@ -425,9 +404,7 @@ def test_strip_sse_data_from_chunk(): @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_streaming_handler_with_usage( - sync_mode: bool, final_usage_block: Optional[Usage] = None -): +async def test_streaming_handler_with_usage(sync_mode: bool, final_usage_block: Optional[Usage] = None): import time final_usage_block = final_usage_block or Usage( @@ -542,32 +519,20 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool): with ( patch.object(mock_callback, "log_success_event") as mock_log_success_event, patch.object(mock_callback, "log_stream_event") as mock_log_stream_event, - patch.object( - mock_callback, "async_log_success_event" - ) as mock_async_log_success_event, - patch.object( - mock_callback, "async_log_stream_event" - ) as mock_async_log_stream_event, + patch.object(mock_callback, "async_log_success_event") as mock_async_log_success_event, + patch.object(mock_callback, "async_log_stream_event") as mock_async_log_stream_event, ): - await test_streaming_handler_with_usage( - sync_mode=sync_mode, final_usage_block=final_usage_block - ) + await test_streaming_handler_with_usage(sync_mode=sync_mode, final_usage_block=final_usage_block) if sync_mode: time.sleep(1) mock_log_success_event.assert_called_once() # mock_log_stream_event.assert_called() - assert ( - mock_log_success_event.call_args.kwargs["response_obj"].usage - == final_usage_block - ) + assert mock_log_success_event.call_args.kwargs["response_obj"].usage == final_usage_block else: await asyncio.sleep(1) mock_async_log_success_event.assert_called_once() # mock_async_log_stream_event.assert_called() - assert ( - mock_async_log_success_event.call_args.kwargs["response_obj"].usage - == final_usage_block - ) + assert mock_async_log_success_event.call_args.kwargs["response_obj"].usage == final_usage_block def test_streaming_handler_with_stop_chunk( @@ -693,17 +658,11 @@ def test_set_response_id_propagation_empty_to_valid( """Test that response_id is properly set when first chunk has empty ID and second chunk has valid ID""" model_response1 = ModelResponseStream(id="", created=1742056047, model=None) - model_response1 = initialized_custom_stream_wrapper.set_model_id( - model_response1.id, model_response1 - ) + model_response1 = initialized_custom_stream_wrapper.set_model_id(model_response1.id, model_response1) assert model_response1.id == "" - model_response2 = ModelResponseStream( - id="valid-id-123", created=1742056048, model=None - ) - model_response2 = initialized_custom_stream_wrapper.set_model_id( - "valid-id-123", model_response2 - ) + model_response2 = ModelResponseStream(id="valid-id-123", created=1742056048, model=None) + model_response2 = initialized_custom_stream_wrapper.set_model_id("valid-id-123", model_response2) assert model_response2.id == "valid-id-123" assert initialized_custom_stream_wrapper.response_id == "valid-id-123" @@ -713,19 +672,13 @@ def test_set_response_id_propagation_valid_to_invalid( ): """Test that response_id is maintained when first chunk has valid ID and second chunk has invalid ID""" - model_response1 = ModelResponseStream( - id="first-valid-id", created=1742056049, model=None - ) - model_response1 = initialized_custom_stream_wrapper.set_model_id( - "first-valid-id", model_response1 - ) + model_response1 = ModelResponseStream(id="first-valid-id", created=1742056049, model=None) + model_response1 = initialized_custom_stream_wrapper.set_model_id("first-valid-id", model_response1) assert model_response1.id == "first-valid-id" assert initialized_custom_stream_wrapper.response_id == "first-valid-id" model_response2 = ModelResponseStream(id="", created=1742056050, model=None) - model_response2 = initialized_custom_stream_wrapper.set_model_id( - "", model_response2 - ) + model_response2 = initialized_custom_stream_wrapper.set_model_id("", model_response2) assert model_response2.id == "first-valid-id" assert initialized_custom_stream_wrapper.response_id == "first-valid-id" @@ -741,9 +694,7 @@ async def test_streaming_completion_start_time(logging_obj: Logging): mock_callback = MockCallback() litellm.success_callback = [mock_callback, "langfuse"] - completion_stream = ModelResponseListIterator( - model_responses=bedrock_chunks, delay=0.1 - ) + completion_stream = ModelResponseListIterator(model_responses=bedrock_chunks, delay=0.1) response = CustomStreamWrapper( completion_stream=completion_stream, @@ -757,10 +708,7 @@ async def test_streaming_completion_start_time(logging_obj: Logging): await asyncio.sleep(2) assert logging_obj.model_call_details["completion_start_time"] is not None - assert ( - logging_obj.model_call_details["completion_start_time"] - < logging_obj.model_call_details["end_time"] - ) + assert logging_obj.model_call_details["completion_start_time"] < logging_obj.model_call_details["end_time"] @pytest.mark.asyncio @@ -769,9 +717,7 @@ async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging): from litellm.llms.vertex_ai.common_utils import VertexAIError async def _raise_bad_request(**kwargs): - raise VertexAIError( - status_code=400, message="invalid maxOutputTokens", headers=None - ) + raise VertexAIError(status_code=400, message="invalid maxOutputTokens", headers=None) response = CustomStreamWrapper( completion_stream=None, @@ -800,9 +746,7 @@ async def test_vertex_streaming_rate_limit_triggers_midstream_fallback( from litellm.llms.vertex_ai.common_utils import VertexAIError async def _raise_rate_limit(**kwargs): - raise VertexAIError( - status_code=429, message="Resource exhausted.", headers=None - ) + raise VertexAIError(status_code=429, message="Resource exhausted.", headers=None) response = CustomStreamWrapper( completion_stream=None, @@ -830,9 +774,7 @@ def test_sync_streaming_rate_limit_triggers_midstream_fallback(logging_obj: Logg from litellm.llms.vertex_ai.common_utils import VertexAIError def _raise_rate_limit(**kwargs): - raise VertexAIError( - status_code=429, message="Resource exhausted.", headers=None - ) + raise VertexAIError(status_code=429, message="Resource exhausted.", headers=None) response = CustomStreamWrapper( completion_stream=None, @@ -857,9 +799,7 @@ def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging): from litellm.llms.vertex_ai.common_utils import VertexAIError def _raise_bad_request(**kwargs): - raise VertexAIError( - status_code=400, message="invalid maxOutputTokens", headers=None - ) + raise VertexAIError(status_code=400, message="invalid maxOutputTokens", headers=None) response = CustomStreamWrapper( completion_stream=None, @@ -912,9 +852,7 @@ async def test_bedrock_midstream_internal_server_error_wraps_for_fallback( decoder = AWSEventStreamDecoder(model="anthropic.claude-3-sonnet-20240229-v1:0") async def _bedrock_stream(): - decoder._parse_message_from_event( - _bedrock_error_event("internalServerException") - ) + decoder._parse_message_from_event(_bedrock_error_event("internalServerException")) yield # unreachable; the line above raises async def _make_call(**kwargs): @@ -978,7 +916,7 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): make_call=_raise_400, ) - with pytest.raises(Exception, match='litellm\\.BadRequestError: BedrockException') as excinfo: + with pytest.raises(Exception, match="litellm\\.BadRequestError: BedrockException") as excinfo: await response.__anext__() assert not isinstance(excinfo.value, MidStreamFallbackError) assert getattr(excinfo.value, "status_code", None) == 400 @@ -996,9 +934,7 @@ def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> Cu yield f"data: {json.dumps(error_payload)}" yield "data: [DONE]" - completion_stream = OpenAIChatCompletionStreamingHandler( - streaming_response=_stream(), sync_stream=False - ) + completion_stream = OpenAIChatCompletionStreamingHandler(streaming_response=_stream(), sync_stream=False) return CustomStreamWrapper( completion_stream=completion_stream, model="qwen-vl", @@ -1103,13 +1039,9 @@ def test_streaming_handler_with_created_time_propagation( """Test that the created time is consistent across chunks""" import time - bad_chunk = ModelResponseStream( - choices=[], created=int(time.time()) - ) # chunk with different created time + bad_chunk = ModelResponseStream(choices=[], created=int(time.time())) # chunk with different created time - completion_stream = ModelResponseListIterator( - model_responses=bedrock_chunks + [bad_chunk] - ) + completion_stream = ModelResponseListIterator(model_responses=bedrock_chunks + [bad_chunk]) response = CustomStreamWrapper( completion_stream=completion_stream, @@ -1196,29 +1128,20 @@ def test_optional_combine_thinking_block_with_none_content( # Process first chunk - should not raise TypeError first_response = ModelResponseStream(**first_chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - first_response - ) - assert ( - first_response.choices[0].delta.content - == "Let me think about this problem" - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(first_response) + assert first_response.choices[0].delta.content == "Let me think about this problem" assert not hasattr(first_response.choices[0].delta, "reasoning_content") assert initialized_custom_stream_wrapper.sent_first_thinking_block is True # Process second chunk - should work with continued reasoning second_response = ModelResponseStream(**second_chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - second_response - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(second_response) assert second_response.choices[0].delta.content == " step by step" assert not hasattr(second_response.choices[0].delta, "reasoning_content") # Process final chunk - should add tag final_response = ModelResponseStream(**final_chunk) - initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices( - final_response - ) + initialized_custom_stream_wrapper._optional_combine_thinking_block_in_choices(final_response) assert final_response.choices[0].delta.content == "The answer is 42" assert initialized_custom_stream_wrapper.sent_last_thinking_block is True assert not hasattr(final_response.choices[0].delta, "reasoning_content") @@ -1230,12 +1153,8 @@ def test_has_special_delta_content( """Test the _has_special_delta_content helper method""" # Test empty choices - empty_response = ModelResponseStream( - id="test", created=1742056047, model=None, choices=[] - ) - assert not initialized_custom_stream_wrapper._has_special_delta_content( - empty_response - ) + empty_response = ModelResponseStream(id="test", created=1742056047, model=None, choices=[]) + assert not initialized_custom_stream_wrapper._has_special_delta_content(empty_response) # Test with tool_calls (simulate with mock object) tool_call_response = ModelResponseStream( @@ -1258,9 +1177,7 @@ def test_has_special_delta_content( ) ], ) - assert initialized_custom_stream_wrapper._has_special_delta_content( - tool_call_response - ) + assert initialized_custom_stream_wrapper._has_special_delta_content(tool_call_response) # Test with function_call (simulate with mock object) function_call_response = ModelResponseStream( @@ -1271,24 +1188,18 @@ def test_has_special_delta_content( StreamingChoices( finish_reason=None, index=0, - delta=Delta( - content=None, function_call={"name": "test_func", "arguments": "{}"} - ), + delta=Delta(content=None, function_call={"name": "test_func", "arguments": "{}"}), ) ], ) - assert initialized_custom_stream_wrapper._has_special_delta_content( - function_call_response - ) + assert initialized_custom_stream_wrapper._has_special_delta_content(function_call_response) # Test with audio (simulate by adding audio attribute) audio_response = ModelResponseStream( id="test", created=1742056047, model=None, - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content=None)) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content=None))], ) # Manually add audio attribute to delta audio_response.choices[0].delta.audio = {"transcript": "test"} @@ -1299,9 +1210,7 @@ def test_has_special_delta_content( id="test", created=1742056047, model=None, - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content=None)) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content=None))], ) # Manually add image attribute to delta image_response.choices[0].delta.images = [{"url": "test.jpg"}] @@ -1312,15 +1221,9 @@ def test_has_special_delta_content( id="test", created=1742056047, model=None, - choices=[ - StreamingChoices( - finish_reason=None, index=0, delta=Delta(content="Hello world") - ) - ], - ) - assert not initialized_custom_stream_wrapper._has_special_delta_content( - regular_response + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hello world"))], ) + assert not initialized_custom_stream_wrapper._has_special_delta_content(regular_response) def test_handle_special_delta_content( @@ -1341,9 +1244,7 @@ def test_handle_special_delta_content( ) # The method should call strip_role_from_delta - result = initialized_custom_stream_wrapper._handle_special_delta_content( - test_response - ) + result = initialized_custom_stream_wrapper._handle_special_delta_content(test_response) # Should return the same response object (modified) assert result is test_response @@ -1363,9 +1264,7 @@ def test_has_any_special_delta_attributes( self.audio = {"transcript": "Hello world"} audio_delta = MockDelta() - result = initialized_custom_stream_wrapper._has_any_special_delta_attributes( - audio_delta - ) + result = initialized_custom_stream_wrapper._has_any_special_delta_attributes(audio_delta) assert result is True # Test with delta that has image attribute @@ -1374,9 +1273,7 @@ def test_has_any_special_delta_attributes( self.images = [{"url": "test.jpg"}] image_delta = MockDeltaImage() - result = initialized_custom_stream_wrapper._has_any_special_delta_attributes( - image_delta - ) + result = initialized_custom_stream_wrapper._has_any_special_delta_attributes(image_delta) assert result is True # Test with delta that has no special attributes @@ -1385,9 +1282,7 @@ def test_has_any_special_delta_attributes( self.content = "regular content" regular_delta = MockDeltaRegular() - result = initialized_custom_stream_wrapper._has_any_special_delta_attributes( - regular_delta - ) + result = initialized_custom_stream_wrapper._has_any_special_delta_attributes(regular_delta) assert result is False @@ -1399,22 +1294,16 @@ def test_calculate_total_usage_with_cost(): id="test-1", created=1745513206, model="openrouter/test", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=chunk1_usage, ) - chunk2_usage = Usage( - completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 - ) + chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk2 = ModelResponseStream( id="test-1", created=1745513207, model="openrouter/test", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=chunk2_usage, ) @@ -1456,27 +1345,21 @@ def test_calculate_total_usage_preserves_prompt_cache_token_details(): prompt_tokens=6017, completion_tokens=4, total_tokens=6021, - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=6004, cache_write_tokens=10 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=6004, cache_write_tokens=10), completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=2), ) chunk_with_details = ModelResponseStream( id="chatcmpl-1", created=1745513206, model="openai/gpt-5.6-sol", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=usage_with_details, ) chunk_without_details = ModelResponseStream( id="chatcmpl-1", created=1745513207, model="openai/gpt-5.6-sol", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=Usage(prompt_tokens=6017, completion_tokens=4, total_tokens=6021), ) @@ -1501,9 +1384,7 @@ def test_calculate_total_usage_preserves_anthropic_cache_creation_ttl_breakdown( id="chatcmpl-1", created=1745513206, model="claude-sonnet-5", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=Usage( prompt_tokens=120, completion_tokens=1, @@ -1521,16 +1402,12 @@ def test_calculate_total_usage_preserves_anthropic_cache_creation_ttl_breakdown( id="chatcmpl-1", created=1745513207, model="claude-sonnet-5", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=Usage( prompt_tokens=120, completion_tokens=4, total_tokens=124, - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=0, cache_creation_tokens=100 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_creation_tokens=100), ), ) @@ -1552,38 +1429,26 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin id="chatcmpl-or", created=1742056047, model="openrouter/claude", - choices=[ - StreamingChoices( - finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant") - ) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant"))], usage=None, ) chunk2 = ModelResponseStream( id="chatcmpl-or", created=1742056048, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=None, ) - chunk3_usage = Usage( - completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 - ) + chunk3_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk3 = ModelResponseStream( id="chatcmpl-or", created=1742056049, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))], usage=chunk3_usage, ) - completion_stream = ModelResponseListIterator( - model_responses=[chunk1, chunk2, chunk3] - ) + completion_stream = ModelResponseListIterator(model_responses=[chunk1, chunk2, chunk3]) response = CustomStreamWrapper( completion_stream=completion_stream, model="openrouter/claude", @@ -1618,20 +1483,14 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(): id="chatcmpl-or", created=1742056047, model="openrouter/claude", - choices=[ - StreamingChoices( - finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant") - ) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant"))], usage=None, ) chunk2 = ModelResponseStream( id="chatcmpl-or", created=1742056048, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=None, ) usage_only_chunk = ModelResponseStream( @@ -1639,9 +1498,7 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(): created=1742056049, model="openrouter/claude", choices=[], - usage=Usage( - completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 - ), + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025), ) class MockCallback(CustomLogger): @@ -1670,9 +1527,7 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(): ) response = CustomStreamWrapper( - completion_stream=ModelResponseListIterator( - model_responses=[chunk1, chunk2, usage_only_chunk] - ), + completion_stream=ModelResponseListIterator(model_responses=[chunk1, chunk2, usage_only_chunk]), model="openrouter/claude", custom_llm_provider="openrouter", logging_obj=stream_logging_obj, @@ -1712,32 +1567,22 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): id="chatcmpl-or", created=1742056047, model="openrouter/claude", - choices=[ - StreamingChoices( - finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant") - ) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant"))], usage=None, ) chunk2 = ModelResponseStream( id="chatcmpl-or", created=1742056048, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=None, ) chunk3 = ModelResponseStream( id="chatcmpl-or", created=1742056049, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) - ], - usage=Usage( - completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 - ), + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025), ) # Build the complete response as stream_chunk_builder does @@ -1754,19 +1599,12 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openrouter") assert "additional_headers" in complete_response._hidden_params - assert ( - complete_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] - == 0.00025 - ) + assert complete_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.00025 # Verify the cost calculator would pick this up from litellm.cost_calculator import get_response_cost_from_hidden_params - provider_cost = get_response_cost_from_hidden_params( - complete_response._hidden_params - ) + provider_cost = get_response_cost_from_hidden_params(complete_response._hidden_params) assert provider_cost == 0.00025 @@ -1795,18 +1633,14 @@ def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator(): id="chatcmpl-pplx", created=1742056048, model="perplexity/sonar", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=None, ), ModelResponseStream( id="chatcmpl-pplx", created=1742056049, model="perplexity/sonar", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))], usage=Usage( completion_tokens=18, prompt_tokens=12, @@ -1821,9 +1655,7 @@ def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator(): ), ] - complete_response = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "test"}] - ) + complete_response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "test"}]) assert complete_response is not None @@ -1920,9 +1752,7 @@ def test_handle_special_delta_attributes( id="test", created=1742056047, model=None, - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="test")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="test"))], ) # Test with delta that has audio attribute @@ -1931,9 +1761,7 @@ def test_handle_special_delta_attributes( self.audio = {"transcript": "Hello world"} audio_delta = MockDelta() - initialized_custom_stream_wrapper._handle_special_delta_attributes( - audio_delta, model_response - ) + initialized_custom_stream_wrapper._handle_special_delta_attributes(audio_delta, model_response) # Should copy the audio attribute assert hasattr(model_response.choices[0].delta, "audio") @@ -1949,14 +1777,10 @@ def test_handle_special_delta_attributes( id="test", created=1742056047, model=None, - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="test")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="test"))], ) - initialized_custom_stream_wrapper._handle_special_delta_attributes( - image_delta, model_response2 - ) + initialized_custom_stream_wrapper._handle_special_delta_attributes(image_delta, model_response2) # Should copy the image attribute print(f"delta: {model_response2.choices[0].delta}") @@ -1971,9 +1795,7 @@ def test_has_special_delta_attribute( """Test the _has_special_delta_attribute helper method""" # Test with None delta - assert not initialized_custom_stream_wrapper._has_special_delta_attribute( - None, "audio" - ) + assert not initialized_custom_stream_wrapper._has_special_delta_attribute(None, "audio") # Test with delta that has the attribute class MockDelta: @@ -1981,9 +1803,7 @@ def test_has_special_delta_attribute( self.audio = {"transcript": "test"} delta_with_audio = MockDelta() - assert initialized_custom_stream_wrapper._has_special_delta_attribute( - delta_with_audio, "audio" - ) + assert initialized_custom_stream_wrapper._has_special_delta_attribute(delta_with_audio, "audio") # Test with delta that doesn't have the attribute class MockDeltaNoAudio: @@ -1991,9 +1811,7 @@ def test_has_special_delta_attribute( self.content = "test" delta_without_audio = MockDeltaNoAudio() - assert not initialized_custom_stream_wrapper._has_special_delta_attribute( - delta_without_audio, "audio" - ) + assert not initialized_custom_stream_wrapper._has_special_delta_attribute(delta_without_audio, "audio") # Test with delta that has the attribute but it's None class MockDeltaNone: @@ -2001,9 +1819,7 @@ def test_has_special_delta_attribute( self.audio = None delta_with_none = MockDeltaNone() - assert not initialized_custom_stream_wrapper._has_special_delta_attribute( - delta_with_none, "audio" - ) + assert not initialized_custom_stream_wrapper._has_special_delta_attribute(delta_with_none, "audio") def test_is_chunk_non_empty_with_empty_tool_calls( @@ -2171,12 +1987,7 @@ _REPETITION_TEST_CASES = [ pytest.param( ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1) + ["different_mid"] - + ["same"] - * ( - litellm.REPEATED_STREAMING_CHUNK_LIMIT - - litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 - + 1 - ), + + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1), False, id="middle_chunk_different_no_raise", ), @@ -2205,6 +2016,7 @@ def test_raise_on_model_repetition( chunks = _build_chunks(chunks_pattern, len(chunks_pattern)) if should_raise: + def _feed(): for chunk in chunks: wrapper.chunks.append(chunk) @@ -2338,12 +2150,12 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): last_chunk = collected[-1] hidden_usage = last_chunk._hidden_params.get("usage") assert hidden_usage is not None, "Expected usage in _hidden_params" - assert ( - hidden_usage.prompt_tokens == 20 - ), f"Expected prompt_tokens=20 from provider, got {hidden_usage.prompt_tokens}" - assert ( - hidden_usage.completion_tokens == 135 - ), f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" + assert hidden_usage.prompt_tokens == 20, ( + f"Expected prompt_tokens=20 from provider, got {hidden_usage.prompt_tokens}" + ) + assert hidden_usage.completion_tokens == 135, ( + f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" + ) @pytest.mark.asyncio @@ -2418,9 +2230,7 @@ def test_content_not_dropped_when_finish_reason_already_set( result = initialized_custom_stream_wrapper.chunk_creator(chunk=content_chunk) - assert ( - result is not None - ), "chunk_creator() returned None — content was dropped (issue #22098)" + assert result is not None, "chunk_creator() returned None — content was dropped (issue #22098)" assert result.choices[0].delta.content == "world!" @@ -2472,14 +2282,10 @@ def test_tool_use_not_dropped_when_finish_reason_already_set( result = initialized_custom_stream_wrapper.chunk_creator(chunk=tool_chunk) - assert ( - result is not None - ), "chunk_creator() returned None — tool_use data was dropped" + assert result is not None, "chunk_creator() returned None — tool_use data was dropped" tool_calls = result.choices[0].delta.tool_calls - assert ( - tool_calls is not None and len(tool_calls) > 0 - ), "tool_calls should contain at least one tool call" + assert tool_calls is not None and len(tool_calls) > 0, "tool_calls should contain at least one tool call" assert tool_calls[0].id == "call_1" assert tool_calls[0].function.name == "get_weather" @@ -2535,9 +2341,7 @@ def test_dispatch_vllm_extracts_output_text( class _VLLMChunk: outputs = [_Output()] - result, _, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, [_VLLMChunk()] - ) + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, [_VLLMChunk()]) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "hello from vllm" @@ -2551,9 +2355,7 @@ def test_dispatch_petals_slices_completion_stream( initialized_custom_stream_wrapper.custom_llm_provider = "petals" initialized_custom_stream_wrapper.completion_stream = "A" * 50 - result, _, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, chunk=None - ) + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk=None) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "A" * 30 @@ -2568,9 +2370,7 @@ def test_dispatch_petals_empty_stream_sets_stop( initialized_custom_stream_wrapper.custom_llm_provider = "petals" initialized_custom_stream_wrapper.completion_stream = "" - result, _, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, chunk=None - ) + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk=None) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "" @@ -2596,9 +2396,7 @@ def test_dispatch_palm_slices_completion_stream( initialized_custom_stream_wrapper.custom_llm_provider = "palm" initialized_custom_stream_wrapper.completion_stream = "B" * 40 - result, _, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, chunk=None - ) + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk=None) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "B" * 30 @@ -2622,9 +2420,7 @@ def test_dispatch_cached_response_extracts_delta( ], ) - result, model_response, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, chunk - ) + result, model_response, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "cached text" @@ -2650,9 +2446,7 @@ def test_dispatch_vertex_ai_legacy_text_and_finish_reason( candidates = [_Candidate()] text = "vertex content" - result, _, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, _VertexChunk() - ) + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, _VertexChunk()) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "vertex content" @@ -2669,9 +2463,7 @@ def test_dispatch_vertex_ai_legacy_without_candidates_stringifies_chunk( def __str__(self) -> str: return "raw vertex blob" - result, _, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, _RawChunk() - ) + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, _RawChunk()) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "raw vertex blob" @@ -2708,9 +2500,7 @@ def test_dispatch_vertex_ai_legacy_function_call( def text(self): raise RuntimeError("Part has no text.") - result, _, _ = _run_dispatch( - initialized_custom_stream_wrapper, _VertexFunctionChunk() - ) + result, _, _ = _run_dispatch(initialized_custom_stream_wrapper, _VertexFunctionChunk()) assert isinstance(result, _ProviderChunkParsed) tool_calls = result.response_obj["original_chunk"].choices[0].delta.tool_calls @@ -2727,11 +2517,7 @@ def test_dispatch_custom_provider_returns_chunk_early( straight through as an early return rather than re-parsing it.""" monkeypatch.setattr(litellm, "_custom_providers", ["my-custom-llm"]) initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-llm" - chunk = ModelResponseStream( - choices=[ - StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason=None) - ] - ) + chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason=None)]) result, _, _ = _run_dispatch(initialized_custom_stream_wrapper, chunk) @@ -2747,11 +2533,7 @@ def test_dispatch_custom_provider_finish_only_returns_none_early( records the reason and returns None so no empty delta is emitted.""" monkeypatch.setattr(litellm, "_custom_providers", ["my-custom-llm"]) initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-llm" - chunk = ModelResponseStream( - choices=[ - StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop") - ] - ) + chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")]) result, _, _ = _run_dispatch(initialized_custom_stream_wrapper, chunk) @@ -2766,9 +2548,7 @@ def test_dispatch_text_completion_codestral_parses_chunk( """text-completion-codestral streams raw SSE JSON strings that the dispatch routes through CodestralTextCompletionConfig to extract content/finish.""" initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-codestral" - chunk = json.dumps( - {"choices": [{"delta": {"content": "codestral text"}, "finish_reason": "stop"}]} - ) + chunk = json.dumps({"choices": [{"delta": {"content": "codestral text"}, "finish_reason": "stop"}]}) result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk) @@ -2807,9 +2587,7 @@ def test_dispatch_ai21_decodes_completion( ): """ai21 does fake streaming over a single byte-encoded JSON completion.""" initialized_custom_stream_wrapper.custom_llm_provider = "ai21" - chunk = json.dumps({"completions": [{"data": {"text": "ai21 text"}}]}).encode( - "utf-8" - ) + chunk = json.dumps({"completions": [{"data": {"text": "ai21 text"}}]}).encode("utf-8") result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk) @@ -2838,9 +2616,7 @@ def test_dispatch_text_completion_openai_with_usage( choices = [_Choice()] usage = _Usage() - result, model_response, completion_obj = _run_dispatch( - initialized_custom_stream_wrapper, _TextChunk() - ) + result, model_response, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, _TextChunk()) assert isinstance(result, _ProviderChunkParsed) assert completion_obj["content"] == "oai text" @@ -3091,9 +2867,7 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool - sent_first_chunk is only marked for chunks with real choices - Chunks with role in delta are not discarded as empty """ - completion_stream = ModelResponseListIterator( - model_responses=_AZURE_CHUNKS_WITH_PROMPT_FILTER - ) + completion_stream = ModelResponseListIterator(model_responses=_AZURE_CHUNKS_WITH_PROMPT_FILTER) response = CustomStreamWrapper( completion_stream=completion_stream, @@ -3120,19 +2894,14 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool chunks.append(chunk) # The prompt_filter chunk should be forwarded with choices=[] - assert ( - len(chunks[0].choices) == 0 - ), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" + assert len(chunks[0].choices) == 0, ( + f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" + ) # At least one chunk must have role='assistant' in its delta - has_role = any( - len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant" - for c in chunks - ) - assert has_role, ( - "No chunk contained role='assistant' in delta (issue #24221). " - "Chunk deltas: " - + str([c.choices[0].delta if c.choices else "no choices" for c in chunks]) + has_role = any(len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant" for c in chunks) + assert has_role, "No chunk contained role='assistant' in delta (issue #24221). Chunk deltas: " + str( + [c.choices[0].delta if c.choices else "no choices" for c in chunks] ) @@ -3206,9 +2975,7 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum(): ) -@pytest.mark.parametrize( - "finish_reason", ["stop", "tool_calls", "length", "content_filter"] -) +@pytest.mark.parametrize("finish_reason", ["stop", "tool_calls", "length", "content_filter"]) def test_chunk_creator_passes_through_model_response_stream( initialized_custom_stream_wrapper: CustomStreamWrapper, finish_reason: str, @@ -3325,9 +3092,9 @@ def test_chunk_creator_strips_finish_reason_from_content_chunk( litellm._custom_providers.remove("my-custom-provider") assert result is not None - assert ( - result.choices[0].finish_reason is None - ), "finish_reason must be stripped from content chunks to avoid double terminal chunks" + assert result.choices[0].finish_reason is None, ( + "finish_reason must be stripped from content chunks to avoid double terminal chunks" + ) assert initialized_custom_stream_wrapper.received_finish_reason == "stop" @@ -3354,9 +3121,7 @@ def test_chunk_creator_tool_calls_not_dropped_on_finish( tool_calls=[ ChatCompletionDeltaToolCall( id="call_abc", - function=Function( - name="get_weather", arguments='{"city":"NYC"}' - ), + function=Function(name="get_weather", arguments='{"city":"NYC"}'), type="function", index=0, ) @@ -3409,9 +3174,7 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost(): 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=30, completion_tokens=1, total_tokens=31), @@ -3486,9 +3249,7 @@ def _wrapper_with_partial_chunks( 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, @@ -3558,9 +3319,7 @@ def test_record_partial_usage_for_failure_carries_up_openai_style_cached_tokens( total_tokens=1010, prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), ) - wrapper, logging_obj = _wrapper_with_partial_chunks( - chunk_model="gpt-4o-mini", usage=recovered - ) + wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="gpt-4o-mini", usage=recovered) wrapper._record_partial_usage_for_failure() @@ -3577,9 +3336,7 @@ def test_record_partial_usage_for_failure_keeps_cache_values_recovered_from_chun cache_read_input_tokens=7, cache_creation_input_tokens=3, ) - wrapper, logging_obj = _wrapper_with_partial_chunks( - chunk_model="gpt-4o-mini", usage=recovered - ) + wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="gpt-4o-mini", usage=recovered) wrapper._record_partial_usage_for_failure() @@ -3600,9 +3357,7 @@ async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage( before the fix it escaped __next__/__anext__ and the request was dropped from SpendLogs while the provider billed the tokens. The wrapper must catch it and recover usage from the raw chunks so cost is still tracked.""" - final_usage_block = Usage( - completion_tokens=392, prompt_tokens=1799, total_tokens=2191 - ) + final_usage_block = Usage(completion_tokens=392, prompt_tokens=1799, total_tokens=2191) final_chunk = ModelResponseStream( id="chatcmpl-raise-test", created=1742056047, @@ -3653,9 +3408,9 @@ async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage( if getattr(chunk, "usage", None) is not None: seen_usage.append(chunk.usage) - assert any( - u.total_tokens == final_usage_block.total_tokens for u in seen_usage - ), "usage recovered from raw chunks was not emitted after stream_chunk_builder raised" + assert any(u.total_tokens == final_usage_block.total_tokens for u in seen_usage), ( + "usage recovered from raw chunks was not emitted after stream_chunk_builder raised" + ) @pytest.mark.parametrize("sync_mode", [True, False]) @@ -3684,9 +3439,7 @@ async def test_stream_chunk_builder_raise_and_usage_recovery_failure_does_not_cr ) response = CustomStreamWrapper( - completion_stream=ModelResponseListIterator( - model_responses=bedrock_chunks + [final_chunk] - ), + completion_stream=ModelResponseListIterator(model_responses=bedrock_chunks + [final_chunk]), model="bedrock/claude-haiku-4-5-20251001-v1:0", custom_llm_provider="bedrock", logging_obj=Logging( @@ -3702,12 +3455,8 @@ async def test_stream_chunk_builder_raise_and_usage_recovery_failure_does_not_cr ) with ( - patch.object( - litellm, "stream_chunk_builder", side_effect=Exception("assembly failed") - ), - patch.object( - sh_module, "calculate_total_usage", side_effect=Exception("recovery failed") - ), + patch.object(litellm, "stream_chunk_builder", side_effect=Exception("assembly failed")), + patch.object(sh_module, "calculate_total_usage", side_effect=Exception("recovery failed")), ): # must not raise even though both assembly and recovery fail if sync_mode: @@ -3781,9 +3530,7 @@ async def test_transport_read_error_after_finish_reason_ends_stream_gracefully( chunks = [chunk async for chunk in response] finish_reasons = [ - chunk.choices[0].finish_reason - for chunk in chunks - if chunk.choices and chunk.choices[0].finish_reason + chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason ] contents = [ chunk.choices[0].delta.content @@ -3818,6 +3565,7 @@ async def test_transport_read_error_before_finish_reason_raises(logging_obj: Log ) received = [] + async def _drain(): async for chunk in response: received.append(chunk) @@ -3826,9 +3574,7 @@ async def test_transport_read_error_before_finish_reason_raises(logging_obj: Log await _drain() fabricated_finish_reasons = [ - chunk.choices[0].finish_reason - for chunk in received - if chunk.choices and chunk.choices[0].finish_reason + chunk.choices[0].finish_reason for chunk in received if chunk.choices and chunk.choices[0].finish_reason ] assert fabricated_finish_reasons == [] @@ -4609,22 +4355,13 @@ def test_chunk_creator_keeps_provider_model_private_across_stream(): assert terminal_result is not None assert first_result.model == "requested-route" assert terminal_result.model == "requested-route" - assert ( - get_hidden_params_dict(first_result)["provider_response_model"] - == "selected-model" - ) - assert ( - get_hidden_params_dict(terminal_result)["provider_response_model"] - == "selected-model" - ) + assert get_hidden_params_dict(first_result)["provider_response_model"] == "selected-model" + assert get_hidden_params_dict(terminal_result)["provider_response_model"] == "selected-model" assembled = litellm.stream_chunk_builder(chunks=[first_result, terminal_result]) assert assembled is not None assert assembled.model == "requested-route" - assert ( - get_hidden_params_dict(assembled)["provider_response_model"] - == "selected-model" - ) + assert get_hidden_params_dict(assembled)["provider_response_model"] == "selected-model" def test_assembled_stream_uses_later_provider_model_for_cost( @@ -4702,35 +4439,20 @@ def test_assembled_stream_uses_later_provider_model_for_cost( assert router_result is not None assert selected_result is not None assert terminal_result is not None - assert ( - get_hidden_params_dict(router_result)["provider_response_model"] - == "azure-model-router" - ) - assert ( - get_hidden_params_dict(selected_result)["provider_response_model"] - == "gpt-4.1-nano-2025-04-14" - ) - assert ( - get_hidden_params_dict(terminal_result)["provider_response_model"] - == "azure-model-router" - ) + assert get_hidden_params_dict(router_result)["provider_response_model"] == "azure-model-router" + assert get_hidden_params_dict(selected_result)["provider_response_model"] == "gpt-4.1-nano-2025-04-14" + assert get_hidden_params_dict(terminal_result)["provider_response_model"] == "azure-model-router" - assembled = litellm.stream_chunk_builder( - chunks=[router_result, selected_result, terminal_result] - ) + assembled = litellm.stream_chunk_builder(chunks=[router_result, selected_result, terminal_result]) assert assembled is not None assert assembled.model == "gpt-4.1-nano-2025-04-14" - assert ( - get_hidden_params_dict(assembled)["provider_response_model"] - == "gpt-4.1-nano-2025-04-14" - ) + assert get_hidden_params_dict(assembled)["provider_response_model"] == "gpt-4.1-nano-2025-04-14" assembled.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) assert litellm.completion_cost( completion_response=assembled, custom_llm_provider="azure", ) == pytest.approx( - 10 * selected_model_info["input_cost_per_token"] - + 5 * selected_model_info["output_cost_per_token"] + 10 * selected_model_info["input_cost_per_token"] + 5 * selected_model_info["output_cost_per_token"] ) @@ -4879,10 +4601,9 @@ class TestStableStreamingResponseId: wrapper.response_id = "chatcmpl-from-provider" assert wrapper.model_response_creator().id == "chatcmpl-from-provider" + @pytest.mark.asyncio async def test_clean_eof_without_finish_reason_raises_midstream_error_issue_40260(): - """OpenAI/Azure clean EOF with partial content and no finish_reason must not - synthesize finish_reason=stop (#40260).""" from litellm.exceptions import MidStreamFallbackError async def source(): @@ -4929,7 +4650,6 @@ async def test_clean_eof_without_finish_reason_raises_midstream_error_issue_4026 def test_clean_eof_without_finish_reason_raises_sync_issue_40260(): - """Sync sibling of #40260: empty/exhausted stream with no provider finish_reason.""" from litellm.exceptions import MidStreamFallbackError log = Logging( @@ -4952,7 +4672,6 @@ def test_clean_eof_without_finish_reason_raises_sync_issue_40260(): def test_provider_finish_reason_still_synthesizes_terminal_chunk_issue_40260(): - """Valid provider finish_reason (length) must still produce a successful terminal chunk.""" log = Logging( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}], @@ -4971,3 +4690,43 @@ def test_provider_finish_reason_still_synthesizes_terminal_chunk_issue_40260(): wrapper.received_finish_reason = "length" chunk = next(wrapper) assert chunk.choices[0].finish_reason == "length" + + +def test_baseten_eof_without_finish_reason_still_succeeds_issue_40260(): + log = Logging( + model="baseten/model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="baseten-eof", + function_id="baseten-eof", + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="baseten/model", + custom_llm_provider="baseten", + logging_obj=log, + ) + chunk = next(wrapper) + assert chunk.choices[0].finish_reason == "stop" + + +def test_vllm_eof_without_finish_reason_still_succeeds_issue_40260(): + log = Logging( + model="vllm/model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="vllm-eof", + function_id="vllm-eof", + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="vllm/model", + custom_llm_provider="vllm", + logging_obj=log, + ) + chunk = next(wrapper) + assert chunk.choices[0].finish_reason == "stop"