Skip to content

Commit 9e4927b

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: 889205531
1 parent 2f2a211 commit 9e4927b

File tree

10 files changed

+3737
-71
lines changed

10 files changed

+3737
-71
lines changed

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

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
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_list = client.agent_engines.runtime_revisions.list(
83+
name=updated_agent_engine.api_resource.name,
84+
)
85+
assert len(runtime_revisions_list) == 2
86+
revision_to_delete = runtime_revisions_list[1]
87+
operation = revision_to_delete.delete()
88+
assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation)
89+
assert operation.done
90+
runtime_revisions_list = client.agent_engines.runtime_revisions.list(
91+
name=updated_agent_engine.api_resource.name,
92+
)
93+
assert len(runtime_revisions_list) == 1
94+
assert runtime_revisions_list[0].name != revision_to_delete.name
95+
client.agent_engines.delete(name=updated_agent_engine.api_resource.name, force=True)
96+
97+
98+
pytestmark = pytest_helper.setup(
99+
file=__file__,
100+
globals_for_file=globals(),
101+
test_method="agent_engines.runtime_revisions.delete",
102+
)
103+
104+
105+
pytest_plugins = ("pytest_asyncio",)
106+
107+
108+
@pytest.mark.asyncio
109+
async def test_delete_runtime_revision_async(
110+
client,
111+
mock_agent_engine_create_base64_encoded_tarball,
112+
mock_agent_engine_create_path_exists,
113+
):
114+
client._api_client._http_options.base_url = (
115+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
116+
)
117+
client._api_client._http_options.api_version = "v1beta1"
118+
119+
with (
120+
mock_agent_engine_create_base64_encoded_tarball,
121+
mock_agent_engine_create_path_exists,
122+
):
123+
agent_engine = client.agent_engines.create(
124+
config={
125+
"display_name": "test-agent-engine-delete-runtime-revision",
126+
"source_packages": [
127+
"test_module.py",
128+
"requirements.txt",
129+
],
130+
"entrypoint_module": "test_module",
131+
"entrypoint_object": "test_object",
132+
"class_methods": _TEST_CLASS_METHODS,
133+
"http_options": {
134+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
135+
"api_version": "v1beta1",
136+
},
137+
},
138+
)
139+
# Create a second runtime revision,
140+
# since it's not possible to delete if there is only one runtime revision.
141+
with (
142+
mock_agent_engine_create_base64_encoded_tarball,
143+
mock_agent_engine_create_path_exists,
144+
):
145+
updated_agent_engine = client.agent_engines.update(
146+
name=agent_engine.api_resource.name,
147+
config={
148+
"display_name": "test-agent-engine-update-traffic-with-agent-after-update",
149+
"source_packages": [
150+
"test_module.py",
151+
"requirements.txt",
152+
],
153+
"entrypoint_module": "test_module",
154+
"entrypoint_object": "test_object",
155+
"class_methods": _TEST_CLASS_METHODS,
156+
"http_options": {
157+
"base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com",
158+
"api_version": "v1beta1",
159+
},
160+
},
161+
)
162+
163+
runtime_revisions_list = await client.aio.agent_engines.runtime_revisions.list(
164+
name=updated_agent_engine.api_resource.name,
165+
)
166+
assert len(runtime_revisions_list) == 2
167+
revision_to_delete = runtime_revisions_list[1]
168+
operation = await revision_to_delete.delete()
169+
assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation)
170+
client.agent_engines.delete(name=updated_agent_engine.api_resource.name, force=True)

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

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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.runtime_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.runtime_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.runtime_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 = await client.aio.agent_engines.runtime_revisions.list(
123+
name=agent_engine.api_resource.name,
124+
)
125+
runtime_revisions_list = list(runtime_revisions_iter)
126+
assert len(runtime_revisions_list) == 1
127+
assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision)
128+
runtime_revision_name = runtime_revisions_list[0].api_resource.name
129+
runtime_revision = await client.aio.agent_engines.runtime_revisions.get(
130+
name=runtime_revision_name,
131+
)
132+
assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision)
133+
assert runtime_revision.api_resource.name == runtime_revision_name
134+
# Clean up resources.
135+
await client.aio.agent_engines.delete(
136+
name=agent_engine.api_resource.name, force=True
137+
)

0 commit comments

Comments
 (0)