From aa2195a995f67cffe63286aafe58c1bdefa0ee95 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:05:38 -0700 Subject: [PATCH] fix(gemini): keep non-chat routes tolerant of a malformed thought_signature_in_tool_call_id, match tool results by exact id first, and keep the deployment's value when drop_params drops a malformed override --- .../litellm_core_utils/get_litellm_params.py | 2 + .../prompt_templates/factory.py | 39 ++--- litellm/main.py | 13 +- litellm/router.py | 26 ++- litellm/types/router.py | 32 +++- .../router_code_coverage.py | 1 + .../test_thought_signature_in_tool_call_id.py | 158 ++++++++++++++++-- 7 files changed, 232 insertions(+), 39 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 8ff374bb55f..121df45d8b6 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -96,6 +96,8 @@ OPTIONAL_KWARGS_KEYS: Final = ( # Backward-compatible alias for existing imports/tests. _OPTIONAL_KWARGS_KEYS: Final = OPTIONAL_KWARGS_KEYS +SUPPLEMENTAL_KWARGS_KEYS: Final = OPTIONAL_KWARGS_KEYS | {"provider_specific_header", "extra_headers"} + _CONTROL_OPTIONS: Final = TypeAdapter(ControlOptions) _CONTROL_OPTION_NAMES: Final = tuple(field.name for field in fields(ControlOptions)) _MAX_SHOWN_INT_BITS: Final = 64 diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 88bdba6d909..2f971fcb789 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1266,11 +1266,20 @@ def encode_tool_call_id_with_signature(tool_call_id: str, thought_signature: str _encode_tool_call_id_with_signature = encode_tool_call_id_with_signature -def _same_tool_call(tool_call_id: str, other_tool_call_id: str) -> bool: - return ( - tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] - == other_tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] +def _tool_call_id_prefix(tool_call_id: object) -> str | None: + return tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] if isinstance(tool_call_id, str) else None + + +def _called_tool_name(tool_call_id: str, tool_calls: Sequence[Mapping[str, object]]) -> str | None: + exact_matches: Final = tuple(call for call in tool_calls if call.get("id") == tool_call_id) + prefix: Final = _tool_call_id_prefix(tool_call_id) + prefix_matches: Final = tuple( + call for call in tool_calls if prefix and _tool_call_id_prefix(call.get("id")) == prefix ) + matches: Final = exact_matches or (prefix_matches if len(prefix_matches) == 1 else ()) + function: Final = matches[0].get("function") if matches else None + name: Final = function.get("name") if isinstance(function, Mapping) else None + return name if isinstance(name, str) else None def get_thought_signature_from_tool(tool: Mapping[str, object]) -> str | None: @@ -1712,22 +1721,14 @@ def convert_to_gemini_tool_call_result( ) ) is not None: media_parts.append(media_part) - name: str | None = message.get("name", "") - # Recover name from last message with tool calls - if last_message_with_tool_calls: - tools: Final = last_message_with_tool_calls.get("tool_calls", []) - msg_tool_call_id: Final = message.get("tool_call_id", None) - for tool in tools: - prev_tool_call_id = tool.get("id", None) - if ( - isinstance(msg_tool_call_id, str) - and isinstance(prev_tool_call_id, str) - and msg_tool_call_id - and prev_tool_call_id - and _same_tool_call(msg_tool_call_id, prev_tool_call_id) - ): - name = tool.get("function", {}).get("name", "") + msg_tool_call_id: Final = message.get("tool_call_id", None) + called_tool_name: Final = ( + _called_tool_name(msg_tool_call_id, last_message_with_tool_calls.get("tool_calls") or ()) + if last_message_with_tool_calls and isinstance(msg_tool_call_id, str) and msg_tool_call_id + else None + ) + name: Final[str | None] = called_tool_name if called_tool_name is not None else message.get("name", "") # Echo the OpenAI tool_call_id on functionResponse (strip thought-signature suffix). gemini_call_id: str | None = None diff --git a/litellm/main.py b/litellm/main.py index 9321dd482e3..f1c6920fee3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -83,8 +83,8 @@ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, - OPTIONAL_KWARGS_KEYS, PROVIDER_AFFINITY_HEADER_KWARG_KEY, + SUPPLEMENTAL_KWARGS_KEYS, InvalidControlOption, parse_control_options, with_control_options, @@ -154,7 +154,7 @@ from litellm.types.completion import ( ) from litellm.types.litellm_params import ControlOptions, RetryStrategy from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import GenericLiteLLMParams +from litellm.types.router import STRICT_LITELLM_PARAMS, GenericLiteLLMParams from litellm.types.utils import ( BaseUrlRoutingHint, CustomPricingLiteLLMParams, @@ -398,13 +398,12 @@ class _SupplementalProviderParams: dropped_keys: frozenset[str] = frozenset() -_SUPPLEMENTAL_KWARGS_KEYS: Final = OPTIONAL_KWARGS_KEYS | {"provider_specific_header", "extra_headers"} - - def _supplemental_provider_params(kwargs: Mapping[str, object], model: str) -> _SupplementalProviderParams: - supplemental: Final = {key: kwargs[key] for key in _SUPPLEMENTAL_KWARGS_KEYS if key in kwargs} + supplemental: Final = {key: kwargs[key] for key in SUPPLEMENTAL_KWARGS_KEYS if key in kwargs} try: - return _SupplementalProviderParams(params=GenericLiteLLMParams.model_validate(supplemental)) + return _SupplementalProviderParams( + params=GenericLiteLLMParams.model_validate(supplemental, context=STRICT_LITELLM_PARAMS) + ) except ValidationError as error: invalid: Final = tuple(detail for detail in error.errors() if detail["loc"]) invalid_keys: Final = frozenset(str(detail["loc"][0]) for detail in invalid) diff --git a/litellm/router.py b/litellm/router.py index f0bb6c3a7a9..61b0d2c5f00 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -96,6 +96,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, inferred_provider, @@ -387,6 +388,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.router import ( CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, + STRICT_LITELLM_PARAMS, VALID_LITELLM_ENVIRONMENTS, AlertingConfig, AllowedFailsPolicy, @@ -666,6 +668,24 @@ def _drops_malformed_request_tools( return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools") +def _malformed_request_params_the_deployment_sets( + deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] +) -> frozenset[str]: + overrides: Final = { + key: request[key] + for key in SUPPLEMENTAL_KWARGS_KEYS + if key in request and deployment_params.get(key) is not None + } + try: + GenericLiteLLMParams.model_validate(overrides, context=STRICT_LITELLM_PARAMS) + except ValidationError as error: + set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} + policy: Final = drop_policy({**deployment_params, **set_defaults, **request}) + malformed: Final = (str(detail["loc"][0]) for detail in error.errors() if detail["loc"]) + return frozenset(key for key in malformed if policy.drops(key)) + return frozenset() + + def _merged_tools(deployment_tools: Sequence[object], request_tools: object) -> list[object] | None: if not isinstance(request_tools, (list, tuple)): return None @@ -4826,6 +4846,8 @@ class Router: if isinstance(raw_deployment_params, LiteLLM_Params) else raw_deployment_params ) + for key in _malformed_request_params_the_deployment_sets(deployment_context, defaults, request_context): + kwargs.pop(key) declared_provider: Final = deployment_context.get("custom_llm_provider") declared_model: Final = deployment_context.get("model") deployment_provider: Final = ( @@ -10462,7 +10484,9 @@ class Router: merged_params: Final[Mapping[str, object]] = ( _litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing}) ) - litellm_params: Final[LiteLLM_Params] = LiteLLM_Params.model_validate(dict(**merged_params)) + litellm_params: Final[LiteLLM_Params] = LiteLLM_Params.model_validate( + dict(**merged_params), context=STRICT_LITELLM_PARAMS + ) declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(_litellm_params) warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=declared_params) warn_on_ignored_model_capability_flags(model_name=_model_name, litellm_params=declared_params) diff --git a/litellm/types/router.py b/litellm/types/router.py index ddf15c79c1b..691bde090f1 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -6,6 +6,7 @@ import datetime import enum from collections.abc import Container, Mapping, Sequence from dataclasses import dataclass, field +from types import MappingProxyType from typing import ( TYPE_CHECKING, Annotated, @@ -22,7 +23,19 @@ from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import httpx from openai import Timeout as SDKTimeout -from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictInt, field_validator, model_validator +from pydantic import ( + BeforeValidator, + ConfigDict, + Field, + JsonValue, + StrictInt, + ValidationError, + ValidationInfo, + ValidatorFunctionWrapHandler, + WrapValidator, + field_validator, + model_validator, +) from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable from litellm._logging import verbose_logger @@ -517,6 +530,19 @@ def holds_secret_pointer(param_name: str) -> bool: _RESERVED_INIT_KEYS: Final = frozenset({"self", "params", "__class__"}) +STRICT_LITELLM_PARAMS: Final = MappingProxyType({"strict_litellm_params": True}) + + +def _none_when_malformed_outside_strict_validation( + value: object, handler: ValidatorFunctionWrapHandler, info: ValidationInfo +) -> object: + try: + return handler(value) + except ValidationError: + if info.context is STRICT_LITELLM_PARAMS: + raise + return None + class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): """ @@ -582,7 +608,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): ) model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: bool | None = False - thought_signature_in_tool_call_id: bool | None = None + thought_signature_in_tool_call_id: Annotated[ + bool | None, WrapValidator(_none_when_malformed_outside_strict_validation) + ] = None model_info: dict | None = None mock_response: str | ModelResponse | Exception | object | None = None diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index f4a295e697d..5f98d80ca63 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -122,6 +122,7 @@ ignored_function_names = [ "_is_fallback_hop", # Tested through the order fallback hop tests in test_router_order_fallback.py "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py "_merged_tools", # Tested through the _update_kwargs_with_deployment tools merge tests in test_router.py + "_malformed_request_params_the_deployment_sets", # Run by _update_kwargs_with_deployment, tested through the drop_params deployment test in test_thought_signature_in_tool_call_id.py "_drops_malformed_request_tools", # Tested through the _update_kwargs_with_deployment tools merge tests in test_router.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py diff --git a/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py index 77682722dfa..4f3b6fa4793 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py +++ b/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py @@ -11,7 +11,7 @@ enable_preview_features=True to be enabled. import json import re -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from itertools import chain from types import MappingProxyType from typing import Final @@ -28,9 +28,11 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( get_thought_signature_from_tool, convert_to_gemini_tool_call_invoke, ) +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.llms.vertex_ai import HttpxPartType @@ -300,13 +302,11 @@ def test_parallel_tool_calls_with_signatures(enable_preview_features): GEMINI_URL_PATTERN: Final = r"https://generativelanguage\.googleapis\.com/v1[a-z]*/models/gemini-3\.5-flash:.*" BOUNDED_TOOL_CALL_ID: Final = re.compile(r"[A-Za-z0-9_-]{1,64}") LONG_SIGNATURE: Final = "Ab+/" * 200 -WEATHER_TOOL: Final = { - "type": "function", - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, - }, +WEATHER_FUNCTION: Final = { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, } +WEATHER_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} USER_TURN: Final = {"role": "user", "content": "What is the weather in Paris?"} @@ -385,9 +385,10 @@ def _returned_tool_calls( api_key="test-gemini-key", **request_params, ) + assert isinstance(response, CustomStreamWrapper) == stream returned_calls: Final = ( chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in response) - if stream + if isinstance(response, CustomStreamWrapper) else response.choices[0].message.tool_calls ) tool_calls: Final = [call.model_dump() for call in returned_calls] @@ -512,6 +513,7 @@ async def _messages_tool_use_blocks(stream: bool) -> list[Mapping[str, object]]: ) if not stream: return [block for block in response["content"] if block["type"] == "tool_use"] + assert isinstance(response, AsyncIterator) relayed: Final = b"".join([chunk async for chunk in response]) events: Final = [json.loads(line.removeprefix(b"data: ")) for line in relayed.splitlines() if line.startswith(b"data:")] return [event["content_block"] for event in events if _is_tool_use_start(event)] @@ -529,7 +531,7 @@ async def test_messages_tool_use_keeps_thought_signature_kept_out_of_its_id( assert tool_uses[0]["provider_specific_fields"]["signature"] == LONG_SIGNATURE -RESPONSES_WEATHER_TOOL: Final = {"type": "function", **WEATHER_TOOL["function"]} +RESPONSES_WEATHER_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} SHORT_IDS: Final = MappingProxyType({"thought_signature_in_tool_call_id": False}) @@ -544,9 +546,10 @@ def _responses_function_call_items(stream: bool, request_params: Mapping[str, ob api_key="test-gemini-key", **request_params, ) + assert isinstance(response, Iterator) == stream items: Final = ( [event.item for event in response if event.type == "response.output_item.done"] - if stream + if isinstance(response, Iterator) and not isinstance(response, ResponsesAPIResponse) else response.output ) return [item for item in items if getattr(item, "type", None) == "function_call"] @@ -724,3 +727,138 @@ def test_deployment_setting_rejects_a_value_that_is_not_a_boolean(): } ] ) + + +def test_deployment_setting_survives_a_request_value_that_drop_params_drops(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gemini-agentcore", + "litellm_params": { + "model": "gemini/gemini-3.5-flash", + "api_key": "test-gemini-key", + "thought_signature_in_tool_call_id": False, + }, + } + ] + ) + with respx.mock: + respx.post(url__regex=GEMINI_URL_PATTERN).mock(return_value=_gemini_response(None, False)) + response: Final = router.completion( + model="gemini-agentcore", + messages=[USER_TURN], + tools=[WEATHER_TOOL], + thought_signature_in_tool_call_id="maybe", + drop_params=True, + ) + + assert _is_bounded(response.choices[0].message.tool_calls[0].id) + + +def test_responses_request_setting_that_is_not_a_boolean_answers_400(): + with pytest.raises(litellm.BadRequestError, match="thought_signature_in_tool_call_id") as raised: + _responses_function_call_items(False, {"thought_signature_in_tool_call_id": "maybe"}) + + assert raised.value.status_code == 400 + + +def test_responses_request_setting_that_is_not_a_boolean_is_dropped_under_drop_params(): + function_calls: Final = _responses_function_call_items( + False, {"thought_signature_in_tool_call_id": "maybe", "drop_params": True} + ) + + assert len(function_calls) == 1 + assert get_thought_signature_from_tool({"id": function_calls[0].call_id}) == LONG_SIGNATURE + + +@pytest.mark.parametrize("drop_params", [False, True]) +async def test_messages_request_setting_that_is_not_a_boolean_follows_drop_params( + drop_params: bool, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=False) as mocked: + mocked.post(url__regex=GEMINI_URL_PATTERN).mock(return_value=_gemini_response(None, False)) + request: Final = litellm.anthropic.messages.acreate( + max_tokens=256, + messages=[USER_TURN], + model="gemini/gemini-3.5-flash", + tools=[ANTHROPIC_WEATHER_TOOL], + api_key="test-gemini-key", + thought_signature_in_tool_call_id="maybe", + drop_params=drop_params, + ) + if not drop_params: + with pytest.raises(litellm.BadRequestError, match="thought_signature_in_tool_call_id"): + await request + return + response: Final = await request + + tool_uses: Final = [block for block in response["content"] if block["type"] == "tool_use"] + assert len(tool_uses) == 1 + + +def test_setting_that_is_not_a_boolean_is_ignored_where_nothing_reads_it(): + embedding_vector: Final = [0.25, 0.5] + with respx.mock: + route: Final = respx.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": embedding_vector}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, + ) + ) + response: Final = litellm.embedding( + model="openai/text-embedding-3-small", + input=["hi"], + api_key="test-openai-key", + thought_signature_in_tool_call_id="maybe", + ) + sent_body: Final = json.loads(route.calls.last.request.content) + + assert response.data[0]["embedding"] == embedding_vector + assert "thought_signature_in_tool_call_id" not in sent_body + + +TIME_TOOL: Final = { + "type": "function", + "function": {"name": "get_time", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, +} + + +@pytest.mark.parametrize( + "tool_call_ids", + [ + ("call_1__thought__S1", "call_1__thought__S2"), + ("__thought__S1", "__thought__S2"), + ], +) +def test_tool_results_keep_their_own_names_when_tool_call_ids_share_a_prefix(tool_call_ids: tuple[str, str]): + called_tools: Final = (("get_weather", tool_call_ids[0]), ("get_time", tool_call_ids[1])) + with respx.mock: + route: Final = respx.post(url__regex=GEMINI_URL_PATTERN).mock(return_value=_gemini_answer()) + litellm.completion( + model="gemini/gemini-3.5-flash", + messages=[ + USER_TURN, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": call_id, "type": "function", "function": {"name": name, "arguments": '{"city": "Paris"}'}} + for name, call_id in called_tools + ], + }, + *({"role": "tool", "tool_call_id": call_id, "content": f'{{"{name}": 1}}'} for name, call_id in called_tools), + ], + tools=[WEATHER_TOOL, TIME_TOOL], + api_key="test-gemini-key", + ) + sent_body: Final = json.loads(route.calls.last.request.content) + + parts: Final = chain.from_iterable(turn["parts"] for turn in sent_body["contents"]) + sent_names: Final = [part["function_response"]["name"] for part in parts if "function_response" in part] + assert sent_names == ["get_weather", "get_time"]