From 2427ee59e99804145c0b7cf46acf8b777a774acf Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 8 May 2026 01:29:12 -0700 Subject: [PATCH] feat: Rename modules for runtime versioning support PiperOrigin-RevId: 912392675 --- .../test_delete_ae_runtime_revision.py | 12 +-- .../replays/test_get_ae_runtime_revision.py | 10 +-- .../replays/test_list_ae_runtime_revisions.py | 6 +- .../test_update_traffic_agent_engine.py | 6 +- .../test_agent_engine_runtime_revisions.py | 53 +++++++------ vertexai/_genai/agent_engines.py | 36 +++++---- vertexai/_genai/runtimes.py | 78 ------------------- 7 files changed, 63 insertions(+), 138 deletions(-) delete mode 100644 vertexai/_genai/runtimes.py diff --git a/tests/unit/vertexai/genai/replays/test_delete_ae_runtime_revision.py b/tests/unit/vertexai/genai/replays/test_delete_ae_runtime_revision.py index 5054503d70..98ff25da54 100644 --- a/tests/unit/vertexai/genai/replays/test_delete_ae_runtime_revision.py +++ b/tests/unit/vertexai/genai/replays/test_delete_ae_runtime_revision.py @@ -79,18 +79,18 @@ def test_delete_runtime_revision( }, ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=updated_agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 2 revision_to_delete = runtime_revisions_list[1] - operation = client.agent_engines.runtimes.revisions.delete( + operation = client.agent_engines.revisions.delete( name=revision_to_delete.api_resource.name, ) assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation) assert operation.done - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=updated_agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -105,7 +105,7 @@ def test_delete_runtime_revision( pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.runtimes.revisions.delete", + test_method="agent_engines.revisions.delete", ) @@ -167,7 +167,7 @@ async def test_delete_runtime_revision_async( }, ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.aio.agent_engines.revisions.list( name=updated_agent_engine.api_resource.name, ) runtime_revisions_list = [] @@ -175,7 +175,7 @@ async def test_delete_runtime_revision_async( runtime_revisions_list.append(revision) assert len(runtime_revisions_list) == 2 revision_to_delete = runtime_revisions_list[1] - operation = await client.aio.agent_engines.runtimes.revisions.delete( + operation = await client.aio.agent_engines.revisions.delete( name=revision_to_delete.api_resource.name, ) assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation) diff --git a/tests/unit/vertexai/genai/replays/test_get_ae_runtime_revision.py b/tests/unit/vertexai/genai/replays/test_get_ae_runtime_revision.py index b356b613e4..3c2bc912ac 100644 --- a/tests/unit/vertexai/genai/replays/test_get_ae_runtime_revision.py +++ b/tests/unit/vertexai/genai/replays/test_get_ae_runtime_revision.py @@ -58,7 +58,7 @@ def test_get_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-get-runtime-revisions" ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -66,7 +66,7 @@ def test_get_runtime_revisions( assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) runtime_revision_name = runtime_revisions_list[0].api_resource.name - runtime_revision = client.agent_engines.runtimes.revisions.get( + runtime_revision = client.agent_engines.revisions.get( name=runtime_revision_name, ) assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision) @@ -78,7 +78,7 @@ def test_get_runtime_revisions( pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.runtimes.revisions.get", + test_method="agent_engines.revisions.get", ) pytest_plugins = ("pytest_asyncio",) @@ -119,7 +119,7 @@ async def test_async_get_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-get-runtime-revisions" ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.aio.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = [] @@ -128,7 +128,7 @@ async def test_async_get_runtime_revisions( assert len(runtime_revisions_list) == 1 assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) runtime_revision_name = runtime_revisions_list[0].api_resource.name - runtime_revision = await client.aio.agent_engines.runtimes.revisions.get( + runtime_revision = await client.aio.agent_engines.revisions.get( name=runtime_revision_name, ) assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision) diff --git a/tests/unit/vertexai/genai/replays/test_list_ae_runtime_revisions.py b/tests/unit/vertexai/genai/replays/test_list_ae_runtime_revisions.py index 87c276cec0..3203369c99 100644 --- a/tests/unit/vertexai/genai/replays/test_list_ae_runtime_revisions.py +++ b/tests/unit/vertexai/genai/replays/test_list_ae_runtime_revisions.py @@ -58,7 +58,7 @@ def test_list_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-list-runtime-revisions" ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -74,7 +74,7 @@ def test_list_runtime_revisions( pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.runtimes.revisions.list", + test_method="agent_engines.revisions.list", ) pytest_plugins = ("pytest_asyncio",) @@ -115,7 +115,7 @@ async def test_async_list_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-list-runtime-revisions" ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.aio.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = [] diff --git a/tests/unit/vertexai/genai/replays/test_update_traffic_agent_engine.py b/tests/unit/vertexai/genai/replays/test_update_traffic_agent_engine.py index 41438e5e95..859c9d39ad 100644 --- a/tests/unit/vertexai/genai/replays/test_update_traffic_agent_engine.py +++ b/tests/unit/vertexai/genai/replays/test_update_traffic_agent_engine.py @@ -80,7 +80,7 @@ def test_agent_engines_update_traffic_to_manual_split( }, ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -148,7 +148,7 @@ def test_agent_engines_update_traffic_with_agent_update( == "test-agent-engine-update-traffic-with-agent-before-update" ) assert agent_engine.api_resource.traffic_config is None - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -197,7 +197,7 @@ def test_agent_engines_update_traffic_with_agent_update( == "test-agent-engine-update-traffic-with-agent-after-update" ) assert updated_agent_engine.api_resource.traffic_config == traffic_config - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.agent_engines.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) diff --git a/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py b/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py index c5ab408d10..16c8dcb86a 100644 --- a/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py +++ b/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py @@ -844,7 +844,7 @@ def teardown_method(self): def test_get_delete_runtime_revision_operation(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.agent_engines.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -854,7 +854,7 @@ def test_get_delete_runtime_revision_operation(self): } ), ) - operation = self.client.agent_engines.runtimes.revisions._get_delete_runtime_revision_operation( + operation = self.client.agent_engines.revisions._get_delete_runtime_revision_operation( operation_name=_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, ) request_mock.assert_called_with( @@ -870,7 +870,7 @@ def test_get_delete_runtime_revision_operation(self): def test_await_operation(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.agent_engines.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -882,7 +882,7 @@ def test_await_operation(self): ) operation = _agent_engines_utils._await_operation( operation_name=_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, - get_operation_fn=self.client.agent_engines.runtimes.revisions._get_delete_runtime_revision_operation, + get_operation_fn=self.client.agent_engines.revisions._get_delete_runtime_revision_operation, ) request_mock.assert_called_with( "get", @@ -895,9 +895,9 @@ def test_await_operation(self): ) def test_register_api_methods(self): - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( + agent = self.client.agent_engines.revisions._register_api_methods( agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions._api_client, + api_client=self.client.agent_engines.revisions._api_client, api_resource=_genai_types.ReasoningEngineRuntimeRevision( spec=_genai_types.ReasoningEngineSpec( class_methods=[ @@ -932,10 +932,10 @@ def teardown_method(self): def test_get_runtime_revision(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.agent_engines.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.runtimes.revisions.get( + self.client.agent_engines.revisions.get( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) request_mock.assert_called_with( @@ -947,12 +947,12 @@ def test_get_runtime_revision(self): def test_list_runtime_revisions(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.agent_engines.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") expected_query_params = {"filter": _TEST_LIST_FILTER} list( - self.client.agent_engines.runtimes.revisions.list( + self.client.agent_engines.revisions.list( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config=expected_query_params ) ) @@ -968,7 +968,7 @@ def test_list_runtime_revisions(self): def test_delete_runtime_revision(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.agent_engines.revisions._api_client, "request" ) as request_mock: request_mock.side_effect = [ # First call: response to delete. @@ -984,7 +984,7 @@ def test_delete_runtime_revision(self): ), ] - self.client.agent_engines.runtimes.revisions.delete( + self.client.agent_engines.revisions.delete( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) request_mock.call_args_list[0].assert_called_with( @@ -1003,12 +1003,12 @@ def test_delete_runtime_revision(self): def test_query_runtime_revision(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.agent_engines.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( + agent = self.client.agent_engines.revisions._register_api_methods( agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions, + api_client=self.client.agent_engines.revisions, api_resource=_genai_types.ReasoningEngineRuntimeRevision( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1032,10 +1032,10 @@ def test_query_runtime_revision(self): ) def test_query_agent_engine_async(self): - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( + agent = self.client.agent_engines.revisions._register_api_methods( agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( api_async_client=runtime_revisions.AsyncRuntimeRevisions( - api_client_=self.client.agent_engines.runtimes.revisions._api_client + api_client_=self.client.agent_engines.revisions._api_client ), api_resource=_genai_types.ReasoningEngineRuntimeRevision( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, @@ -1048,7 +1048,7 @@ def test_query_agent_engine_async(self): ) ) with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "async_request" + self.client.agent_engines.revisions._api_client, "async_request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") asyncio.run(agent.async_query(query=_TEST_QUERY_PROMPT)) @@ -1065,11 +1065,11 @@ def test_query_agent_engine_async(self): def test_query_agent_engine_stream(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request_streamed" + self.client.agent_engines.revisions._api_client, "request_streamed" ) as request_mock: - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( + agent = self.client.agent_engines.revisions._register_api_methods( agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions, + api_client=self.client.agent_engines.revisions, api_resource=_genai_types.ReasoningEngineRuntimeRevision( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1097,13 +1097,13 @@ async def mock_async_generator(): yield genai_types.HttpResponse(body=b"") with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, + self.client.agent_engines.revisions._api_client, "async_request_streamed", ) as request_mock: request_mock.return_value = mock_async_generator() - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( + agent = self.client.agent_engines.revisions._register_api_methods( agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions, + api_client=self.client.agent_engines.revisions, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1154,8 +1154,7 @@ def teardown_method(self): def test_delete_runtime_revision(self): with mock.patch.object( - self.client.aio.agent_engines.runtimes.revisions._api_client, - "async_request", + self.client.aio.agent_engines.revisions._api_client, "async_request" ) as request_mock: request_mock.side_effect = [ # First call: response to delete. @@ -1171,7 +1170,7 @@ def test_delete_runtime_revision(self): ), ] asyncio.run( - self.client.aio.agent_engines.runtimes.revisions.delete( + self.client.aio.agent_engines.revisions.delete( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) ) diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 4ffbcd9e5c..396a0b6f1d 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -38,12 +38,12 @@ from . import sessions as sessions_module from . import memories as memories_module from . import a2a_tasks as a2a_tasks_module - from . import runtimes as runtimes_module + from . import runtime_revisions as runtime_revisions_module _ = sessions_module __ = memories_module ___ = a2a_tasks_module - ____ = runtimes_module + ____ = runtime_revisions_module logger = logging.getLogger("vertexai_genai.agentengines") @@ -1515,22 +1515,24 @@ def _update( _memories = None _sandboxes = None _sessions = None - _runtimes = None + _runtime_revisions = None @property - def runtimes(self) -> "runtimes_module.Runtimes": - if self._runtimes is None: + def revisions(self) -> "runtime_revisions_module.RuntimeRevisions": + if self._runtime_revisions is None: try: - # We need to lazy load the runtimes module to handle the + # We need to lazy load the runtime_revisions module to handle the # possibility of ImportError when dependencies are not installed. - self._runtimes = importlib.import_module(".runtimes", __package__) + self._runtime_revisions = importlib.import_module( + ".runtime_revisions", __package__ + ) except ImportError as e: raise ImportError( - "The 'agent_engines.runtimes' module requires additional " + "The 'agent_engines.runtime_revisions' module requires additional " "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._runtimes.Runtimes(self._api_client) # type: ignore[no-any-return] + return self._runtime_revisions.RuntimeRevisions(self._api_client) # type: ignore[no-any-return] @property def a2a_tasks(self) -> "a2a_tasks_module.A2aTasks": @@ -3929,7 +3931,7 @@ async def _update( _a2a_tasks = None _memories = None _sessions = None - _runtimes = None + _runtime_revisions = None async def delete( self, @@ -3960,19 +3962,21 @@ async def delete( return operation @property - def runtimes(self) -> "runtimes_module.AsyncRuntimes": - if self._runtimes is None: + def revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": + if self._runtime_revisions is None: try: - # We need to lazy load the runtimes module to handle the + # We need to lazy load the runtime_revisions module to handle the # possibility of ImportError when dependencies are not installed. - self._runtimes = importlib.import_module(".runtimes", __package__) + self._runtime_revisions = importlib.import_module( + ".runtime_revisions", __package__ + ) except ImportError as e: raise ImportError( - "The 'agent_engines.runtimes' module requires additional " + "The 'agent_engines.runtime_revisions' module requires additional " "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._runtimes.AsyncRuntimes(self._api_client) # type: ignore[no-any-return] + return self._runtime_revisions.AsyncRuntimeRevisions(self._api_client) # type: ignore[no-any-return] @property def a2a_tasks(self) -> "a2a_tasks_module.AsyncA2aTasks": diff --git a/vertexai/_genai/runtimes.py b/vertexai/_genai/runtimes.py deleted file mode 100644 index d292a1137a..0000000000 --- a/vertexai/_genai/runtimes.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -# Handwritten placeholder code for the runtimes.py file. -# Should be replaced by the generated file. - -import importlib -import logging -import typing - -from google.genai import _api_module - - -if typing.TYPE_CHECKING: - from . import runtime_revisions as runtime_revisions_module - - _ = runtime_revisions_module - - -logger = logging.getLogger("vertexai_genai.runtimes") - -logger.setLevel(logging.INFO) - - -class Runtimes(_api_module.BaseModule): - - _revisions = None - - @property - def revisions(self) -> "runtime_revisions_module.RuntimeRevisions": - if self._revisions is None: - try: - # We need to lazy load the revisions module to handle the - # possibility of ImportError when dependencies are not installed. - self._revisions = importlib.import_module( - ".runtime_revisions", __package__ - ) - except ImportError as e: - raise ImportError( - "The 'agent_engines.runtimes.revisions' module requires " - "additional packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._revisions.RuntimeRevisions(self._api_client) # type: ignore[no-any-return] - - -class AsyncRuntimes(_api_module.BaseModule): - - _revisions = None - - @property - def revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": - if self._revisions is None: - try: - # We need to lazy load the revisions module to handle the - # possibility of ImportError when dependencies are not installed. - self._revisions = importlib.import_module( - ".runtime_revisions", __package__ - ) - except ImportError as e: - raise ImportError( - "The 'agent_engines.runtimes.revisions' module requires " - "additional packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._revisions.AsyncRuntimeRevisions(self._api_client) # type: ignore[no-any-return]