From 28499a92466994232191e5aaf7745180abb4a640 Mon Sep 17 00:00:00 2001 From: Shawn Yang Date: Mon, 17 Nov 2025 12:16:51 -0800 Subject: [PATCH] feat: GenAI SDK client - Support `build options` in Agent Engine GCS Deployment. PiperOrigin-RevId: 833444035 --- .../unit/vertexai/genai/test_agent_engines.py | 27 ++++++++ vertexai/_genai/agent_engines.py | 4 ++ vertexai/_genai/types/common.py | 63 +++++++++++++++++++ 3 files changed, 94 insertions(+) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index afd9acadb4..f61f4bed84 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/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", diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index ed29cf5092..81629cb291 100644 --- a/vertexai/_genai/agent_engines.py +++ b/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( diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 1e9ec94c93..da5e9d47ff 100644 --- a/vertexai/_genai/types/common.py +++ b/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]