diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index 3c0637626c..f9f60233cb 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -1002,6 +1002,7 @@ def test_create_agent_engine_config_full(self, mock_prepare): container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY, encryption_spec=_TEST_AGENT_ENGINE_ENCRYPTION_SPEC, python_version=_TEST_PYTHON_VERSION_OVERRIDE, + dedicated_ingress_endpoint_enabled=True, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION @@ -1032,6 +1033,7 @@ def test_create_agent_engine_config_full(self, mock_prepare): "max_instances": _TEST_AGENT_ENGINE_MAX_INSTANCES, "resource_limits": _TEST_AGENT_ENGINE_RESOURCE_LIMITS, "container_concurrency": _TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY, + "dedicated_ingress_endpoint_enabled": True, } assert config["encryption_spec"] == _TEST_AGENT_ENGINE_ENCRYPTION_SPEC assert config["spec"]["class_methods"] == [_TEST_AGENT_ENGINE_CLASS_METHOD_1] @@ -1486,6 +1488,7 @@ def test_update_agent_engine_config_full(self, mock_prepare): service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, python_version=_TEST_PYTHON_VERSION_OVERRIDE, + dedicated_ingress_endpoint_enabled=True, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION @@ -1510,6 +1513,7 @@ def test_update_agent_engine_config_full(self, mock_prepare): }, }, ], + "dedicated_ingress_endpoint_enabled": True, } assert config["spec"]["class_methods"] == [_TEST_AGENT_ENGINE_CLASS_METHOD_1] assert ( @@ -1530,6 +1534,7 @@ def test_update_agent_engine_config_full(self, mock_prepare): "spec.class_methods", "spec.deployment_spec.env", "spec.deployment_spec.secret_env", + "spec.deployment_spec.dedicated_ingress_endpoint_enabled", "spec.agent_framework", "spec.identity_type", "spec.service_account", @@ -2190,6 +2195,7 @@ def test_create_agent_engine_with_env_vars_dict( image_spec=None, agent_config_source=None, container_spec=None, + dedicated_ingress_endpoint_enabled=None, keep_alive_probe=None, ) request_mock.assert_called_with( @@ -2213,6 +2219,116 @@ def test_create_agent_engine_with_env_vars_dict( None, ) + @mock.patch.object(agent_engines.AgentEngines, "_create_config") + @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) + def test_create_agent_engine_with_dedicated_ingress_endpoint( + self, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_create_config, + ): + mock_create_config.return_value = { + "display_name": _TEST_AGENT_ENGINE_DISPLAY_NAME, + "description": _TEST_AGENT_ENGINE_DESCRIPTION, + "spec": { + "package_spec": { + "python_version": _TEST_PYTHON_VERSION, + "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, + "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, + }, + "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], + "deployment_spec": { + "dedicated_ingress_endpoint_enabled": True, + }, + "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, + }, + } + mock_await_operation.return_value = _genai_types.AgentEngineOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.agent_engines._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.agent_engines.create( + agent=self.test_agent, + config=_genai_types.AgentEngineConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + staging_bucket=_TEST_STAGING_BUCKET, + dedicated_ingress_endpoint_enabled=True, + ), + ) + mock_create_config.assert_called_with( + mode="create", + agent=self.test_agent, + staging_bucket=_TEST_STAGING_BUCKET, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + description=None, + gcs_dir_name=None, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + env_vars=None, + service_account=None, + identity_type=None, + context_spec=None, + psc_interface_config=None, + agent_gateway_config=None, + min_instances=None, + max_instances=None, + resource_limits=None, + container_concurrency=None, + encryption_spec=None, + agent_server_mode=None, + labels=None, + class_methods=None, + source_packages=None, + developer_connect_source=None, + entrypoint_module=None, + entrypoint_object=None, + requirements_file=None, + agent_framework=None, + python_version=None, + build_options=None, + image_spec=None, + agent_config_source=None, + container_spec=None, + dedicated_ingress_endpoint_enabled=True, + keep_alive_probe=None, + ) + request_mock.assert_called_with( + "post", + "reasoningEngines", + { + "displayName": _TEST_AGENT_ENGINE_DISPLAY_NAME, + "description": _TEST_AGENT_ENGINE_DESCRIPTION, + "spec": { + "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, + "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], + "deployment_spec": { + "dedicated_ingress_endpoint_enabled": True, + }, + "package_spec": { + "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, + "python_version": _TEST_PYTHON_VERSION, + "requirements_gcs_uri": ( + _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI + ), + }, + }, + }, + None, + ) + @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") @mock.patch.object( @@ -2296,6 +2412,7 @@ def test_create_agent_engine_with_custom_service_account( image_spec=None, agent_config_source=None, container_spec=None, + dedicated_ingress_endpoint_enabled=None, keep_alive_probe=None, ) request_mock.assert_called_with( @@ -2401,6 +2518,7 @@ def test_create_agent_engine_with_experimental_mode( image_spec=None, agent_config_source=None, container_spec=None, + dedicated_ingress_endpoint_enabled=None, keep_alive_probe=None, ) request_mock.assert_called_with( @@ -2575,6 +2693,7 @@ def test_create_agent_engine_with_class_methods( image_spec=None, agent_config_source=None, container_spec=None, + dedicated_ingress_endpoint_enabled=None, keep_alive_probe=None, ) request_mock.assert_called_with( @@ -2675,6 +2794,7 @@ def test_create_agent_engine_with_agent_framework( image_spec=None, agent_config_source=None, container_spec=None, + dedicated_ingress_endpoint_enabled=None, keep_alive_probe=None, ) request_mock.assert_called_with( diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 4ffbcd9e5c..9a01e2a8c0 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, + dedicated_ingress_endpoint_enabled=config.dedicated_ingress_endpoint_enabled, ) operation = self._create(config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( @@ -2316,6 +2317,7 @@ def _create_config( types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict ] = None, container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, + dedicated_ingress_endpoint_enabled: Optional[bool] = None, keep_alive_probe: Optional[dict[str, Any]] = None, traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, ) -> types.UpdateAgentEngineConfigDict: @@ -2452,12 +2454,13 @@ def _create_config( or resource_limits is not None or container_concurrency is not None or keep_alive_probe is not None + or dedicated_ingress_endpoint_enabled is not None ) if agent_engine_spec is None and is_deployment_spec_updated: raise ValueError( "To update `env_vars`, `psc_interface_config`, `min_instances`, " - "`max_instances`, `resource_limits`, `container_concurrency`, or " - "`keep_alive_probe`, you must also provide the `agent` variable or " + "`max_instances`, `resource_limits`, `container_concurrency`, " + "`keep_alive_probe`, or `dedicated_ingress_endpoint_enabled`, you must also provide the `agent` variable or " "the source code options (`source_packages`, " "`developer_connect_source` or `agent_config_source`)." ) @@ -2476,6 +2479,7 @@ def _create_config( resource_limits=resource_limits, container_concurrency=container_concurrency, keep_alive_probe=keep_alive_probe, + dedicated_ingress_endpoint_enabled=dedicated_ingress_endpoint_enabled, ) update_masks.extend(deployment_update_masks) agent_engine_spec["deployment_spec"] = deployment_spec @@ -2544,6 +2548,7 @@ def _generate_deployment_spec_or_raise( resource_limits: Optional[dict[str, str]] = None, container_concurrency: Optional[int] = None, keep_alive_probe: Optional[dict[str, Any]] = None, + dedicated_ingress_endpoint_enabled: Optional[bool] = None, ) -> Tuple[dict[str, Any], Sequence[str]]: deployment_spec: dict[str, Any] = {} update_masks = [] @@ -2597,6 +2602,13 @@ def _generate_deployment_spec_or_raise( if keep_alive_probe is not None: deployment_spec["keep_alive_probe"] = keep_alive_probe update_masks.append("spec.deployment_spec.keep_alive_probe") + if dedicated_ingress_endpoint_enabled is not None: + deployment_spec["dedicated_ingress_endpoint_enabled"] = ( + dedicated_ingress_endpoint_enabled + ) + update_masks.append( + "spec.deployment_spec.dedicated_ingress_endpoint_enabled" + ) return deployment_spec, update_masks def _update_deployment_spec_with_env_vars_dict_or_raise( @@ -2794,6 +2806,7 @@ def update( agent_config_source=agent_config_source, container_spec=container_spec, keep_alive_probe=keep_alive_probe, + dedicated_ingress_endpoint_enabled=config.dedicated_ingress_endpoint_enabled, traffic_config=traffic_config, ) operation = self._update(name=name, config=api_config) diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index e33f01470d..52958c818d 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -7629,6 +7629,10 @@ class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): default=None, description="""Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""", ) + dedicated_ingress_endpoint_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""", + ) class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): @@ -7666,6 +7670,9 @@ class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): keep_alive_probe: Optional[KeepAliveProbeDict] """Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" + dedicated_ingress_endpoint_enabled: Optional[bool] + """Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""" + ReasoningEngineSpecDeploymentSpecOrDict = Union[ ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict @@ -8476,6 +8483,10 @@ class CreateAgentEngineConfig(_common.BaseModel): Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""", ) + dedicated_ingress_endpoint_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""", + ) class CreateAgentEngineConfigDict(TypedDict, total=False): @@ -8617,6 +8628,9 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" + dedicated_ingress_endpoint_enabled: Optional[bool] + """Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""" + CreateAgentEngineConfigOrDict = Union[ CreateAgentEngineConfig, CreateAgentEngineConfigDict @@ -9143,6 +9157,10 @@ class UpdateAgentEngineConfig(_common.BaseModel): Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""", ) + dedicated_ingress_endpoint_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""", + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -9293,6 +9311,9 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" + dedicated_ingress_endpoint_enabled: Optional[bool] + """Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -19574,6 +19595,10 @@ class AgentEngineConfig(_common.BaseModel): traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( default=None, description="""The traffic config for the Agent Engine.""" ) + dedicated_ingress_endpoint_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""", + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -19761,6 +19786,9 @@ class AgentEngineConfigDict(TypedDict, total=False): traffic_config: Optional[ReasoningEngineTrafficConfigDict] """The traffic config for the Agent Engine.""" + dedicated_ingress_endpoint_enabled: Optional[bool] + """Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]