scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #35826 from BerriAI/litellm_fix_toolcall_stream_linear_assembly

Browse source 
perf(streaming): assemble streamed tool-call arguments in linear time
This commit is contained in:
Mateo Wang 2026-08-04 18:06:23 -07:00 • committed by GitHub
commit cb27d998f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 171 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

70
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -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(

116
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -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)