Skip to content

Commit b207217

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add support for agent gateway in agent engine
PiperOrigin-RevId: 896753301
1 parent 7142c62 commit b207217

File tree

3 files changed

+81
-0
lines changed

3 files changed

+81
-0
lines changed

‎tests/unit/vertexai/genai/test_agent_engines.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -526,6 +526,20 @@ def register_operations(self) -> Dict[str, List[str]]:
526526
}
527527
],
528528
}
529+
_TEST_AGENT_ENGINE_AGENT_GATEWAY_CONFIG = {
530+
"client_to_agent_config": {
531+
"agent_gateway": (
532+
"projects/test-project/locations/us-central1/agentGateways/"
533+
"test-client-to-agent-gateway"
534+
),
535+
},
536+
"agent_to_anywhere_config": {
537+
"agent_gateway": (
538+
"projects/test-project/locations/us-central1/agentGateways/"
539+
"test-agent-to-anywhere-gateway"
540+
),
541+
},
542+
}
529543
_TEST_AGENT_ENGINE_MIN_INSTANCES = 2
530544
_TEST_AGENT_ENGINE_MAX_INSTANCES = 4
531545
_TEST_AGENT_ENGINE_RESOURCE_LIMITS = {
@@ -975,6 +989,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
975989
service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
976990
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
977991
psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG,
992+
agent_gateway_config=_TEST_AGENT_ENGINE_AGENT_GATEWAY_CONFIG,
978993
min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES,
979994
max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES,
980995
resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS,
@@ -1006,6 +1021,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
10061021
},
10071022
],
10081023
"psc_interface_config": _TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG,
1024+
"agent_gateway_config": _TEST_AGENT_ENGINE_AGENT_GATEWAY_CONFIG,
10091025
"min_instances": _TEST_AGENT_ENGINE_MIN_INSTANCES,
10101026
"max_instances": _TEST_AGENT_ENGINE_MAX_INSTANCES,
10111027
"resource_limits": _TEST_AGENT_ENGINE_RESOURCE_LIMITS,
@@ -2097,6 +2113,7 @@ def test_create_agent_engine_with_env_vars_dict(
20972113
identity_type=None,
20982114
context_spec=None,
20992115
psc_interface_config=None,
2116+
agent_gateway_config=None,
21002117
min_instances=None,
21012118
max_instances=None,
21022119
resource_limits=None,
@@ -2201,6 +2218,7 @@ def test_create_agent_engine_with_custom_service_account(
22012218
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
22022219
context_spec=None,
22032220
psc_interface_config=None,
2221+
agent_gateway_config=None,
22042222
min_instances=None,
22052223
max_instances=None,
22062224
resource_limits=None,
@@ -2304,6 +2322,7 @@ def test_create_agent_engine_with_experimental_mode(
23042322
identity_type=None,
23052323
context_spec=None,
23062324
psc_interface_config=None,
2325+
agent_gateway_config=None,
23072326
min_instances=None,
23082327
max_instances=None,
23092328
resource_limits=None,
@@ -2476,6 +2495,7 @@ def test_create_agent_engine_with_class_methods(
24762495
identity_type=None,
24772496
context_spec=None,
24782497
psc_interface_config=None,
2498+
agent_gateway_config=None,
24792499
min_instances=None,
24802500
max_instances=None,
24812501
resource_limits=None,
@@ -2573,6 +2593,7 @@ def test_create_agent_engine_with_agent_framework(
25732593
service_account=None,
25742594
context_spec=None,
25752595
psc_interface_config=None,
2596+
agent_gateway_config=None,
25762597
min_instances=None,
25772598
max_instances=None,
25782599
resource_limits=None,

‎vertexai/_genai/agent_engines.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,13 @@ def _CreateAgentEngineConfig_to_vertex(
163163
if getv(from_object, ["python_version"]) is not None:
164164
setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"]))
165165

166+
if getv(from_object, ["agent_gateway_config"]) is not None:
167+
setv(
168+
parent_object,
169+
["agentGatewayConfig"],
170+
getv(from_object, ["agent_gateway_config"]),
171+
)
172+
166173
return to_object
167174

168175

@@ -372,6 +379,13 @@ def _UpdateAgentEngineConfig_to_vertex(
372379
if getv(from_object, ["python_version"]) is not None:
373380
setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"]))
374381

382+
if getv(from_object, ["agent_gateway_config"]) is not None:
383+
setv(
384+
parent_object,
385+
["agentGatewayConfig"],
386+
getv(from_object, ["agent_gateway_config"]),
387+
)
388+
375389
if getv(from_object, ["update_mask"]) is not None:
376390
setv(
377391
parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"])
@@ -1485,6 +1499,7 @@ def create(
14851499
service_account=config.service_account,
14861500
context_spec=context_spec,
14871501
psc_interface_config=config.psc_interface_config,
1502+
agent_gateway_config=config.agent_gateway_config,
14881503
min_instances=config.min_instances,
14891504
max_instances=config.max_instances,
14901505
resource_limits=config.resource_limits,
@@ -1787,6 +1802,9 @@ def _create_config(
17871802
service_account: Optional[str] = None,
17881803
context_spec: Optional[types.ReasoningEngineContextSpecDict] = None,
17891804
psc_interface_config: Optional[types.PscInterfaceConfigDict] = None,
1805+
agent_gateway_config: Optional[
1806+
types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
1807+
] = None,
17901808
min_instances: Optional[int] = None,
17911809
max_instances: Optional[int] = None,
17921810
resource_limits: Optional[dict[str, str]] = None,
@@ -1937,6 +1955,7 @@ def _create_config(
19371955
is_deployment_spec_updated = (
19381956
env_vars is not None
19391957
or psc_interface_config is not None
1958+
or agent_gateway_config is not None
19401959
or min_instances is not None
19411960
or max_instances is not None
19421961
or resource_limits is not None
@@ -1959,6 +1978,7 @@ def _create_config(
19591978
) = self._generate_deployment_spec_or_raise(
19601979
env_vars=env_vars,
19611980
psc_interface_config=psc_interface_config,
1981+
agent_gateway_config=agent_gateway_config,
19621982
min_instances=min_instances,
19631983
max_instances=max_instances,
19641984
resource_limits=resource_limits,
@@ -2021,6 +2041,9 @@ def _generate_deployment_spec_or_raise(
20212041
*,
20222042
env_vars: Optional[dict[str, Union[str, Any]]] = None,
20232043
psc_interface_config: Optional[types.PscInterfaceConfigDict] = None,
2044+
agent_gateway_config: Optional[
2045+
types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
2046+
] = None,
20242047
min_instances: Optional[int] = None,
20252048
max_instances: Optional[int] = None,
20262049
resource_limits: Optional[dict[str, str]] = None,
@@ -2045,6 +2068,9 @@ def _generate_deployment_spec_or_raise(
20452068
if psc_interface_config:
20462069
deployment_spec["psc_interface_config"] = psc_interface_config
20472070
update_masks.append("spec.deployment_spec.psc_interface_config")
2071+
if agent_gateway_config:
2072+
deployment_spec["agent_gateway_config"] = agent_gateway_config
2073+
update_masks.append("spec.deployment_spec.agent_gateway_config")
20482074
if min_instances is not None:
20492075
if not 0 <= min_instances <= 10:
20502076
raise ValueError(
@@ -2242,6 +2268,7 @@ def update(
22422268
service_account=config.service_account,
22432269
context_spec=context_spec,
22442270
psc_interface_config=config.psc_interface_config,
2271+
agent_gateway_config=config.agent_gateway_config,
22452272
min_instances=config.min_instances,
22462273
max_instances=config.max_instances,
22472274
resource_limits=config.resource_limits,

‎vertexai/_genai/types/common.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7993,6 +7993,12 @@ class CreateAgentEngineConfig(_common.BaseModel):
79937993
subdirectory and the path must be added to `extra_packages`.
79947994
""",
79957995
)
7996+
agent_gateway_config: Optional[
7997+
ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
7998+
] = Field(
7999+
default=None,
8000+
description="""Agent Gateway configuration for a Reasoning Engine deployment.""",
8001+
)
79968002

79978003

79988004
class CreateAgentEngineConfigDict(TypedDict, total=False):
@@ -8124,6 +8130,11 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
81248130
subdirectory and the path must be added to `extra_packages`.
81258131
"""
81268132

8133+
agent_gateway_config: Optional[
8134+
ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
8135+
]
8136+
"""Agent Gateway configuration for a Reasoning Engine deployment."""
8137+
81278138

81288139
CreateAgentEngineConfigOrDict = Union[
81298140
CreateAgentEngineConfig, CreateAgentEngineConfigDict
@@ -8640,6 +8651,12 @@ class UpdateAgentEngineConfig(_common.BaseModel):
86408651
subdirectory and the path must be added to `extra_packages`.
86418652
""",
86428653
)
8654+
agent_gateway_config: Optional[
8655+
ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
8656+
] = Field(
8657+
default=None,
8658+
description="""Agent Gateway configuration for a Reasoning Engine deployment.""",
8659+
)
86438660
update_mask: Optional[str] = Field(
86448661
default=None,
86458662
description="""The update mask to apply. For the `FieldMask` definition, see
@@ -8776,6 +8793,11 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
87768793
subdirectory and the path must be added to `extra_packages`.
87778794
"""
87788795

8796+
agent_gateway_config: Optional[
8797+
ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
8798+
]
8799+
"""Agent Gateway configuration for a Reasoning Engine deployment."""
8800+
87798801
update_mask: Optional[str]
87808802
"""The update mask to apply. For the `FieldMask` definition, see
87818803
https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
@@ -16412,6 +16434,12 @@ class AgentEngineConfig(_common.BaseModel):
1641216434
container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field(
1641316435
default=None, description="""The container spec for the Agent Engine."""
1641416436
)
16437+
agent_gateway_config: Optional[
16438+
ReasoningEngineSpecDeploymentSpecAgentGatewayConfig
16439+
] = Field(
16440+
default=None,
16441+
description="""Agent Gateway configuration for a Reasoning Engine deployment.""",
16442+
)
1641516443

1641616444

1641716445
class AgentEngineConfigDict(TypedDict, total=False):
@@ -16586,6 +16614,11 @@ class AgentEngineConfigDict(TypedDict, total=False):
1658616614
container_spec: Optional[ReasoningEngineSpecContainerSpecDict]
1658716615
"""The container spec for the Agent Engine."""
1658816616

16617+
agent_gateway_config: Optional[
16618+
ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
16619+
]
16620+
"""Agent Gateway configuration for a Reasoning Engine deployment."""
16621+
1658916622

1659016623
AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]
1659116624

0 commit comments

Comments
 (0)