Merge afcf221919 into aeede53112

litellm/responses/litellm_completion_transformation/streaming_iterator.py

@@ -18,6 +18,7 @@ from litellm.responses.litellm_completion_transformation.custom_tools import (
     extract_custom_tool_names,
     is_custom_tool_call,
     serialize_tool_call_arguments,
+    unwrap_custom_tool_arguments,
 )
 from litellm.responses.litellm_completion_transformation.reasoning_items import mint_reasoning_item_id
 from litellm.responses.litellm_completion_transformation.tool_search import (
@@ -44,6 +45,7 @@ from litellm.types.llms.openai import (
     ContentPartDonePartRefusal,
     FunctionCallArgumentsDeltaEvent,
     FunctionCallArgumentsDoneEvent,
+    GenericEvent,
     OutputItemAddedEvent,
     OutputItemDoneEvent,
     OutputTextAnnotationAddedEvent,
@@ -225,6 +227,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
         self._accumulated_provider_specific_fields: dict[str, object] = {}
         request_tools: Final = self.responses_api_request.get("tools")
         self._custom_tool_names: set[str] = extract_custom_tool_names(request_tools)
+        self._custom_tool_call_ids: frozenset[str] = frozenset()
+        self._unnamed_tool_args_by_call_id: Mapping[str, str] = MappingProxyType({})
         self._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools)
         self._web_search_calls: dict[str, object] = {}  # mutable-ok: latest call by provider id
         self._queued_web_search_call_ids: set[str] = set()  # mutable-ok: emitted call ids
@@ -486,36 +490,90 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
             else:
                 fn_name = str(getattr(fn, "name", "") or "")
                 fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
-            if call_id not in self._tool_output_index_by_call_id:
-                self._queue_held_annotations()
-            output_index = self._get_or_assign_tool_output_index(call_id)
+            self._queue_named_tool_call_delta(call_id, fn_name, fn_args_delta)
 
-            if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None:
-                self._queue_tool_search_call_added(call_id, output_index, fn_args_delta, search_execution)
-                continue
-            self._started_tool_call = True
+    def _queue_named_tool_call_delta(self, call_id: str, fn_name: str, fn_args_delta: str) -> None:
+        if call_id not in self._tool_output_index_by_call_id:
+            self._queue_held_annotations()
+        output_index: Final = self._get_or_assign_tool_output_index(call_id)
+        if not fn_name and call_id not in self._tool_item_id_by_call_id:
+            self._unnamed_tool_args_by_call_id = MappingProxyType(
+                {
+                    **self._unnamed_tool_args_by_call_id,
+                    call_id: self._unnamed_tool_args_by_call_id.get(call_id, "") + fn_args_delta,
+                }
+            )
+            return
+        args_delta: Final = self._unnamed_tool_args_by_call_id.get(call_id, "") + fn_args_delta
+        if call_id in self._unnamed_tool_args_by_call_id:
+            self._unnamed_tool_args_by_call_id = MappingProxyType(
+                {key: value for key, value in self._unnamed_tool_args_by_call_id.items() if key != call_id}
+            )
+        if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None:
+            self._queue_tool_search_call_added(call_id, output_index, args_delta, search_execution)
+            return
+        if is_custom_tool_call(fn_name, self._custom_tool_names):
+            self._custom_tool_call_ids = self._custom_tool_call_ids | {call_id}
+        self._started_tool_call = True
 
-            if call_id not in self._tool_args_by_call_id:
-                self._tool_args_by_call_id[call_id] = ""
-                item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
-                self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
-                self._queue_output_item_added(output_index, BaseLiteLLMOpenAIResponseObject(**item_kwargs))
+        if call_id not in self._tool_args_by_call_id:
+            self._tool_args_by_call_id[call_id] = ""
+            item_kwargs: Final = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
+            self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
+            self._queue_output_item_added(output_index, BaseLiteLLMOpenAIResponseObject(**item_kwargs))
 
-            if fn_args_delta:
-                self._tool_args_by_call_id[call_id] += fn_args_delta
+        self._tool_args_by_call_id[call_id] += args_delta
+        if call_id in self._custom_tool_call_ids:
+            return
+        for offset in range(0, len(args_delta), 10):
+            self._queue_function_tool_argument_delta(call_id, output_index, args_delta[offset : offset + 10])
 
-                # Split large argument deltas into smaller chunks to match OpenAI's streaming behavior
-                # This is especially important for providers like Bedrock that send complete arguments at once
-                chunk_size = 10  # Match typical OpenAI delta size
-                for i in range(0, len(fn_args_delta), chunk_size):
-                    delta_chunk = fn_args_delta[i : i + chunk_size]
-                    delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent(
-                        type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
-                        item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
-                        output_index=output_index,
-                        delta=delta_chunk,
-                    )
-                    self._pending_tool_events.append(delta_event)
+    def _queue_function_tool_argument_delta(self, call_id: str, output_index: int, delta: str) -> None:
+        delta_event: Final = FunctionCallArgumentsDeltaEvent(
+            type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
+            item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
+            output_index=output_index,
+            delta=delta,
+        )
+        self._pending_tool_events.append(delta_event)
+
+    def _queue_custom_tool_input_done(self, call_id: str, output_index: int, arguments: str) -> None:
+        input_text: Final = unwrap_custom_tool_arguments(arguments)
+        item_id: Final = self._tool_item_id_by_call_id.get(call_id, call_id)
+        if input_text:
+            self._pending_tool_events.append(
+                GenericEvent.model_validate(
+                    {
+                        "type": ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA.value,
+                        "item_id": item_id,
+                        "output_index": output_index,
+                        "delta": input_text,
+                    }
+                )
+            )
+        self._pending_tool_events.append(
+            GenericEvent.model_validate(
+                {
+                    "type": ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DONE.value,
+                    "item_id": item_id,
+                    "output_index": output_index,
+                    "input": input_text,
+                }
+            )
+        )
+
+    def _queue_function_tool_arguments_done(self, call_id: str, output_index: int, final_args: str) -> None:
+        already_streamed: Final = self._tool_args_by_call_id.get(call_id, "")
+        remaining_args: Final = final_args[len(already_streamed) :] if final_args else ""
+        for offset in range(0, len(remaining_args), 10):
+            self._queue_function_tool_argument_delta(call_id, output_index, remaining_args[offset : offset + 10])
+        done_event: Final = FunctionCallArgumentsDoneEvent(
+            type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
+            item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
+            output_index=output_index,
+            arguments=final_args,
+        )
+        self._pending_tool_events.append(done_event)
 
     def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None:
         """
@@ -576,32 +634,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
                 self._queue_output_item_added(output_index, BaseLiteLLMOpenAIResponseObject(**item_kwargs))
 
             final_args = fn_args or self._tool_args_by_call_id.get(call_id, "")
-
-            # Emit delta events for arguments that weren't streamed yet
-            # This handles cases where Bedrock sends the complete tool call at the end
-            already_streamed = self._tool_args_by_call_id.get(call_id, "")
-            remaining_args = final_args[len(already_streamed) :] if final_args else ""
-
-            if remaining_args:
-                # Split into smaller chunks to match OpenAI's streaming behavior
-                chunk_size = 10  # Match typical OpenAI delta size
-                for i in range(0, len(remaining_args), chunk_size):
-                    delta_chunk = remaining_args[i : i + chunk_size]
-                    delta_event = FunctionCallArgumentsDeltaEvent(
-                        type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
-                        item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
-                        output_index=output_index,
-                        delta=delta_chunk,
-                    )
-                    self._pending_tool_events.append(delta_event)
-
-            done_event = FunctionCallArgumentsDoneEvent(
-                type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
-                item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
-                output_index=output_index,
-                arguments=final_args,
-            )
-            self._pending_tool_events.append(done_event)
+            if is_custom_tool_call(fn_name, self._custom_tool_names):
+                self._queue_custom_tool_input_done(call_id, output_index, final_args)
+            else:
+                self._queue_function_tool_arguments_done(call_id, output_index, final_args)
 
             item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed")
             item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])

tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py

@@ -2,7 +2,6 @@ import asyncio
 import json
 from collections.abc import Mapping, Sequence
 from copy import deepcopy
-from types import MappingProxyType
 from typing import Final, Literal, cast
 
 import pytest
@@ -27,7 +26,6 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     anthropic_messages_pt,
 )
 from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks
-from litellm.responses.litellm_completion_transformation.tool_search import tool_search_routing
 from litellm.responses.litellm_completion_transformation.transformation import (
     TOOL_CALLS_CACHE,
     ChatCompletionSession,
@@ -4264,40 +4262,18 @@ class TestEnsureOutputItemContentPartAdded:
     output_item.added for message items."""
 
     def _make_iterator(self):
-        """Create a minimal LiteLLMCompletionStreamingIterator for testing."""
+        from unittest.mock import Mock
+
         from litellm.responses.litellm_completion_transformation.streaming_iterator import (
             LiteLLMCompletionStreamingIterator,
         )
 
-        iterator = LiteLLMCompletionStreamingIterator.__new__(LiteLLMCompletionStreamingIterator)
-        iterator.sent_output_item_added_event = False
-        iterator.sent_message_item_added_event = False
-        iterator.sent_content_part_added_event = False
-        iterator._sequence_number = 0
-        iterator._cached_item_id = None
-        iterator._cached_reasoning_item_id = None
-        iterator._reasoning_active = False
-        iterator._pending_response_events = []
-        iterator._pending_tool_events = []
-        iterator._annotations_awaiting_message = ()
-        iterator._tool_output_index_by_call_id = {}
-        iterator._tool_args_by_call_id = {}
-        iterator._tool_item_id_by_call_id = {}
-        iterator._tool_call_id_by_index = {}
-        iterator._ambiguous_tool_call_indexes = set()
-        iterator._next_output_index = 0
-        iterator._message_output_index = None
-        iterator._reasoning_output_index = None
-        iterator._final_tool_events_queued = False
-        iterator._custom_tool_names = set()
-        iterator.responses_api_request = {}
-        iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(None)
-        iterator._web_search_calls = {}
-        iterator._queued_web_search_call_ids = set()
-        iterator._tool_search = tool_search_routing(None)
-        iterator._tool_search_executions = MappingProxyType({})
-        iterator.tool_call_provider_fields = False
-        return iterator
+        return LiteLLMCompletionStreamingIterator(
+            model="test-model",
+            litellm_custom_stream_wrapper=Mock(spec=litellm.CustomStreamWrapper, logging_obj=Mock()),
+            request_input="",
+            responses_api_request={},
+        )
 
     def _make_text_chunk(self):
         """Create a mock ModelResponseStream with a text delta."""

tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py

@@ -100,6 +100,198 @@ def _response_ids(events) -> list[str]:
     ]
 
 
