diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index 2f954e709a..518d371e5e 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -526,6 +526,20 @@ def register_operations(self) -> Dict[str, List[str]]: } ], } +_TEST_AGENT_ENGINE_AGENT_GATEWAY_CONFIG = { + "client_to_agent_config": { + "agent_gateway": ( + "projects/test-project/locations/us-central1/agentGateways/" + "test-client-to-agent-gateway" + ), + }, + "agent_to_anywhere_config": { + "agent_gateway": ( + "projects/test-project/locations/us-central1/agentGateways/" + "test-agent-to-anywhere-gateway" + ), + }, +} _TEST_AGENT_ENGINE_MIN_INSTANCES = 2 _TEST_AGENT_ENGINE_MAX_INSTANCES = 4 _TEST_AGENT_ENGINE_RESOURCE_LIMITS = { @@ -975,6 +989,7 @@ def test_create_agent_engine_config_full(self, mock_prepare): service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG, + agent_gateway_config=_TEST_AGENT_ENGINE_AGENT_GATEWAY_CONFIG, min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES, max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES, resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS, @@ -1006,6 +1021,7 @@ def test_create_agent_engine_config_full(self, mock_prepare): }, ], "psc_interface_config": _TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG, + "agent_gateway_config": _TEST_AGENT_ENGINE_AGENT_GATEWAY_CONFIG, "min_instances": _TEST_AGENT_ENGINE_MIN_INSTANCES, "max_instances": _TEST_AGENT_ENGINE_MAX_INSTANCES, "resource_limits": _TEST_AGENT_ENGINE_RESOURCE_LIMITS, @@ -2097,6 +2113,7 @@ def test_create_agent_engine_with_env_vars_dict( identity_type=None, context_spec=None, psc_interface_config=None, + agent_gateway_config=None, min_instances=None, max_instances=None, resource_limits=None, @@ -2201,6 +2218,7 @@ def test_create_agent_engine_with_custom_service_account( identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, context_spec=None, psc_interface_config=None, + agent_gateway_config=None, min_instances=None, max_instances=None, resource_limits=None, @@ -2304,6 +2322,7 @@ def test_create_agent_engine_with_experimental_mode( identity_type=None, context_spec=None, psc_interface_config=None, + agent_gateway_config=None, min_instances=None, max_instances=None, resource_limits=None, @@ -2476,6 +2495,7 @@ def test_create_agent_engine_with_class_methods( identity_type=None, context_spec=None, psc_interface_config=None, + agent_gateway_config=None, min_instances=None, max_instances=None, resource_limits=None, @@ -2573,6 +2593,7 @@ def test_create_agent_engine_with_agent_framework( service_account=None, context_spec=None, psc_interface_config=None, + agent_gateway_config=None, min_instances=None, max_instances=None, resource_limits=None, diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 4272e5b34e..b72880ce22 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -163,6 +163,13 @@ def _CreateAgentEngineConfig_to_vertex( if getv(from_object, ["python_version"]) is not None: setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + return to_object @@ -372,6 +379,13 @@ def _UpdateAgentEngineConfig_to_vertex( if getv(from_object, ["python_version"]) is not None: setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + if getv(from_object, ["update_mask"]) is not None: setv( parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) @@ -1485,6 +1499,7 @@ def create( service_account=config.service_account, context_spec=context_spec, psc_interface_config=config.psc_interface_config, + agent_gateway_config=config.agent_gateway_config, min_instances=config.min_instances, max_instances=config.max_instances, resource_limits=config.resource_limits, @@ -1787,6 +1802,9 @@ def _create_config( service_account: Optional[str] = None, context_spec: Optional[types.ReasoningEngineContextSpecDict] = None, psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, + agent_gateway_config: Optional[ + types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, resource_limits: Optional[dict[str, str]] = None, @@ -1937,6 +1955,7 @@ def _create_config( is_deployment_spec_updated = ( env_vars is not None or psc_interface_config is not None + or agent_gateway_config is not None or min_instances is not None or max_instances is not None or resource_limits is not None @@ -1959,6 +1978,7 @@ def _create_config( ) = self._generate_deployment_spec_or_raise( env_vars=env_vars, psc_interface_config=psc_interface_config, + agent_gateway_config=agent_gateway_config, min_instances=min_instances, max_instances=max_instances, resource_limits=resource_limits, @@ -2021,6 +2041,9 @@ def _generate_deployment_spec_or_raise( *, env_vars: Optional[dict[str, Union[str, Any]]] = None, psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, + agent_gateway_config: Optional[ + types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, resource_limits: Optional[dict[str, str]] = None, @@ -2045,6 +2068,9 @@ def _generate_deployment_spec_or_raise( if psc_interface_config: deployment_spec["psc_interface_config"] = psc_interface_config update_masks.append("spec.deployment_spec.psc_interface_config") + if agent_gateway_config: + deployment_spec["agent_gateway_config"] = agent_gateway_config + update_masks.append("spec.deployment_spec.agent_gateway_config") if min_instances is not None: if not 0 <= min_instances <= 10: raise ValueError( @@ -2242,6 +2268,7 @@ def update( service_account=config.service_account, context_spec=context_spec, psc_interface_config=config.psc_interface_config, + agent_gateway_config=config.agent_gateway_config, min_instances=config.min_instances, max_instances=config.max_instances, resource_limits=config.resource_limits, diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index cc115e03cd..97ecc9c3eb 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -7993,6 +7993,12 @@ class CreateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + ) class CreateAgentEngineConfigDict(TypedDict, total=False): @@ -8124,6 +8130,11 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Reasoning Engine deployment.""" + CreateAgentEngineConfigOrDict = Union[ CreateAgentEngineConfig, CreateAgentEngineConfigDict @@ -8640,6 +8651,12 @@ class UpdateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -8776,6 +8793,11 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Reasoning Engine deployment.""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -16412,6 +16434,12 @@ class AgentEngineConfig(_common.BaseModel): container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( default=None, description="""The container spec for the Agent Engine.""" ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -16586,6 +16614,11 @@ class AgentEngineConfigDict(TypedDict, total=False): container_spec: Optional[ReasoningEngineSpecContainerSpecDict] """The container spec for the Agent Engine.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Reasoning Engine deployment.""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]