mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
Make thought sign in tool call id as a beta feat
This commit is contained in:
Sameer Kankute
parent
427074ac6e
commit
40c203e32b
2 changed files with 188 additions and 101 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue