Merge 8810cabdfd into e9b1c11098

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -168,8 +168,17 @@ class _ToolCallFragment(TypedDict, total=False):
     provider_specific_fields: ReadOnly[dict[str, object] | None]
 
 
+_ToolCallEntry: TypeAlias = _ToolCallFragment | ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall
+
+
 class _ToolCallDelta(TypedDict, total=False):
-    tool_calls: ReadOnly[Sequence[_ToolCallFragment | ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]]
+    tool_calls: ReadOnly[Sequence[_ToolCallEntry] | _ToolCallEntry]
+
+
+def _normalize_tool_calls(tool_calls: Sequence[_ToolCallEntry] | _ToolCallEntry | None) -> Sequence[_ToolCallEntry]:
+    if isinstance(tool_calls, (dict, ChatCompletionDeltaToolCall, ChatCompletionDeltaCustomToolCall)):
+        return (tool_calls,)
+    return tool_calls or ()
 
 
 class _ToolCallChoice(TypedDict, total=False):
@@ -554,7 +563,7 @@ class ChunkProcessor:
                 if not delta:
                     continue
                 choice_index = choice.get("index", 0)
-                for tool_call in delta.get("tool_calls") or ():
+                for tool_call in _normalize_tool_calls(delta.get("tool_calls")):
                     if not tool_call:
                         continue
                     if isinstance(tool_call, dict):
@@ -601,10 +610,9 @@ class ChunkProcessor:
             choices = chunk["choices"]
             for choice in choices:
                 delta = choice.get("delta", {})
-                tool_calls = delta.get("tool_calls") or ()
                 choice_index = choice.get("index", 0)
 
-                for tool_call in tool_calls:
+                for tool_call in _normalize_tool_calls(delta.get("tool_calls")):
                     # Handle both dict and object formats
                     if not tool_call:
                         continue

tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -16,6 +16,10 @@ from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
 )
 from litellm.llms.anthropic.chat.handler import ModelResponseIterator
 from litellm.types.utils import (
+    ChatCompletionDeltaCustomToolCall,
+    ChatCompletionDeltaCustomToolCallPayload,
+    ChatCompletionMessageCustomToolCall,
+    Choices,
     ChatCompletionDeltaToolCall,
     ChatCompletionMessageToolCall,
     Delta,
@@ -31,6 +35,53 @@ from litellm.types.utils import (
 )
 
 
+@pytest.mark.parametrize("custom", (False, True))
+@pytest.mark.parametrize("as_dict", (False, True))
+@pytest.mark.parametrize("single", (False, True))
+def test_stream_chunk_builder_preserves_single_tool_call(custom: bool, as_dict: bool, single: bool) -> None:
+    tool_call: Final = (
+        ChatCompletionDeltaCustomToolCall(
+            index=0,
+            id="call_one",
+            type="custom",
+            custom=ChatCompletionDeltaCustomToolCallPayload(name="lookup", input='{"key": "value"}'),
+        )
+        if custom
+        else ChatCompletionDeltaToolCall(
+            index=0,
+            id="call_one",
+            type="function",
+            function=Function(name="lookup", arguments='{"key": "value"}'),
+        )
+    )
+    entry: Final = tool_call.model_dump() if as_dict else tool_call
+    chunk: Final = ModelResponseStream(
+        id="chatcmpl-single",
+        model="test-model",
+        created=0,
+        choices=[
+            StreamingChoices(index=0, finish_reason="tool_calls", delta=Delta(tool_calls=entry if single else [entry]))
+        ],
+    )
+    response: Final = stream_chunk_builder([chunk])
+    assert response is not None
+    choice: Final = response.choices[0]
+    assert isinstance(choice, Choices)
+    combined: Final = choice.message.tool_calls
+    assert combined is not None and len(combined) == 1
+    assert combined[0].id == "call_one"
+    if custom:
+        assert isinstance(combined[0], ChatCompletionMessageCustomToolCall)
+        assert combined[0].type == "custom"
+        assert combined[0].custom.name == "lookup"
+        assert combined[0].custom.input == '{"key": "value"}'
+    else:
+        assert isinstance(combined[0], ChatCompletionMessageToolCall)
+        assert combined[0].type == "function"
+        assert combined[0].function.name == "lookup"
+        assert combined[0].function.arguments == '{"key": "value"}'
+
+
 def test_get_combined_tool_content():
     chunks = [
         ModelResponseStream(