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]