chore: Update langchain agent system test to include tool config

PiperOrigin-RevId: 634821368
googleapis · May 17, 2024 · 932a245 · 932a245
1 parent f8e76a6
commit 932a245
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/tests/system/vertexai/test_reasoning_engines.py b/tests/system/vertexai/test_reasoning_engines.py
@@ -22,6 +22,7 @@
 import vertexai
 from tests.system.aiplatform import e2e_base
 from vertexai.preview import reasoning_engines
+from vertexai.preview.generative_models import ToolConfig
 
 
 _BLOB_FILENAME = vertexai.reasoning_engines._reasoning_engines._BLOB_FILENAME
@@ -48,7 +49,16 @@ def test_langchain_template(self, shared_state):
         )
         # Test prebuilt langchain_template
         created_app = reasoning_engines.ReasoningEngine.create(
-            reasoning_engines.LangchainAgent(model="gemini-1.0-pro"),
+            reasoning_engines.LangchainAgent(
+                model="gemini-1.5-pro-preview-0409",
+                model_tool_kwargs={
+                    "tool_config": {
+                        "function_calling_config": {
+                            "mode": ToolConfig.FunctionCallingConfig.Mode.AUTO,
+                        },
+                    },
+                },
+            ),
             requirements=["google-cloud-aiplatform[reasoningengine,langchain]"],
             display_name="test-display-name",
             description="test-description",