Skip to content

Commit 4c0cc66

Browse files
committed
fix(litellm): preserve thought_signature across tool call roundtrip
1 parent 4010716 commit 4c0cc66

File tree

1 file changed

+50
-0
lines changed

1 file changed

+50
-0
lines changed

tests/unittests/models/test_litellm.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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+
22262276
def test_message_to_generate_content_response_inline_tool_call_text():
22272277
message = ChatCompletionAssistantMessage(
22282278
role="assistant",

0 commit comments

Comments
 (0)