feat: add support for container_spec in AgentEngines

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1146,6 +1146,47 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises(
                 )
             assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
 
+    def test_create_agent_engine_config_with_container_spec(self):
+        container_spec = {"image_uri": "gcr.io/test-project/test-image"}
+        config = self.client.agent_engines._create_config(
+            mode="create",
+            display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+            description=_TEST_AGENT_ENGINE_DESCRIPTION,
+            container_spec=container_spec,
+            class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+            identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
+        )
+        assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
+        assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
+        assert config["spec"]["container_spec"] == container_spec
+        assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
+        assert (
+            config["spec"]["identity_type"]
+            == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+        )
+
+    def test_create_agent_engine_config_with_container_spec_and_others_raises(self):
+        container_spec = {"image_uri": "gcr.io/test-project/test-image"}
+        with pytest.raises(ValueError) as excinfo:
+            self.client.agent_engines._create_config(
+                mode="create",
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                container_spec=container_spec,
+                agent=self.test_agent,
+            )
+        assert "please do not specify `agent`" in str(excinfo.value)
+
+        with pytest.raises(ValueError) as excinfo:
+            self.client.agent_engines._create_config(
+                mode="create",
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                container_spec=container_spec,
+                source_packages=["."],
+            )
+        assert "please do not specify `source_packages`" in str(excinfo.value)
+
     @mock.patch.object(
         _agent_engines_utils,
         "_create_base64_encoded_tarball",
@@ -1916,6 +1957,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                container_spec=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2018,6 +2060,7 @@ def test_create_agent_engine_with_custom_service_account(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                container_spec=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2119,6 +2162,7 @@ def test_create_agent_engine_with_experimental_mode(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                container_spec=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2289,6 +2333,7 @@ def test_create_agent_engine_with_class_methods(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                container_spec=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2385,6 +2430,7 @@ def test_create_agent_engine_with_agent_framework(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                container_spec=None,
             )
             request_mock.assert_called_with(
                 "post",

vertexai/_genai/agent_engines.py

@@ -786,7 +786,12 @@ def _list_pager(
     def _is_lightweight_creation(
         self, agent: Any, config: types.AgentEngineConfig
     ) -> bool:
-        if agent or config.source_packages or config.developer_connect_source:
+        if (
+            agent
+            or config.source_packages
+            or config.developer_connect_source
+            or config.container_spec
+        ):
             return False
         return True
 
@@ -975,6 +980,7 @@ def create(
             python_version=config.python_version,
             build_options=config.build_options,
             image_spec=config.image_spec,
+            container_spec=config.container_spec,
         )
         operation = self._create(config=api_config)
         reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
@@ -1235,6 +1241,7 @@ def _create_config(
         image_spec: Optional[
             types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
         ] = None,
+        container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None,
     ) -> types.UpdateAgentEngineConfigDict:
         import sys
 
@@ -1289,6 +1296,19 @@ def _create_config(
                 "Please specify only one of `source_packages` or `developer_connect_source` in `config`."
             )
 
+        if container_spec:
+            if agent:
+                raise ValueError(
+                    "If you have provided `container_spec` in `config`, please "
+                    "do not specify `agent` in `agent_engines.create()` or "
+                    "`agent_engines.update()`."
+                )
+            if source_packages or developer_connect_source:
+                raise ValueError(
+                    "If you have provided `container_spec` in `config`, please "
+                    "do not specify `source_packages` or `developer_connect_source` in `config`."
+                )
+
         agent_engine_spec: Any = None
         if agent:
             agent_engine_spec = {}
@@ -1322,6 +1342,24 @@ def _create_config(
                 build_options=build_options,
                 image_spec=image_spec,
             )
+        elif container_spec:
+            agent_engine_spec = {}
+            if class_methods is None:
+                raise ValueError(
+                    "`class_methods` must be specified if `container_spec` is specified."
+                )
+            update_masks.append("spec.class_methods")
+            class_methods_spec_list = (
+                _agent_engines_utils._class_methods_to_class_methods_spec(
+                    class_methods=class_methods
+                )
+            )
+            agent_engine_spec["class_methods"] = [
+                _agent_engines_utils._to_dict(class_method_spec)
+                for class_method_spec in class_methods_spec_list
+            ]
+            update_masks.append("spec.container_spec")
+            agent_engine_spec["container_spec"] = container_spec
 
         is_deployment_spec_updated = (
             env_vars is not None
@@ -1335,8 +1373,9 @@ def _create_config(
             raise ValueError(
                 "To update `env_vars`, `psc_interface_config`, `min_instances`, "
                 "`max_instances`, `resource_limits`, or `container_concurrency`, "
-                "you must also provide the `agent` variable or the source code "
-                "options (`source_packages` or `developer_connect_source`)."
+                "you must also provide the `agent` variable, the source code "
+                "options (`source_packages` or `developer_connect_source`), "
+                "or the container spec (`container_spec`)."
             )
 
         if agent_engine_spec is not None:
@@ -1633,6 +1672,7 @@ def update(
             agent_framework=config.agent_framework,
             python_version=config.python_version,
             build_options=config.build_options,
+            container_spec=config.container_spec,
         )
         operation = self._update(name=name, config=api_config)
         reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(

vertexai/_genai/types/__init__.py

@@ -800,6 +800,9 @@
 from .common import ReasoningEngineDict
 from .common import ReasoningEngineOrDict
 from .common import ReasoningEngineSpec
+from .common import ReasoningEngineSpecContainerSpec
+from .common import ReasoningEngineSpecContainerSpecDict
+from .common import ReasoningEngineSpecContainerSpecOrDict
 from .common import ReasoningEngineSpecDeploymentSpec
 from .common import ReasoningEngineSpecDeploymentSpecDict
 from .common import ReasoningEngineSpecDeploymentSpecOrDict
@@ -1579,6 +1582,9 @@
     "VertexBaseConfig",
     "VertexBaseConfigDict",
     "VertexBaseConfigOrDict",
+    "ReasoningEngineSpecContainerSpec",
+    "ReasoningEngineSpecContainerSpecDict",
+    "ReasoningEngineSpecContainerSpecOrDict",
     "SecretRef",
     "SecretRefDict",
     "SecretRefOrDict",

vertexai/_genai/types/common.py

@@ -6030,6 +6030,27 @@ class _GetCustomJobParametersDict(TypedDict, total=False):
 ]
 
 
+class ReasoningEngineSpecContainerSpec(_common.BaseModel):
+    """Specification for deploying from a container image."""
+
+    image_uri: Optional[str] = Field(
+        default=None,
+        description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""",
+    )
+
+
+class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False):
+    """Specification for deploying from a container image."""
+
+    image_uri: Optional[str]
+    """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica."""
+
+
+ReasoningEngineSpecContainerSpecOrDict = Union[
+    ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict
+]
+
+
 class SecretRef(_common.BaseModel):
     """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
 
@@ -6412,6 +6433,10 @@ class ReasoningEngineSpec(_common.BaseModel):
         default=None,
         description="""Optional. Declarations for object class methods in OpenAPI specification format.""",
     )
+    container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field(
+        default=None,
+        description="""Deploy from a container image with a defined entrypoint and commands.""",
+    )
     deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field(
         default=None,
         description="""Optional. The specification of a Reasoning Engine deployment.""",
@@ -6450,6 +6475,9 @@ class ReasoningEngineSpecDict(TypedDict, total=False):
     class_methods: Optional[list[dict[str, Any]]]
     """Optional. Declarations for object class methods in OpenAPI specification format."""
 
+    container_spec: Optional[ReasoningEngineSpecContainerSpecDict]
+    """Deploy from a container image with a defined entrypoint and commands."""
+
     deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict]
     """Optional. The specification of a Reasoning Engine deployment."""
 
@@ -15291,6 +15319,9 @@ class AgentEngineConfig(_common.BaseModel):
     image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field(
         default=None, description="""The image spec for the Agent Engine."""
     )
+    container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field(
+        default=None, description="""The container spec for the Agent Engine."""
+    )
 
 
 class AgentEngineConfigDict(TypedDict, total=False):
@@ -15457,6 +15488,9 @@ class AgentEngineConfigDict(TypedDict, total=False):
     image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict]
     """The image spec for the Agent Engine."""
 
+    container_spec: Optional[ReasoningEngineSpecContainerSpecDict]
+    """The container spec for the Agent Engine."""
+
 
 AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]