+def _custom_tool_stream_chunk(delta: Delta, finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id=CHAT_COMPLETION_ID,
+        created=1748575031,
+        model="test-model",
+        object="chat.completion.chunk",
+        choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)],
+    )
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("input_mode", ["split", "late_name", "terminal_only"])
+@pytest.mark.parametrize("raw_input", ["", "hello", 'line one\n"世界" \\path 😀'])
+@pytest.mark.asyncio
+async def test_custom_tool_input_events_round_trip_with_interleaved_function(
+    sync_mode: bool, input_mode: str, raw_input: str
+):
+    terminal_only: Final = input_mode == "terminal_only"
+    late_name: Final = input_mode == "late_name"
+    arguments: Final = json.dumps({"content": raw_input}, ensure_ascii=True)
+    first: Final = Delta(
+        content=None,
+        tool_calls=[
+            {
+                "index": 0,
+                "id": "call-custom",
+                "type": "function",
+                "function": {
+                    **({} if late_name else {"name": "write_note"}),
+                    "arguments": arguments if terminal_only else arguments[:3] if late_name else "",
+                },
+            },
+            {
+                "index": 1,
+                "id": "call-function",
+                "type": "function",
+                "function": {
+                    **({} if late_name else {"name": "lookup"}),
+                    "arguments": '{"query":' if late_name else '{"query":"ready"}',
+                },
+            },
+        ],
+    )
+    pieces: Final = (
+        ()
+        if terminal_only
+        else tuple(arguments[offset : offset + 3] for offset in range(3 if late_name else 0, len(arguments), 3))
+    )
+    chunks: Final = (
+        _custom_tool_stream_chunk(first, "tool_calls" if terminal_only else None),
+        *(
+            (
+                _custom_tool_stream_chunk(
+                    Delta(
+                        content=None,
+                        tool_calls=[
+                            {"index": 0, "function": {"name": "write_note", "arguments": ""}},
+                            {"index": 1, "function": {"name": "lookup", "arguments": '"ready"}'}},
+                        ],
+                    )
+                ),
+            )
+            if late_name
+            else ()
+        ),
+        *(
+            _custom_tool_stream_chunk(Delta(content=None, tool_calls=[{"index": 0, "function": {"arguments": piece}}]))
+            for piece in pieces
+        ),
+        *(() if terminal_only else (_custom_tool_stream_chunk(Delta(content=None), "tool_calls"),)),
+    )
+    iterator: Final = LiteLLMCompletionStreamingIterator(
+        model="test-model",
+        litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks),
+        request_input="Write a note and look up its status",
+        responses_api_request={
+            "tools": [
+                {"type": "custom", "name": "write_note", "format": {"type": "text"}},
+                {"type": "function", "name": "lookup", "parameters": {"type": "object"}},
+            ]
+        },
+        custom_llm_provider="openai",
+    )
+    events: Final = (
+        tuple(event.model_dump(exclude_none=True) for event in iterator)
+        if sync_mode
+        else tuple([event.model_dump(exclude_none=True) async for event in iterator])
+    )
+    assert tuple(event["sequence_number"] for event in events) == tuple(range(len(events)))
+    completed: Final = next(event["response"] for event in events if event["type"] == "response.completed")
+    custom_item: Final = next(item for item in completed["output"] if item["type"] == "custom_tool_call")
+    function_item: Final = next(item for item in completed["output"] if item["type"] == "function_call")
+    custom_deltas: Final = tuple(event for event in events if event["type"] == "response.custom_tool_call_input.delta")
+    custom_done: Final = tuple(event for event in events if event["type"] == "response.custom_tool_call_input.done")
+    custom_add: Final = next(
+        event
+        for event in events
+        if event["type"] == "response.output_item.added" and event["item"].get("call_id") == "call-custom"
+    )
+    function_add: Final = next(
+        event
+        for event in events
+        if event["type"] == "response.output_item.added" and event["item"].get("call_id") == "call-function"
+    )
+    assert custom_add["item"]["type"] == "custom_tool_call"
+    assert custom_add["item"]["id"] == custom_item["id"]
+    assert function_add["item"]["type"] == "function_call"
+    assert function_add["item"]["name"] == function_item["name"] == "lookup"
+    assert custom_item["input"] == raw_input
+    assert "".join(event["delta"] for event in custom_deltas) == raw_input
+    assert len(custom_done) == 1
+    assert custom_done[0]["input"] == custom_item["input"]
+    assert custom_item["call_id"] == "call-custom"
+    assert all(event["item_id"] == custom_item["id"] for event in (*custom_deltas, *custom_done))
+    assert all(isinstance(event["sequence_number"], int) for event in (*custom_deltas, *custom_done))
+    assert all(event["output_index"] == custom_done[0]["output_index"] for event in custom_deltas)
+    function_events: Final = tuple(
+        event for event in events if event["type"].startswith("response.function_call_arguments.")
+    )
+    assert all(event["item_id"] == function_item["id"] for event in function_events)
+    function_deltas: Final = tuple(event for event in function_events if event["type"].endswith(".delta"))
+    assert json.loads("".join(event["delta"] for event in function_deltas)) == {"query": "ready"}
+    function_done: Final = tuple(event for event in function_events if event["type"].endswith(".done"))
+    assert len(function_done) == 1
+    assert function_done[0]["arguments"] == function_item["arguments"]
+    item_done: Final = next(
+        event
+        for event in events
+        if event["type"] == "response.output_item.done" and event["item"]["id"] == custom_item["id"]
+    )
+    assert item_done["item"] == custom_item
+    custom_sequence: Final = tuple(
+        event["sequence_number"] for event in (*custom_deltas, *custom_done, item_done)
+    )
+    assert all(left < right for left, right in zip(custom_sequence, custom_sequence[1:]))
+    assert item_done["sequence_number"] > custom_done[0]["sequence_number"]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("tool_type", ["custom", "function"])
+@pytest.mark.parametrize("late_name", [True, False])
+@pytest.mark.parametrize("citation_first", [True, False])
+@pytest.mark.asyncio
+async def test_emitted_tool_call_prevents_text_only_continuation(
+    sync_mode: bool, tool_type: str, late_name: bool, citation_first: bool
+) -> None:
+    iterator: Final = LiteLLMCompletionStreamingIterator(
+        model="test-model",
+        litellm_custom_stream_wrapper=_FakeStreamWrapper(
+            [
+                *([_annotation_only_chunk()] if citation_first else []),
+                _custom_tool_stream_chunk(
+                    Delta(
+                        content=None,
+                        tool_calls=[{
+                            "index": 0,
+                            "id": "call-note",
+                            "type": "function",
+                            "function": {**({} if late_name else {"name": "write_note"}), "arguments": ""},
+                        }],
+                    )
+                ),
+                _custom_tool_stream_chunk(
+                    Delta(
+                        content=None,
+                        tool_calls=[{
+                            "index": 0,
+                            "function": {
+                                **({"name": "write_note"} if late_name else {}),
+                                "arguments": '{"content":"note"}',
+                            },
+                        }],
+                    ),
+                    "tool_calls",
+                ),
+            ]
+        ),
+        request_input="Write a note",
+        responses_api_request={"tools": [{"type": tool_type, "name": "write_note"}]},
+        custom_llm_provider="openai",
+    )
+    events: Final = await _collect_events(iterator, sync_mode)
+    expected_type: Final = "custom_tool_call" if tool_type == "custom" else "function_call"
+    added_types: Final = tuple(
+        event.item.type
+        for event in events
+        if event.type == "response.output_item.added"
+    )
+    assert added_types == (("message", expected_type) if citation_first else (expected_type,))
+    assert iterator.started_tool_call
+
+
 def test_tool_call_delta_is_emitted_as_responses_events():
     iterator = LiteLLMCompletionStreamingIterator(
         model="test-model",
@@ -1831,21 +2023,36 @@ async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode: b
 
 
 @pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("execution", ["server", "client"])
+@pytest.mark.parametrize("late_name", [False, True])
 @pytest.mark.asyncio
-async def test_anthropic_tool_search_call_streams_as_tool_search_call_item(sync_mode: bool):
+async def test_anthropic_tool_search_call_streams_as_tool_search_call_item(
+    sync_mode: bool, execution: str, late_name: bool
+):
+    call_id: Final = "srvtoolu_01Search" if execution == "server" else "toolu_01Search"
+    name: Final = "tool_search_tool_regex" if execution == "server" else "tool_search"
     deltas: Final = (
         Delta(
             content=None,
             tool_calls=[
                 {
                     "index": 0,
-                    "id": "srvtoolu_01Search",
+                    "id": call_id,
                     "type": "function",
-                    "function": {"name": "tool_search_tool_regex", "arguments": ""},
+                    "function": {
+                        **({} if late_name else {"name": name}),
+                        "arguments": '{"pattern": ' if late_name else "",
+                    },
                 }
             ],
         ),
-        Delta(content=None, tool_calls=[{"index": 0, "function": {"arguments": '{"pattern": '}}]),
+        Delta(
+            content=None,
+            tool_calls=[{
+                "index": 0,
+                "function": {"name": name, "arguments": ""} if late_name else {"arguments": '{"pattern": '},
+            }],
+        ),
         Delta(content=None, tool_calls=[{"index": 0, "function": {"arguments": '"weather|forecast"}'}}]),
         Delta(
             content=None,
@@ -1879,7 +2086,7 @@ async def test_anthropic_tool_search_call_streams_as_tool_search_call_item(sync_
         request_input="What is the weather in Paris right now?",
         responses_api_request={
             "tools": [
-                {"type": "tool_search", "execution": "server"},
+                {"type": "tool_search", "execution": execution},
                 {"type": "function", "name": "get_weather", "parameters": {"type": "object"}, "defer_loading": True},
             ]
         },
@@ -1902,16 +2109,17 @@ async def test_anthropic_tool_search_call_streams_as_tool_search_call_item(sync_
     ]
 
     assert completed["type"] == "response.completed"
-    assert search_item["id"] == "tsc_srvtoolu_01Search"
-    assert search_item["call_id"] == "srvtoolu_01Search"
+    assert search_item["id"] == f"tsc_{call_id}"
+    assert search_item["call_id"] == call_id
     assert search_item["arguments"] == {"pattern": "weather|forecast"}
-    assert search_item["execution"] == "server"
+    assert search_item["execution"] == execution
     assert search_item["status"] == "completed"
     assert [event["type"] for event in search_events] == [
         "response.output_item.added",
         "response.output_item.done",
     ]
     assert search_events[0]["item"]["status"] == "in_progress"
+    assert search_events[0]["item"]["execution"] == execution
     assert search_events[0]["item"]["arguments"] == {}
     assert search_events[-1]["item"] == search_item
     assert function_item["name"] == "get_weather"