Skip to content

Commit b8eaefb

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add Agent Engine Runtime Versioning support to SDK.
PiperOrigin-RevId: 910749459
1 parent 9fd0fe0 commit b8eaefb

11 files changed

Lines changed: 3788 additions & 1 deletion

‎tests/unit/vertexai/genai/replays/test_delete_ae_runtime_revision.py‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
import pytest
18+
19+
20+
from tests.unit.vertexai.genai.replays import pytest_helper
21+
from vertexai._genai import types
22+
23+
_TEST_CLASS_METHODS = [
24+
{"name": "query", "api_mode": ""},
25+
]
26+
27+
28+
def test_delete_runtime_revision(
29+
client,
30+
mock_agent_engine_create_base64_encoded_tarball,
31+
mock_agent_engine_create_path_exists,
32+
):
33+
client._api_client._http_options.base_url = (
34+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
35+
)
36+
client._api_client._http_options.api_version = "v1beta1"
37+
38+
with (
39+
mock_agent_engine_create_base64_encoded_tarball,
40+
mock_agent_engine_create_path_exists,
41+
):
42+
agent_engine = client.agent_engines.create(
43+
config={
44+
"display_name": "test-agent-engine-delete-runtime-revision",
45+
"source_packages": [
46+
"test_module.py",
47+
"requirements.txt",
48+
],
49+
"entrypoint_module": "test_module",
50+
"entrypoint_object": "test_object",
51+
"class_methods": _TEST_CLASS_METHODS,
52+
"http_options": {
53+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
54+
"api_version": "v1beta1",
55+
},
56+
},
57+
)
58+
# Create a second runtime revision,
59+
# since it's not possible to delete if there is only one runtime revision.
60+
with (
61+
mock_agent_engine_create_base64_encoded_tarball,
62+
mock_agent_engine_create_path_exists,
63+
):
64+
updated_agent_engine = client.agent_engines.update(
65+
name=agent_engine.api_resource.name,
66+
config={
67+
"display_name": "test-agent-engine-update-traffic-with-agent-after-update",
68+
"source_packages": [
69+
"test_module.py",
70+
"requirements.txt",
71+
],
72+
"entrypoint_module": "test_module",
73+
"entrypoint_object": "test_object",
74+
"class_methods": _TEST_CLASS_METHODS,
75+
"http_options": {
76+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
77+
"api_version": "v1beta1",
78+
},
79+
},
80+
)
81+
82+
runtime_revisions_iter = client.agent_engines.runtimes.revisions.list(
83+
name=updated_agent_engine.api_resource.name,
84+
)
85+
runtime_revisions_list = list(runtime_revisions_iter)
86+
assert len(runtime_revisions_list) == 2
87+
revision_to_delete = runtime_revisions_list[1]
88+
operation = client.agent_engines.runtimes.revisions.delete(
89+
name=revision_to_delete.api_resource.name,
90+
)
91+
assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation)
92+
assert operation.done
93+
runtime_revisions_iter = client.agent_engines.runtimes.revisions.list(
94+
name=updated_agent_engine.api_resource.name,
95+
)
96+
runtime_revisions_list = list(runtime_revisions_iter)
97+
assert len(runtime_revisions_list) == 1
98+
assert (
99+
runtime_revisions_list[0].api_resource.name
100+
!= revision_to_delete.api_resource.name
101+
)
102+
client.agent_engines.delete(name=updated_agent_engine.api_resource.name, force=True)
103+
104+
105+
pytestmark = pytest_helper.setup(
106+
file=__file__,
107+
globals_for_file=globals(),
108+
test_method="agent_engines.runtimes.revisions.delete",
109+
)
110+
111+
112+
pytest_plugins = ("pytest_asyncio",)
113+
114+
115+
@pytest.mark.asyncio
116+
async def test_delete_runtime_revision_async(
117+
client,
118+
mock_agent_engine_create_base64_encoded_tarball,
119+
mock_agent_engine_create_path_exists,
120+
):
121+
client._api_client._http_options.base_url = (
122+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
123+
)
124+
client._api_client._http_options.api_version = "v1beta1"
125+
126+
with (
127+
mock_agent_engine_create_base64_encoded_tarball,
128+
mock_agent_engine_create_path_exists,
129+
):
130+
agent_engine = client.agent_engines.create(
131+
config={
132+
"display_name": "test-agent-engine-delete-runtime-revision",
133+
"source_packages": [
134+
"test_module.py",
135+
"requirements.txt",
136+
],
137+
"entrypoint_module": "test_module",
138+
"entrypoint_object": "test_object",
139+
"class_methods": _TEST_CLASS_METHODS,
140+
"http_options": {
141+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
142+
"api_version": "v1beta1",
143+
},
144+
},
145+
)
146+
# Create a second runtime revision,
147+
# since it's not possible to delete if there is only one runtime revision.
148+
with (
149+
mock_agent_engine_create_base64_encoded_tarball,
150+
mock_agent_engine_create_path_exists,
151+
):
152+
updated_agent_engine = client.agent_engines.update(
153+
name=agent_engine.api_resource.name,
154+
config={
155+
"display_name": "test-agent-engine-update-traffic-with-agent-after-update",
156+
"source_packages": [
157+
"test_module.py",
158+
"requirements.txt",
159+
],
160+
"entrypoint_module": "test_module",
161+
"entrypoint_object": "test_object",
162+
"class_methods": _TEST_CLASS_METHODS,
163+
"http_options": {
164+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
165+
"api_version": "v1beta1",
166+
},
167+
},
168+
)
169+
170+
runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list(
171+
name=updated_agent_engine.api_resource.name,
172+
)
173+
runtime_revisions_list = []
174+
async for revision in runtime_revisions_iter:
175+
runtime_revisions_list.append(revision)
176+
assert len(runtime_revisions_list) == 2
177+
revision_to_delete = runtime_revisions_list[1]
178+
operation = await client.aio.agent_engines.runtimes.revisions.delete(
179+
name=revision_to_delete.api_resource.name,
180+
)
181+
assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation)
182+
await client.aio.agent_engines.delete(
183+
name=updated_agent_engine.api_resource.name, force=True
184+
)

