feat: Support for explicitly providing class_methods in Agent Engine config.

PiperOrigin-RevId: 819409737

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -667,6 +667,16 @@ def register_operations(self) -> Dict[str, List[str]]:
 }
 
 
+_TEST_AGENT_ENGINE_CLASS_METHODS = [
+    {
+        "name": "query",
+        "description": "Simple query method",
+        "parameters": {"type": "object", "properties": {"input": {"type": "string"}}},
+        "api_mode": "",
+    }
+]
+
+
 def _create_empty_fake_package(package_name: str) -> str:
     """Creates a temporary directory structure representing an empty fake Python package.
 
@@ -1354,6 +1364,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 encryption_spec=None,
                 agent_server_mode=None,
                 labels=None,
+                class_methods=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1435,6 +1446,7 @@ def test_create_agent_engine_with_custom_service_account(
                 encryption_spec=None,
                 labels=None,
                 agent_server_mode=None,
+                class_methods=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1518,6 +1530,7 @@ def test_create_agent_engine_with_experimental_mode(
                 encryption_spec=None,
                 labels=None,
                 agent_server_mode=_genai_types.AgentServerMode.EXPERIMENTAL,
+                class_methods=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1540,6 +1553,85 @@ def test_create_agent_engine_with_experimental_mode(
                 None,
             )
 
+    @mock.patch.object(agent_engines.AgentEngines, "_create_config")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
+    def test_create_agent_engine_with_class_methods(
+        self,
+        mock_await_operation,
+        mock_create_config,
+    ):
+        mock_create_config.return_value = {
+            "display_name": _TEST_AGENT_ENGINE_DISPLAY_NAME,
+            "description": _TEST_AGENT_ENGINE_DESCRIPTION,
+            "spec": {
+                "package_spec": {
+                    "python_version": _TEST_PYTHON_VERSION,
+                    "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
+                    "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
+                },
+                "class_methods": _TEST_AGENT_ENGINE_CLASS_METHODS,
+            },
+        }
+        mock_await_operation.return_value = _genai_types.AgentEngineOperation(
+            response=_genai_types.ReasoningEngine(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                spec=_TEST_AGENT_ENGINE_SPEC,
+            )
+        )
+        with mock.patch.object(
+            self.client.agent_engines._api_client, "request"
+        ) as request_mock:
+            request_mock.return_value = genai_types.HttpResponse(body="")
+            self.client.agent_engines.create(
+                agent=self.test_agent,
+                config=_genai_types.AgentEngineConfig(
+                    display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                    requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+                    extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
+                    staging_bucket=_TEST_STAGING_BUCKET,
+                    class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+                ),
+            )
+            mock_create_config.assert_called_with(
+                mode="create",
+                agent=self.test_agent,
+                staging_bucket=_TEST_STAGING_BUCKET,
+                requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=None,
+                gcs_dir_name=None,
+                extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
+                env_vars=None,
+                service_account=None,
+                context_spec=None,
+                psc_interface_config=None,
+                min_instances=None,
+                max_instances=None,
+                resource_limits=None,
+                container_concurrency=None,
+                encryption_spec=None,
+                labels=None,
+                agent_server_mode=None,
+                class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+            )
+            request_mock.assert_called_with(
+                "post",
+                "reasoningEngines",
+                {
+                    "displayName": _TEST_AGENT_ENGINE_DISPLAY_NAME,
+                    "description": _TEST_AGENT_ENGINE_DESCRIPTION,
+                    "spec": {
+                        "class_methods": _TEST_AGENT_ENGINE_CLASS_METHODS,
+                        "package_spec": {
+                            "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
+                            "python_version": _TEST_PYTHON_VERSION,
+                            "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
+                        },
+                    },
+                },
+                None,
+            )
+
     @pytest.mark.usefixtures("caplog")
     @mock.patch.object(_agent_engines_utils, "_prepare")
     @mock.patch.object(_agent_engines_utils, "_await_operation")

vertexai/_genai/_agent_engines_utils.py

@@ -566,6 +566,13 @@ def _generate_class_methods_spec_or_raise(
     return class_methods_spec
 
 
+def _class_methods_to_class_methods_spec(
+    class_methods: List[dict[str, Any]],
+) -> List[proto.Message]:
+    """Converts a list of class methods to a list of ReasoningEngineSpec.ClassMethod messages."""
+    return [_to_proto(class_method) for class_method in class_methods]
+
+
 def _is_pydantic_serializable(param: inspect.Parameter) -> bool:
     """Checks if the parameter is pydantic serializable."""
 

vertexai/_genai/agent_engines.py

@@ -870,6 +870,7 @@ def create(
             encryption_spec=config.encryption_spec,
             agent_server_mode=config.agent_server_mode,
             labels=config.labels,
+            class_methods=config.class_methods,
         )
         operation = self._create(config=api_config)
         # TODO: Use a more specific link.
@@ -928,6 +929,7 @@ def _create_config(
         encryption_spec: Optional[genai_types.EncryptionSpecDict] = None,
         labels: Optional[dict[str, str]] = None,
         agent_server_mode: Optional[types.AgentServerMode] = None,
+        class_methods: Optional[Sequence[dict[str, Any]]] = None,
     ) -> types.UpdateAgentEngineConfigDict:
         import sys
 
@@ -975,9 +977,6 @@ def _create_config(
             extra_packages = _agent_engines_utils._validate_extra_packages_or_raise(
                 extra_packages=extra_packages,
             )
-            extra_packages = _agent_engines_utils._validate_extra_packages_or_raise(
-                extra_packages=extra_packages,
-            )
             # Prepares the Agent Engine for creation/update in Vertex AI. This
             # involves packaging and uploading the artifacts for agent_engine,
             # requirements and extra_packages to `staging_bucket/gcs_dir_name`.
@@ -1041,13 +1040,27 @@ def _create_config(
             if service_account is not None:
                 agent_engine_spec["service_account"] = service_account
                 update_masks.append("spec.service_account")
-            class_methods = _agent_engines_utils._generate_class_methods_spec_or_raise(
-                agent=agent,
-                operations=_agent_engines_utils._get_registered_operations(agent=agent),
-            )
+
+            update_masks.append("spec.class_methods")
+            class_methods_spec = []
+            if class_methods is not None:
+                class_methods_spec = (
+                    _agent_engines_utils._class_methods_to_class_methods_spec(
+                        class_methods=class_methods
+                    )
+                )
+            else:
+                class_methods_spec = (
+                    _agent_engines_utils._generate_class_methods_spec_or_raise(
+                        agent=agent,
+                        operations=_agent_engines_utils._get_registered_operations(
+                            agent=agent
+                        ),
+                    )
+                )
             agent_engine_spec["class_methods"] = [
-                _agent_engines_utils._to_dict(class_method)
-                for class_method in class_methods
+                _agent_engines_utils._to_dict(class_method_spec)
+                for class_method_spec in class_methods_spec
             ]
 
             if agent_server_mode:
@@ -1059,7 +1072,6 @@ def _create_config(
                     "agent_server_mode"
                 ] = agent_server_mode
 
-            update_masks.append("spec.class_methods")
             agent_engine_spec["agent_framework"] = (
                 _agent_engines_utils._get_agent_framework(agent=agent)
             )
@@ -1283,6 +1295,7 @@ def update(
             resource_limits=config.resource_limits,
             container_concurrency=config.container_concurrency,
             labels=config.labels,
+            class_methods=config.class_methods,
         )
         operation = self._update(name=name, config=api_config)
         logger.info(

vertexai/_genai/types.py

@@ -4630,6 +4630,15 @@ class CreateAgentEngineConfig(_common.BaseModel):
     labels: Optional[dict[str, str]] = Field(
         default=None, description="""The labels to be used for the Agent Engine."""
     )
+    class_methods: Optional[list[dict[str, Any]]] = Field(
+        default=None,
+        description="""The class methods to be used for the Agent Engine.
+      If specified, they'll override the class methods that are autogenerated by
+      default. By default, methods are generated by inspecting the agent object
+      and generating a corresponding method for each method defined on the
+      agent class.
+      """,
+    )
 
 
 class CreateAgentEngineConfigDict(TypedDict, total=False):
@@ -4687,6 +4696,14 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
     labels: Optional[dict[str, str]]
     """The labels to be used for the Agent Engine."""
 
+    class_methods: Optional[list[dict[str, Any]]]
+    """The class methods to be used for the Agent Engine.
+      If specified, they'll override the class methods that are autogenerated by
+      default. By default, methods are generated by inspecting the agent object
+      and generating a corresponding method for each method defined on the
+      agent class.
+      """
+
 
 CreateAgentEngineConfigOrDict = Union[
     CreateAgentEngineConfig, CreateAgentEngineConfigDict
@@ -5244,6 +5261,15 @@ class UpdateAgentEngineConfig(_common.BaseModel):
     labels: Optional[dict[str, str]] = Field(
         default=None, description="""The labels to be used for the Agent Engine."""
     )
+    class_methods: Optional[list[dict[str, Any]]] = Field(
+        default=None,
+        description="""The class methods to be used for the Agent Engine.
+      If specified, they'll override the class methods that are autogenerated by
+      default. By default, methods are generated by inspecting the agent object
+      and generating a corresponding method for each method defined on the
+      agent class.
+      """,
+    )
     update_mask: Optional[str] = Field(
         default=None,
         description="""The update mask to apply. For the `FieldMask` definition, see
