Fix ollama issues with ChatCompletionAssistantToolCall

PiperOrigin-RevId: 761732316
This commit is contained in:
Google Team Member 2025-05-21 17:26:31 -07:00 committed by Copybara-Service
parent 505d936007
commit 4542af5650
2 changed files with 9 additions and 9 deletions

View File

@ -30,6 +30,7 @@ from typing import Union
from google.genai import types
from litellm import acompletion
from litellm import ChatCompletionAssistantMessage
from litellm import ChatCompletionAssistantToolCall
from litellm import ChatCompletionDeveloperMessage
from litellm import ChatCompletionImageUrlObject
from litellm import ChatCompletionMessageToolCall
@ -180,12 +181,12 @@ def _content_to_message_param(
for part in content.parts:
if part.function_call:
tool_calls.append(
ChatCompletionMessageToolCall(
ChatCompletionAssistantToolCall(
type="function",
id=part.function_call.id,
function=Function(
name=part.function_call.name,
arguments=part.function_call.args,
arguments=json.dumps(part.function_call.args),
),
)
)

View File

@ -700,13 +700,12 @@ def test_content_to_message_param_function_call():
message = _content_to_message_param(content)
assert message["role"] == "assistant"
assert message["content"] == None
assert message["tool_calls"][0].type == "function"
assert message["tool_calls"][0].id == "test_tool_call_id"
assert message["tool_calls"][0].function.name == "test_function"
assert (
message["tool_calls"][0].function.arguments
== '{"test_arg": "test_value"}'
)
tool_call = message["tool_calls"][0]
assert tool_call["type"] == "function"
assert tool_call["id"] == "test_tool_call_id"
assert tool_call["function"]["name"] == "test_function"
assert tool_call["function"]["arguments"] == '{"test_arg": "test_value"}'
def test_message_to_generate_content_response_text():