From 770eadae4df97ee41bd7ffaa03a751dd67a6ca78 Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Mon, 2 Nov 2026 13:20:11 +0000 Subject: [PATCH] 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: