feat: Add support for agent gateway in agent engine

PiperOrigin-RevId: 896753301

Author: vertex-sdk-bot

vertexai/_genai/agent_engines.py

@@ -1474,6 +1474,7 @@ def create(
         api_config = self._create_config(
             mode="create",
             agent=agent,
+            agent_gateway_config=config.agent_gateway_config,
             identity_type=config.identity_type,
             staging_bucket=config.staging_bucket,
             requirements=config.requirements,
@@ -1776,6 +1777,9 @@ def _create_config(
         *,
         mode: str,
         agent: Any = None,
+        agent_gateway_config: Optional[
+            types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
+        ] = None,
         identity_type: Optional[types.IdentityType] = None,
         staging_bucket: Optional[str] = None,
         requirements: Optional[Union[str, Sequence[str]]] = None,
@@ -1996,6 +2000,12 @@ def _create_config(
                         ) from e
             update_masks.append("spec.agent_framework")
 
+        if agent_gateway_config is not None:
+            if agent_engine_spec is None:
+                agent_engine_spec = {}
+            agent_engine_spec["agent_gateway_config"] = agent_gateway_config
+            update_masks.append("spec.agent_gateway_config")
+
         if identity_type is not None or service_account is not None:
             if agent_engine_spec is None:
                 agent_engine_spec = {}
@@ -2231,6 +2241,7 @@ def update(
         api_config = self._create_config(
             mode="update",
             agent=agent,
+            agent_gateway_config=config.agent_gateway_config,
             identity_type=config.identity_type,
             staging_bucket=config.staging_bucket,
             requirements=config.requirements,

vertexai/_genai/types/common.py

@@ -16383,6 +16383,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):
@@ -16557,6 +16563,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]