From 427bd75b6c5924738d8bca66dc214fa019d7eba9 Mon Sep 17 00:00:00 2001
From: Alexey Volkov <avolkov@google.com>
Date: Mon, 30 Sep 2024 20:49:37 -0700
Subject: [PATCH] chore: GenAI - `Part.function_call.args` is now a proper
 `dict`

Fixes https://github.com/googleapis/python-aiplatform/issues/4079

PiperOrigin-RevId: 680833836
---
 .../generative_models/_generative_models.py   | 44 +++++++++++++++----
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
index fc26c92ace..d969be6d44 100644
--- a/vertexai/generative_models/_generative_models.py
+++ b/vertexai/generative_models/_generative_models.py
@@ -2298,13 +2298,13 @@ def text(self) -> str:
             ) from e
 
     @property
-    def function_calls(self) -> Sequence[gapic_tool_types.FunctionCall]:
+    def function_calls(self) -> Sequence["FunctionCall"]:
         if not self.content or not self.content.parts:
             return []
         return [
             part.function_call
             for part in self.content.parts
-            if part and part.function_call
+            if part._raw_part._pb.WhichOneof("data") == "function_call"
         ]
 
 
@@ -2479,8 +2479,8 @@ def file_data(self) -> gapic_content_types.FileData:
         return self._raw_part.file_data
 
     @property
-    def function_call(self) -> gapic_tool_types.FunctionCall:
-        return self._raw_part.function_call
+    def function_call(self) -> "FunctionCall":
+        return FunctionCall._from_gapic(self._raw_part.function_call)
 
     @property
     def function_response(self) -> gapic_tool_types.FunctionResponse:
@@ -2491,6 +2491,35 @@ def _image(self) -> "Image":
         return Image.from_bytes(data=self._raw_part.inline_data.data)
 
 
+class FunctionCall:
+    """Function call."""
+
+    def __init__(self):
+        self._raw_message = aiplatform_types.FunctionCall()
+
+    @classmethod
+    def _from_gapic(cls, raw_message: aiplatform_types.FunctionCall) -> "FunctionCall":
+        response = cls()
+        response._raw_message = raw_message
+        return response
+
+    def to_dict(self) -> Dict[str, Any]:
+        return _proto_to_dict(self._raw_message)
+
+    def __repr__(self) -> str:
+        return self._raw_message.__repr__()
+
+    @property
+    def name(self) -> str:
+        return self._raw_message.name
+
+    @property
+    def args(self) -> Dict[str, Any]:
+        # We cannot use `type(self.args).to_dict(self.args)`
+        # due to: AttributeError: type object 'MapComposite' has no attribute 'to_dict'
+        return self.to_dict().get("args")
+
+
 class SafetySetting:
     """Parameters for the generation."""
 
@@ -2949,10 +2978,9 @@ def respond_to_model_response(
                     )
 
                 try:
-                    # We cannot use `function_args = type(function_call.args).to_dict(function_call.args)`
-                    # due to: AttributeError: type object 'MapComposite' has no attribute 'to_dict'
-                    function_args = type(function_call).to_dict(function_call)["args"]
-                    function_call_result = callable_function._function(**function_args)
+                    function_call_result = callable_function._function(
+                        **function_call.args
+                    )
                     if not isinstance(function_call_result, Mapping):
                         # If the function returns a single value, wrap it in the
                         # format that Part.from_function_response can accept.