diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 0f8ba6ef98a..c1e5ca5042f 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/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"]) diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index d5860b7519b..42f2bb3d726 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/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.""" diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 1a6e3cc7832..92dd1cfc0b5 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/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"