feat: add build configuration support to Vertex AI Reasoning Engine

PiperOrigin-RevId: 963119503

Author: vertex-sdk-bot

agentplatform/_genai/agent_engines.py

@@ -2006,6 +2006,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(
@@ -2319,6 +2320,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
 
@@ -2525,6 +2527,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
 
@@ -2796,6 +2817,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(

agentplatform/_genai/types/__init__.py

@@ -1444,6 +1444,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
@@ -2588,6 +2591,9 @@
     "ReasoningEngineSpecContainerSpec",
     "ReasoningEngineSpecContainerSpecDict",
     "ReasoningEngineSpecContainerSpecOrDict",
+    "ReasoningEngineSpecBuildSpec",
+    "ReasoningEngineSpecBuildSpecDict",
+    "ReasoningEngineSpecBuildSpecOrDict",
     "ReasoningEngineSpec",
     "ReasoningEngineSpecDict",
     "ReasoningEngineSpecOrDict",

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]
 
@@ -26372,6 +26407,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):
@@ -26559,6 +26598,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]
 

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",

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(

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",