@@ -2223,6 +2223,56 @@ def test_message_to_generate_content_response_tool_call():
22232223 assert response .content .parts [0 ].function_call .id == "test_tool_call_id"
22242224
22252225
2226+ def test_message_to_generate_content_response_tool_call_with_thought_signature ():
2227+ signature = b"gemini_signature"
2228+ encoded_signature = base64 .b64encode (signature ).decode ("utf-8" )
2229+ message = ChatCompletionAssistantMessage (
2230+ role = "assistant" ,
2231+ content = None ,
2232+ tool_calls = [
2233+ ChatCompletionMessageToolCall (
2234+ type = "function" ,
2235+ id = f"test_tool_call_id__thought__{ encoded_signature } " ,
2236+ function = Function (
2237+ name = "test_function" ,
2238+ arguments = '{"test_arg": "test_value"}' ,
2239+ ),
2240+ )
2241+ ],
2242+ )
2243+
2244+ response = _message_to_generate_content_response (message )
2245+ assert response .content .role == "model"
2246+ assert response .content .parts [0 ].function_call .name == "test_function"
2247+ assert response .content .parts [0 ].function_call .args == {
2248+ "test_arg" : "test_value"
2249+ }
2250+ assert response .content .parts [0 ].function_call .id == "test_tool_call_id"
2251+ assert response .content .parts [0 ].thought_signature == signature
2252+
2253+
2254+ @pytest .mark .asyncio
2255+ async def test_content_to_message_param_embeds_thought_signature_in_tool_call ():
2256+ part = types .Part .from_function_call (
2257+ name = "test_function" ,
2258+ args = {"test_arg" : "test_value" },
2259+ )
2260+ part .function_call .id = "test_tool_call_id"
2261+ part .thought_signature = b"gemini_signature"
2262+ content = types .Content (role = "model" , parts = [part ])
2263+
2264+ message = await _content_to_message_param (content )
2265+
2266+ tool_calls = message ["tool_calls" ]
2267+ assert tool_calls is not None
2268+ assert len (tool_calls ) == 1
2269+ assert (
2270+ tool_calls [0 ]["id" ]
2271+ == "test_tool_call_id__thought__"
2272+ + base64 .b64encode (b"gemini_signature" ).decode ("utf-8" )
2273+ )
2274+
2275+
22262276def test_message_to_generate_content_response_inline_tool_call_text ():
22272277 message = ChatCompletionAssistantMessage (
22282278 role = "assistant" ,
0 commit comments