From 51bc802c1fef1f81ae419979f24bfdcc7dedb19a Mon Sep 17 00:00:00 2001 From: Dan Loftus Date: Tue, 1 Sep 2026 21:22:50 -0400 Subject: [PATCH] fix(responses): preserve identity across replacement deltas --- .../streaming_iterator.py | 49 ++++++------- .../test_streaming_iterator_transformation.py | 70 ++++++++++++++----- 2 files changed, 76 insertions(+), 43 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 3e138429f99..d0ec1114511 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -119,8 +119,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_call_id_by_index: dict[int, str] = {} self._streamed_tool_call_ids_in_order: list[str] = [] # mutable-ok: accumulates call ids across stream chunks self._resolved_tool_call_id_by_position: dict[int, str] = {} # mutable-ok: terminal correlation state - self._streamed_call_id_by_terminal_id: dict[str, str] = {} # mutable-ok: terminal identity correlation - self._ambiguous_tool_call_indexes: set[int] = set() + self._streamed_call_id_by_provider_id: dict[str, str] = {} # mutable-ok: provider identity correlation self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item self._final_tool_events_queued: bool = False self._sequence_number: int = 0 @@ -159,8 +158,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return None def _streamed_tool_call_id_at_position(self, position: int) -> str | None: - if position in self._ambiguous_tool_call_indexes: - return None indexed_call_id: Final = self._tool_call_id_by_index.get(position) if indexed_call_id is not None: return indexed_call_id @@ -177,8 +174,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """Match a terminal aggregate tool call to the identity emitted while streaming.""" tool_call_index: Final = self._normalize_tool_call_index(tool_call) if tool_call_index is not None: - if tool_call_index in self._ambiguous_tool_call_indexes: - return None indexed_call_id: Final = self._tool_call_id_by_index.get(tool_call_index) if indexed_call_id is not None: return indexed_call_id @@ -186,6 +181,23 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return None return self._streamed_tool_call_id_at_position(position) + def _resolve_streamed_tool_call_id(self, tool_call_index: int | None, call_id_raw: object) -> str | None: + if tool_call_index is None: + return str(call_id_raw) if call_id_raw else None + + indexed_call_id: Final = self._tool_call_id_by_index.get(tool_call_index) + if indexed_call_id is not None: + if call_id_raw: + self._streamed_call_id_by_provider_id[str(call_id_raw)] = indexed_call_id + return indexed_call_id + + if not call_id_raw: + return None + + call_id: Final = str(call_id_raw) + self._tool_call_id_by_index[tool_call_index] = call_id + return call_id + def _responses_namespace_tool_call_fields(self, fn_name: str) -> tuple[str, str | None]: mapped: Final = self._namespace_tool_names.get(fn_name) if mapped: @@ -255,25 +267,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): for tc in tool_calls: tc_index = self._normalize_tool_call_index(tc) call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) - call_id = "" - - if call_id_raw: - call_id = str(call_id_raw) - if tc_index is not None: - existing_call_id = self._tool_call_id_by_index.get(tc_index) - if existing_call_id is not None and existing_call_id != call_id: - # Reusing the same index for multiple call_ids is ambiguous for id-less deltas. - # Guard against silent misrouting by disabling index fallback for this index. - self._ambiguous_tool_call_indexes.add(tc_index) - self._tool_call_id_by_index[tc_index] = call_id - elif tc_index is not None: - if tc_index in self._ambiguous_tool_call_indexes: - continue - mapped_call_id = self._tool_call_id_by_index.get(tc_index) - if mapped_call_id: - call_id = mapped_call_id - - if not call_id: + call_id = self._resolve_streamed_tool_call_id(tc_index, call_id_raw) + if call_id is None: continue fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) @@ -338,7 +333,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): terminal_call_id = str(call_id_raw) call_id = self._streamed_tool_call_id_for_terminal_call(tc, position) or terminal_call_id self._resolved_tool_call_id_by_position[position] = call_id - self._streamed_call_id_by_terminal_id[terminal_call_id] = call_id + self._streamed_call_id_by_provider_id[terminal_call_id] = call_id output_index = self._get_or_assign_tool_output_index(call_id) fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) @@ -1216,7 +1211,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _output_item_with_streamed_tool_identity(self, item: object, tool_position: int) -> object: terminal_call_id: Final = getattr(item, "call_id", None) resolved_by_call_id: Final = ( - self._streamed_call_id_by_terminal_id.get(terminal_call_id) if isinstance(terminal_call_id, str) else None + self._streamed_call_id_by_provider_id.get(terminal_call_id) if isinstance(terminal_call_id, str) else None ) resolved_by_position: Final = self._resolved_tool_call_id_by_position.get(tool_position) streamed_call_id: Final = ( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 855612a5eda..4ed8b478e1c 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -804,7 +804,7 @@ def test_completed_snapshot_correlates_function_after_server_tool_replacement(): assert function_call.arguments == '{"city":"Paris"}' -def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): +def test_reused_index_with_new_call_id_preserves_first_streamed_identity(): iterator = LiteLLMCompletionStreamingIterator( model="test-model", litellm_custom_stream_wrapper=AsyncMock(), @@ -828,39 +828,77 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): "index": 0, "id": "call_b", "type": "function", - "function": {"name": "tool_b", "arguments": '{"b":'}, + "function": {"name": "tool_a", "arguments": "1"}, } ] ) - # Ambiguous chunk: index reused and id missing. We should skip fallback rather than misroute. iterator._queue_tool_call_delta_events( [ { "index": 0, "type": "function", - "function": {"arguments": "1}"}, + "function": {"arguments": "}"}, } ] ) + streamed_argument_events = [ + event + for event in iterator._pending_tool_events + if event.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] + + terminal_response = ModelResponse( + id="chatcmpl-terminal", + created=123, + model="test-model", + object="chat.completion", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "index": 0, + "id": "call_b", + "type": "function", + "function": {"name": "tool_a", "arguments": '{"a":1}'}, + } + ], + }, + } + ], + ) + iterator._queue_final_tool_call_done_events(terminal_response) + completed = iterator._emit_response_completed_event(terminal_response) all_events = [] while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + added_items = [event.item for event in all_events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] + argument_events = [ + event + for event in all_events + if event.type + in { + ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, + } ] - arguments_by_call_id = {} - for evt in delta_events: - arguments_by_call_id.setdefault(evt.item_id, "") - arguments_by_call_id[evt.item_id] += evt.delta + done_item = next(event.item for event in all_events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) - assert arguments_by_call_id["fc_call_a"] == '{"a":' - assert arguments_by_call_id["fc_call_b"] == '{"b":' - assert arguments_by_call_id["fc_call_a"] != '{"a":1}' - assert arguments_by_call_id["fc_call_b"] != '{"b":1}' + assert completed is not None + assert [(item.id, item.call_id) for item in added_items] == [("fc_call_a", "call_a")] + assert {event.item_id for event in streamed_argument_events} == {"fc_call_a"} + assert "".join(event.delta for event in streamed_argument_events) == '{"a":1}' + assert {event.item_id for event in argument_events} == {"fc_call_a"} + assert argument_events[-1].arguments == '{"a":1}' + assert (done_item.id, done_item.call_id) == ("fc_call_a", "call_a") + completed_call = next(item for item in completed.response.output if item.type == "function_call") + assert (completed_call.id, completed_call.call_id) == ("fc_call_a", "call_a") @pytest.mark.asyncio