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 e85fb8da86a..18b61a82ea6 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -173,7 +173,9 @@ 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,
@@ -247,7 +249,9 @@ 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
@@ -257,14 +261,18 @@ 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
@@ -375,15 +383,19 @@ 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
@@ -394,9 +406,18 @@ 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
@@ -404,7 +425,9 @@ 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(
@@ -519,20 +542,32 @@ 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(
@@ -658,11 +693,17 @@ 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"
@@ -672,13 +713,19 @@ 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"
@@ -694,7 +741,9 @@ 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,
@@ -708,7 +757,10 @@ 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
@@ -717,7 +769,9 @@ 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,
@@ -746,7 +800,9 @@ 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,
@@ -774,7 +830,9 @@ 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,
@@ -799,7 +857,9 @@ 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,
@@ -852,7 +912,9 @@ 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):
@@ -916,7 +978,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
@@ -934,7 +996,9 @@ 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",
@@ -1039,9 +1103,13 @@ 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,
@@ -1128,20 +1196,29 @@ 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")
@@ -1153,8 +1230,12 @@ 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(
@@ -1177,7 +1258,9 @@ 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(
@@ -1188,18 +1271,24 @@ 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"}
@@ -1210,7 +1299,9 @@ 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"}]
@@ -1221,9 +1312,15 @@ def test_has_special_delta_content(
id="test",
created=1742056047,
model=None,
- choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hello world"))],
+ choices=[
+ StreamingChoices(
+ finish_reason=None, index=0, delta=Delta(content="Hello world")
+ )
+ ],
+ )
+ assert not initialized_custom_stream_wrapper._has_special_delta_content(
+ regular_response
)
- assert not initialized_custom_stream_wrapper._has_special_delta_content(regular_response)
def test_handle_special_delta_content(
@@ -1244,7 +1341,9 @@ 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
@@ -1264,7 +1363,9 @@ 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
@@ -1273,7 +1374,9 @@ 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
@@ -1282,7 +1385,9 @@ 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
@@ -1294,16 +1399,22 @@ 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,
)
@@ -1345,21 +1456,27 @@ 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),
)
@@ -1384,7 +1501,9 @@ 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,
@@ -1402,12 +1521,16 @@ 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
+ ),
),
)
@@ -1429,26 +1552,38 @@ 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",
@@ -1483,14 +1618,20 @@ 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(
@@ -1498,7 +1639,9 @@ 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):
@@ -1527,7 +1670,9 @@ 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,
@@ -1567,22 +1712,32 @@ 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
@@ -1599,12 +1754,19 @@ 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
@@ -1633,14 +1795,18 @@ 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,
@@ -1655,7 +1821,9 @@ 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
@@ -1752,7 +1920,9 @@ 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
@@ -1761,7 +1931,9 @@ 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")
@@ -1777,10 +1949,14 @@ 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}")
@@ -1795,7 +1971,9 @@ 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:
@@ -1803,7 +1981,9 @@ 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:
@@ -1811,7 +1991,9 @@ 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:
@@ -1819,7 +2001,9 @@ 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(
@@ -1987,7 +2171,12 @@ _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",
),
@@ -2016,7 +2205,6 @@ 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)
@@ -2150,12 +2338,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
@@ -2230,7 +2418,9 @@ 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!"
@@ -2282,10 +2472,14 @@ 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"
@@ -2341,7 +2535,9 @@ 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"
@@ -2355,7 +2551,9 @@ 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
@@ -2370,7 +2568,9 @@ 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"] == ""
@@ -2396,7 +2596,9 @@ 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
@@ -2420,7 +2622,9 @@ 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"
@@ -2446,7 +2650,9 @@ 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"
@@ -2463,7 +2669,9 @@ 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"
@@ -2500,7 +2708,9 @@ 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
@@ -2517,7 +2727,11 @@ 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)
@@ -2533,7 +2747,11 @@ 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)
@@ -2548,7 +2766,9 @@ 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)
@@ -2587,7 +2807,9 @@ 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)
@@ -2616,7 +2838,9 @@ 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"
@@ -2867,7 +3091,9 @@ 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,
@@ -2894,14 +3120,19 @@ 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])
)
@@ -2975,7 +3206,9 @@ 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,
@@ -3092,9 +3325,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"
@@ -3121,7 +3354,9 @@ 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,
)
@@ -3174,7 +3409,9 @@ 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),
@@ -3249,7 +3486,9 @@ 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,
@@ -3319,7 +3558,9 @@ 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()
@@ -3336,7 +3577,9 @@ 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()
@@ -3357,7 +3600,9 @@ 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,
@@ -3408,9 +3653,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])
@@ -3439,7 +3684,9 @@ 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(
@@ -3455,8 +3702,12 @@ 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:
@@ -3530,7 +3781,9 @@ 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
@@ -3565,7 +3818,6 @@ 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)
@@ -3574,7 +3826,9 @@ 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 == []
@@ -4355,13 +4609,22 @@ 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(
@@ -4439,20 +4702,35 @@ 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"]
)
@@ -4601,7 +4879,6 @@ 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():
from litellm.exceptions import MidStreamFallbackError
@@ -4730,3 +5007,4 @@ def test_vllm_eof_without_finish_reason_still_succeeds_issue_40260():
)
chunk = next(wrapper)
assert chunk.choices[0].finish_reason == "stop"
+