From e5ba4d3227a37661efb2bf1e18f504af279e29df Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 20 Feb 2026 16:03:34 -0800 Subject: [PATCH] perf: reduce responses streaming CPU for text-only streams --- .../streaming_chunk_builder_utils.py | 27 +++++- litellm/main.py | 74 ++++++++++++++- .../streaming_iterator.py | 32 ++++++- .../test_anthropic_responses_api.py | 2 +- .../test_streaming_chunk_builder_utils.py | 91 +++++++++++++++++++ 5 files changed, 214 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 76c7246b87e..143d87ebf34 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -41,10 +41,29 @@ class ChunkProcessor: def _sort_chunks(self, chunks: list) -> list: if not chunks: return [] - if chunks[0]._hidden_params.get("created_at"): - return sorted( - chunks, key=lambda x: x._hidden_params.get("created_at", float("inf")) - ) + + first_chunk = chunks[0] + first_hidden_params: Dict[str, Any] = {} + if isinstance(first_chunk, dict): + candidate = first_chunk.get("_hidden_params", {}) + if isinstance(candidate, dict): + first_hidden_params = candidate + else: + candidate = getattr(first_chunk, "_hidden_params", {}) + if isinstance(candidate, dict): + first_hidden_params = candidate + + if first_hidden_params.get("created_at"): + def _created_at(chunk: Any) -> Union[int, float]: + if isinstance(chunk, dict): + params = chunk.get("_hidden_params", {}) + else: + params = getattr(chunk, "_hidden_params", {}) + if isinstance(params, dict): + return cast(Union[int, float], params.get("created_at", float("inf"))) + return float("inf") + + return sorted(chunks, key=_created_at) return chunks def update_model_response_with_hidden_params( diff --git a/litellm/main.py b/litellm/main.py index 80a2f74c571..356ca7ecf13 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2506,10 +2506,10 @@ def completion( # type: ignore # noqa: PLR0915 # Add GitHub Copilot headers (same as /responses endpoint does) if custom_llm_provider == "github_copilot": + from litellm.llms.github_copilot.authenticator import Authenticator from litellm.llms.github_copilot.common_utils import ( get_copilot_default_headers, ) - from litellm.llms.github_copilot.authenticator import Authenticator copilot_auth = Authenticator() copilot_api_key = copilot_auth.get_api_key() @@ -7230,6 +7230,71 @@ def stream_chunk_builder( # noqa: PLR0915 # Initialize the response dictionary response = processor.build_base_response(chunks) + # Fast path for the common text-only streaming case: + # avoid repeated multi-pass list scans over chunks. + simple_content_parts: List[str] = [] + is_simple_text_stream = True + for chunk in chunks: + if len(chunk["choices"]) == 0: + continue + + choice = chunk["choices"][0] + delta_obj = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) + if isinstance(delta_obj, dict): + delta = delta_obj + elif hasattr(delta_obj, "model_dump"): + delta = cast(Dict[str, Any], delta_obj.model_dump()) + else: + delta = {} + + if ( + delta.get("tool_calls") is not None + or delta.get("function_call") is not None + or delta.get("reasoning_content") is not None + or delta.get("thinking_blocks") is not None + or delta.get("annotations") is not None + or delta.get("audio") is not None + or delta.get("images") is not None + or delta.get("provider_specific_fields") is not None + ): + is_simple_text_stream = False + break + + content = delta.get("content") + if isinstance(content, str) and content: + simple_content_parts.append(content) + + if is_simple_text_stream: + if simple_content_parts: + response["choices"][0]["message"]["content"] = "".join(simple_content_parts) + completion_output = get_content_from_model_response(response) + usage = processor.calculate_usage( + chunks=chunks, + model=model, + completion_output=completion_output, + messages=messages, + reasoning_tokens=0, + ) + setattr(response, "usage", usage) + + # Propagate provider_specific_fields from chunk hidden params when present. + for chunk in reversed(chunks): + if isinstance(chunk, dict): + hidden = chunk.get("_hidden_params") + else: + hidden = getattr(chunk, "_hidden_params", None) + if isinstance(hidden, dict) and "provider_specific_fields" in hidden: + response._hidden_params.setdefault( + "provider_specific_fields", {} + ).update(hidden["provider_specific_fields"]) + break + + if litellm.include_cost_in_streaming_usage and logging_obj is not None: + setattr( + usage, "cost", logging_obj._response_cost_calculator(result=response) + ) + return response + tool_call_chunks = [ chunk for chunk in chunks @@ -7386,8 +7451,11 @@ def stream_chunk_builder( # noqa: PLR0915 # Propagate provider_specific_fields from the last chunk (contains provider # metadata like traffic_type set during streaming) for chunk in reversed(chunks): - hidden = getattr(chunk, "_hidden_params", None) - if hidden and "provider_specific_fields" in hidden: + if isinstance(chunk, dict): + hidden = chunk.get("_hidden_params") + else: + hidden = getattr(chunk, "_hidden_params", None) + if isinstance(hidden, dict) and "provider_specific_fields" in hidden: response._hidden_params.setdefault( "provider_specific_fields", {} ).update(hidden["provider_specific_fields"]) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 817575f30d3..6e32a0d48d7 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1,6 +1,6 @@ import time import uuid -from typing import List, Optional, Union, cast +from typing import Any, Dict, List, Optional, Union, cast import litellm from litellm.main import stream_chunk_builder @@ -68,7 +68,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) self.custom_llm_provider: Optional[str] = custom_llm_provider self.litellm_metadata: Optional[dict] = litellm_metadata or {} - self.collected_chat_completion_chunks: List[ModelResponseStream] = [] + # Store lightweight dict snapshots for stream_chunk_builder to reduce + # repeated Pydantic attribute access in end-of-stream assembly. + self.collected_chat_completion_chunks: List[Dict[str, Any]] = [] self.finished: bool = False self.litellm_logging_obj = litellm_custom_stream_wrapper.logging_obj self.sent_response_created_event: bool = False @@ -465,6 +467,22 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ), ) + @staticmethod + def _snapshot_chunk_for_stream_chunk_builder( + chunk: ModelResponseStream, + ) -> Dict[str, Any]: + """ + Convert a streaming chunk into a plain dict for end-of-stream assembly. + Keep _hidden_params so downstream usage/header behavior is preserved. + """ + chunk_dict = chunk.model_dump() + hidden_params = getattr(chunk, "_hidden_params", None) + if hidden_params is not None: + chunk_dict["_hidden_params"] = ( + dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params + ) + return chunk_dict + def create_reasoning_summary_text_done_event( self, reasoning_item_id: str, @@ -811,7 +829,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk = cast(ModelResponseStream, chunk) self._ensure_output_item_for_chunk(chunk) # Proceed to transformation - self.collected_chat_completion_chunks.append(chunk) + self.collected_chat_completion_chunks.append( + self._snapshot_chunk_for_stream_chunk_builder(chunk) + ) if self._reasoning_active and not self._reasoning_done_emitted: # Incrementally accumulate reasoning content instead of # calling stream_chunk_builder on every chunk (O(n²)) @@ -902,7 +922,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Emit any just-queued output_item event if self._pending_response_events: return self._pending_response_events.pop(0) - self.collected_chat_completion_chunks.append(chunk) + self.collected_chat_completion_chunks.append( + self._snapshot_chunk_for_stream_chunk_builder( + cast(ModelResponseStream, chunk) + ) + ) response_api_chunk = ( self._transform_chat_completion_chunk_to_response_api_chunk( chunk diff --git a/tests/llm_responses_api_testing/test_anthropic_responses_api.py b/tests/llm_responses_api_testing/test_anthropic_responses_api.py index 5df1045b7c0..9957f2c342b 100644 --- a/tests/llm_responses_api_testing/test_anthropic_responses_api.py +++ b/tests/llm_responses_api_testing/test_anthropic_responses_api.py @@ -30,7 +30,7 @@ class TestAnthropicResponsesAPITest(BaseResponsesAPITest): def get_base_completion_call_args(self): #litellm._turn_on_debug() return { - "model": "anthropic/claude-sonnet-4-5-20250929", + "model": "anthropic/claude-sonnet-4-5", } async def test_basic_openai_responses_delete_endpoint(self, sync_mode=False): diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index da6d8027921..c86e146b0ef 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -8,6 +8,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +from litellm import stream_chunk_builder from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor from litellm.types.utils import ( ChatCompletionDeltaToolCall, @@ -512,3 +513,93 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.completion_tokens == 27 assert usage.total_tokens == 77 assert usage.server_tool_use['web_search_requests'] == 2 + + +def test_sort_chunks_handles_dict_hidden_params_created_at(): + chunks = [ + { + "id": "chunk_2", + "object": "chat.completion.chunk", + "created": 2, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "b"}}], + "_hidden_params": {"created_at": 2}, + }, + { + "id": "chunk_1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "a"}}], + "_hidden_params": {"created_at": 1}, + }, + ] + + processor = ChunkProcessor(chunks=chunks) + assert processor.chunks[0]["id"] == "chunk_1" + assert processor.chunks[1]["id"] == "chunk_2" + + +def test_stream_chunk_builder_accepts_dict_snapshot_chunks(): + chunk1 = ModelResponseStream( + id="chatcmpl-123", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello ", role="assistant"), + ) + ], + ) + chunk2 = ModelResponseStream( + id="chatcmpl-123", + created=2, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="world", role=None), + ) + ], + ) + chunk1._hidden_params = {"created_at": 1} + chunk2._hidden_params = {"created_at": 2} + + chunks = [] + for chunk in [chunk2, chunk1]: + chunk_dict = chunk.model_dump() + chunk_dict["_hidden_params"] = chunk._hidden_params + chunks.append(chunk_dict) + + response = stream_chunk_builder(chunks=chunks) + assert response is not None + assert response.choices[0].message.content == "Hello world" + + +def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): + chunk = ModelResponseStream( + id="chatcmpl-123", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hi", role="assistant"), + ) + ], + ) + chunk_dict = chunk.model_dump() + chunk_dict["_hidden_params"] = { + "provider_specific_fields": {"traffic_type": "default"} + } + + response = stream_chunk_builder(chunks=[chunk_dict]) + assert response is not None + assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"