diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index b21703509cb..a2f9c80f577 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 Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence +from itertools import groupby +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from litellm._logging import verbose_logger @@ -205,6 +207,52 @@ 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 + + @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[ @@ -250,9 +298,7 @@ class ChunkProcessor: "id": None, "name": None, "type": None, - "arguments": (), "custom_name": None, - "custom_input": (), "provider_specific_fields": None, } @@ -267,21 +313,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 +331,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 +360,10 @@ class ChunkProcessor: if isinstance(provider_fields, dict): tool_call_map[index]["provider_specific_fields"].update(provider_fields) + joined_fragments: Final = 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()): tool_call_data = tool_call_map[index] @@ -333,12 +373,12 @@ class ChunkProcessor: id=tool_call_data["id"], custom=ChatCompletionCustomToolCallPayload( name=tool_call_data["custom_name"], - input="".join(tool_call_data["custom_input"]), + input=joined_fragments.get((index, "custom_input"), ""), ), ) ) elif tool_call_data["id"] and tool_call_data["name"]: - combined_arguments = "".join(tool_call_data["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 2db5461702a..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 @@ -1064,3 +1064,119 @@ 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: dict[str, object] | ChatCompletionDeltaToolCall) -> dict[str, object]: + 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_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)] + 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)