From 1ec5c7bc026d2c591f181166a7c8d264e16c15fe Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Mon, 2 Nov 2026 12:01:09 +0000 Subject: [PATCH 1/2] fix(google_genai): preserve and disambiguate tool call IDs in generateContent adapter Preserve explicit functionCall.id and functionResponse.id values across inbound, non-streaming outbound, and streaming outbound generateContent translations, and assign unique suffixed IDs (call__2, ...) to repeated tool calls while pairing responses per turn without colliding with explicit IDs. --- .../google_genai/adapters/transformation.py | 169 +++++++-- .../google_genai/test_google_genai_adapter.py | 342 +++++++++++++++++- 2 files changed, 486 insertions(+), 25 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index c1895376675..a8d8d035d71 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,11 +1,11 @@ import json -from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence -from itertools import chain +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from itertools import chain, count, groupby from types import MappingProxyType from typing import Any, Final, TypeAlias, TypeVar, cast from pydantic import JsonValue, TypeAdapter, ValidationError -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NamedTuple, NotRequired, ReadOnly, Required, TypedDict from litellm import verbose_logger from litellm.exceptions import BadRequestError @@ -49,27 +49,42 @@ _JsonDictList: TypeAlias = list[_JsonDict] class _ToolCallAccumulator(TypedDict): + id: NotRequired[ReadOnly[str | None]] name: ReadOnly[str] arguments: ReadOnly[str] class _GenAIFunctionCall(TypedDict): + id: NotRequired[ReadOnly[str]] name: ReadOnly[str] args: ReadOnly[Mapping[str, object]] +class _PendingToolCall(NamedTuple): + call_id: str + explicit_id: str | None + default_id: str + + class _GenAIPart(TypedDict, total=False): text: ReadOnly[str] functionCall: ReadOnly[_GenAIFunctionCall] class _GenAIFunctionResponse(TypedDict, total=False): + id: ReadOnly[str | None] name: ReadOnly[str] response: ReadOnly[object] parts: ReadOnly[object] +class _ContentEntry(NamedTuple): + message: AllMessageValues + response: _GenAIFunctionResponse | None + + class _GenAIRequestFunctionCall(TypedDict, total=False): + id: ReadOnly[str | None] name: ReadOnly[str] args: ReadOnly[object] @@ -208,13 +223,30 @@ def _genai_message_parts(parts: object) -> tuple[_GenAIMessagePart, ...]: return tuple(converted for part in items if (converted := _genai_message_part(part)) is not None) -def _genai_tool_response(part: object) -> ChatCompletionToolMessage | None: +def _default_tool_call_id(name: str | None) -> str: + return f"call_{name or 'unknown'}" + + +def _genai_function_call_part(call_id: str | None, name: str, args: Mapping[str, object]) -> _GenAIPart: + # Gemini clients echo functionCall.id back as functionResponse.id so the next turn can pair each + # result with its call; dropping the provider's ID here would force pairing by function name. + function_call: Final[_GenAIFunctionCall] = ( + {"id": call_id, "name": name, "args": args} if call_id else {"name": name, "args": args} + ) + return {"functionCall": function_call} + + +def _genai_function_response(part: object) -> _GenAIFunctionResponse | None: fields: Final = _validated(_CONFIG_FIELDS, part) if fields is None: return None - response: Final = _validated(_GENAI_FUNCTION_RESPONSE, fields.get("functionResponse")) or _validated( + return _validated(_GENAI_FUNCTION_RESPONSE, fields.get("functionResponse")) or _validated( _GENAI_FUNCTION_RESPONSE, fields.get("function_response") ) + + +def _genai_tool_response(part: object) -> ChatCompletionToolMessage | None: + response: Final = _genai_function_response(part) if response is None: return None result: Final = json.dumps(response.get("response", {})) @@ -226,7 +258,7 @@ def _genai_tool_response(part: object) -> ChatCompletionToolMessage | None: media: Final = tuple(converted for fields in media_fields if (converted := _genai_media_part(fields)) is not None) return ChatCompletionToolMessage( role="tool", - tool_call_id=f"call_{response.get('name', 'unknown')}", + tool_call_id=response.get("id") or _default_tool_call_id(response.get("name")), content=[ChatCompletionTextObject(type="text", text=result), *media] if media else result, ) @@ -241,7 +273,7 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None: if call is None: return None return ChatCompletionAssistantToolCall( - id=f"call_{call.get('name', 'unknown')}", + id=call.get("id"), type="function", function=ChatCompletionToolCallFunctionChunk( name=call.get("name", ""), arguments=json.dumps(call.get("args", {})) @@ -249,6 +281,86 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None: ) +def _genai_content_entries(content: object) -> Iterator[_ContentEntry]: + # A converted tool message can no longer tell an explicit functionResponse.id from the call_ + # fallback, so the raw part is carried alongside it. _genai_content_messages emits tool messages in + # part order, which makes pairing by position safe. + fields: Final = _validated(_CONFIG_FIELDS, content) + parts: Final = _genai_content_parts(fields.get("parts") if fields is not None else None) + responses: Final = iter(resp for part in parts if (resp := _genai_function_response(part)) is not None) + return ( + _ContentEntry(message, next(responses, None) if message["role"] == "tool" else None) + for message in _genai_content_messages(content) + ) + + +def _reconcile_tool_call_ids(contents: Iterable[object]) -> Iterator[AllMessageValues]: + # Providers such as Anthropic and OpenAI reject a request whose tool call IDs repeat, but Gemini + # clients may omit functionCall.id / functionResponse.id or reuse the same ID in later turns. + used_ids: Final = set[str]() + pending: Final = list[_PendingToolCall]() + + def next_unique_id(base_id: str) -> str: + candidates: Final = chain((base_id,), (f"{base_id}_{n}" for n in count(2))) + unique_id: Final = next(candidate for candidate in candidates if candidate not in used_ids) + used_ids.add(unique_id) + return unique_id + + def register_call(call: ChatCompletionAssistantToolCall) -> ChatCompletionAssistantToolCall: + default_id: Final = _default_tool_call_id(call["function"].get("name")) + explicit_id: Final = call["id"] or None + call_id: Final = next_unique_id(explicit_id or default_id) + pending.append(_PendingToolCall(call_id=call_id, explicit_id=explicit_id, default_id=default_id)) + return ChatCompletionAssistantToolCall(id=call_id, type=call["type"], function=call["function"]) + + def pop_pending(*predicates: Callable[[_PendingToolCall], bool]) -> str | None: + for predicate in predicates: + for index, item in enumerate(pending): + if predicate(item): + return pending.pop(index).call_id + return None + + def resolve_response_id(response: _GenAIFunctionResponse) -> str: + explicit_id: Final = response.get("id") or None + default_id: Final = _default_tool_call_id(response.get("name")) + if explicit_id is not None: + # A client ID can coincide with a suffix generated for another call, so the client's own call wins. + return ( + pop_pending(lambda item: item.explicit_id == explicit_id, lambda item: item.call_id == explicit_id) + or explicit_id + ) + # ID-less calls go first so that calls with an ID stay available to the responses echoing it. + return ( + pop_pending( + lambda item: item.explicit_id is None and item.default_id == default_id, + lambda item: item.default_id == default_id, + ) + or default_id + ) + + entries: Final = chain.from_iterable(_genai_content_entries(content) for content in contents) + # Consecutive model contents form one turn (e.g. streamed parallel calls recorded chunk by chunk), + # so unanswered calls are dropped only when a new run of assistant messages starts. + for role, run in groupby(entries, key=lambda entry: entry.message["role"]): + if role == "assistant": + pending.clear() + for message, response in run: + if message["role"] == "assistant" and (tool_calls := message.get("tool_calls")): + yield ChatCompletionAssistantMessage( + role="assistant", + content=message.get("content"), + tool_calls=[register_call(call) for call in tool_calls], + ) + elif message["role"] == "tool" and response is not None: + yield ChatCompletionToolMessage( + role="tool", + tool_call_id=resolve_response_id(response), + content=message["content"], + ) + else: + yield message + + def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]: fields: Final = _validated(_CONFIG_FIELDS, content) if fields is None: @@ -366,13 +478,13 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): parsed_args: Mapping[str, object] = self._parse_accumulated_args( tool_call_data["arguments"] or "{}" ) - function_call_part: _GenAIPart = { - "functionCall": { - "name": tool_call_data["name"] or "undefined_tool_name", - "args": parsed_args, - } - } - parts.append(function_call_part) + parts.append( + _genai_function_call_part( + tool_call_data.get("id"), + tool_call_data["name"] or "undefined_tool_name", + parsed_args, + ) + ) except json.JSONDecodeError: # This can happen if the stream is abruptly cut off mid-argument string. verbose_logger.warning( @@ -638,7 +750,7 @@ class GoogleGenAIAdapter: if system_parts and "text" in system_parts[0] else () ) - return [*system_messages, *chain.from_iterable(_genai_content_messages(content) for content in contents)] + return [*system_messages, *_reconcile_tool_call_ids(contents)] def translate_completion_to_generate_content( self, @@ -799,13 +911,13 @@ class GoogleGenAIAdapter: except json.JSONDecodeError: args = {} - function_call_part: _GenAIPart = { - "functionCall": { - "name": function.name or "undefined_tool_name", - "args": args, - } - } - parts.append(function_call_part) + parts.append( + _genai_function_call_part( + getattr(tool_call, "id", None), + function.name or "undefined_tool_name", + args, + ) + ) return parts if parts else [{"text": ""}] @@ -837,9 +949,12 @@ class GoogleGenAIAdapter: if tool_call_index is None: continue # Index is essential for tracking streaming tool calls + # Providers usually send the ID only on the first delta of each tool call + tool_call_id: str | None = getattr(tool_call, "id", None) # Initialize accumulator for this index if it's new if tool_call_index not in wrapper.accumulated_tool_calls: wrapper.accumulated_tool_calls[tool_call_index] = { + "id": tool_call_id, "name": "", "arguments": "", } @@ -856,6 +971,7 @@ class GoogleGenAIAdapter: previous_data: _ToolCallAccumulator = wrapper.accumulated_tool_calls[tool_call_index] wrapper.accumulated_tool_calls[tool_call_index] = _ToolCallAccumulator( + id=tool_call_id or previous_data.get("id"), name=function_name or previous_data["name"], arguments=previous_data["arguments"] + (args_chunk or ""), ) @@ -874,8 +990,13 @@ class GoogleGenAIAdapter: # The part will be created by a later chunk that brings the name. if accumulated_name: # If successful, create the part and clean up - function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}} - parts.append(function_call_part) + parts.append( + _genai_function_call_part( + accumulated_data.get("id"), + accumulated_name, + parsed_args, + ) + ) # Remove the completed tool call from the accumulator del wrapper.accumulated_tool_calls[tool_call_index] diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 4ab642dc99d..33b2fb40dcb 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -5,14 +5,29 @@ Test to verify the Google GenAI generate_content adapter functionality import json import unittest +from collections.abc import AsyncIterator, Sequence +from itertools import chain from typing import Final import pytest import litellm -from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper from litellm.google_genai.main import agenerate_content from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + ChatCompletionMessageToolCall, + Choices, + Delta, + Function, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + Usage, +) @pytest.mark.parametrize(("data_key", "mime_key"), [("inline_data", "mime_type"), ("inlineData", "mimeType")]) @@ -1820,3 +1835,328 @@ async def test_generate_content_sends_response_schema_and_tool_parameters_to_the } assert response["candidates"][0]["content"]["parts"] == [{"text": '{"park_name": "EPCOT"}'}] assert "text" not in response + + +def _assistant_tool_call_ids(messages: Sequence[AllMessageValues]) -> tuple[tuple[str | None, ...], ...]: + return tuple( + tuple(call["id"] for call in msg.get("tool_calls") or ()) for msg in messages if msg["role"] == "assistant" + ) + + +def _tool_response_ids(messages: Sequence[AllMessageValues]) -> tuple[str, ...]: + return tuple(msg["tool_call_id"] for msg in messages if msg["role"] == "tool") + + +def _tool_call_stream_chunk( + call_id: str | None, + name: str | None, + arguments: str | None, + index: int, +) -> ModelResponseStream: + return ModelResponseStream( + id="stream-1", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=call_id, + type="function", + function=Function(name=name, arguments=arguments), + index=index, + ) + ], + ), + ) + ], + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion.chunk", + ) + + +@pytest.mark.parametrize( + ("call_key", "response_key"), + [("functionCall", "functionResponse"), ("function_call", "function_response")], +) +def test_native_function_calls_preserve_explicit_ids_and_disambiguate_duplicates_fifo( + call_key: str, + response_key: str, +) -> None: + contents: Final = [ + {"role": "user", "parts": [{"text": "Compare London, Paris, and Tokyo weather"}]}, + { + "role": "model", + "parts": [ + {call_key: {"name": "lookup", "args": {"city": "London"}}}, + {call_key: {"id": None, "name": "lookup", "args": {"city": "Paris"}}}, + {call_key: {"id": "toolu_1", "name": "forecast", "args": {"city": "London"}}}, + {call_key: {"id": "toolu_2", "name": "forecast", "args": {"city": "Paris"}}}, + ], + }, + { + "role": "user", + "parts": [ + {response_key: {"name": "lookup", "response": {"temp": 15}}}, + {response_key: {"id": None, "name": "lookup", "response": {"temp": 18}}}, + {response_key: {"id": "toolu_2", "name": "forecast", "response": {"rain": False}}}, + {response_key: {"id": "toolu_1", "name": "forecast", "response": {"rain": True}}}, + ], + }, + { + "role": "model", + "parts": [ + {call_key: {"name": "lookup", "args": {"city": "Tokyo"}}}, + {call_key: {"id": "toolu_tokyo", "name": "lookup", "args": {"city": "Tokyo"}}}, + {call_key: {"name": "forecast", "args": {"city": "Tokyo"}}}, + ], + }, + { + "role": "user", + "parts": [ + {response_key: {"name": "lookup", "response": {"temp": 22}}}, + {response_key: {"name": "lookup", "response": {"temp": 23}}}, + {response_key: {"name": "forecast", "response": {"rain": False}}}, + ], + }, + ] + messages: Final = GoogleGenAIAdapter()._transform_contents_to_messages(contents) + assert [msg["role"] for msg in messages] == [ + "user", + "assistant", + "tool", + "tool", + "tool", + "tool", + "assistant", + "tool", + "tool", + "tool", + ] + assert _assistant_tool_call_ids(messages) == ( + ("call_lookup", "call_lookup_2", "toolu_1", "toolu_2"), + ("call_lookup_3", "toolu_tokyo", "call_forecast"), + ) + assert _tool_response_ids(messages) == ( + "call_lookup", + "call_lookup_2", + "toolu_2", + "toolu_1", + "call_lookup_3", + "toolu_tokyo", + "call_forecast", + ) + + +def test_native_function_call_ids_resolve_collisions_between_generated_and_explicit_ids() -> None: + contents: Final = [ + { + "role": "model", + "parts": [ + {"functionCall": {"id": "call_lookup", "name": "lookup", "args": {"q": "explicit_base"}}}, + {"functionCall": {"name": "lookup", "args": {"q": "generated_first"}}}, + {"functionCall": {"id": "call_lookup_2", "name": "lookup", "args": {"q": "collides_with_generated"}}}, + {"functionCall": {"name": "lookup", "args": {"q": "generated_second"}}}, + ], + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "call_lookup_2", "name": "lookup", "response": {"hit": "collides"}}}, + {"functionResponse": {"id": "call_lookup_2", "name": "lookup", "response": {"hit": "echoed_gen"}}}, + {"functionResponse": {"name": "lookup", "response": {"hit": "omitted_for_generated_second"}}}, + {"functionResponse": {"id": "call_lookup", "name": "lookup", "response": {"hit": "explicit_base"}}}, + ], + }, + ] + messages: Final = GoogleGenAIAdapter()._transform_contents_to_messages(contents) + assert _assistant_tool_call_ids(messages) == (("call_lookup", "call_lookup_2", "call_lookup_2_2", "call_lookup_3"),) + assert _tool_response_ids(messages) == ( + "call_lookup_2_2", + "call_lookup_2", + "call_lookup_3", + "call_lookup", + ) + + +def test_native_duplicate_explicit_ids_and_missing_names_are_disambiguated() -> None: + contents: Final = [ + { + "role": "model", + "parts": [ + {"functionCall": {"id": "dup_id", "name": "search", "args": {"q": "first"}}}, + {"functionCall": {"id": "dup_id", "name": "search", "args": {"q": "second"}}}, + {"functionCall": {"args": {"q": "nameless"}}}, + ], + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "dup_id", "name": "search", "response": {"hit": "first"}}}, + {"functionResponse": {"id": "dup_id", "name": "search", "response": {"hit": "second"}}}, + {"functionResponse": {"response": {"hit": "nameless"}}}, + {"functionResponse": {"id": "unmatched_explicit", "name": "search", "response": {"hit": "orphan"}}}, + ], + }, + ] + messages: Final = GoogleGenAIAdapter()._transform_contents_to_messages(contents) + assert _assistant_tool_call_ids(messages) == (("dup_id", "dup_id_2", "call_unknown"),) + assert _tool_response_ids(messages) == ("dup_id", "dup_id_2", "call_unknown", "unmatched_explicit") + + +def test_native_abandoned_turn_reserves_ids_without_matching_later_responses() -> None: + contents: Final = [ + { + "role": "model", + "parts": [ + {"functionCall": {"id": "call_lookup_2", "name": "lookup", "args": {"q": "abandoned_explicit"}}}, + {"functionCall": {"name": "lookup", "args": {"q": "abandoned_omitted"}}}, + ], + }, + {"role": "user", "parts": [{"text": "Cancel those lookups and run a new one"}]}, + { + "role": "model", + "parts": [{"functionCall": {"name": "lookup", "args": {"q": "fresh"}}}], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "lookup", "response": {"hit": "fresh"}}}], + }, + ] + messages: Final = GoogleGenAIAdapter()._transform_contents_to_messages(contents) + assert _assistant_tool_call_ids(messages) == ( + ("call_lookup_2", "call_lookup"), + ("call_lookup_3",), + ) + assert _tool_response_ids(messages) == ("call_lookup_3",) + + +def test_native_function_calls_split_across_consecutive_model_contents_share_one_turn() -> None: + contents: Final = [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "lookup", "args": {"q": "answered"}}}, + {"functionCall": {"name": "lookup", "args": {"q": "abandoned"}}}, + ], + }, + {"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"hit": "answered"}}}]}, + {"role": "user", "parts": [{"text": "Never mind, compare Tokyo and Osaka instead"}]}, + {"role": "model", "parts": [{"text": "Checking both cities."}]}, + {"role": "model", "parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Tokyo"}}}]}, + {"role": "model", "parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Osaka"}}}]}, + { + "role": "user", + "parts": [ + {"functionResponse": {"name": "get_weather", "response": {"city": "Tokyo"}}}, + {"functionResponse": {"name": "get_weather", "response": {"city": "Osaka"}}}, + {"functionResponse": {"name": "lookup", "response": {"hit": "late"}}}, + ], + }, + ] + messages: Final = GoogleGenAIAdapter()._transform_contents_to_messages(contents) + assert _assistant_tool_call_ids(messages) == ( + ("call_lookup", "call_lookup_2"), + (), + ("call_get_weather",), + ("call_get_weather_2",), + ) + tool_results: Final = [(msg["tool_call_id"], msg["content"]) for msg in messages if msg["role"] == "tool"] + assert tool_results == [ + ("call_lookup", '{"hit": "answered"}'), + ("call_get_weather", '{"city": "Tokyo"}'), + ("call_get_weather_2", '{"city": "Osaka"}'), + ("call_lookup", '{"hit": "late"}'), + ] + + +def test_outbound_generate_content_preserves_and_roundtrips_tool_call_ids() -> None: + adapter: Final = GoogleGenAIAdapter() + completion_response: Final = ModelResponse( + id="resp-1", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="toolu_01abc", + type="function", + function=Function(name="get_weather", arguments='{"city": "Tokyo"}'), + ), + ChatCompletionMessageToolCall( + id="", + type="function", + function=Function(name="ping", arguments="{}"), + ), + ], + ), + ) + ], + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion", + usage=Usage(prompt_tokens=10, completion_tokens=12, total_tokens=22), + ) + outbound: Final = adapter.translate_completion_to_generate_content(completion_response) + expected_content: Final = { + "role": "model", + "parts": [ + {"functionCall": {"id": "toolu_01abc", "name": "get_weather", "args": {"city": "Tokyo"}}}, + {"functionCall": {"name": "ping", "args": {}}}, + ], + } + assert outbound["candidates"] == [ + { + "content": expected_content, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [], + } + ] + + roundtrip_messages: Final = adapter._transform_contents_to_messages( + [ + expected_content, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "toolu_01abc", "name": "get_weather", "response": {"temp": "20C"}}}, + {"functionResponse": {"name": "ping", "response": {"ok": True}}}, + ], + }, + ] + ) + assert _assistant_tool_call_ids(roundtrip_messages) == (("toolu_01abc", "call_ping"),) + assert _tool_response_ids(roundtrip_messages) == ("toolu_01abc", "call_ping") + + +@pytest.mark.asyncio +async def test_streaming_generate_content_accumulates_and_flushes_tool_call_ids() -> None: + stream_chunks: Final = ( + _tool_call_stream_chunk("toolu_stream_1", "read_file", '{"pa', 0), + _tool_call_stream_chunk(None, None, 'th": "a.py"}', 0), + _tool_call_stream_chunk(None, "no_id_tool", "{}", 1), + # Empty arguments stay buffered until stream end, exercising the __anext__ flush path. + _tool_call_stream_chunk("toolu_stream_2", None, "", 2), + _tool_call_stream_chunk(None, "list_dir", "", 2), + ) + + async def completion_stream() -> AsyncIterator[ModelResponseStream]: + for item in stream_chunks: + yield item + + events: Final = [event async for event in GoogleGenAIStreamWrapper(completion_stream=completion_stream())] + streamed_parts: Final = chain.from_iterable(event["candidates"][0]["content"]["parts"] for event in events) + streamed_calls: Final = [part["functionCall"] for part in streamed_parts if "functionCall" in part] + assert streamed_calls == [ + {"id": "toolu_stream_1", "name": "read_file", "args": {"path": "a.py"}}, + {"name": "no_id_tool", "args": {}}, + {"id": "toolu_stream_2", "name": "list_dir", "args": {}}, + ] From 770eadae4df97ee41bd7ffaa03a751dd67a6ca78 Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Mon, 2 Nov 2026 13:20:11 +0000 Subject: [PATCH 2/2] perf(google_genai): resolve tool call IDs in linear time Resume a per-base suffix counter instead of rescanning earlier suffixes for every repeated tool call, and index unanswered calls by ID and function name so each functionResponse is paired without scanning the whole turn. Pairing order and precedence are unchanged. --- .../google_genai/adapters/transformation.py | 95 ++++++++++++------- 1 file changed, 60 insertions(+), 35 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index a8d8d035d71..16d5e1be12f 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,5 +1,7 @@ import json +from collections import defaultdict, deque from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from enum import Enum from itertools import chain, count, groupby from types import MappingProxyType from typing import Any, Final, TypeAlias, TypeVar, cast @@ -60,12 +62,6 @@ class _GenAIFunctionCall(TypedDict): args: ReadOnly[Mapping[str, object]] -class _PendingToolCall(NamedTuple): - call_id: str - explicit_id: str | None - default_id: str - - class _GenAIPart(TypedDict, total=False): text: ReadOnly[str] functionCall: ReadOnly[_GenAIFunctionCall] @@ -294,14 +290,68 @@ def _genai_content_entries(content: object) -> Iterator[_ContentEntry]: ) +class _PendingQueueKind(Enum): + EXPLICIT = "explicit" + OMITTED = "omitted" + DEFAULT = "default" + + +class _PendingToolCalls: + # Each call is queued under its own ID (or as ID-less) and under its function name, so a response + # resolves in O(1) without scanning every unanswered call. Matching through one queue leaves a stale + # entry in the other; _unmatched lets _pop skip those lazily instead of removing them eagerly. + def __init__(self) -> None: + self._unmatched: Final = set[str]() + self._queues: Final = defaultdict[tuple[_PendingQueueKind, str], deque[str]](deque) + + def clear(self) -> None: + self._unmatched.clear() + self._queues.clear() + + def add(self, call_id: str, explicit_id: str | None, default_id: str) -> None: + self._unmatched.add(call_id) + self._queues[_PendingQueueKind.DEFAULT, default_id].append(call_id) + bucket: Final = _PendingQueueKind.EXPLICIT if explicit_id is not None else _PendingQueueKind.OMITTED + self._queues[bucket, explicit_id if explicit_id is not None else default_id].append(call_id) + + def _claim(self, call_id: str) -> str | None: + if call_id not in self._unmatched: + return None + self._unmatched.remove(call_id) + return call_id + + def _pop(self, kind: _PendingQueueKind, key: str) -> str | None: + queue: Final = self._queues.get((kind, key)) + while queue: + if (claimed := self._claim(queue.popleft())) is not None: + return claimed + return None + + def resolve(self, response: _GenAIFunctionResponse) -> str: + explicit_id: Final = response.get("id") or None + default_id: Final = _default_tool_call_id(response.get("name")) + if explicit_id is not None: + # A client ID can coincide with a suffix generated for another call, so the client's own call wins. + return self._pop(_PendingQueueKind.EXPLICIT, explicit_id) or self._claim(explicit_id) or explicit_id + # ID-less calls go first so that calls with an ID stay available to the responses echoing it. + return ( + self._pop(_PendingQueueKind.OMITTED, default_id) + or self._pop(_PendingQueueKind.DEFAULT, default_id) + or default_id + ) + + def _reconcile_tool_call_ids(contents: Iterable[object]) -> Iterator[AllMessageValues]: # Providers such as Anthropic and OpenAI reject a request whose tool call IDs repeat, but Gemini # clients may omit functionCall.id / functionResponse.id or reuse the same ID in later turns. used_ids: Final = set[str]() - pending: Final = list[_PendingToolCall]() + # used_ids only grows, so resuming each base ID's suffix counter never skips a free suffix and keeps + # repeated tool calls from rescanning every earlier suffix. + suffix_counters: Final = defaultdict[str, Iterator[int]](lambda: count(2)) + pending: Final = _PendingToolCalls() def next_unique_id(base_id: str) -> str: - candidates: Final = chain((base_id,), (f"{base_id}_{n}" for n in count(2))) + candidates: Final = (f"{base_id}_{n}" for n in suffix_counters[base_id]) if base_id in used_ids else (base_id,) unique_id: Final = next(candidate for candidate in candidates if candidate not in used_ids) used_ids.add(unique_id) return unique_id @@ -310,34 +360,9 @@ def _reconcile_tool_call_ids(contents: Iterable[object]) -> Iterator[AllMessageV default_id: Final = _default_tool_call_id(call["function"].get("name")) explicit_id: Final = call["id"] or None call_id: Final = next_unique_id(explicit_id or default_id) - pending.append(_PendingToolCall(call_id=call_id, explicit_id=explicit_id, default_id=default_id)) + pending.add(call_id=call_id, explicit_id=explicit_id, default_id=default_id) return ChatCompletionAssistantToolCall(id=call_id, type=call["type"], function=call["function"]) - def pop_pending(*predicates: Callable[[_PendingToolCall], bool]) -> str | None: - for predicate in predicates: - for index, item in enumerate(pending): - if predicate(item): - return pending.pop(index).call_id - return None - - def resolve_response_id(response: _GenAIFunctionResponse) -> str: - explicit_id: Final = response.get("id") or None - default_id: Final = _default_tool_call_id(response.get("name")) - if explicit_id is not None: - # A client ID can coincide with a suffix generated for another call, so the client's own call wins. - return ( - pop_pending(lambda item: item.explicit_id == explicit_id, lambda item: item.call_id == explicit_id) - or explicit_id - ) - # ID-less calls go first so that calls with an ID stay available to the responses echoing it. - return ( - pop_pending( - lambda item: item.explicit_id is None and item.default_id == default_id, - lambda item: item.default_id == default_id, - ) - or default_id - ) - entries: Final = chain.from_iterable(_genai_content_entries(content) for content in contents) # Consecutive model contents form one turn (e.g. streamed parallel calls recorded chunk by chunk), # so unanswered calls are dropped only when a new run of assistant messages starts. @@ -354,7 +379,7 @@ def _reconcile_tool_call_ids(contents: Iterable[object]) -> Iterator[AllMessageV elif message["role"] == "tool" and response is not None: yield ChatCompletionToolMessage( role="tool", - tool_call_id=resolve_response_id(response), + tool_call_id=pending.resolve(response), content=message["content"], ) else: