From 57e613282ade7c37bc57545ecf6fb8f4bf221d13 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 23 Dec 2025 14:55:15 +0530 Subject: [PATCH 1/6] Re add thought signature as normal feat --- .../gemini/vertex_and_google_ai_studio_gemini.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index b1810b40cf9..a5cc3dca8c1 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1318,13 +1318,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tool_response_chunk["provider_specific_fields"] = { # type: ignore "thought_signature": thought_signature } - # Only embed in ID if preview features are enabled - if litellm.enable_preview_features: - _tool_response_chunk[ - "id" - ] = _encode_tool_call_id_with_signature( - _tool_response_chunk["id"] or "", thought_signature - ) + _tool_response_chunk[ + "id" + ] = _encode_tool_call_id_with_signature( + _tool_response_chunk["id"] or "", thought_signature + ) _tools.append(_tool_response_chunk) cumulative_tool_call_idx += 1 if len(_tools) == 0: From de99e30d0d619a4b03939219521a5f051ae49751 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 23 Dec 2025 14:56:00 +0530 Subject: [PATCH 2/6] Add: a new pre call hook for checking tool thought signature in id --- litellm/utils.py | 157 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) diff --git a/litellm/utils.py b/litellm/utils.py index 805fbafcfce..7a1b0a34361 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -568,6 +568,111 @@ def get_dynamic_callbacks( return returned_callbacks +def _is_gemini_model(model: Optional[str], custom_llm_provider: Optional[str]) -> bool: + """ + Check if the target model is a Gemini or Vertex AI Gemini model. + """ + if custom_llm_provider in ["gemini", "vertex_ai", "vertex_ai_beta"]: + # For vertex_ai, check if it's actually a Gemini model + if custom_llm_provider in ["vertex_ai", "vertex_ai_beta"]: + return model is not None and "gemini" in model.lower() + return True + + # Check if model name contains gemini + return model is not None and "gemini" in model.lower() + + +def _remove_thought_signature_from_id(tool_call_id: str, separator: str) -> str: + """ + Remove thought signature from a tool call ID. + """ + if separator in tool_call_id: + return tool_call_id.split(separator, 1)[0] + return tool_call_id + + +def _process_assistant_message_tool_calls( + msg_copy: dict, thought_signature_separator: str +) -> dict: + """ + Process assistant message to remove thought signatures from tool call IDs. + """ + role = msg_copy.get("role") + tool_calls = msg_copy.get("tool_calls") + + if role == "assistant" and isinstance(tool_calls, list): + new_tool_calls = [] + for tc in tool_calls: + # Handle both dict and Pydantic model tool calls + if hasattr(tc, "model_dump"): + # It's a Pydantic model, convert to dict + tc_dict = tc.model_dump() + elif isinstance(tc, dict): + tc_dict = tc.copy() + else: + new_tool_calls.append(tc) + continue + + # Remove thought signature from ID if present + if isinstance(tc_dict.get("id"), str): + if thought_signature_separator in tc_dict["id"]: + tc_dict["id"] = _remove_thought_signature_from_id( + tc_dict["id"], thought_signature_separator + ) + + new_tool_calls.append(tc_dict) + msg_copy["tool_calls"] = new_tool_calls + + return msg_copy + + +def _process_tool_message_id(msg_copy: dict, thought_signature_separator: str) -> dict: + """ + Process tool message to remove thought signature from tool_call_id. + """ + if msg_copy.get("role") == "tool" and isinstance( + msg_copy.get("tool_call_id"), str + ): + if thought_signature_separator in msg_copy["tool_call_id"]: + msg_copy["tool_call_id"] = _remove_thought_signature_from_id( + msg_copy["tool_call_id"], thought_signature_separator + ) + + return msg_copy + + +def _remove_thought_signatures_from_messages( + messages: List, thought_signature_separator: str +) -> List: + """ + Remove thought signatures from tool call IDs in all messages. + """ + processed_messages = [] + + for msg in messages: + # Handle Pydantic models (convert to dict) + if hasattr(msg, "model_dump"): + msg_dict = msg.model_dump() + elif isinstance(msg, dict): + msg_dict = msg.copy() + else: + # Unknown type, keep as is + processed_messages.append(msg) + continue + + # Process assistant messages with tool_calls + msg_dict = _process_assistant_message_tool_calls( + msg_dict, thought_signature_separator + ) + + # Process tool messages with tool_call_id + msg_dict = _process_tool_message_id(msg_dict, thought_signature_separator) + + processed_messages.append(msg_dict) + + return processed_messages + + def function_setup( # noqa: PLR0915 original_function: str, rules_obj, start_time, *args, **kwargs ): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc. @@ -779,6 +884,58 @@ def function_setup( # noqa: PLR0915 input=buffer.getvalue(), model=model, ) + + ### REMOVE THOUGHT SIGNATURES FROM TOOL CALL IDS FOR NON-GEMINI MODELS ### + # Gemini models embed thought signatures in tool call IDs. When sending + # messages with tool calls to non-Gemini providers, we need to remove these + # signatures to ensure compatibility. + if isinstance(messages, list) and len(messages) > 0: + try: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + from litellm.litellm_core_utils.prompt_templates.factory import ( + THOUGHT_SIGNATURE_SEPARATOR, + ) + + # Get custom_llm_provider to determine target provider + custom_llm_provider = kwargs.get("custom_llm_provider") + + # If custom_llm_provider not in kwargs, try to determine it from the model + if not custom_llm_provider and model: + try: + _, custom_llm_provider, _, _ = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + ) + except Exception: + # If we can't determine the provider, skip this processing + pass + + # Only process if target is NOT a Gemini model + if not _is_gemini_model(model, custom_llm_provider): + verbose_logger.debug( + f"Removing thought signatures from tool call IDs for non-Gemini model: {model}" + ) + + # Process messages to remove thought signatures + processed_messages = _remove_thought_signatures_from_messages( + messages, THOUGHT_SIGNATURE_SEPARATOR + ) + + # Update messages in kwargs or args + if "messages" in kwargs: + kwargs["messages"] = processed_messages + elif len(args) > 1: + args = list(args) + args[1] = processed_messages + args = tuple(args) + + except Exception as e: + # Log the error but don't fail the request + verbose_logger.warning( + f"Error removing thought signatures from tool call IDs: {str(e)}" + ) elif ( call_type == CallTypes.embedding.value or call_type == CallTypes.aembedding.value From 0addab4f63088592fc1b4f40088e33b3270533d6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 23 Dec 2025 14:56:42 +0530 Subject: [PATCH 3/6] Add test for removal of thought signature --- tests/local_testing/test_function_setup.py | 176 ++++++++++++++++++++- 1 file changed, 175 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_function_setup.py b/tests/local_testing/test_function_setup.py index 5cc3ce12304..23a82fd7a6e 100644 --- a/tests/local_testing/test_function_setup.py +++ b/tests/local_testing/test_function_setup.py @@ -9,9 +9,10 @@ import os, io sys.path.insert( 0, os.path.abspath("../..") -) # Adds the parent directory to the, system path +) # Adds the parent directory to the system path import pytest, uuid from litellm.utils import function_setup, Rules +from litellm.litellm_core_utils.prompt_templates.factory import THOUGHT_SIGNATURE_SEPARATOR from datetime import datetime @@ -31,3 +32,176 @@ def test_empty_content(): messages=[], litellm_call_id=str(uuid.uuid4()), ) + + +def test_thought_signature_removal_for_non_gemini(): + """ + Test that thought signatures are removed from tool call IDs when sending to non-Gemini models + """ + rules_obj = Rules() + + # Create messages with thought signatures (as would come from Gemini) + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": f"call_123{THOUGHT_SIGNATURE_SEPARATOR}sig1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "SF"}' + } + } + ] + }, + { + "role": "tool", + "tool_call_id": f"call_123{THOUGHT_SIGNATURE_SEPARATOR}sig1", + "content": "Sunny, 72°F" + } + ] + + # Call function_setup with OpenAI model (non-Gemini) + logging_obj, kwargs = function_setup( + original_function="acompletion", + rules_obj=rules_obj, + start_time=datetime.now(), + model="gpt-4", + messages=messages, + litellm_call_id=str(uuid.uuid4()), + custom_llm_provider="openai" + ) + + # Verify thought signatures were removed + processed_messages = kwargs["messages"] + assert processed_messages[1]["tool_calls"][0]["id"] == "call_123" + assert processed_messages[2]["tool_call_id"] == "call_123" + assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[1]["tool_calls"][0]["id"] + assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[2]["tool_call_id"] + + +def test_thought_signature_preserved_for_gemini(): + """ + Test that thought signatures are preserved when sending to Gemini models + """ + rules_obj = Rules() + + # Create messages with thought signatures + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": f"call_456{THOUGHT_SIGNATURE_SEPARATOR}sig2", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}' + } + } + ] + }, + { + "role": "tool", + "tool_call_id": f"call_456{THOUGHT_SIGNATURE_SEPARATOR}sig2", + "content": "Rainy, 65°F" + } + ] + + # Call function_setup with Gemini model + logging_obj, kwargs = function_setup( + original_function="acompletion", + rules_obj=rules_obj, + start_time=datetime.now(), + model="gemini-1.5-pro", + messages=messages, + litellm_call_id=str(uuid.uuid4()), + custom_llm_provider="vertex_ai" + ) + + # Verify thought signatures were preserved (messages should be unchanged) + processed_messages = kwargs["messages"] + assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[1]["tool_calls"][0]["id"] + assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[2]["tool_call_id"] + + +def test_thought_signature_removal_with_multiple_tool_calls(): + """ + Test that thought signatures are removed from multiple tool calls + """ + rules_obj = Rules() + + messages = [ + {"role": "user", "content": "Get weather and time"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": f"call_1{THOUGHT_SIGNATURE_SEPARATOR}sig1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"} + }, + { + "id": f"call_2{THOUGHT_SIGNATURE_SEPARATOR}sig2", + "type": "function", + "function": {"name": "get_time", "arguments": "{}"} + } + ] + }, + { + "role": "tool", + "tool_call_id": f"call_1{THOUGHT_SIGNATURE_SEPARATOR}sig1", + "content": "Sunny" + }, + { + "role": "tool", + "tool_call_id": f"call_2{THOUGHT_SIGNATURE_SEPARATOR}sig2", + "content": "3:00 PM" + } + ] + + logging_obj, kwargs = function_setup( + original_function="acompletion", + rules_obj=rules_obj, + start_time=datetime.now(), + model="claude-3-opus", + messages=messages, + litellm_call_id=str(uuid.uuid4()), + custom_llm_provider="anthropic" + ) + + processed_messages = kwargs["messages"] + + # Check all tool call IDs are cleaned + assert processed_messages[1]["tool_calls"][0]["id"] == "call_1" + assert processed_messages[1]["tool_calls"][1]["id"] == "call_2" + assert processed_messages[2]["tool_call_id"] == "call_1" + assert processed_messages[3]["tool_call_id"] == "call_2" + + +def test_messages_without_tool_calls_unchanged(): + """ + Test that messages without tool calls pass through unchanged + """ + rules_obj = Rules() + + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"} + ] + + logging_obj, kwargs = function_setup( + original_function="acompletion", + rules_obj=rules_obj, + start_time=datetime.now(), + model="gpt-4", + messages=messages, + litellm_call_id=str(uuid.uuid4()), + custom_llm_provider="openai" + ) + + # Messages should be unchanged + assert kwargs["messages"] == messages From 16985ab2f3ea021383cbee4e0e64ae8e9c3be2a1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 23 Dec 2025 17:13:48 +0530 Subject: [PATCH 4/6] Potential fix for code scanning alert no. 3942: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- litellm/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 7a1b0a34361..2ee2296ab6e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -915,7 +915,7 @@ def function_setup( # noqa: PLR0915 # Only process if target is NOT a Gemini model if not _is_gemini_model(model, custom_llm_provider): verbose_logger.debug( - f"Removing thought signatures from tool call IDs for non-Gemini model: {model}" + "Removing thought signatures from tool call IDs for non-Gemini model" ) # Process messages to remove thought signatures From dce95d3a03cedc562bce182343b9dfb818fb03d0 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 23 Dec 2025 17:16:11 +0530 Subject: [PATCH 5/6] fix mypy issue --- litellm/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 2ee2296ab6e..9216e1e53e0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -927,9 +927,9 @@ def function_setup( # noqa: PLR0915 if "messages" in kwargs: kwargs["messages"] = processed_messages elif len(args) > 1: - args = list(args) - args[1] = processed_messages - args = tuple(args) + args_list = list(args) + args_list[1] = processed_messages + args = tuple(args_list) except Exception as e: # Log the error but don't fail the request From b2b0604173c7e4709da7d3fdfd64abe7d8deaa61 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 23 Dec 2025 21:36:21 +0530 Subject: [PATCH 6/6] Fix : mypy error --- .../test_thought_signature_in_tool_call_id.py | 276 +++++++----------- 1 file changed, 105 insertions(+), 171 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py index 46bb8930a7a..5fe51ed23b9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py @@ -10,15 +10,16 @@ enable_preview_features=True to be enabled. """ import pytest + import litellm -from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, -) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, - convert_to_gemini_tool_call_invoke, _encode_tool_call_id_with_signature, _get_thought_signature_from_tool, + convert_to_gemini_tool_call_invoke, +) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, ) from litellm.types.llms.vertex_ai import HttpxPartType @@ -71,52 +72,36 @@ def test_tool_call_id_includes_signature_in_response(enable_preview_features): """Test that tool call IDs in responses include embedded thought signatures only when preview features are enabled""" test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - # Save original state - original_flag = litellm.enable_preview_features - litellm.enable_preview_features = enable_preview_features - - try: - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, - ) - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=False, + parts_with_signature = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, ) + ] - # Verify tool call exists - assert tools is not None - assert len(tools) == 1 - tool_call_id = tools[0]["id"] - - # Verify signature is always in provider_specific_fields - assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == test_signature + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_with_signature, + cumulative_tool_call_idx=0, + is_function_call=False, + ) - if enable_preview_features: - # When preview features enabled, signature should be embedded in ID - assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id - # Verify we can decode it using the factory function - tool_obj = {"id": tool_call_id, "type": "function"} - decoded_sig = _get_thought_signature_from_tool(tool_obj) - assert decoded_sig == test_signature - else: - # When preview features disabled, signature should NOT be embedded in ID - assert THOUGHT_SIGNATURE_SEPARATOR not in tool_call_id - # But we can still extract from provider_specific_fields - tool_obj = {"id": tool_call_id, "type": "function", "provider_specific_fields": {"thought_signature": test_signature}} - decoded_sig = _get_thought_signature_from_tool(tool_obj) - assert decoded_sig == test_signature - finally: - # Restore original state - litellm.enable_preview_features = original_flag + # Verify tool call exists + assert tools is not None + assert len(tools) == 1 + tool_call_id = tools[0]["id"] + + # Verify signature is always in provider_specific_fields + assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == test_signature + + # When preview features enabled, signature should be embedded in ID + assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id + # Verify we can decode it using the factory function + tool_obj = {"id": tool_call_id, "type": "function"} + decoded_sig = _get_thought_signature_from_tool(tool_obj) + assert decoded_sig == test_signature def test_get_thought_signature_backward_compatibility(): @@ -204,90 +189,57 @@ def test_openai_client_e2e_flow(enable_preview_features): """ test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - # Save original state - original_flag = litellm.enable_preview_features - litellm.enable_preview_features = enable_preview_features + # Step 1: Gemini returns function call with thought signature + gemini_parts = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, + ) + ] - try: - # Step 1: Gemini returns function call with thought signature - gemini_parts = [ - HttpxPartType( - functionCall={ + # Step 2: LiteLLM transforms to OpenAI format + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=gemini_parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + assert tools is not None + assert len(tools) == 1 + tool_call_id = tools[0]["id"] + + assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id + + # Step 3: OpenAI client sends back assistant message + # For the disabled case, we simulate that the client might have provider_specific_fields + # or we use the embedded ID if preview features were enabled + openai_assistant_message = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, # Preserved from response (with embedded signature) + "type": "function", + "function": { "name": "get_current_temperature", - "args": {"location": "Paris"}, + "arguments": '{"location": "Paris"}', }, - thoughtSignature=test_signature, - ) - ] - - # Step 2: LiteLLM transforms to OpenAI format - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=gemini_parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - assert tools is not None - assert len(tools) == 1 - tool_call_id = tools[0]["id"] - - if enable_preview_features: - # When preview features enabled, signature should be embedded in ID - assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id - else: - # When preview features disabled, signature should NOT be embedded in ID - assert THOUGHT_SIGNATURE_SEPARATOR not in tool_call_id - - # Step 3: OpenAI client sends back assistant message - # For the disabled case, we simulate that the client might have provider_specific_fields - # or we use the embedded ID if preview features were enabled - if enable_preview_features: - openai_assistant_message = { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": tool_call_id, # Preserved from response (with embedded signature) - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - } - ], - } - else: - # When preview features disabled, simulate that provider_specific_fields might be preserved - # (though in real OpenAI client usage, this might not happen) - # For this test, we'll use provider_specific_fields to show extraction still works - openai_assistant_message = { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": tool_call_id, # ID without embedded signature - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "provider_specific_fields": {"thought_signature": test_signature}, - } - ], } + ], + } + # Step 4: LiteLLM converts back to Gemini format, extracting signature + gemini_parts_converted = convert_to_gemini_tool_call_invoke( + openai_assistant_message + ) - # Step 4: LiteLLM converts back to Gemini format, extracting signature - gemini_parts_converted = convert_to_gemini_tool_call_invoke( - openai_assistant_message - ) + # Verify signature is preserved through the round trip + assert len(gemini_parts_converted) == 1 + assert "thoughtSignature" in gemini_parts_converted[0] + assert gemini_parts_converted[0]["thoughtSignature"] == test_signature - # Verify signature is preserved through the round trip - assert len(gemini_parts_converted) == 1 - assert "thoughtSignature" in gemini_parts_converted[0] - assert gemini_parts_converted[0]["thoughtSignature"] == test_signature - finally: - # Restore original state - litellm.enable_preview_features = original_flag @pytest.mark.parametrize("enable_preview_features", [True, False]) @@ -296,54 +248,36 @@ def test_parallel_tool_calls_with_signatures(enable_preview_features): signature1 = "signature_for_first_call" # Only first call has signature (Gemini behavior for parallel calls) - # Save original state - original_flag = litellm.enable_preview_features - litellm.enable_preview_features = enable_preview_features + gemini_parts = [ + HttpxPartType( + functionCall={"name": "get_temperature", "args": {"location": "Paris"}}, + thoughtSignature=signature1, + ), + HttpxPartType( + functionCall={"name": "get_temperature", "args": {"location": "London"}}, + # No signature for second parallel call + ), + ] - try: - gemini_parts = [ - HttpxPartType( - functionCall={"name": "get_temperature", "args": {"location": "Paris"}}, - thoughtSignature=signature1, - ), - HttpxPartType( - functionCall={"name": "get_temperature", "args": {"location": "London"}}, - # No signature for second parallel call - ), - ] + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=gemini_parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=gemini_parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) + assert tools is not None + assert len(tools) == 2 - assert tools is not None - assert len(tools) == 2 + # First tool call should have signature in provider_specific_fields + assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == signature1 + + # When preview features enabled, first tool call has signature in ID + assert THOUGHT_SIGNATURE_SEPARATOR in tools[0]["id"] + sig1 = _get_thought_signature_from_tool({"id": tools[0]["id"], "type": "function"}) + assert sig1 == signature1 - # First tool call should have signature in provider_specific_fields - assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == signature1 - - if enable_preview_features: - # When preview features enabled, first tool call has signature in ID - assert THOUGHT_SIGNATURE_SEPARATOR in tools[0]["id"] - sig1 = _get_thought_signature_from_tool({"id": tools[0]["id"], "type": "function"}) - assert sig1 == signature1 - else: - # When preview features disabled, signature should NOT be in ID - assert THOUGHT_SIGNATURE_SEPARATOR not in tools[0]["id"] - # But we can extract from provider_specific_fields - sig1 = _get_thought_signature_from_tool({ - "id": tools[0]["id"], - "type": "function", - "provider_specific_fields": {"thought_signature": signature1} - }) - assert sig1 == signature1 - # Second tool call has no signature in ID (regardless of flag) - assert THOUGHT_SIGNATURE_SEPARATOR not in tools[1]["id"] - sig2 = _get_thought_signature_from_tool({"id": tools[1]["id"], "type": "function"}) - assert sig2 is None - finally: - # Restore original state - litellm.enable_preview_features = original_flag + # Second tool call has no signature in ID (regardless of flag) + assert THOUGHT_SIGNATURE_SEPARATOR not in tools[1]["id"] + sig2 = _get_thought_signature_from_tool({"id": tools[1]["id"], "type": "function"}) + assert sig2 is None