diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index c87357122ce..9db951b022e 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/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 diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 527c83d5760..52f994aa7fe 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/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(