feat: GenAI SDK client - Support build options in Agent Engine GCS Deployment.

PiperOrigin-RevId: 833444035

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1062,6 +1062,28 @@ def test_create_agent_engine_config_with_source_packages(
                 == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
             )
 
+    @mock.patch.object(_agent_engines_utils, "_prepare")
+    @mock.patch.object(_agent_engines_utils, "_validate_extra_packages_or_raise")
+    def test_create_agent_engine_config_with_build_options(
+        self, mock_validate_extra_packages, mock_prepare
+    ):
+        build_options = {"installation_scripts": ["install.sh"]}
+        extra_packages = ["install.sh"]
+
+        self.client.agent_engines._create_config(
+            mode="create",
+            agent=self.test_agent,
+            staging_bucket=_TEST_STAGING_BUCKET,
+            display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+            extra_packages=extra_packages,
+            build_options=build_options,
+        )
+
+        mock_validate_extra_packages.assert_called_once_with(
+            extra_packages=extra_packages,
+            build_options=build_options,
+        )
+
     @mock.patch.object(_agent_engines_utils, "_prepare")
     def test_update_agent_engine_config_full(self, mock_prepare):
         config = self.client.agent_engines._create_config(
@@ -1598,6 +1620,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 requirements_file=None,
                 agent_framework=None,
                 python_version=None,
+                build_options=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1691,6 +1714,7 @@ def test_create_agent_engine_with_custom_service_account(
                 requirements_file=None,
                 agent_framework=None,
                 python_version=None,
+                build_options=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1783,6 +1807,7 @@ def test_create_agent_engine_with_experimental_mode(
                 requirements_file=None,
                 agent_framework=None,
                 python_version=None,
+                build_options=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1938,6 +1963,7 @@ def test_create_agent_engine_with_class_methods(
                 requirements_file=None,
                 agent_framework=None,
                 python_version=None,
+                build_options=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2025,6 +2051,7 @@ def test_create_agent_engine_with_agent_framework(
                 agent_framework=_TEST_AGENT_FRAMEWORK,
                 identity_type=None,
                 python_version=None,
+                build_options=None,
             )
             request_mock.assert_called_with(
                 "post",

vertexai/_genai/agent_engines.py

@@ -938,6 +938,7 @@ def create(
             requirements_file=config.requirements_file,
             agent_framework=config.agent_framework,
             python_version=config.python_version,
+            build_options=config.build_options,
         )
         operation = self._create(config=api_config)
         # TODO: Use a more specific link.
@@ -1004,6 +1005,7 @@ def _create_config(
         requirements_file: Optional[str] = None,
         agent_framework: Optional[str] = None,
         python_version: Optional[str] = None,
+        build_options: Optional[dict[str, list[str]]] = None,
     ) -> types.UpdateAgentEngineConfigDict:
         import sys
 
@@ -1067,6 +1069,7 @@ def _create_config(
             )
             extra_packages = _agent_engines_utils._validate_extra_packages_or_raise(
                 extra_packages=extra_packages,
+                build_options=build_options,
             )
             # Prepares the Agent Engine for creation/update in Vertex AI. This
             # involves packaging and uploading the artifacts for agent_engine,
@@ -1465,6 +1468,7 @@ def update(
             requirements_file=config.requirements_file,
             agent_framework=config.agent_framework,
             python_version=config.python_version,
+            build_options=config.build_options,
         )
         operation = self._update(name=name, config=api_config)
         logger.info(

vertexai/_genai/types/common.py

@@ -5471,6 +5471,17 @@ class CreateAgentEngineConfig(_common.BaseModel):
       Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
       """,
     )
+    build_options: Optional[dict[str, list[str]]] = Field(
+        default=None,
+        description="""The build options for the Agent Engine.
+      The following keys are supported:
+      - installation_scripts:
+          Optional. The paths to the installation scripts to be
+          executed in the Docker image.
+          The scripts must be located in the `installation_scripts`
+          subdirectory and the path must be added to `extra_packages`.
+      """,
+    )
 
 
 class CreateAgentEngineConfigDict(TypedDict, total=False):
@@ -5587,6 +5598,16 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
       Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
       """
 
+    build_options: Optional[dict[str, list[str]]]
+    """The build options for the Agent Engine.
+      The following keys are supported:
+      - installation_scripts:
+          Optional. The paths to the installation scripts to be
+          executed in the Docker image.
+          The scripts must be located in the `installation_scripts`
+          subdirectory and the path must be added to `extra_packages`.
+      """
+
 
 CreateAgentEngineConfigOrDict = Union[
     CreateAgentEngineConfig, CreateAgentEngineConfigDict
@@ -6210,6 +6231,17 @@ class UpdateAgentEngineConfig(_common.BaseModel):
       Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
       """,
     )
+    build_options: Optional[dict[str, list[str]]] = Field(
+        default=None,
+        description="""The build options for the Agent Engine.
+      The following keys are supported:
+      - installation_scripts:
+          Optional. The paths to the installation scripts to be
+          executed in the Docker image.
+          The scripts must be located in the `installation_scripts`
+          subdirectory and the path must be added to `extra_packages`.
+      """,
+    )
     update_mask: Optional[str] = Field(
         default=None,
         description="""The update mask to apply. For the `FieldMask` definition, see
@@ -6331,6 +6363,16 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
       Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
       """
 
+    build_options: Optional[dict[str, list[str]]]
+    """The build options for the Agent Engine.
+      The following keys are supported:
+      - installation_scripts:
+          Optional. The paths to the installation scripts to be
+          executed in the Docker image.
+          The scripts must be located in the `installation_scripts`
+          subdirectory and the path must be added to `extra_packages`.
+      """
+
     update_mask: Optional[str]
     """The update mask to apply. For the `FieldMask` definition, see
       https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
@@ -13251,6 +13293,17 @@ class AgentEngineConfig(_common.BaseModel):
       Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
       """,
     )
+    build_options: Optional[dict[str, list[str]]] = Field(
+        default=None,
+        description="""The build options for the Agent Engine.
+      The following keys are supported:
+      - installation_scripts:
+          Optional. The paths to the installation scripts to be
+          executed in the Docker image.
+          The scripts must be located in the `installation_scripts`
+          subdirectory and the path must be added to `extra_packages`.
+      """,
+    )
 
 
 class AgentEngineConfigDict(TypedDict, total=False):
@@ -13399,6 +13452,16 @@ class AgentEngineConfigDict(TypedDict, total=False):
       Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
       """
 
+    build_options: Optional[dict[str, list[str]]]
+    """The build options for the Agent Engine.
+      The following keys are supported:
+      - installation_scripts:
+          Optional. The paths to the installation scripts to be
+          executed in the Docker image.
+          The scripts must be located in the `installation_scripts`
+          subdirectory and the path must be added to `extra_packages`.
+      """
+
 
 AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]