Merge pull request #18374 from BerriAI/litellm_add_gemini_thought_signature_support

Add gemini thought signature support via tool call id

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:

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(
+                            "Removing thought signatures from tool call IDs for non-Gemini 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 = list(args)
+                            args_list[1] = processed_messages
+                            args = tuple(args_list)
+
+                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

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

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