@@ -5306,6 +5332,14 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
     labels: Optional[dict[str, str]]
     """The labels to be used for the Agent Engine."""
 
+    class_methods: Optional[list[dict[str, Any]]]
+    """The class methods to be used for the Agent Engine.
+      If specified, they'll override the class methods that are autogenerated by
+      default. By default, methods are generated by inspecting the agent object
+      and generating a corresponding method for each method defined on the
+      agent class.
+      """
+
     update_mask: Optional[str]
     """The update mask to apply. For the `FieldMask` definition, see
       https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
@@ -11806,6 +11840,15 @@ class AgentEngineConfig(_common.BaseModel):
     agent_server_mode: Optional[AgentServerMode] = Field(
         default=None, description="""The agent server mode to use for deployment."""
     )
+    class_methods: Optional[list[dict[str, Any]]] = Field(
+        default=None,
+        description="""The class methods to be used for the Agent Engine.
+      If specified, they'll override the class methods that are autogenerated by
+      default. By default, methods are generated by inspecting the agent object
+      and generating a corresponding method for each method defined on the
+      agent class.
+      """,
+    )
 
 
 class AgentEngineConfigDict(TypedDict, total=False):
@@ -11892,6 +11935,14 @@ class AgentEngineConfigDict(TypedDict, total=False):
     agent_server_mode: Optional[AgentServerMode]
     """The agent server mode to use for deployment."""
 
+    class_methods: Optional[list[dict[str, Any]]]
+    """The class methods to be used for the Agent Engine.
+      If specified, they'll override the class methods that are autogenerated by
+      default. By default, methods are generated by inspecting the agent object
+      and generating a corresponding method for each method defined on the
+      agent class.
+      """
+
 
 AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]
 

vertexai/agent_engines/_agent_engines.py

@@ -1903,3 +1903,10 @@ def _generate_class_methods_spec_or_raise(
             class_methods_spec.append(class_method)
 
     return class_methods_spec
+
+
+def _class_methods_to_class_methods_spec(
+    class_methods: List[dict[str, Any]],
+) -> List[proto.Message]:
+    """Converts a list of class methods to a list of ReasoningEngineSpec.ClassMethod messages."""
+    return [_utils.to_proto(class_method) for class_method in class_methods]