‎tests/unit/vertexai/genai/replays/test_get_ae_runtime_revision.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
import pytest
18+
19+
from tests.unit.vertexai.genai.replays import pytest_helper
20+
from vertexai._genai import types
21+
22+
_TEST_CLASS_METHODS = [
23+
{"name": "query", "api_mode": ""},
24+
]
25+
26+
27+
def test_get_runtime_revisions(
28+
client,
29+
mock_agent_engine_create_base64_encoded_tarball,
30+
mock_agent_engine_create_path_exists,
31+
):
32+
client._api_client._http_options.base_url = (
33+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
34+
)
35+
client._api_client._http_options.api_version = "v1beta1"
36+
37+
with (
38+
mock_agent_engine_create_base64_encoded_tarball,
39+
mock_agent_engine_create_path_exists,
40+
):
41+
agent_engine = client.agent_engines.create(
42+
config={
43+
"display_name": "test-agent-engine-get-runtime-revisions",
44+
"source_packages": [
45+
"test_module.py",
46+
"requirements.txt",
47+
],
48+
"entrypoint_module": "test_module",
49+
"entrypoint_object": "test_object",
50+
"class_methods": _TEST_CLASS_METHODS,
51+
"http_options": {
52+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
53+
"api_version": "v1beta1",
54+
},
55+
},
56+
)
57+
assert (
58+
agent_engine.api_resource.display_name
59+
== "test-agent-engine-get-runtime-revisions"
60+
)
61+
runtime_revisions_iter = client.agent_engines.runtimes.revisions.list(
62+
name=agent_engine.api_resource.name,
63+
)
64+
runtime_revisions_list = list(runtime_revisions_iter)
65+
assert len(runtime_revisions_list) == 1
66+
67+
assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision)
68+
runtime_revision_name = runtime_revisions_list[0].api_resource.name
69+
runtime_revision = client.agent_engines.runtimes.revisions.get(
70+
name=runtime_revision_name,
71+
)
72+
assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision)
73+
assert runtime_revision.api_resource.name == runtime_revision_name
74+
# Clean up resources.
75+
agent_engine.delete(force=True)
76+
77+
78+
pytestmark = pytest_helper.setup(
79+
file=__file__,
80+
globals_for_file=globals(),
81+
test_method="agent_engines.runtimes.revisions.get",
82+
)
83+
84+
pytest_plugins = ("pytest_asyncio",)
85+
86+
87+
@pytest.mark.asyncio
88+
async def test_async_get_runtime_revisions(
89+
client,
90+
mock_agent_engine_create_base64_encoded_tarball,
91+
mock_agent_engine_create_path_exists,
92+
):
93+
client._api_client._http_options.base_url = (
94+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
95+
)
96+
client._api_client._http_options.api_version = "v1beta1"
97+
98+
with (
99+
mock_agent_engine_create_base64_encoded_tarball,
100+
mock_agent_engine_create_path_exists,
101+
):
102+
agent_engine = client.agent_engines.create(
103+
config={
104+
"display_name": "test-agent-engine-get-runtime-revisions",
105+
"source_packages": [
106+
"test_module.py",
107+
"requirements.txt",
108+
],
109+
"entrypoint_module": "test_module",
110+
"entrypoint_object": "test_object",
111+
"class_methods": _TEST_CLASS_METHODS,
112+
"http_options": {
113+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
114+
"api_version": "v1beta1",
115+
},
116+
},
117+
)
118+
assert (
119+
agent_engine.api_resource.display_name
120+
== "test-agent-engine-get-runtime-revisions"
121+
)
122+
runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list(
123+
name=agent_engine.api_resource.name,
124+
)
125+
runtime_revisions_list = []
126+
async for revision in runtime_revisions_iter:
127+
runtime_revisions_list.append(revision)
128+
assert len(runtime_revisions_list) == 1
129+
assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision)
130+
runtime_revision_name = runtime_revisions_list[0].api_resource.name
131+
runtime_revision = await client.aio.agent_engines.runtimes.revisions.get(
132+
name=runtime_revision_name,
133+
)
134+
assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision)
135+
assert runtime_revision.api_resource.name == runtime_revision_name
136+
# Clean up resources.
137+
await client.aio.agent_engines.delete(
138+
name=agent_engine.api_resource.name, force=True
139+
)

0 commit comments

Comments
 (0)