From 4542af5650c13817cf85d080e2b280433c883cfa Mon Sep 17 00:00:00 2001
From: Google Team Member <genai-sdk-bot@google.com>
Date: Wed, 21 May 2025 17:26:31 -0700
Subject: [PATCH] Fix ollama issues with ChatCompletionAssistantToolCall

PiperOrigin-RevId: 761732316
---
 src/google/adk/models/lite_llm.py      |  5 +++--
 tests/unittests/models/test_litellm.py | 13 ++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py
index 8826432..07f27f2 100644
--- a/src/google/adk/models/lite_llm.py
+++ b/src/google/adk/models/lite_llm.py
@@ -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),
                 ),
             )
         )
diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py
index 9e74d2b..44a14e1 100644
--- a/tests/unittests/models/test_litellm.py
+++ b/tests/unittests/models/test_litellm.py
@@ -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():