feat: Support Agent Gateway in AE

PiperOrigin-RevId: 886194959

Author: vertex-sdk-bot

vertexai/_genai/agent_engines.py

@@ -110,6 +110,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
 
 
@@ -284,6 +291,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"])

vertexai/_genai/types/__init__.py

@@ -776,6 +776,27 @@
 from .common import ReasoningEngineOrDict
 from .common import ReasoningEngineSpec
 from .common import ReasoningEngineSpecDeploymentSpec
+from .common import ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
+from .common import (
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig,
+)
+from .common import (
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict,
+)
+from .common import (
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict,
+)
+from .common import (
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig,
+)
+from .common import (
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict,
+)
+from .common import (
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict,
+)
+from .common import ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
+from .common import ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict
 from .common import ReasoningEngineSpecDeploymentSpecDict
 from .common import ReasoningEngineSpecDeploymentSpecOrDict
 from .common import ReasoningEngineSpecDict
@@ -1526,6 +1547,15 @@
     "SecretEnvVar",
     "SecretEnvVarDict",
     "SecretEnvVarOrDict",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfig",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict",
+    "ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict",
     "ReasoningEngineSpecDeploymentSpec",
     "ReasoningEngineSpecDeploymentSpecDict",
     "ReasoningEngineSpecDeploymentSpecOrDict",

vertexai/_genai/types/common.py

@@ -5857,6 +5857,95 @@ class SecretEnvVarDict(TypedDict, total=False):
 SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict]
 
 
+class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig(
+    _common.BaseModel
+):
+    """Configuration for traffic targeting a Reasoning Engine."""
+
+    agent_gateway: Optional[str] = Field(
+        default=None,
+        description="""Required. The resource name of the Agent Gateway to use for inbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `CLIENT_TO_AGENT`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""",
+    )
+
+
+class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict(
+    TypedDict, total=False
+):
+    """Configuration for traffic targeting a Reasoning Engine."""
+
+    agent_gateway: Optional[str]
+    """Required. The resource name of the Agent Gateway to use for inbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `CLIENT_TO_AGENT`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`"""
+
+
+ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig,
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict,
+]
+
+
+class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig(
+    _common.BaseModel
+):
+    """Configuration for traffic originating from a Reasoning Engine."""
+
+    agent_gateway: Optional[str] = Field(
+        default=None,
+        description="""Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""",
+    )
+
+
+class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict(
+    TypedDict, total=False
+):
+    """Configuration for traffic originating from a Reasoning Engine."""
+
+    agent_gateway: Optional[str]
+    """Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`"""
+
+
+ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig,
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict,
+]
+
+
+class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel):
+    """FUCK."""
+
+    client_to_agent_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig
+    ] = Field(
+        default=None,
+        description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""",
+    )
+    agent_to_anywhere_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig
+    ] = Field(
+        default=None,
+        description="""Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""",
+    )
+
+
+class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False):
+    """FUCK."""
+
+    client_to_agent_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict
+    ]
+    """Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway."""
+
+    agent_to_anywhere_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict
+    ]
+    """Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway."""
+
+
+ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfig,
+    ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict,
+]
+
+
 class ReasoningEngineSpecDeploymentSpec(_common.BaseModel):
     """The specification of a Reasoning Engine deployment."""
 
@@ -5890,6 +5979,12 @@ class ReasoningEngineSpecDeploymentSpec(_common.BaseModel):
         default=None,
         description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""",
     )
+    agent_gateway_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
+    ] = Field(
+        default=None,
+        description="""Optional. Agent Gateway configuration for the Reasoning Engine deployment.""",
+    )
 
 
 class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False):
@@ -5919,6 +6014,11 @@ class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False):
     secret_env: Optional[list[SecretEnvVarDict]]
     """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent."""
 
+    agent_gateway_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
+    ]
+    """Optional. Agent Gateway configuration for the Reasoning Engine deployment."""
+
 
 ReasoningEngineSpecDeploymentSpecOrDict = Union[
     ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict
@@ -6953,6 +7053,9 @@ class CreateAgentEngineConfig(_common.BaseModel):
           subdirectory and the path must be added to `extra_packages`.
       """,
     )
+    agent_gateway_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
+    ] = Field(default=None, description="""Optional. FUCK.""")
 
 
 class CreateAgentEngineConfigDict(TypedDict, total=False):
@@ -7084,6 +7187,11 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
           subdirectory and the path must be added to `extra_packages`.
       """
 
+    agent_gateway_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
+    ]
+    """Optional. FUCK."""
+
 
 CreateAgentEngineConfigOrDict = Union[
     CreateAgentEngineConfig, CreateAgentEngineConfigDict
@@ -7726,6 +7834,9 @@ class UpdateAgentEngineConfig(_common.BaseModel):
           subdirectory and the path must be added to `extra_packages`.
       """,
     )
+    agent_gateway_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
+    ] = Field(default=None, description="""Optional. FUCK.""")
     update_mask: Optional[str] = Field(
         default=None,
         description="""The update mask to apply. For the `FieldMask` definition, see
@@ -7862,6 +7973,11 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
           subdirectory and the path must be added to `extra_packages`.
       """
 
+    agent_gateway_config: Optional[
+        ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
+    ]
+    """Optional. FUCK."""
+
     update_mask: Optional[str]
     """The update mask to apply. For the `FieldMask` definition, see
       https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""