From 95aa057a72234558d6e825d57e12ee79fdfaf893 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Aug 2026 19:34:53 -0700 Subject: [PATCH] feat: add build configuration support to Vertex AI Reasoning Engine PiperOrigin-RevId: 958722459 --- agentplatform/_genai/agent_engines.py | 22 +++++++ agentplatform/_genai/types/__init__.py | 6 ++ agentplatform/_genai/types/common.py | 42 +++++++++++++ .../agentplatform/genai/test_agent_engines.py | 60 +++++++++++++++++++ vertexai/_genai/agent_engines.py | 22 +++++++ vertexai/_genai/types/__init__.py | 6 ++ vertexai/_genai/types/common.py | 42 +++++++++++++ 7 files changed, 200 insertions(+) diff --git a/agentplatform/_genai/agent_engines.py b/agentplatform/_genai/agent_engines.py index 2fd26fda80..60ca138c06 100644 --- a/agentplatform/_genai/agent_engines.py +++ b/agentplatform/_genai/agent_engines.py @@ -1988,6 +1988,7 @@ def create( agent_config_source=agent_config_source, container_spec=config.container_spec, keep_alive_probe=keep_alive_probe, + build_config=config.build_config, ) operation = self._create(config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( @@ -2301,6 +2302,7 @@ def _create_config( container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, keep_alive_probe: Optional[dict[str, Any]] = None, traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, + build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = None, ) -> types.UpdateAgentEngineConfigDict: import sys @@ -2507,6 +2509,25 @@ def _create_config( agent_engine_spec["service_account"] = service_account update_masks.append("spec.service_account") + if build_config is not None: + if agent_engine_spec is None: + agent_engine_spec = {} + build_spec: dict[str, Any] = {} + if isinstance(build_config, dict): + worker_pool = build_config.get("worker_pool") + build_service_account = build_config.get("service_account") + else: + worker_pool = getattr(build_config, "worker_pool", None) + build_service_account = getattr(build_config, "service_account", None) + if worker_pool is not None: + build_spec["worker_pool"] = worker_pool + update_masks.append("spec.build_spec.worker_pool") + if build_service_account is not None: + build_spec["service_account"] = build_service_account + update_masks.append("spec.build_spec.service_account") + if build_spec: + agent_engine_spec["build_spec"] = build_spec + if agent_engine_spec is not None: config["spec"] = agent_engine_spec @@ -2778,6 +2799,7 @@ def update( container_spec=container_spec, keep_alive_probe=keep_alive_probe, traffic_config=traffic_config, + build_config=config.build_config, ) operation = self._update(name=name, config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 9e9278030a..cf0dae337e 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -1462,6 +1462,9 @@ from .common import ReasoningEngineRuntimeRevisionDict from .common import ReasoningEngineRuntimeRevisionOrDict from .common import ReasoningEngineSpec +from .common import ReasoningEngineSpecBuildSpec +from .common import ReasoningEngineSpecBuildSpecDict +from .common import ReasoningEngineSpecBuildSpecOrDict from .common import ReasoningEngineSpecContainerSpec from .common import ReasoningEngineSpecContainerSpecDict from .common import ReasoningEngineSpecContainerSpecOrDict @@ -2609,6 +2612,9 @@ "ReasoningEngineSpecContainerSpec", "ReasoningEngineSpecContainerSpecDict", "ReasoningEngineSpecContainerSpecOrDict", + "ReasoningEngineSpecBuildSpec", + "ReasoningEngineSpecBuildSpecDict", + "ReasoningEngineSpecBuildSpecOrDict", "ReasoningEngineSpec", "ReasoningEngineSpecDict", "ReasoningEngineSpecOrDict", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index b0343429c1..ff34bbfd82 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -8852,6 +8852,34 @@ class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): ] +class ReasoningEngineSpecBuildSpec(_common.BaseModel): + """Specification for building container image.""" + + service_account: Optional[str] = Field( + default=None, + description="""Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent.""", + ) + worker_pool: Optional[str] = Field( + default=None, + description="""Optional. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""", + ) + + +class ReasoningEngineSpecBuildSpecDict(TypedDict, total=False): + """Specification for building container image.""" + + service_account: Optional[str] + """Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent.""" + + worker_pool: Optional[str] + """Optional. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""" + + +ReasoningEngineSpecBuildSpecOrDict = Union[ + ReasoningEngineSpecBuildSpec, ReasoningEngineSpecBuildSpecDict +] + + class ReasoningEngineSpec(_common.BaseModel): """The specification of an agent engine.""" @@ -8895,6 +8923,10 @@ class ReasoningEngineSpec(_common.BaseModel): default=None, description="""Deploy from a container image with a defined entrypoint and commands.""", ) + build_spec: Optional[ReasoningEngineSpecBuildSpec] = Field( + default=None, + description="""Optional. Configuration for building container image.""", + ) class ReasoningEngineSpecDict(TypedDict, total=False): @@ -8930,6 +8962,9 @@ class ReasoningEngineSpecDict(TypedDict, total=False): container_spec: Optional[ReasoningEngineSpecContainerSpecDict] """Deploy from a container image with a defined entrypoint and commands.""" + build_spec: Optional[ReasoningEngineSpecBuildSpecDict] + """Optional. Configuration for building container image.""" + ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] @@ -26547,6 +26582,10 @@ class AgentEngineConfig(_common.BaseModel): traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( default=None, description="""The traffic config for the Agent Engine.""" ) + build_config: Optional[ReasoningEngineSpecBuildSpec] = Field( + default=None, + description="""The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""", + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -26734,6 +26773,9 @@ class AgentEngineConfigDict(TypedDict, total=False): traffic_config: Optional[ReasoningEngineTrafficConfigDict] """The traffic config for the Agent Engine.""" + build_config: Optional[ReasoningEngineSpecBuildSpecDict] + """The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index e92f121e8d..a37deb56fd 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -554,6 +554,16 @@ def register_operations(self) -> Dict[str, List[str]]: _genai_types.IdentityType.SERVICE_ACCOUNT ) _TEST_AGENT_ENGINE_ENCRYPTION_SPEC = {"kms_key_name": "test-kms-key"} +_TEST_AGENT_ENGINE_BUILD_WORKER_POOL = ( + "projects/test-project/locations/us-central1/workerPools/test-pool" +) +_TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT = ( + "test-build-sa@test-project.iam.gserviceaccount.com" +) +_TEST_AGENT_ENGINE_BUILD_CONFIG = _genai_types.ReasoningEngineSpecBuildSpecDict( + worker_pool=_TEST_AGENT_ENGINE_BUILD_WORKER_POOL, + service_account=_TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT, +) _TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE = { "http_get": { "path": "/health", @@ -1046,6 +1056,51 @@ def test_create_agent_engine_config_full(self, mock_prepare): == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT ) + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_agent_engine_config_with_build_config(self, mock_prepare): + config = self.client.agent_engines._create_config( + mode="create", + agent=self.test_agent, + staging_bucket=_TEST_STAGING_BUCKET, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + build_config=_TEST_AGENT_ENGINE_BUILD_CONFIG, + ) + assert config["spec"]["build_spec"] == { + "worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL, + "service_account": _TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT, + } + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_agent_engine_config_with_build_config_worker_pool_only( + self, mock_prepare + ): + config = self.client.agent_engines._create_config( + mode="create", + agent=self.test_agent, + staging_bucket=_TEST_STAGING_BUCKET, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + build_config={"worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL}, + ) + assert config["spec"]["build_spec"] == { + "worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL, + } + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_update_agent_engine_config_with_build_config(self, mock_prepare): + config = self.client.agent_engines._create_config( + mode="update", + build_config=_TEST_AGENT_ENGINE_BUILD_CONFIG, + ) + assert config["spec"]["build_spec"] == { + "worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL, + "service_account": _TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT, + } + update_mask = config["update_mask"].split(",") + assert "spec.build_spec.worker_pool" in update_mask + assert "spec.build_spec.service_account" in update_mask + @mock.patch.object( _agent_engines_utils, "_create_base64_encoded_tarball", @@ -2236,6 +2291,7 @@ def test_create_agent_engine_with_env_vars_dict( agent_config_source=None, container_spec=None, keep_alive_probe=None, + build_config=None, ) request_mock.assert_called_with( "post", @@ -2342,6 +2398,7 @@ def test_create_agent_engine_with_custom_service_account( agent_config_source=None, container_spec=None, keep_alive_probe=None, + build_config=None, ) request_mock.assert_called_with( "post", @@ -2447,6 +2504,7 @@ def test_create_agent_engine_with_experimental_mode( agent_config_source=None, container_spec=None, keep_alive_probe=None, + build_config=None, ) request_mock.assert_called_with( "post", @@ -2621,6 +2679,7 @@ def test_create_agent_engine_with_class_methods( agent_config_source=None, container_spec=None, keep_alive_probe=None, + build_config=None, ) request_mock.assert_called_with( "post", @@ -2721,6 +2780,7 @@ def test_create_agent_engine_with_agent_framework( agent_config_source=None, container_spec=None, keep_alive_probe=None, + build_config=None, ) request_mock.assert_called_with( "post", diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index f642e7da51..d93ebcdc47 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -2005,6 +2005,7 @@ def create( agent_config_source=agent_config_source, container_spec=config.container_spec, keep_alive_probe=keep_alive_probe, + build_config=config.build_config, ) operation = self._create(config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( @@ -2318,6 +2319,7 @@ def _create_config( container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, keep_alive_probe: Optional[dict[str, Any]] = None, traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, + build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = None, ) -> types.UpdateAgentEngineConfigDict: import sys @@ -2524,6 +2526,25 @@ def _create_config( agent_engine_spec["service_account"] = service_account update_masks.append("spec.service_account") + if build_config is not None: + if agent_engine_spec is None: + agent_engine_spec = {} + build_spec: dict[str, Any] = {} + if isinstance(build_config, dict): + worker_pool = build_config.get("worker_pool") + build_service_account = build_config.get("service_account") + else: + worker_pool = getattr(build_config, "worker_pool", None) + build_service_account = getattr(build_config, "service_account", None) + if worker_pool is not None: + build_spec["worker_pool"] = worker_pool + update_masks.append("spec.build_spec.worker_pool") + if build_service_account is not None: + build_spec["service_account"] = build_service_account + update_masks.append("spec.build_spec.service_account") + if build_spec: + agent_engine_spec["build_spec"] = build_spec + if agent_engine_spec is not None: config["spec"] = agent_engine_spec @@ -2795,6 +2816,7 @@ def update( container_spec=container_spec, keep_alive_probe=keep_alive_probe, traffic_config=traffic_config, + build_config=config.build_config, ) operation = self._update(name=name, config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index a75584f31b..99efe8771b 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -1031,6 +1031,9 @@ from .common import ReasoningEngineRuntimeRevisionDict from .common import ReasoningEngineRuntimeRevisionOrDict from .common import ReasoningEngineSpec +from .common import ReasoningEngineSpecBuildSpec +from .common import ReasoningEngineSpecBuildSpecDict +from .common import ReasoningEngineSpecBuildSpecOrDict from .common import ReasoningEngineSpecContainerSpec from .common import ReasoningEngineSpecContainerSpecDict from .common import ReasoningEngineSpecContainerSpecOrDict @@ -2082,6 +2085,9 @@ "ReasoningEngineSpecContainerSpec", "ReasoningEngineSpecContainerSpecDict", "ReasoningEngineSpecContainerSpecOrDict", + "ReasoningEngineSpecBuildSpec", + "ReasoningEngineSpecBuildSpecDict", + "ReasoningEngineSpecBuildSpecOrDict", "ReasoningEngineSpec", "ReasoningEngineSpecDict", "ReasoningEngineSpecOrDict", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 17d95148d1..5bb3b8811b 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -8414,6 +8414,34 @@ class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): ] +class ReasoningEngineSpecBuildSpec(_common.BaseModel): + """Specification for building container image.""" + + worker_pool: Optional[str] = Field( + default=None, + description="""Optional. Identifier. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""", + ) + service_account: Optional[str] = Field( + default=None, + description="""Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent. Format: `projects/{project}/serviceAccounts/{service_account}` or `{service_account}@{project}.iam.gserviceaccount.com`""", + ) + + +class ReasoningEngineSpecBuildSpecDict(TypedDict, total=False): + """Specification for building container image.""" + + worker_pool: Optional[str] + """Optional. Identifier. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""" + + service_account: Optional[str] + """Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent. Format: `projects/{project}/serviceAccounts/{service_account}` or `{service_account}@{project}.iam.gserviceaccount.com`""" + + +ReasoningEngineSpecBuildSpecOrDict = Union[ + ReasoningEngineSpecBuildSpec, ReasoningEngineSpecBuildSpecDict +] + + class ReasoningEngineSpec(_common.BaseModel): """The specification of an agent engine.""" @@ -8457,6 +8485,10 @@ class ReasoningEngineSpec(_common.BaseModel): default=None, description="""Deploy from a container image with a defined entrypoint and commands.""", ) + build_spec: Optional[ReasoningEngineSpecBuildSpec] = Field( + default=None, + description="""Optional. Configuration for building container image.""", + ) class ReasoningEngineSpecDict(TypedDict, total=False): @@ -8492,6 +8524,9 @@ class ReasoningEngineSpecDict(TypedDict, total=False): container_spec: Optional[ReasoningEngineSpecContainerSpecDict] """Deploy from a container image with a defined entrypoint and commands.""" + build_spec: Optional[ReasoningEngineSpecBuildSpecDict] + """Optional. Configuration for building container image.""" + ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] @@ -20373,6 +20408,10 @@ class AgentEngineConfig(_common.BaseModel): traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( default=None, description="""The traffic config for the Agent Engine.""" ) + build_config: Optional[ReasoningEngineSpecBuildSpec] = Field( + default=None, + description="""The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""", + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -20560,6 +20599,9 @@ class AgentEngineConfigDict(TypedDict, total=False): traffic_config: Optional[ReasoningEngineTrafficConfigDict] """The traffic config for the Agent Engine.""" + build_config: Optional[ReasoningEngineSpecBuildSpecDict] + """The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]