diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index c705f93753..7b17825e2b 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/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", diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 23fcf2ecf9..44dad62c7b 100644 --- a/vertexai/_genai/agent_engines.py +++ b/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( diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index e1d4e73d23..8da56fb7fb 100644 --- a/vertexai/_genai/types/__init__.py +++ b/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", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 0903178153..9ca1e225d2 100644 --- a/vertexai/_genai/types/common.py +++ b/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]