Skip to content

chore: Agent Gateway Config exposed to open source public GitHub repo. #6557

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

chore: Agent Gateway Config exposed to open source public GitHub repo. #6557

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions tests/unit/vertexai/genai/test_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -538,6 +538,16 @@ def register_operations(self) -> Dict[str, List[str]]:
_genai_types.IdentityType.SERVICE_ACCOUNT
)
_TEST_AGENT_ENGINE_ENCRYPTION_SPEC = {"kms_key_name": "test-kms-key"}
_TEST_AGENT_GATEWAY_CONFIG = {
"client_to_agent_config": {
"agent_gateway": "projects/test-target-project/locations/us-central1/agentGateways/"
"test-client-to-agent-gateway",
},
"agent_to_anywhere_config": {
"agent_gateway": "projects/test-target-project/locations/us-central1/agentGateways/"
"test-agent-to-anywhere-gateway",
},
}
_TEST_AGENT_ENGINE_SPEC = _genai_types.ReasoningEngineSpecDict(
agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK,
class_methods=[_TEST_AGENT_ENGINE_CLASS_METHOD_1],
Expand Down Expand Up @@ -981,6 +991,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,
agent_gateway_config=_TEST_AGENT_GATEWAY_CONFIG,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
Expand Down Expand Up @@ -1021,6 +1032,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
config["spec"]["identity_type"]
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
)
assert config["spec"]["agent_gateway_config"] == _TEST_AGENT_GATEWAY_CONFIG

@mock.patch.object(
_agent_engines_utils,
Expand Down Expand Up @@ -1413,6 +1425,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,
agent_gateway_config=_TEST_AGENT_GATEWAY_CONFIG,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
Expand Down Expand Up @@ -1447,6 +1460,7 @@ def test_update_agent_engine_config_full(self, mock_prepare):
config["spec"]["identity_type"]
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
)
assert config["spec"]["agent_gateway_config"] == _TEST_AGENT_GATEWAY_CONFIG
assert config["update_mask"] == ",".join(
[
"display_name",
Expand All @@ -1458,6 +1472,7 @@ def test_update_agent_engine_config_full(self, mock_prepare):
"spec.deployment_spec.env",
"spec.deployment_spec.secret_env",
"spec.agent_framework",
"spec.agent_gateway_config",
"spec.identity_type",
"spec.service_account",
]
Expand Down Expand Up @@ -2116,6 +2131,7 @@ def test_create_agent_engine_with_env_vars_dict(
image_spec=None,
agent_config_source=None,
container_spec=None,
agent_gateway_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2220,6 +2236,7 @@ def test_create_agent_engine_with_custom_service_account(
image_spec=None,
agent_config_source=None,
container_spec=None,
agent_gateway_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2323,6 +2340,7 @@ def test_create_agent_engine_with_experimental_mode(
image_spec=None,
agent_config_source=None,
container_spec=None,
agent_gateway_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2495,6 +2513,7 @@ def test_create_agent_engine_with_class_methods(
image_spec=None,
agent_config_source=None,
container_spec=None,
agent_gateway_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2593,6 +2612,7 @@ def test_create_agent_engine_with_agent_framework(
image_spec=None,
agent_config_source=None,
container_spec=None,
agent_gateway_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down
66 changes: 66 additions & 0 deletions vertexai/_genai/types/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -536,6 +536,12 @@
from .common import IntermediateExtractedMemoryDict
from .common import IntermediateExtractedMemoryOrDict
from .common import JobState
from .common import KeepAliveProbe
from .common import KeepAliveProbeDict
from .common import KeepAliveProbeHttpGet
from .common import KeepAliveProbeHttpGetDict
from .common import KeepAliveProbeHttpGetOrDict
from .common import KeepAliveProbeOrDict
from .common import Language
from .common import ListAgentEngineConfig
from .common import ListAgentEngineConfigDict
Expand Down Expand Up @@ -825,6 +831,27 @@
from .common import ReasoningEngineSpecContainerSpecDict
from .common import ReasoningEngineSpecContainerSpecOrDict
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
Expand Down Expand Up @@ -856,6 +883,18 @@
from .common import ReasoningEngineSpecSourceCodeSpecPythonSpec
from .common import ReasoningEngineSpecSourceCodeSpecPythonSpecDict
from .common import ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict
from .common import ReasoningEngineTrafficConfig
from .common import ReasoningEngineTrafficConfigDict
from .common import ReasoningEngineTrafficConfigOrDict
from .common import ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest
from .common import ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict
from .common import ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict
from .common import ReasoningEngineTrafficConfigTrafficSplitManual
from .common import ReasoningEngineTrafficConfigTrafficSplitManualDict
from .common import ReasoningEngineTrafficConfigTrafficSplitManualOrDict
from .common import ReasoningEngineTrafficConfigTrafficSplitManualTarget
from .common import ReasoningEngineTrafficConfigTrafficSplitManualTargetDict
from .common import ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict
from .common import ReservationAffinity
from .common import ReservationAffinityDict
from .common import ReservationAffinityOrDict
Expand Down Expand Up @@ -1665,6 +1704,21 @@
"SecretEnvVar",
"SecretEnvVarDict",
"SecretEnvVarOrDict",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfig",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict",
"ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict",
"KeepAliveProbeHttpGet",
"KeepAliveProbeHttpGetDict",
"KeepAliveProbeHttpGetOrDict",
"KeepAliveProbe",
"KeepAliveProbeDict",
"KeepAliveProbeOrDict",
"ReasoningEngineSpecDeploymentSpec",
"ReasoningEngineSpecDeploymentSpecDict",
"ReasoningEngineSpecDeploymentSpecOrDict",
Expand Down Expand Up @@ -1701,6 +1755,18 @@
"ReasoningEngineSpec",
"ReasoningEngineSpecDict",
"ReasoningEngineSpecOrDict",
"ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest",
"ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict",
"ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict",
"ReasoningEngineTrafficConfigTrafficSplitManualTarget",
"ReasoningEngineTrafficConfigTrafficSplitManualTargetDict",
"ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict",
"ReasoningEngineTrafficConfigTrafficSplitManual",
"ReasoningEngineTrafficConfigTrafficSplitManualDict",
"ReasoningEngineTrafficConfigTrafficSplitManualOrDict",
"ReasoningEngineTrafficConfig",
"ReasoningEngineTrafficConfigDict",
"ReasoningEngineTrafficConfigOrDict",
"ReasoningEngine",
"ReasoningEngineDict",
"ReasoningEngineOrDict",
Expand Down
Loading
Loading