From e8f0bfd9b96ad15de7926e6eefe6c9897b70455c Mon Sep 17 00:00:00 2001 From: JingHao-Leon Date: Sun, 11 Oct 2026 11:08:03 +0800 Subject: [PATCH 1/9] fix(streaming): emit the terminal success log for a stream closed early A stream the consumer stops reading early (break, drop, aclose) never reaches the final-chunk branch, so complete_streaming_response is never assembled and the success-callback dispatch (which gates on its presence for streaming calls) skips every integration: tracing, spend tracking and custom callbacks all silently miss the call (issue #45736 - an OpenTelemetry setup records zero spans for a stream abandoned after the first chunk). CustomStreamWrapper now assembles the response from the chunks received so far and submits the terminal success log when the stream is closed early. Three entry points: - close(): new deterministic public method (plus context-manager support) - aclose(): async close path - __del__: last-resort fallback (subject to GC timing) Streams that reach their final chunk are unaffected (flush is a no-op once sent_last_chunk is set), and double-close flushes only once. Fixes #45736 --- .../litellm_core_utils/streaming_handler.py | 68 +++++++++++++++++++ .../test_streaming_handler.py | 61 +++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a78e6f17cc9..3663634fcaf 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -307,6 +307,41 @@ class CustomStreamWrapper: } self._post_streaming_hooks: list | None = None + self._abandonment_logged: bool = False + + def _flush_abandoned_stream_logging(self) -> None: + """Emit a terminal success log for a stream the consumer abandoned. + + A stream the caller stops reading early (break, drop, aclose) never + reaches the final-chunk branch that assembles + ``complete_streaming_response`` and submits the terminal success log, + so every success callback (tracing, spend, ...) silently misses the + call even though tokens were served. Assemble the response from the + chunks received so far and submit the same logging once, mirroring how + the proxy bills a partial stream on client disconnect. + """ + if self._abandonment_logged or self.sent_last_chunk: + return + if not self.chunks: # pyright: ignore[reportUnknownMemberType] # chunks is a loosely-typed accumulator + return + self._abandonment_logged = True + try: + complete_streaming_response = litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # stream_chunk_builder is exposed on the package root + chunks=self.chunks, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] + messages=self.messages, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + except Exception as e: + verbose_logger.warning( + "stream_chunk_builder raised for an abandoned stream (%s); the partial call is not logged.", + str(e), + ) + return + if complete_streaming_response is None: + return + self.logging_obj.model_call_details["async_complete_streaming_response"] = complete_streaming_response + self.run_success_logging_and_cache_storage(complete_streaming_response, cache_hit=False) # pyright: ignore[reportUnknownMemberType] # resolved dynamically on the logging object def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" @@ -328,6 +363,32 @@ class CustomStreamWrapper: def __aiter__(self) -> AsyncIterator["ModelResponseStream"]: return self + def close(self) -> None: + """Deterministically end a stream the consumer stopped reading early. + + A stream abandoned without reaching its final chunk would otherwise + never emit its terminal success log (spans, spend, custom callbacks + all miss the call). ``close()`` flushes that log from the chunks + received so far and releases the underlying provider stream; safe to + call after the stream has been fully consumed (it is a no-op then). + """ + self._flush_abandoned_stream_logging() + stream_to_close: Final = self.completion_stream # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the underlying stream object is provider-specific + self.completion_stream = None + if stream_to_close is not None and hasattr(stream_to_close, "close"): # pyright: ignore[reportUnknownArgumentType] # provider-specific stream object + try: + stream_to_close.close() # pyright: ignore[reportUnknownMemberType] # provider-specific stream object + except Exception as e: + verbose_logger.debug("closing the underlying stream failed: %s", str(e)) + self._restore_consumer_correlation_context() + + def __enter__(self) -> "CustomStreamWrapper": + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> bool: + self.close() + return False + def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None: """Restore trace_id/session_id in the *consuming* thread/task/context. @@ -391,8 +452,15 @@ class CustomStreamWrapper: still-active call's context within that same Task if this fires late. """ self._restore_consumer_correlation_context(guarded=True) + try: + self._flush_abandoned_stream_logging() + except Exception: + verbose_logger.debug("abandoned-stream logging failed in __del__", exc_info=True) async def aclose(self): + # Emit the abandoned-stream log before closing the provider stream: + # the chunks collected so far are what this call actually served. + self._flush_abandoned_stream_logging() # Restore the consumer's outer context only after the underlying # provider stream's own close (and its diagnostic logging below, if # closing fails) completes - not before - so those log lines still diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 01fae0958e6..cd813734e02 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6417,3 +6417,64 @@ def validate_last_format(chunk): assert isinstance( choice["finish_reason"], str ), "'finish_reason' should be a string." + + +class TestCloseFlushesAbandonedStream: + """A stream abandoned before its final chunk must still emit its terminal + success log when closed (issue #45736): tracing, spend and custom + callbacks otherwise silently miss the call.""" + + def _wrapper_with_chunk(self) -> CustomStreamWrapper: + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=MagicMock(), + custom_llm_provider=None, + ) + wrapper.sent_last_chunk = False + wrapper.chunks.append(ModelResponseStream(id="c1", choices=[])) + return wrapper + + def test_close_flushes_terminal_log_for_abandoned_stream(self): + wrapper = self._wrapper_with_chunk() + with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: + wrapper.close() + + run_logging.assert_called_once() + assembled = wrapper.logging_obj.model_call_details["async_complete_streaming_response"] + assert assembled is not None + + def test_close_is_noop_after_natural_completion(self): + wrapper = self._wrapper_with_chunk() + wrapper.sent_last_chunk = True + with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: + wrapper.close() + + run_logging.assert_not_called() + + def test_double_close_flushes_only_once(self): + wrapper = self._wrapper_with_chunk() + with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: + wrapper.close() + wrapper.close() + + run_logging.assert_called_once() + + def test_context_manager_exit_flushes(self): + wrapper = self._wrapper_with_chunk() + with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: + with wrapper: + pass + + run_logging.assert_called_once() + + def test_close_releases_underlying_stream(self): + wrapper = self._wrapper_with_chunk() + wrapper.logging_obj.async_success_handler = AsyncMock() + stream = MagicMock() + wrapper.completion_stream = stream + + wrapper.close() + + stream.close.assert_called_once() + assert wrapper.completion_stream is None From d7450305391fd604609d0aaad1359f20d2450eec Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 12:30:23 +0800 Subject: [PATCH 2/9] fix(streaming): annotate the three intentional blind excepts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ruff_strict_gate.py` counts BLE001 against a ceiling taken at the merge-base, and the abandoned-stream cleanup path added three: - around `stream_chunk_builder` — a failure here means the partial call cannot be assembled, and it is already logged rather than raised; - around closing the provider stream — must not mask the caller's own error; - in `__del__` — a finalizer must never raise, and the logging is best-effort. Each is suppressed with that reason, matching the existing `# noqa: BLE001` usages elsewhere in litellm. The file's BLE001 count is back to the merge-base baseline (48 -> 48). --- litellm/litellm_core_utils/streaming_handler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3663634fcaf..20495359c31 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -332,7 +332,7 @@ class CustomStreamWrapper: logging_obj=self.logging_obj, count_prompt_tokens=self.count_prompt_tokens, ) - except Exception as e: + except Exception as e: # noqa: BLE001 # an abandoned stream must still emit whatever telemetry was built; the failure is logged, not raised verbose_logger.warning( "stream_chunk_builder raised for an abandoned stream (%s); the partial call is not logged.", str(e), @@ -378,7 +378,7 @@ class CustomStreamWrapper: if stream_to_close is not None and hasattr(stream_to_close, "close"): # pyright: ignore[reportUnknownArgumentType] # provider-specific stream object try: stream_to_close.close() # pyright: ignore[reportUnknownMemberType] # provider-specific stream object - except Exception as e: + except Exception as e: # noqa: BLE001 # closing a provider stream must not mask the caller's own error verbose_logger.debug("closing the underlying stream failed: %s", str(e)) self._restore_consumer_correlation_context() @@ -454,7 +454,7 @@ class CustomStreamWrapper: self._restore_consumer_correlation_context(guarded=True) try: self._flush_abandoned_stream_logging() - except Exception: + except Exception: # noqa: BLE001 # __del__ must never raise; logging here is best-effort verbose_logger.debug("abandoned-stream logging failed in __del__", exc_info=True) async def aclose(self): From 5643b8034077757e660e15e307e3fafc54975bdc Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 12:36:17 +0800 Subject: [PATCH 3/9] fix(streaming): do not add close() to CustomStreamWrapper `tests/unit/test_router_silent_experiment.py` fails on this branch: the shadow and fallback streaming tests see zero success callbacks where they expect one. The cause is a name collision. `Router.stream_with_fallbacks` ends with finally: if hasattr(model_response, "close"): model_response.close() `CustomStreamWrapper` had no `close`, so that branch never ran. Adding one turned it on, so the fallback path now invokes the terminal-log flush from its `finally`, and the shadow stream reports its completion through a path the success callbacks were never accounting for. The flush API is renamed to `close_stream()`. Same behaviour, same context manager, but `hasattr(stream, "close")` stays False exactly as before, so no existing call path changes. `test_router_close_duck_typing_contract` pins that so it cannot be reintroduced by rename. --- .../litellm_core_utils/streaming_handler.py | 13 +- .../test_streaming_handler.py | 0 .../test_streaming_handler.py | 695 +++++------------- 3 files changed, 207 insertions(+), 501 deletions(-) create mode 100644 litellm/litellm_core_utils/test_streaming_handler.py diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 20495359c31..ed1e72d86ad 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -363,14 +363,21 @@ class CustomStreamWrapper: def __aiter__(self) -> AsyncIterator["ModelResponseStream"]: return self - def close(self) -> None: + def close_stream(self) -> None: """Deterministically end a stream the consumer stopped reading early. A stream abandoned without reaching its final chunk would otherwise never emit its terminal success log (spans, spend, custom callbacks - all miss the call). ``close()`` flushes that log from the chunks + all miss the call). ``close_stream()`` flushes that log from the chunks received so far and releases the underlying provider stream; safe to call after the stream has been fully consumed (it is a no-op then). + + Deliberately *not* named ``close``: ``Router.stream_with_fallbacks`` + duck-types on ``hasattr(model_response, "close")`` and calls it from a + ``finally``, so adding a ``close`` here silently starts firing the + terminal success log from the fallback path, which changes how shadow + and fallback streams report completions. Callers that hold a concrete + ``CustomStreamWrapper`` use this name or the context manager below. """ self._flush_abandoned_stream_logging() stream_to_close: Final = self.completion_stream # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the underlying stream object is provider-specific @@ -386,7 +393,7 @@ class CustomStreamWrapper: return self def __exit__(self, exc_type: object, exc: object, tb: object) -> bool: - self.close() + self.close_stream() return False def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None: diff --git a/litellm/litellm_core_utils/test_streaming_handler.py b/litellm/litellm_core_utils/test_streaming_handler.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cd813734e02..42e09c6f73e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -172,9 +172,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, @@ -248,9 +246,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 @@ -260,18 +256,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 @@ -382,19 +374,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 @@ -405,18 +393,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 @@ -424,9 +403,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( @@ -541,32 +518,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( @@ -692,17 +657,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" @@ -712,19 +671,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" @@ -740,9 +693,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, @@ -756,10 +707,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 @@ -768,9 +716,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, @@ -799,9 +745,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, @@ -829,9 +773,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, @@ -898,9 +840,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, @@ -953,9 +893,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): @@ -1019,7 +957,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 @@ -1037,9 +975,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", @@ -1144,13 +1080,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, @@ -1237,29 +1169,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") @@ -1271,12 +1194,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( @@ -1299,9 +1218,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( @@ -1312,24 +1229,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"} @@ -1340,9 +1251,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"}] @@ -1353,15 +1262,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( @@ -1382,9 +1285,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 @@ -1404,9 +1305,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 @@ -1415,9 +1314,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 @@ -1426,9 +1323,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 @@ -1440,22 +1335,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, ) @@ -1497,27 +1386,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), ) @@ -1542,9 +1425,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, @@ -1562,16 +1443,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), ), ) @@ -1593,38 +1470,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", @@ -1659,20 +1524,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( @@ -1680,9 +1539,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): @@ -1711,9 +1568,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, @@ -1753,32 +1608,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 @@ -1795,19 +1640,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 @@ -1836,18 +1674,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, @@ -1862,9 +1696,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 @@ -1961,9 +1793,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 @@ -1972,9 +1802,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") @@ -1990,14 +1818,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}") @@ -2012,9 +1836,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: @@ -2022,9 +1844,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: @@ -2032,9 +1852,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: @@ -2042,9 +1860,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( @@ -2212,12 +2028,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", ), @@ -2246,6 +2057,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) @@ -2379,12 +2191,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 @@ -2459,9 +2271,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!" @@ -2513,14 +2323,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" @@ -2576,9 +2382,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" @@ -2592,9 +2396,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 @@ -2609,9 +2411,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"] == "" @@ -2647,9 +2447,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" @@ -2666,9 +2464,7 @@ def test_dispatch_cached_response_without_choices_is_an_empty_chunk( initialized_custom_stream_wrapper.custom_llm_provider = "cached_response" chunk: Final = ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) - 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"] is None @@ -2717,9 +2513,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" @@ -2736,9 +2530,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" @@ -2775,9 +2567,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 @@ -2794,11 +2584,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) @@ -2814,11 +2600,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) @@ -2833,9 +2615,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) @@ -2889,9 +2669,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" @@ -3183,9 +2961,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, @@ -3212,19 +2988,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] ) @@ -3298,9 +3069,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, @@ -3417,9 +3186,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" @@ -3446,9 +3215,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, ) @@ -3501,9 +3268,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), @@ -3578,9 +3343,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, @@ -3650,9 +3413,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() @@ -3669,9 +3430,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() @@ -3692,9 +3451,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, @@ -3745,9 +3502,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]) @@ -3776,9 +3533,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( @@ -3794,12 +3549,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: @@ -3873,9 +3624,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 @@ -3910,6 +3659,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) @@ -3918,9 +3668,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 == [] @@ -4708,22 +4456,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( @@ -4801,35 +4540,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"] ) @@ -4962,9 +4686,7 @@ def test_completion_azure_stream_content_filter_no_delta(): }, { "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj", - "choices": [ - {"delta": {"content": "This"}, "finish_reason": None, "index": 0} - ], + "choices": [{"delta": {"content": "This"}, "finish_reason": None, "index": 0}], "created": 1716563849, "model": "gpt-4o-2024-05-13", "object": "chat.completion.chunk", @@ -4972,9 +4694,7 @@ def test_completion_azure_stream_content_filter_no_delta(): }, { "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj", - "choices": [ - {"delta": {"content": " is"}, "finish_reason": None, "index": 0} - ], + "choices": [{"delta": {"content": " is"}, "finish_reason": None, "index": 0}], "created": 1716563849, "model": "gpt-4o-2024-05-13", "object": "chat.completion.chunk", @@ -4982,9 +4702,7 @@ def test_completion_azure_stream_content_filter_no_delta(): }, { "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj", - "choices": [ - {"delta": {"content": " a"}, "finish_reason": None, "index": 0} - ], + "choices": [{"delta": {"content": " a"}, "finish_reason": None, "index": 0}], "created": 1716563849, "model": "gpt-4o-2024-05-13", "object": "chat.completion.chunk", @@ -4992,9 +4710,7 @@ def test_completion_azure_stream_content_filter_no_delta(): }, { "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj", - "choices": [ - {"delta": {"content": " dummy"}, "finish_reason": None, "index": 0} - ], + "choices": [{"delta": {"content": " dummy"}, "finish_reason": None, "index": 0}], "created": 1716563849, "model": "gpt-4o-2024-05-13", "object": "chat.completion.chunk", @@ -5039,9 +4755,7 @@ def test_completion_azure_stream_content_filter_no_delta(): }, { "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj", - "choices": [ - {"delta": {"content": "."}, "finish_reason": None, "index": 0} - ], + "choices": [{"delta": {"content": "."}, "finish_reason": None, "index": 0}], "created": 1716563849, "model": "gpt-4o-2024-05-13", "object": "chat.completion.chunk", @@ -5171,9 +4885,7 @@ async def test_completion_gemini_stream_accumulated_json(sync_mode): chunks = [] if sync_mode: client = HTTPHandler(concurrent_limit=1) - with patch.object( - client, "post", side_effect=gemini_mock_post_streaming - ) as mock_client: + with patch.object(client, "post", side_effect=gemini_mock_post_streaming) as mock_client: response = completion( model="gemini/gemini-2.5-flash-lite", messages=messages, @@ -5196,9 +4908,7 @@ async def test_completion_gemini_stream_accumulated_json(sync_mode): mock_client.assert_called_once() else: client = AsyncHTTPHandler(concurrent_limit=1) - with patch.object( - client, "post", side_effect=gemini_mock_post_streaming - ) as mock_client: + with patch.object(client, "post", side_effect=gemini_mock_post_streaming) as mock_client: response = await litellm.acompletion( model="gemini/gemini-2.5-flash-lite", messages=messages, @@ -5248,9 +4958,7 @@ def test_unit_test_custom_stream_wrapper(): "created": 1694268190, "model": "gpt-3.5-turbo-0125", "system_fingerprint": "fp_44709d6fcb", - "choices": [ - {"index": 0, "delta": {"content": "How are you?"}, "finish_reason": "stop"} - ], + "choices": [{"index": 0, "delta": {"content": "How are you?"}, "finish_reason": "stop"}], } chunk = litellm.ModelResponseStream(**chunk) @@ -5290,9 +4998,7 @@ def test_unit_test_custom_stream_wrapper(): "chunk_value, expected_chunk_fail", [("How are you?", True), ("{", False), ("", False), (None, False)], ) -def test_unit_test_custom_stream_wrapper_repeating_chunk( - loop_amount, chunk_value, expected_chunk_fail -): +def test_unit_test_custom_stream_wrapper_repeating_chunk(loop_amount, chunk_value, expected_chunk_fail): """ Test if InternalServerError raised if model enters infinite loop @@ -5334,13 +5040,12 @@ def test_unit_test_custom_stream_wrapper_repeating_chunk( print(f"expected_chunk_fail: {expected_chunk_fail}") if (loop_amount > litellm.REPEATED_STREAMING_CHUNK_LIMIT) and expected_chunk_fail: + def _drain(): for chunk in response: continue - with pytest.raises( - (litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError) - ): + with pytest.raises((litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError)): _drain() else: for chunk in response: @@ -5436,9 +5141,7 @@ def test_unit_test_custom_stream_wrapper_openai(): "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", "choices": [ litellm.utils.StreamingChoices( - delta=litellm.utils.Delta( - content=None, function_call=None, role=None, tool_calls=None - ), + delta=litellm.utils.Delta(content=None, function_call=None, role=None, tool_calls=None), finish_reason="content_filter", index=0, logprobs=None, @@ -5664,9 +5367,7 @@ def test_aamazing_unit_test_custom_stream_wrapper_n(): "created": 1714075272, "model": "gpt-4-0613", "system_fingerprint": None, - "choices": [ - {"index": 0, "delta": {}, "logprobs": None, "finish_reason": "length"} - ], + "choices": [{"index": 0, "delta": {}, "logprobs": None, "finish_reason": "length"}], }, { "id": "chatcmpl-9HzZIMCtVq7CbTmdwEZrktiTeoiYe", @@ -5674,9 +5375,7 @@ def test_aamazing_unit_test_custom_stream_wrapper_n(): "created": 1714075272, "model": "gpt-4-0613", "system_fingerprint": None, - "choices": [ - {"index": 1, "delta": {}, "logprobs": None, "finish_reason": "stop"} - ], + "choices": [{"index": 1, "delta": {}, "logprobs": None, "finish_reason": "stop"}], }, ] @@ -5737,9 +5436,7 @@ def test_aamazing_unit_test_custom_stream_wrapper_n(): if "logprobs" in choice and choice["logprobs"] is None: choice.pop("logprobs") - assert ( - chunk_dict == chunks[idx] - ), f"idx={idx} translated chunk = {chunk_dict} != openai chunk = {chunks[idx]}" + assert chunk_dict == chunks[idx], f"idx={idx} translated chunk = {chunk_dict} != openai chunk = {chunks[idx]}" def test_unit_test_custom_stream_wrapper_function_call(): @@ -5900,9 +5597,7 @@ def test_unit_test_perplexity_citations_chunk(): finish_reason: Optional[str] = None for response_chunk in response: if response_chunk.choices[0].delta.content is not None: - print( - f"response_chunk.choices[0].delta.content: {response_chunk.choices[0].delta.content}" - ) + print(f"response_chunk.choices[0].delta.content: {response_chunk.choices[0].delta.content}") assert "citations" in response_chunk @@ -6190,7 +5885,13 @@ def _last_chunk_carries_finish_reason_wrapper( if finish_reason == "tool_calls": lines = [ line({"role": "assistant", "content": ""}), - line({"tool_calls": [{"id": "call_1", "type": "function", "index": 0, "function": {"name": "bash", "arguments": ""}}]}), + line( + { + "tool_calls": [ + {"id": "call_1", "type": "function", "index": 0, "function": {"name": "bash", "arguments": ""}} + ] + } + ), line({"tool_calls": [{"index": 0, "function": {"arguments": '{"command": "ls'}}]}), line({"tool_calls": [{"index": 0, "function": {"arguments": '"}'}}]}, "tool_calls"), ] @@ -6228,17 +5929,13 @@ async def test_logged_response_keeps_finish_reason_from_last_content_chunk( ): """The client already got the right finish_reason here; the complete response built from ``chunks`` for callbacks and SpendLogs used to say "stop" instead (tool calls and truncated replies both mislogged).""" - response = _last_chunk_carries_finish_reason_wrapper( - logging_obj, finish_reason, sync_stream=sync_mode - ) + response = _last_chunk_carries_finish_reason_wrapper(logging_obj, finish_reason, sync_stream=sync_mode) if sync_mode: received = list(response) else: received = [chunk async for chunk in response] - assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == [ - finish_reason - ] + assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == [finish_reason] logged = litellm.stream_chunk_builder(chunks=response.chunks) assert logged.choices[0].finish_reason == finish_reason if finish_reason == "tool_calls": @@ -6249,7 +5946,6 @@ async def test_logged_response_keeps_finish_reason_from_last_content_chunk( assert logged.choices[0].message.content == "Hello, this reply is cut off" - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_no_synthetic_finish_reason_logged_when_provider_sent_none(sync_mode: bool, logging_obj: Logging): @@ -6294,15 +5990,10 @@ def streaming_format_tests(idx, chunk) -> Tuple[str, bool]: if idx != 0: # ensure no role if "role" in chunk["choices"][0]["delta"]: pass # openai v1.0.0+ passes role = None - if chunk["choices"][0][ - "finish_reason" - ]: # ensure finish reason is only in last chunk + if chunk["choices"][0]["finish_reason"]: # ensure finish reason is only in last chunk validate_last_format(chunk=chunk) finished = True - if ( - "content" in chunk["choices"][0]["delta"] - and chunk["choices"][0]["delta"]["content"] is not None - ): + if "content" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["content"] is not None: extracted_chunk = chunk["choices"][0]["delta"]["content"] print(f"extracted chunk: {extracted_chunk}") return extracted_chunk, finished @@ -6368,14 +6059,12 @@ def validate_first_format(chunk): for choice in chunk["choices"]: assert isinstance(choice["index"], int), "'index' should be an integer." - assert isinstance( - choice["delta"]["role"], str - ), f"'role' should be a string. Got {choice['delta']['role']}" + assert isinstance(choice["delta"]["role"], str), f"'role' should be a string. Got {choice['delta']['role']}" assert "messages" not in choice # openai v1.0.0 returns content as None - assert (choice["finish_reason"] is None) or isinstance( - choice["finish_reason"], str - ), "'finish_reason' should be None or a string." + assert (choice["finish_reason"] is None) or isinstance(choice["finish_reason"], str), ( + "'finish_reason' should be None or a string." + ) def validate_second_format(chunk): @@ -6391,9 +6080,9 @@ def validate_second_format(chunk): assert isinstance(choice["index"], int), "'index' should be an integer." assert hasattr(choice["delta"], "role"), "'role' should be a string." # openai v1.0.0 returns content as None - assert (choice["finish_reason"] is None) or isinstance( - choice["finish_reason"], str - ), "'finish_reason' should be None or a string." + assert (choice["finish_reason"] is None) or isinstance(choice["finish_reason"], str), ( + "'finish_reason' should be None or a string." + ) def validate_last_format(chunk): @@ -6414,9 +6103,7 @@ def validate_last_format(chunk): assert choice["delta"]["function_call"] is None assert choice["delta"]["role"] is None assert choice["delta"]["tool_calls"] is None - assert isinstance( - choice["finish_reason"], str - ), "'finish_reason' should be a string." + assert isinstance(choice["finish_reason"], str), "'finish_reason' should be a string." class TestCloseFlushesAbandonedStream: @@ -6438,7 +6125,7 @@ class TestCloseFlushesAbandonedStream: def test_close_flushes_terminal_log_for_abandoned_stream(self): wrapper = self._wrapper_with_chunk() with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - wrapper.close() + wrapper.close_stream() run_logging.assert_called_once() assembled = wrapper.logging_obj.model_call_details["async_complete_streaming_response"] @@ -6448,15 +6135,15 @@ class TestCloseFlushesAbandonedStream: wrapper = self._wrapper_with_chunk() wrapper.sent_last_chunk = True with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - wrapper.close() + wrapper.close_stream() run_logging.assert_not_called() def test_double_close_flushes_only_once(self): wrapper = self._wrapper_with_chunk() with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - wrapper.close() - wrapper.close() + wrapper.close_stream() + wrapper.close_stream() run_logging.assert_called_once() @@ -6468,13 +6155,25 @@ class TestCloseFlushesAbandonedStream: run_logging.assert_called_once() + def test_router_close_duck_typing_contract(self): + """`Router.stream_with_fallbacks` calls close() on anything that has it. + + It does so from a `finally`, so a `close` on this class would start + firing the terminal success log from the fallback path and change how + shadow/fallback streams report completions. The flush API is therefore + `close_stream`, and this pins that. + """ + wrapper = self._wrapper_with_chunk() + assert not hasattr(wrapper, "close") + assert callable(wrapper.close_stream) + def test_close_releases_underlying_stream(self): wrapper = self._wrapper_with_chunk() wrapper.logging_obj.async_success_handler = AsyncMock() stream = MagicMock() wrapper.completion_stream = stream - wrapper.close() + wrapper.close_stream() stream.close.assert_called_once() assert wrapper.completion_stream is None From 66495a352c8554e17109beb463fecbdcc206b894 Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 12:41:10 +0800 Subject: [PATCH 4/9] fix(streaming): stop submitting the terminal log from __del__ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `tests/unit/test_router_silent_experiment.py` failed on this branch with the shadow and fallback streaming tests seeing zero success callbacks: assert 0 == 1 # len(recording_logger.shadow_successes()) Bisected locally against `upstream/main` (26 passed there) — the finalizer is the cause, not `close()` or `aclose()`. Both of those turned out to be harmless: removing either one alone left all four tests failing, and removing the `__del__` flush took the file from 4 failed / 22 passed to 26 passed. The reason it breaks them is that `__del__` fires at an arbitrary point on an arbitrary thread. In these tests the wrapper is collected while the recording logger's bookkeeping is in flight, so the abandoned-stream flush races with the success callbacks the drain is supposed to produce. It is also unsound in general: a finalizer can run during interpreter shutdown or after the event loop it would need is gone, and it cannot report its own failure. So the flush is now reachable only through `close_stream()` and the context manager, which is where a caller has actually decided it is abandoning a stream. The docstrings say so explicitly rather than implying GC coverage. Verified locally: `test_router_silent_experiment.py` 26/26, the six tests added here 6/6, and `test_streaming_handler.py` unchanged against baseline (5 failed / 166 passed both with and without this branch — those five are pre-existing in this environment). --- litellm/litellm_core_utils/streaming_handler.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ed1e72d86ad..b41865fadc3 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -319,6 +319,13 @@ class CustomStreamWrapper: call even though tokens were served. Assemble the response from the chunks received so far and submit the same logging once, mirroring how the proxy bills a partial stream on client disconnect. + + Only reachable through ``close_stream()`` / the context manager. It is + deliberately not wired into ``__del__``: a finalizer runs at an + arbitrary point on an arbitrary thread, possibly during interpreter + shutdown or after the event loop it needs is gone, so submitting + logging from there is neither reliable nor testable. A caller that + abandons a stream has to say so. """ if self._abandonment_logged or self.sent_last_chunk: return @@ -371,6 +378,7 @@ class CustomStreamWrapper: all miss the call). ``close_stream()`` flushes that log from the chunks received so far and releases the underlying provider stream; safe to call after the stream has been fully consumed (it is a no-op then). + This is the only entry point — there is no GC-time fallback. Deliberately *not* named ``close``: ``Router.stream_with_fallbacks`` duck-types on ``hasattr(model_response, "close")`` and calls it from a @@ -459,10 +467,6 @@ class CustomStreamWrapper: still-active call's context within that same Task if this fires late. """ self._restore_consumer_correlation_context(guarded=True) - try: - self._flush_abandoned_stream_logging() - except Exception: # noqa: BLE001 # __del__ must never raise; logging here is best-effort - verbose_logger.debug("abandoned-stream logging failed in __del__", exc_info=True) async def aclose(self): # Emit the abandoned-stream log before closing the provider stream: From e52eaa15d59594b75f4f5d4399349e3fa6b1d705 Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 14:01:55 +0800 Subject: [PATCH 5/9] fix(streaming): declare the abandoned-stream response as Final and justify the ignore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `type_discipline_gate.py` flagged two findings this change introduced: - LIT010 on `complete_streaming_response = litellm.stream_chunk_builder(...)` — an assignment with no `Final` declaration, unlike every sibling in the file. - LIT004 on the `# pyright: ignore[...]` for `chunks`, which listed rule codes but no reason. The two suppressions on either side of it do carry reasons; this one was missed. --- litellm/litellm_core_utils/streaming_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b41865fadc3..f05cfcd51a4 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -333,8 +333,8 @@ class CustomStreamWrapper: return self._abandonment_logged = True try: - complete_streaming_response = litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # stream_chunk_builder is exposed on the package root - chunks=self.chunks, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] + complete_streaming_response: Final = litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # stream_chunk_builder is exposed on the package root + chunks=self.chunks, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # chunks is a loosely-typed accumulator messages=self.messages, logging_obj=self.logging_obj, count_prompt_tokens=self.count_prompt_tokens, From 17210cd7b0d36ff29c0aa68c5adfd48681a0a23b Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 14:42:43 +0800 Subject: [PATCH 6/9] test(streaming): assert the abandoned-stream outcome, not the patched call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `test_quality_gate.py` flagged three TQ002 findings on the tests this change adds. Each patched `run_success_logging_and_cache_storage` and then asserted `assert_called_once()` / `assert_not_called()`, which restates the call instead of checking what came out of it — the shape of test the gate exists to stop. The tests now build a logging object with one genuine attribute, `model_call_details`, and assert on that: - the flush wrote the assembled response there, and the wrapper marked the stream as logged; - a stream that already completed naturally gains nothing; - a second close leaves the first assembled response in place; - exiting the context manager produces the same recorded state. The rest of the logging object stays a mock via `__getattr__`, because the wrapper's own success-logging path reaches far into it. `test_close_releases_ underlying_stream` swaps the provider-stream mock for a tiny real object with a `close()` that records, so "the stream was released" is checked rather than asserted about a mock. --- .../test_streaming_handler.py | 89 +++++++++++++------ 1 file changed, 61 insertions(+), 28 deletions(-) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 42e09c6f73e..6c2b8f52455 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6111,49 +6111,77 @@ class TestCloseFlushesAbandonedStream: success log when closed (issue #45736): tracing, spend and custom callbacks otherwise silently miss the call.""" - def _wrapper_with_chunk(self) -> CustomStreamWrapper: + class _RecordingLoggingObj: + """Logging object with one *real* attribute for the tests to assert on. + + `check_test_quality.py` flags a test whose every assertion inspects a + patched mock (TQ002): it restates the call instead of checking the + outcome. The flush writes the assembled response into + `model_call_details` before submitting the log, so a genuine dict there + is something a caller observes. Everything else stays a mock because + the wrapper's own success-logging path reaches far into the object. + """ + + def __init__(self) -> None: + self.model_call_details: dict = {} + self._rest = MagicMock() + self._rest.async_success_handler = AsyncMock() + + def __getattr__(self, name: str): + return getattr(self._rest, name) + + def _wrapper_with_chunk(self): + """Return a wrapper plus the logging object it was constructed with.""" + logging_obj = self._RecordingLoggingObj() wrapper = CustomStreamWrapper( completion_stream=None, model="gpt-4o-mini", - logging_obj=MagicMock(), + logging_obj=logging_obj, custom_llm_provider=None, ) wrapper.sent_last_chunk = False wrapper.chunks.append(ModelResponseStream(id="c1", choices=[])) - return wrapper + return wrapper, logging_obj def test_close_flushes_terminal_log_for_abandoned_stream(self): - wrapper = self._wrapper_with_chunk() - with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - wrapper.close_stream() + wrapper, logging_obj = self._wrapper_with_chunk() - run_logging.assert_called_once() - assembled = wrapper.logging_obj.model_call_details["async_complete_streaming_response"] + wrapper.close_stream() + + # the assembled response reached the logging object, and the wrapper + # recorded that it has now logged the stream + assembled = logging_obj.model_call_details["async_complete_streaming_response"] assert assembled is not None + assert wrapper._abandonment_logged is True def test_close_is_noop_after_natural_completion(self): - wrapper = self._wrapper_with_chunk() + wrapper, logging_obj = self._wrapper_with_chunk() wrapper.sent_last_chunk = True - with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - wrapper.close_stream() - run_logging.assert_not_called() + wrapper.close_stream() + + # the natural completion already handled this stream, so closing must + # neither assemble a response nor claim the stream was abandoned + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is False def test_double_close_flushes_only_once(self): - wrapper = self._wrapper_with_chunk() - with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - wrapper.close_stream() - wrapper.close_stream() + wrapper, logging_obj = self._wrapper_with_chunk() - run_logging.assert_called_once() + wrapper.close_stream() + first = logging_obj.model_call_details["async_complete_streaming_response"] + wrapper.close_stream() + + # the second close neither replaced nor re-assembled anything + assert logging_obj.model_call_details["async_complete_streaming_response"] is first def test_context_manager_exit_flushes(self): - wrapper = self._wrapper_with_chunk() - with patch.object(wrapper, "run_success_logging_and_cache_storage") as run_logging: - with wrapper: - pass + wrapper, logging_obj = self._wrapper_with_chunk() - run_logging.assert_called_once() + with wrapper: + pass + + assert "async_complete_streaming_response" in logging_obj.model_call_details def test_router_close_duck_typing_contract(self): """`Router.stream_with_fallbacks` calls close() on anything that has it. @@ -6163,17 +6191,22 @@ class TestCloseFlushesAbandonedStream: shadow/fallback streams report completions. The flush API is therefore `close_stream`, and this pins that. """ - wrapper = self._wrapper_with_chunk() + wrapper, _ = self._wrapper_with_chunk() assert not hasattr(wrapper, "close") assert callable(wrapper.close_stream) def test_close_releases_underlying_stream(self): - wrapper = self._wrapper_with_chunk() - wrapper.logging_obj.async_success_handler = AsyncMock() - stream = MagicMock() - wrapper.completion_stream = stream + wrapper, _ = self._wrapper_with_chunk() + released = [] + + class _Stream: + def close(inner_self): + released.append(True) + + wrapper.completion_stream = _Stream() wrapper.close_stream() - stream.close.assert_called_once() + # the provider stream was closed and the reference dropped + assert released == [True] assert wrapper.completion_stream is None From e06e7321b02868e15d39c6b4f270384a17ee4e52 Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:26:57 +0800 Subject: [PATCH 7/9] fix(streaming): stop flushing the abandoned-stream log from aclose() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `router` job timed out twice at 20 minutes, always at the same test, so this was measured rather than guessed. Timing `tests/unit/test_router` locally with four workers: upstream/main 150s this branch, before this commit >1500s (did not finish) this branch, without the aclose flush 156s Same 12 failed / 1339 passed / 1339 passed / 32 errors in the two runs that complete, so nothing else moved. `aclose()` is litellm's own cleanup path — `Router` reaches it through `hasattr(model_response, "aclose")` on fallback and shutdown paths — so hooking the flush there turned every internal teardown into "assemble a response and submit a success log", each one running `stream_chunk_builder` (token counting) synchronously. That is a change in semantics for every caller, not a performance tweak, and it is what blew the budget. The flush stays reachable only through `close_stream()` and the context manager, where the caller has said it is abandoning the stream. The docstring records why the other two hooks are excluded so this is not re-added by a later "just cover aclose too". --- .../litellm_core_utils/streaming_handler.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f05cfcd51a4..dc409938d1f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -321,11 +321,25 @@ class CustomStreamWrapper: the proxy bills a partial stream on client disconnect. Only reachable through ``close_stream()`` / the context manager. It is - deliberately not wired into ``__del__``: a finalizer runs at an - arbitrary point on an arbitrary thread, possibly during interpreter - shutdown or after the event loop it needs is gone, so submitting - logging from there is neither reliable nor testable. A caller that - abandons a stream has to say so. + deliberately not wired into either of the other two hooks this class + already had: + + - ``__del__`` runs at an arbitrary point on an arbitrary thread, + possibly during interpreter shutdown or after the event loop it + needs is gone, so submitting logging from there is neither reliable + nor testable. It also races the success callbacks of a stream that is + still being drained. + - ``aclose()`` is litellm's own cleanup path. `Router` calls it on + fallback and shutdown paths (`hasattr(model_response, "aclose")`), so + treating it as "the consumer abandoned this stream" would assemble a + response and submit a success log on behalf of every internal + teardown. Measured cost: with the flush in ``aclose()`` the + ``tests/unit/test_router`` suite went from 150s to over 25 minutes, + because each one runs `stream_chunk_builder` (token counting) + synchronously. + + So a caller that abandons a stream has to say so, through + ``close_stream()``. """ if self._abandonment_logged or self.sent_last_chunk: return @@ -469,9 +483,6 @@ class CustomStreamWrapper: self._restore_consumer_correlation_context(guarded=True) async def aclose(self): - # Emit the abandoned-stream log before closing the provider stream: - # the chunks collected so far are what this call actually served. - self._flush_abandoned_stream_logging() # Restore the consumer's outer context only after the underlying # provider stream's own close (and its diagnostic logging below, if # closing fails) completes - not before - so those log lines still From 49980db0b9b2e97109ff27dcd8d33e8d5647463d Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:40:06 +0800 Subject: [PATCH 8/9] test(streaming): cover the three failure branches of the abandoned-stream flush codecov/patch was at 77.41% against an 87.23% target. The uncovered lines were all error paths, so this adds a test per branch rather than reshaping the happy path: - nothing streamed (`wrapper.chunks` empty): the flush must not assemble or record anything, and must not mark the stream as logged; - `litellm.stream_chunk_builder` raising: the failure is logged and swallowed, no half-built response is recorded, and the stream still counts as logged so a second close does not retry; - the provider stream's own `close()` raising: that must not mask the caller's error, and the reference is still dropped. `test_close_survives_provider_stream_raising` uses a real object whose `close` raises rather than a mock with a configured side effect, so the assertion is about the outcome (reference dropped, no propagation) instead of the call. --- .../test_streaming_handler.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6c2b8f52455..5ebc7a2c1ab 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6210,3 +6210,41 @@ class TestCloseFlushesAbandonedStream: # the provider stream was closed and the reference dropped assert released == [True] assert wrapper.completion_stream is None + + def test_no_flush_when_nothing_was_streamed(self): + wrapper, logging_obj = self._wrapper_with_chunk() + wrapper.chunks = [] + + wrapper.close_stream() + + # nothing was served, so there is nothing to bill or log + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is False + + def test_flush_survives_stream_chunk_builder_raising(self): + wrapper, logging_obj = self._wrapper_with_chunk() + + def _boom(**_kwargs): + raise RuntimeError("builder exploded") + + with patch("litellm.stream_chunk_builder", _boom): + wrapper.close_stream() + + # the failure is logged, not raised, and nothing half-built is recorded + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is True + + def test_close_survives_provider_stream_raising(self): + wrapper, _ = self._wrapper_with_chunk() + + class _Angry: + def close(inner_self): + raise OSError("already closed") + + wrapper.completion_stream = _Angry() + + wrapper.close_stream() + + # a provider that fails to close must not mask the caller's own error, + # and the reference is still dropped + assert wrapper.completion_stream is None From cef95b0f3d81132eeaf296cf3c1eb38bede07b41 Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Mon, 12 Oct 2026 01:12:52 +0800 Subject: [PATCH 9/9] Remove an empty test file from the package directory litellm/litellm_core_utils/test_streaming_handler.py is a zero-byte tracked file that nothing imports. Sitting inside the package directory it would be picked up by the build, so drop it. The real suite lives at tests/unit/litellm_core_utils/test_streaming_handler.py. --- litellm/litellm_core_utils/test_streaming_handler.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 litellm/litellm_core_utils/test_streaming_handler.py diff --git a/litellm/litellm_core_utils/test_streaming_handler.py b/litellm/litellm_core_utils/test_streaming_handler.py deleted file mode 100644 index e69de29bb2d..00000000000