From 40c203e32b5991cefa14d9d7befaac6385ee0538 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 3 Dec 2025 11:04:53 +0530 Subject: [PATCH] Make thought sign in tool call id as a beta feat --- .../vertex_and_google_ai_studio_gemini.py | 12 +- .../test_thought_signature_in_tool_call_id.py | 277 ++++++++++++------ 2 files changed, 188 insertions(+), 101 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 5fef8c1ec49..1c2efe1eaa7 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 @@ -1205,14 +1205,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): } # Embed thought signature in ID for OpenAI client compatibility if thought_signature: - _tool_response_chunk[ - "id" - ] = _encode_tool_call_id_with_signature( - _tool_response_chunk["id"] or "", thought_signature - ) _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 + ) _tools.append(_tool_response_chunk) cumulative_tool_call_idx += 1 if len(_tools) == 0: 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 68c0f3bdbfc..46bb8930a7a 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 @@ -4,9 +4,13 @@ Tests for embedding thought signatures in tool call IDs for OpenAI client compat When using OpenAI clients (instead of LiteLLM SDK), provider_specific_fields are not preserved. This test suite validates that thought signatures can be embedded in tool call IDs and extracted when converting back to Gemini format. + +Note: Embedding signatures in tool call IDs is a beta feature that requires +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, ) @@ -62,36 +66,57 @@ def test_encode_tool_call_id_without_signature(): assert decoded_signature is None -def test_tool_call_id_includes_signature_in_response(): - """Test that tool call IDs in responses include embedded thought signatures""" +@pytest.mark.parametrize("enable_preview_features", [True, False]) +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" - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, + # 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, ) - ] - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=False, - ) + # 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 - # Verify tool call ID includes thought signature - assert tools is not None - assert len(tools) == 1 - tool_call_id = tools[0]["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 + 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 def test_get_thought_signature_backward_compatibility(): @@ -168,97 +193,157 @@ def test_convert_to_gemini_with_embedded_signature(): assert gemini_parts[0]["thoughtSignature"] == test_signature -def test_openai_client_e2e_flow(): +@pytest.mark.parametrize("enable_preview_features", [True, False]) +def test_openai_client_e2e_flow(enable_preview_features): """ End-to-end test simulating OpenAI client usage: 1. LiteLLM receives response from Gemini with thought signature - 2. LiteLLM embeds signature in tool call ID + 2. LiteLLM embeds signature in tool call ID (if preview features enabled) 3. OpenAI client sends message back with same tool call ID - 4. LiteLLM extracts signature from ID and sends to Gemini + 4. LiteLLM extracts signature from ID/provider_specific_fields and sends to Gemini """ test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - # Step 1: Gemini returns function call with thought signature - gemini_parts = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, - ) - ] + # Save original state + original_flag = litellm.enable_preview_features + litellm.enable_preview_features = enable_preview_features - # Step 2: LiteLLM transforms to OpenAI format with embedded signature - 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 (preserves tool_call_id) - openai_assistant_message = { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": tool_call_id, # Preserved from response - "type": "function", - "function": { + try: + # Step 1: Gemini returns function call with thought signature + gemini_parts = [ + HttpxPartType( + functionCall={ "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', + "args": {"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 -def test_parallel_tool_calls_with_signatures(): +@pytest.mark.parametrize("enable_preview_features", [True, False]) +def test_parallel_tool_calls_with_signatures(enable_preview_features): """Test that parallel tool calls preserve signatures correctly""" signature1 = "signature_for_first_call" # Only first call has signature (Gemini behavior for parallel calls) - 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 - ), - ] + # Save original state + original_flag = litellm.enable_preview_features + litellm.enable_preview_features = enable_preview_features - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=gemini_parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) + 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 + ), + ] - assert tools is not None - assert len(tools) == 2 + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=gemini_parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) - # 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 + assert tools is not None + assert len(tools) == 2 - # Second tool call has no signature in ID - 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 + # 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