feat: GenAI - Support FunctionDeclaration.response schema

Ark-kun · copybara-github · commit 4288fec2b922 · 2024-11-26T20:25:45.000-08:00
PiperOrigin-RevId: 700543936
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -126,6 +126,27 @@
     "required": ["location"],
 }
 
+_REQUEST_FUNCTION_RESPONSE_SCHEMA_STRUCT = {
+    "type": "object",
+    "properties": {
+        "location": {
+            "type": "string",
+            "description": "The city and state, e.g. San Francisco, CA",
+        },
+        "unit": {
+            "type": "string",
+            "enum": [
+                "celsius",
+                "fahrenheit",
+            ],
+        },
+        "weather": {
+            "type": "string",
+        },
+    },
+}
+
+
 # Input and expected output schema for renaming tests.
 _RENAMING_INPUT_SCHEMA = {
     "type": "object",
@@ -1085,6 +1106,7 @@ def test_chat_function_calling(self, generative_models: generative_models):
             name="get_current_weather",
             description="Get the current weather in a given location",
             parameters=_REQUEST_FUNCTION_PARAMETER_SCHEMA_STRUCT,
+            response=_REQUEST_FUNCTION_RESPONSE_SCHEMA_STRUCT,
         )
         weather_tool = generative_models.Tool(
             function_declarations=[get_current_weather_func],
@@ -1180,6 +1202,7 @@ def test_chat_forced_function_calling(self, generative_models: generative_models
             name="get_current_weather",
             description="Get the current weather in a given location",
             parameters=_REQUEST_FUNCTION_PARAMETER_SCHEMA_STRUCT,
+            response=_REQUEST_FUNCTION_RESPONSE_SCHEMA_STRUCT,
         )
         weather_tool = generative_models.Tool(
             function_declarations=[get_current_weather_func],
@@ -1239,6 +1262,7 @@ def test_conversion_methods(self, generative_models: generative_models):
             name="get_current_weather",
             description="Get the current weather in a given location",
             parameters=_REQUEST_FUNCTION_PARAMETER_SCHEMA_STRUCT,
+            response=_REQUEST_FUNCTION_RESPONSE_SCHEMA_STRUCT,
         )
         weather_tool = generative_models.Tool(
             function_declarations=[get_current_weather_func],
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -2120,6 +2120,16 @@ class FunctionDeclaration:
                     "location"
                 ]
             },
+            # Optional:
+            response={
+                "type": "object",
+                "properties": {
+                    "weather": {
+                        "type": "string",
+                        "description": "The weather in the city"
+                    },
+                },
+            },
         )
         weather_tool = generative_models.Tool(
             function_declarations=[get_current_weather_func],
@@ -2160,6 +2170,7 @@ def __init__(
         name: str,
         parameters: Dict[str, Any],
         description: Optional[str] = None,
+        response: Dict[str, Any] = None,
     ):
         """Constructs a FunctionDeclaration.
 
@@ -2168,12 +2179,24 @@ def __init__(
             parameters: Describes the parameters to this function in JSON Schema Object format.
             description: Description and purpose of the function.
                 Model uses it to decide how and whether to call the function.
+            response: Describes the response type of this function in JSON Schema format.
         """
         parameters = copy.deepcopy(parameters)
         _fix_schema_dict_for_gapic_in_place(parameters)
         raw_schema = _dict_to_proto(aiplatform_types.Schema, parameters)
+
+        if response:
+            response = copy.deepcopy(response)
+            _fix_schema_dict_for_gapic_in_place(response)
+            raw_response_schema = _dict_to_proto(aiplatform_types.Schema, response)
+        else:
+            raw_response_schema = None
+
         self._raw_function_declaration = gapic_tool_types.FunctionDeclaration(
-            name=name, description=description, parameters=raw_schema
+            name=name,
+            description=description,
+            parameters=raw_schema,
+            response=raw_response_schema,
         )
 
     @classmethod
