Skip to content

Commit 4542af5

Browse files
google-genai-botcopybara-github
authored andcommitted
Fix ollama issues with ChatCompletionAssistantToolCall
PiperOrigin-RevId: 761732316
1 parent 505d936 commit 4542af5

File tree

2 files changed

+9
-9
lines changed

2 files changed

+9
-9
lines changed

src/google/adk/models/lite_llm.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from google.genai import types
3131
from litellm import acompletion
3232
from litellm import ChatCompletionAssistantMessage
33+
from litellm import ChatCompletionAssistantToolCall
3334
from litellm import ChatCompletionDeveloperMessage
3435
from litellm import ChatCompletionImageUrlObject
3536
from litellm import ChatCompletionMessageToolCall
@@ -180,12 +181,12 @@ def _content_to_message_param(
180181
for part in content.parts:
181182
if part.function_call:
182183
tool_calls.append(
183-
ChatCompletionMessageToolCall(
184+
ChatCompletionAssistantToolCall(
184185
type="function",
185186
id=part.function_call.id,
186187
function=Function(
187188
name=part.function_call.name,
188-
arguments=part.function_call.args,
189+
arguments=json.dumps(part.function_call.args),
189190
),
190191
)
191192
)

tests/unittests/models/test_litellm.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -700,13 +700,12 @@ def test_content_to_message_param_function_call():
700700
message = _content_to_message_param(content)
701701
assert message["role"] == "assistant"
702702
assert message["content"] == None
703-
assert message["tool_calls"][0].type == "function"
704-
assert message["tool_calls"][0].id == "test_tool_call_id"
705-
assert message["tool_calls"][0].function.name == "test_function"
706-
assert (
707-
message["tool_calls"][0].function.arguments
708-
== '{"test_arg": "test_value"}'
709-
)
703+
704+
tool_call = message["tool_calls"][0]
705+
assert tool_call["type"] == "function"
706+
assert tool_call["id"] == "test_tool_call_id"
707+
assert tool_call["function"]["name"] == "test_function"
708+
assert tool_call["function"]["arguments"] == '{"test_arg": "test_value"}'
710709

711710

712711
def test_message_to_generate_content_response_text():

0 commit comments

Comments
 (0)