From fb1674923d6ec62887d475054aa8c89b5fd14c62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:11:48 -0700 Subject: [PATCH 1/2] perf(streaming): assemble streamed tool-call arguments in linear time --- .../streaming_chunk_builder_utils.py | 63 ++++++++++--- .../test_streaming_chunk_builder_utils.py | 94 +++++++++++++++++++ 2 files changed, 142 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 6e7ed370294..1f22f241452 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,6 @@ import base64 import time -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Union, cast from litellm._logging import verbose_logger @@ -205,6 +205,38 @@ class ChunkProcessor: response = self.update_model_response_with_hidden_params(model_response=response, chunk=chunk) return response + @staticmethod + def _iter_tool_call_fragments( + tool_call_chunks: Sequence[Mapping[str, Any]], + ) -> Iterator[tuple[int, str, str]]: + for chunk in tool_call_chunks: + for choice in chunk["choices"]: + delta = choice.get("delta") + if not delta: + continue + for tool_call in delta.get("tool_calls", ()): + if not tool_call: + continue + if isinstance(tool_call, dict): + index = tool_call.get("index", 0) + function = tool_call.get("function") + if isinstance(function, dict): + if function.get("arguments"): + yield index, "arguments", function["arguments"] + elif getattr(function, "arguments", None): + yield index, "arguments", function.arguments + custom = tool_call.get("custom") + if isinstance(custom, dict) and custom.get("input"): + yield index, "custom_input", custom["input"] + else: + index = getattr(tool_call, "index", 0) + function = getattr(tool_call, "function", None) + if getattr(function, "arguments", None): + yield index, "arguments", function.arguments + custom = getattr(tool_call, "custom", None) + if getattr(custom, "input", None): + yield index, "custom_input", custom.input + def get_combined_tool_content( self, tool_call_chunks: Sequence[Mapping[str, Any]] ) -> list[ @@ -250,9 +282,7 @@ class ChunkProcessor: "id": None, "name": None, "type": None, - "arguments": (), "custom_name": None, - "custom_input": (), "provider_specific_fields": None, } @@ -267,21 +297,15 @@ class ChunkProcessor: if isinstance(function, dict): if function.get("name"): tool_call_map[index]["name"] = function["name"] - if function.get("arguments"): - tool_call_map[index]["arguments"] += (function["arguments"],) else: # function is an object if hasattr(function, "name") and function.name: tool_call_map[index]["name"] = function.name - if hasattr(function, "arguments") and function.arguments: - tool_call_map[index]["arguments"] += (function.arguments,) custom = tool_call.get("custom") if isinstance(custom, dict): if custom.get("name"): tool_call_map[index]["custom_name"] = custom["name"] - if custom.get("input"): - tool_call_map[index]["custom_input"] += (custom["input"],) else: # tool_call is an object if hasattr(tool_call, "id") and tool_call.id: @@ -291,15 +315,11 @@ class ChunkProcessor: if hasattr(tool_call, "function"): if hasattr(tool_call.function, "name") and tool_call.function.name: tool_call_map[index]["name"] = tool_call.function.name - if hasattr(tool_call.function, "arguments") and tool_call.function.arguments: - tool_call_map[index]["arguments"] += (tool_call.function.arguments,) custom = getattr(tool_call, "custom", None) if custom is not None: if getattr(custom, "name", None): tool_call_map[index]["custom_name"] = custom.name - if getattr(custom, "input", None): - tool_call_map[index]["custom_input"] += (custom.input,) # Preserve provider_specific_fields from streaming chunks provider_fields = None @@ -324,6 +344,8 @@ class ChunkProcessor: if isinstance(provider_fields, dict): tool_call_map[index]["provider_specific_fields"].update(provider_fields) + fragment_records = tuple(self._iter_tool_call_fragments(tool_call_chunks)) + # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): tool_call_data = tool_call_map[index] @@ -333,12 +355,23 @@ class ChunkProcessor: id=tool_call_data["id"], custom=ChatCompletionCustomToolCallPayload( name=tool_call_data["custom_name"], - input="".join(tool_call_data["custom_input"]), + input="".join( + fragment + for fragment_index, field, fragment in fragment_records + if fragment_index == index and field == "custom_input" + ), ), ) ) elif tool_call_data["id"] and tool_call_data["name"]: - combined_arguments = "".join(tool_call_data["arguments"]) or "{}" + combined_arguments = ( + "".join( + fragment + for fragment_index, field, fragment in fragment_records + if fragment_index == index and field == "arguments" + ) + or "{}" + ) # Build function - provider_specific_fields should be on tool_call level, not function level function = Function( 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 2db5461702a..cfa566428d0 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 @@ -1064,3 +1064,97 @@ def test_get_combined_tool_content_custom_tool_call_without_type_field(): "type": "custom", "custom": {"name": "ApplyPatch", "input": "*** Begin Patch"}, } + + +def _tool_call_delta_chunk(tool_call): + return {"choices": [{"delta": {"tool_calls": [tool_call]}}]} + + +def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_order(): + processor = ChunkProcessor.__new__(ChunkProcessor) + first_fragments = [f"a{i};" for i in range(300)] + second_fragments = [f"b{i};" for i in range(300)] + header_chunks = [ + _tool_call_delta_chunk({"index": 0, "id": "call_a", "type": "function", "function": {"name": "tool_a"}}), + _tool_call_delta_chunk({"index": 1, "id": "call_b", "type": "function", "function": {"name": "tool_b"}}), + _tool_call_delta_chunk({"index": 2, "id": "call_c", "type": "function", "function": {"name": "tool_c"}}), + ] + fragment_chunks = [ + _tool_call_delta_chunk({"index": index, "function": {"arguments": fragment}}) + for first, second in zip(first_fragments, second_fragments) + for index, fragment in ((0, first), (1, second)) + ] + + combined = processor.get_combined_tool_content(header_chunks + fragment_chunks) + + assert [tool_call.id for tool_call in combined] == ["call_a", "call_b", "call_c"] + assert combined[0].function.name == "tool_a" + assert combined[0].function.arguments == "".join(first_fragments) + assert combined[1].function.name == "tool_b" + assert combined[1].function.arguments == "".join(second_fragments) + assert combined[2].function.arguments == "{}" + + +def test_get_combined_tool_content_joins_many_object_shaped_argument_fragments_in_order(): + processor = ChunkProcessor.__new__(ChunkProcessor) + first_fragments = [f"x{i}|" for i in range(300)] + second_fragments = [f"y{i}|" for i in range(300)] + header_chunks = [ + _tool_call_delta_chunk( + ChatCompletionDeltaToolCall( + id="call_x", type="function", index=0, function=Function(name="tool_x", arguments="") + ) + ), + _tool_call_delta_chunk( + ChatCompletionDeltaToolCall( + id="call_y", type="function", index=1, function=Function(name="tool_y", arguments="") + ) + ), + ] + fragment_chunks = [ + _tool_call_delta_chunk(ChatCompletionDeltaToolCall(index=index, function=Function(arguments=fragment))) + for first, second in zip(first_fragments, second_fragments) + for index, fragment in ((0, first), (1, second)) + ] + + combined = processor.get_combined_tool_content(header_chunks + fragment_chunks) + + assert [tool_call.id for tool_call in combined] == ["call_x", "call_y"] + assert combined[0].function.name == "tool_x" + assert combined[0].function.arguments == "".join(first_fragments) + assert combined[1].function.name == "tool_y" + assert combined[1].function.arguments == "".join(second_fragments) + + +def test_get_combined_tool_content_joins_many_custom_tool_input_fragments_in_order(): + from types import SimpleNamespace + + from litellm.types.utils import ChatCompletionMessageCustomToolCall + + processor = ChunkProcessor.__new__(ChunkProcessor) + dict_fragments = [f"d{i}," for i in range(200)] + object_fragments = [f"o{i}," for i in range(200)] + header_chunks = [ + _tool_call_delta_chunk({"index": 0, "id": "call_d", "type": "custom", "custom": {"name": "apply_patch"}}), + _tool_call_delta_chunk( + SimpleNamespace(index=1, id="call_o", type="custom", custom=SimpleNamespace(name="run_script", input="")) + ), + ] + fragment_chunks = [ + _tool_call_delta_chunk(tool_call) + for dict_fragment, object_fragment in zip(dict_fragments, object_fragments) + for tool_call in ( + {"index": 0, "custom": {"input": dict_fragment}}, + SimpleNamespace(index=1, custom=SimpleNamespace(input=object_fragment)), + ) + ] + + combined = processor.get_combined_tool_content(header_chunks + fragment_chunks) + + assert [tool_call.id for tool_call in combined] == ["call_d", "call_o"] + assert isinstance(combined[0], ChatCompletionMessageCustomToolCall) + assert combined[0].custom.name == "apply_patch" + assert combined[0].custom.input == "".join(dict_fragments) + assert isinstance(combined[1], ChatCompletionMessageCustomToolCall) + assert combined[1].custom.name == "run_script" + assert combined[1].custom.input == "".join(object_fragments) From 1dad33749c0c2d78b85558e3e0a698e5c44f685c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:38:37 -0700 Subject: [PATCH 2/2] perf(streaming): group tool-call fragments once instead of rescanning per index --- .../streaming_chunk_builder_utils.py | 33 +++++++++++-------- .../test_streaming_chunk_builder_utils.py | 24 +++++++++++++- 2 files changed, 42 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1f22f241452..029a18d7514 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,8 @@ import base64 import time from collections.abc import Iterator, Mapping, Sequence +from itertools import groupby +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Union, cast from litellm._logging import verbose_logger @@ -237,6 +239,20 @@ class ChunkProcessor: if getattr(custom, "input", None): yield index, "custom_input", custom.input + @staticmethod + def _join_fragments_by_index_and_field( + fragment_records: Iterator[tuple[int, str, str]], + ) -> Mapping[tuple[int, str], str]: + def group_key(record: tuple[int, str, str]) -> tuple[int, str]: + return record[0], record[1] + + return MappingProxyType( + { + key: "".join(fragment for _, _, fragment in group) + for key, group in groupby(sorted(fragment_records, key=group_key), key=group_key) + } + ) + def get_combined_tool_content( self, tool_call_chunks: Sequence[Mapping[str, Any]] ) -> list[ @@ -344,7 +360,7 @@ class ChunkProcessor: if isinstance(provider_fields, dict): tool_call_map[index]["provider_specific_fields"].update(provider_fields) - fragment_records = tuple(self._iter_tool_call_fragments(tool_call_chunks)) + joined_fragments = self._join_fragments_by_index_and_field(self._iter_tool_call_fragments(tool_call_chunks)) # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): @@ -355,23 +371,12 @@ class ChunkProcessor: id=tool_call_data["id"], custom=ChatCompletionCustomToolCallPayload( name=tool_call_data["custom_name"], - input="".join( - fragment - for fragment_index, field, fragment in fragment_records - if fragment_index == index and field == "custom_input" - ), + input=joined_fragments.get((index, "custom_input"), ""), ), ) ) elif tool_call_data["id"] and tool_call_data["name"]: - combined_arguments = ( - "".join( - fragment - for fragment_index, field, fragment in fragment_records - if fragment_index == index and field == "arguments" - ) - or "{}" - ) + combined_arguments = joined_fragments.get((index, "arguments"), "") or "{}" # Build function - provider_specific_fields should be on tool_call level, not function level function = Function( 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 cfa566428d0..0114db381cf 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 @@ -1066,7 +1066,7 @@ def test_get_combined_tool_content_custom_tool_call_without_type_field(): } -def _tool_call_delta_chunk(tool_call): +def _tool_call_delta_chunk(tool_call: dict[str, object] | ChatCompletionDeltaToolCall) -> dict[str, object]: return {"choices": [{"delta": {"tool_calls": [tool_call]}}]} @@ -1095,6 +1095,28 @@ def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_ assert combined[2].function.arguments == "{}" +def test_get_combined_tool_content_joins_fragments_across_many_parallel_tool_calls(): + processor = ChunkProcessor.__new__(ChunkProcessor) + indexes = range(40) + header_chunks = [ + _tool_call_delta_chunk( + {"index": index, "id": f"call_{index}", "type": "function", "function": {"name": f"tool_{index}"}} + ) + for index in indexes + ] + fragment_chunks = [ + _tool_call_delta_chunk({"index": index, "function": {"arguments": f"{index}.{position};"}}) + for position in range(5) + for index in indexes + ] + + combined = processor.get_combined_tool_content(header_chunks + fragment_chunks) + + assert [tool_call.id for tool_call in combined] == [f"call_{index}" for index in indexes] + for index, tool_call in zip(indexes, combined): + assert tool_call.function.arguments == "".join(f"{index}.{position};" for position in range(5)) + + def test_get_combined_tool_content_joins_many_object_shaped_argument_fragments_in_order(): processor = ChunkProcessor.__new__(ChunkProcessor) first_fragments = [f"x{i}|" for i in range(300)]