fix(generative-ai): Add tests for function_calling code snippet (GoogleCloudPlatform#11019)

koverholt · gcf-owl-bot[bot] · web-flow · commit bb3b34f531e4 · 2023-12-12T22:02:58.000-06:00
* Add tests for function_calling.py * Add test file * Update test * Move main construct out of region tag * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Update google-cloud-aiplatform to 1.38.0 * Update sample region tag --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/generative_ai/function_calling.py b/generative_ai/function_calling.py
@@ -12,45 +12,48 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# [START aiplatform_function_calling]
+# [START aiplatform_gemini_function_calling]
 from vertexai.preview.generative_models import (
     FunctionDeclaration,
     GenerativeModel,
     Tool,
 )
 
-# Load the Vertex AI Gemini API to use function calling
-model = GenerativeModel("gemini-pro")
-
-# Specify a function declaration and parameters for an API request
-get_current_weather_func = FunctionDeclaration(
-    name="get_current_weather",
-    description="Get the current weather in a given location",
-    # Function parameters are specified in OpenAPI JSON schema format
-    parameters={
-        "type": "object",
-        "properties": {"location": {"type": "string", "description": "Location"}},
-    },
-)
 
-# Define a tool that includes the above get_current_weather_func
-weather_tool = Tool(
-    function_declarations=[get_current_weather_func],
-)
+def generate_function_call(prompt: str) -> str:
+    # Load the Vertex AI Gemini API to use function calling
+    model = GenerativeModel("gemini-pro")
 
-# Prompt to ask the model about weather, which will invoke the Tool
-prompt = "What is the weather like in Boston?"
+    # Specify a function declaration and parameters for an API request
+    get_current_weather_func = FunctionDeclaration(
+        name="get_current_weather",
+        description="Get the current weather in a given location",
+        # Function parameters are specified in OpenAPI JSON schema format
+        parameters={
+            "type": "object",
+            "properties": {"location": {"type": "string", "description": "Location"}},
+        },
+    )
+
+    # Define a tool that includes the above get_current_weather_func
+    weather_tool = Tool(
+        function_declarations=[get_current_weather_func],
+    )
+
+    # Prompt to ask the model about weather, which will invoke the Tool
+    prompt = prompt
+
+    # Instruct the model to generate content using the Tool that you just created:
+    response = model.generate_content(
+        prompt,
+        generation_config={"temperature": 0},
+        tools=[weather_tool],
+    )
+
+    return str(response)
 
-# Instruct the model to generate content using the Tool that you just created:
-response = model.generate_content(
-    prompt,
-    generation_config={"temperature": 0},
-    tools=[weather_tool],
-)
 
-# Print the entire response
-print(response)
+# [END aiplatform_gemini_function_calling]
 
-# Print the part of the response that contains info about the function call
-print(response.candidates[0].content.parts[0].function_call)
-# [END aiplatform_function_calling]
+if __name__ == "__main__":
+    generate_function_call("What is the weather like in Boston?")
diff --git a/generative_ai/function_calling_test.py b/generative_ai/function_calling_test.py
@@ -0,0 +1,42 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import backoff
+from google.api_core.exceptions import ResourceExhausted
+
+import function_calling
+
+
+_PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+_LOCATION = "us-central1"
+
+
+function_expected_responses = [
+    "function_call",
+    "get_current_weather",
+    "args",
+    "fields",
+    "location",
+    "Boston",
+]
+
+
+@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
+def test_interview() -> None:
+    content = function_calling.generate_function_call(
+        prompt="What is the weather like in Boston?"
+    )
+    assert all(x in content for x in function_expected_responses)
