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 new file mode 100644 index 0000000000..8555290657 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_delete_ae_runtime_revision.py @@ -0,0 +1,181 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +import pytest + + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +_TEST_CLASS_METHODS = [ + {"name": "query", "api_mode": ""}, +] + + +def test_delete_runtime_revision( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-delete-runtime-revision", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + # Create a second runtime revision, + # since it's not possible to delete if there is only one runtime revision. + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + updated_agent_engine = client.agent_engines.update( + name=agent_engine.api_resource.name, + config={ + "display_name": "test-agent-engine-update-traffic-with-agent-after-update", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + + runtime_revisions_iter = client.agent_engines.runtime_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.runtime_revisions.delete( + name=revision_to_delete.api_resource.name, + ) + assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation) + assert operation.done + runtime_revisions_iter = client.agent_engines.runtime_revisions.list( + name=updated_agent_engine.api_resource.name, + ) + runtime_revisions_list = list(runtime_revisions_iter) + assert len(runtime_revisions_list) == 1 + assert runtime_revisions_list[0].api_resource.name != revision_to_delete.api_resource.name + client.agent_engines.delete(name=updated_agent_engine.api_resource.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.runtime_revisions.delete", +) + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_delete_runtime_revision_async( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-delete-runtime-revision", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + # Create a second runtime revision, + # since it's not possible to delete if there is only one runtime revision. + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + updated_agent_engine = client.agent_engines.update( + name=agent_engine.api_resource.name, + config={ + "display_name": "test-agent-engine-update-traffic-with-agent-after-update", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + + runtime_revisions_iter = client.aio.agent_engines.runtime_revisions.list( + name=updated_agent_engine.api_resource.name, + ) + runtime_revisions_list = [] + async for revision in runtime_revisions_iter: + runtime_revisions_list.append(revision) + assert len(runtime_revisions_list) == 2 + revision_to_delete = runtime_revisions_list[1] + operation = await client.aio.agent_engines.runtime_revisions.delete( + name=revision_to_delete.api_resource.name, + ) + assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation) + await client.aio.agent_engines.delete( + name=updated_agent_engine.api_resource.name, force=True + ) 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 new file mode 100644 index 0000000000..947a130c7a --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_get_ae_runtime_revision.py @@ -0,0 +1,139 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +import pytest + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +_TEST_CLASS_METHODS = [ + {"name": "query", "api_mode": ""}, +] + + +def test_get_runtime_revisions( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-get-runtime-revisions", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + assert ( + agent_engine.api_resource.display_name + == "test-agent-engine-get-runtime-revisions" + ) + runtime_revisions_iter = client.agent_engines.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = list(runtime_revisions_iter) + 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 = client.agent_engines.runtime_revisions.get( + name=runtime_revision_name, + ) + assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision) + assert runtime_revision.api_resource.name == runtime_revision_name + # Clean up resources. + agent_engine.delete(force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.runtime_revisions.get", +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_async_get_runtime_revisions( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-get-runtime-revisions", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + assert ( + agent_engine.api_resource.display_name + == "test-agent-engine-get-runtime-revisions" + ) + runtime_revisions_iter = client.aio.agent_engines.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = [] + async for revision in runtime_revisions_iter: + runtime_revisions_list.append(revision) + 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.runtime_revisions.get( + name=runtime_revision_name, + ) + assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision) + assert runtime_revision.api_resource.name == runtime_revision_name + # Clean up resources. + await client.aio.agent_engines.delete( + name=agent_engine.api_resource.name, force=True + ) 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 new file mode 100644 index 0000000000..b309c7762f --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_list_ae_runtime_revisions.py @@ -0,0 +1,132 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +import pytest + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +_TEST_CLASS_METHODS = [ + {"name": "query", "api_mode": ""}, +] + + +def test_list_runtime_revisions( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-list-runtime-revisions", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + assert ( + agent_engine.api_resource.display_name + == "test-agent-engine-list-runtime-revisions" + ) + runtime_revisions_iter = client.agent_engines.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = list(runtime_revisions_iter) + assert len(runtime_revisions_list) == 1 + assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) + assert isinstance( + runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision + ) + # Clean up resources. + agent_engine.delete(force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.runtime_revisions.list", +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_async_list_runtime_revisions( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-list-runtime-revisions", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + assert ( + agent_engine.api_resource.display_name + == "test-agent-engine-list-runtime-revisions" + ) + runtime_revisions_iter = client.aio.agent_engines.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = [] + async for revision in runtime_revisions_iter: + runtime_revisions_list.append(revision) + assert len(runtime_revisions_list) == 1 + assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) + assert isinstance( + runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision + ) + # Clean up resources. + await client.aio.agent_engines.delete( + name=agent_engine.api_resource.name, force=True + ) 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 new file mode 100644 index 0000000000..cd7fb2168c --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_update_traffic_agent_engine.py @@ -0,0 +1,213 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +_TEST_CLASS_METHODS = [ + {"name": "query", "api_mode": ""}, +] + + +def test_agent_engines_update_traffic_to_always_latest(client): + + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + agent_engine = client.agent_engines.create() + + traffic_config = types.ReasoningEngineTrafficConfig( + traffic_split_always_latest=types.ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(), + ) + + updated_agent_engine = client.agent_engines.update( + name=agent_engine.api_resource.name, + config=types.AgentEngineConfig( + traffic_config=traffic_config, + ), + ) + + assert updated_agent_engine.api_resource.traffic_config == traffic_config + + agent_engine.delete(force=True) + + +def test_agent_engines_update_traffic_to_manual_split( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-update-traffic-to-manual-split", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + + runtime_revisions_iter = client.agent_engines.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = list(runtime_revisions_iter) + assert len(runtime_revisions_list) == 1 + assert isinstance(runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision) + runtime_revision_name = runtime_revisions_list[0].api_resource.name + + traffic_config = types.ReasoningEngineTrafficConfig( + traffic_split_manual=types.ReasoningEngineTrafficConfigTrafficSplitManual( + targets=[ + types.ReasoningEngineTrafficConfigTrafficSplitManualTarget( + runtime_revision_name=runtime_revision_name, + percent=100, + ), + ], + ), + ) + + updated_agent_engine = client.agent_engines.update( + name=agent_engine.api_resource.name, + config=types.AgentEngineConfig( + traffic_config=traffic_config, + ), + ) + + assert updated_agent_engine.api_resource.traffic_config == traffic_config + agent_engine.delete(force=True) + + +def test_agent_engines_update_traffic_with_agent_update( + client, + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, +): + + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" + ) + client._api_client._http_options.api_version = "v1beta1" + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + agent_engine = client.agent_engines.create( + config={ + "display_name": "test-agent-engine-update-traffic-with-agent-before-update", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + }, + ) + assert ( + agent_engine.api_resource.display_name + == "test-agent-engine-update-traffic-with-agent-before-update" + ) + assert agent_engine.api_resource.traffic_config is None + runtime_revisions_iter = client.agent_engines.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = list(runtime_revisions_iter) + assert len(runtime_revisions_list) == 1 + assert isinstance(runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision) + runtime_revision_name = runtime_revisions_list[0].api_resource.name + + traffic_config = types.ReasoningEngineTrafficConfig( + traffic_split_manual=types.ReasoningEngineTrafficConfigTrafficSplitManual( + targets=[ + types.ReasoningEngineTrafficConfigTrafficSplitManualTarget( + runtime_revision_name=runtime_revision_name, + percent=100, + ), + ], + ), + ) + + with ( + mock_agent_engine_create_base64_encoded_tarball, + mock_agent_engine_create_path_exists, + ): + updated_agent_engine = client.agent_engines.update( + name=agent_engine.api_resource.name, + config={ + "display_name": "test-agent-engine-update-traffic-with-agent-after-update", + "source_packages": [ + "test_module.py", + "requirements.txt", + ], + "entrypoint_module": "test_module", + "entrypoint_object": "test_object", + "class_methods": _TEST_CLASS_METHODS, + "http_options": { + "base_url": "https://us-central1-autopush-aiplatform.sandbox.googleapis.com", + "api_version": "v1beta1", + }, + "traffic_config": traffic_config, + }, + ) + + assert ( + updated_agent_engine.api_resource.display_name + == "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.runtime_revisions.list( + name=agent_engine.api_resource.name, + ) + runtime_revisions_list = list(runtime_revisions_iter) + assert len(runtime_revisions_list) == 2 + assert isinstance(runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision) + assert isinstance(runtime_revisions_list[1].api_resource, types.ReasoningEngineRuntimeRevision) + assert runtime_revisions_list[0].api_resource.name != runtime_revision_name # new revision + assert runtime_revisions_list[1].api_resource.name == runtime_revision_name + + agent_engine.delete(force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.update", +) diff --git a/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py b/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py new file mode 100644 index 0000000000..3bc62bad78 --- /dev/null +++ b/tests/unit/vertexai/genai/test_agent_engine_runtime_revisions.py @@ -0,0 +1,1189 @@ +# 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. +# +import asyncio +import importlib +import json +import os +import sys +import tempfile +from typing import Any, AsyncIterable, Dict, Iterable, List +from unittest import mock +from urllib.parse import urlencode + +from google import auth +from google.auth import credentials as auth_credentials +from google.cloud import aiplatform +import vertexai +from google.cloud.aiplatform import initializer +from vertexai._genai import _agent_engines_utils +from vertexai._genai import runtime_revisions +from vertexai._genai import types as _genai_types +from google.genai import types as genai_types +import pytest + + +_TEST_AGENT_FRAMEWORK = "google-adk" +GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( + "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" +) + + +class CapitalizeEngine: + """A sample Agent Engine.""" + + def set_up(self): + pass + + def query(self, unused_arbitrary_string_name: str) -> str: + """Runs the engine.""" + return unused_arbitrary_string_name.upper() + + def clone(self): + return self + + +class AsyncQueryEngine: + """A sample Agent Engine that implements `async_query`.""" + + def set_up(self): + pass + + async def async_query(self, unused_arbitrary_string_name: str): + """Runs the query asynchronously.""" + return unused_arbitrary_string_name.upper() + + def clone(self): + return self + + +class AsyncStreamQueryEngine: + """A sample Agent Engine that implements `async_stream_query`.""" + + def set_up(self): + pass + + async def async_stream_query( + self, unused_arbitrary_string_name: str + ) -> AsyncIterable[Any]: + """Runs the async stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + def clone(self): + return self + + +class StreamQueryEngine: + """A sample Agent Engine that implements `stream_query`.""" + + def set_up(self): + pass + + def stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]: + """Runs the stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + def clone(self): + return self + + +class OperationRegistrableEngine: + """Add a test class that implements OperationRegistrable.""" + + agent_framework = _TEST_AGENT_FRAMEWORK + + def query(self, unused_arbitrary_string_name: str) -> str: + """Runs the engine.""" + return unused_arbitrary_string_name.upper() + + async def async_query(self, unused_arbitrary_string_name: str) -> str: + """Runs the query asynchronously.""" + return unused_arbitrary_string_name.upper() + + # Add a custom method to test the custom method registration. + def custom_method(self, x: str) -> str: + return x.upper() + + # Add a custom async method to test the custom async method registration. + async def custom_async_method(self, x: str): + return x.upper() + + def stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]: + """Runs the stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + async def async_stream_query( + self, unused_arbitrary_string_name: str + ) -> AsyncIterable[Any]: + """Runs the async stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + # Add a custom method to test the custom stream method registration. + def custom_stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]: + """Runs the stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + # Add a custom method to test the custom stream method registration. + def custom_stream_method(self, unused_arbitrary_string_name: str) -> Iterable[Any]: + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + async def custom_async_stream_method( + self, unused_arbitrary_string_name: str + ) -> AsyncIterable[Any]: + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + async def bidi_stream_query(self, input_queue: asyncio.Queue) -> AsyncIterable[Any]: + """Runs the bidi stream engine.""" + while True: + chunk = await input_queue.get() + yield chunk + + async def custom_bidi_stream_method( + self, input_queue: asyncio.Queue + ) -> AsyncIterable[Any]: + """Runs the async bidi stream engine.""" + while True: + chunk = await input_queue.get() + yield chunk + + def clone(self): + return self + + def register_operations(self) -> Dict[str, List[str]]: + return { + _TEST_STANDARD_API_MODE: [ + _TEST_DEFAULT_METHOD_NAME, + _TEST_CUSTOM_METHOD_NAME, + ], + _TEST_ASYNC_API_MODE: [ + _TEST_DEFAULT_ASYNC_METHOD_NAME, + _TEST_CUSTOM_ASYNC_METHOD_NAME, + ], + _TEST_STREAM_API_MODE: [ + _TEST_DEFAULT_STREAM_METHOD_NAME, + _TEST_CUSTOM_STREAM_METHOD_NAME, + ], + _TEST_ASYNC_STREAM_API_MODE: [ + _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, + _TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, + ], + _TEST_BIDI_STREAM_API_MODE: [ + _TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, + _TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, + ], + } + + +class SameRegisteredOperationsEngine: + """Add a test class that is different from `OperationRegistrableEngine` but has the same registered operations.""" + + def query(self, unused_arbitrary_string_name: str) -> str: + """Runs the engine.""" + return unused_arbitrary_string_name.upper() + + async def async_query(self, unused_arbitrary_string_name: str) -> str: + """Runs the query asynchronously.""" + return unused_arbitrary_string_name.upper() + + # Add a custom method to test the custom method registration + def custom_method(self, x: str) -> str: + return x.upper() + + # Add a custom method that is not registered. + def custom_method_2(self, x: str) -> str: + return x.upper() + + # Add a custom async method to test the custom async method registration. + async def custom_async_method(self, x: str): + return x.upper() + + def stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]: + """Runs the stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + async def async_stream_query( + self, unused_arbitrary_string_name: str + ) -> AsyncIterable[Any]: + """Runs the async stream engine.""" + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + # Add a custom method to test the custom stream method registration. + def custom_stream_method(self, unused_arbitrary_string_name: str) -> Iterable[Any]: + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + async def custom_async_stream_method( + self, unused_arbitrary_string_name: str + ) -> AsyncIterable[Any]: + for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE: + yield chunk + + def clone(self): + return self + + def register_operations(self) -> Dict[str, List[str]]: + return { + _TEST_STANDARD_API_MODE: [ + _TEST_DEFAULT_METHOD_NAME, + _TEST_CUSTOM_METHOD_NAME, + ], + _TEST_ASYNC_API_MODE: [ + _TEST_DEFAULT_ASYNC_METHOD_NAME, + _TEST_CUSTOM_ASYNC_METHOD_NAME, + ], + _TEST_STREAM_API_MODE: [ + _TEST_DEFAULT_STREAM_METHOD_NAME, + _TEST_CUSTOM_STREAM_METHOD_NAME, + ], + _TEST_ASYNC_STREAM_API_MODE: [ + _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, + _TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, + ], + } + + +class OperationNotRegisteredEngine: + """Add a test class that has a method that is not registered.""" + + def query(self, unused_arbitrary_string_name: str) -> str: + """Runs the engine.""" + return unused_arbitrary_string_name.upper() + + def custom_method(self, x: str) -> str: + return x.upper() + + def clone(self): + return self + + def register_operations(self) -> Dict[str, List[str]]: + # `query` method is not exported in registered operations. + return { + _TEST_STANDARD_API_MODE: [ + _TEST_CUSTOM_METHOD_NAME, + ] + } + + +class RegisteredOperationNotExistEngine: + """Add a test class that has a method that is registered but does not exist.""" + + def query(self, unused_arbitrary_string_name: str) -> str: + """Runs the engine.""" + return unused_arbitrary_string_name.upper() + + def custom_method(self, x: str) -> str: + return x.upper() + + def clone(self): + return self + + def register_operations(self) -> Dict[str, List[str]]: + # Registered method `missing_method` is not a method of the AgentEngine. + return { + _TEST_STANDARD_API_MODE: [ + _TEST_DEFAULT_METHOD_NAME, + _TEST_CUSTOM_METHOD_NAME, + "missing_method", + ] + } + + +class MethodToBeUnregisteredEngine: + """An Agent Engine that has a method to be unregistered.""" + + def method_to_be_unregistered(self, unused_arbitrary_string_name: str) -> str: + """Method to be unregistered.""" + return unused_arbitrary_string_name.upper() + + def register_operations(self) -> Dict[str, List[str]]: + # Registered method `missing_method` is not a method of the AgentEngine. + return {_TEST_STANDARD_API_MODE: [_TEST_METHOD_TO_BE_UNREGISTERED_NAME]} + + +_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) +_TEST_STAGING_BUCKET = "gs://test-bucket" +_TEST_LOCATION = "us-central1" +_TEST_PROJECT = "test-project" +_TEST_RESOURCE_ID = "1028944691210842416" +_TEST_RUNTIME_REVISION_ID = "1234567890" +_TEST_OPERATION_ID = "4589432830794137600" +_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" +_TEST_AGENT_ENGINE_RESOURCE_NAME = ( + f"{_TEST_PARENT}/reasoningEngines/{_TEST_RESOURCE_ID}" +) +_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME = ( + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}/runtimeRevisions/{_TEST_RUNTIME_REVISION_ID}" +) +_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME = ( + f"{_TEST_PARENT}/operations/{_TEST_OPERATION_ID}" +) +_TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" +_TEST_AGENT_ENGINE_DESCRIPTION = "Agent Engine Description" +_TEST_LIST_FILTER = f'name="{_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}"' +_TEST_GCS_DIR_NAME = _agent_engines_utils._DEFAULT_GCS_DIR_NAME +_TEST_BLOB_FILENAME = _agent_engines_utils._BLOB_FILENAME +_TEST_REQUIREMENTS_FILE = _agent_engines_utils._REQUIREMENTS_FILE +_TEST_EXTRA_PACKAGES_FILE = _agent_engines_utils._EXTRA_PACKAGES_FILE +_TEST_STANDARD_API_MODE = _agent_engines_utils._STANDARD_API_MODE +_TEST_ASYNC_API_MODE = _agent_engines_utils._ASYNC_API_MODE +_TEST_STREAM_API_MODE = _agent_engines_utils._STREAM_API_MODE +_TEST_ASYNC_STREAM_API_MODE = _agent_engines_utils._ASYNC_STREAM_API_MODE +_TEST_BIDI_STREAM_API_MODE = _agent_engines_utils._BIDI_STREAM_API_MODE +_TEST_DEFAULT_METHOD_NAME = _agent_engines_utils._DEFAULT_METHOD_NAME +_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines_utils._DEFAULT_ASYNC_METHOD_NAME +_TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines_utils._DEFAULT_STREAM_METHOD_NAME +_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = ( + _agent_engines_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME +) +_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME = ( + _agent_engines_utils._DEFAULT_BIDI_STREAM_METHOD_NAME +) +_TEST_CAPITALIZE_ENGINE_METHOD_DOCSTRING = "Runs the engine." +_TEST_STREAM_METHOD_DOCSTRING = "Runs the stream engine." +_TEST_ASYNC_STREAM_METHOD_DOCSTRING = "Runs the async stream engine." +_TEST_BIDI_STREAM_METHOD_DOCSTRING = "Runs the bidi stream engine." +_TEST_MODE_KEY_IN_SCHEMA = _agent_engines_utils._MODE_KEY_IN_SCHEMA +_TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines_utils._METHOD_NAME_KEY_IN_SCHEMA +_TEST_CUSTOM_METHOD_NAME = "custom_method" +_TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method" +_TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method" +_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME = "custom_async_stream_method" +_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME = "custom_bidi_stream_method" +_TEST_CUSTOM_METHOD_DEFAULT_DOCSTRING = """ + Runs the Agent Engine to serve the user request. + + This will be based on the `.custom_method(...)` of the python object that + was passed in when creating the Agent Engine. The method will invoke the + `query` API client of the python object. + + Args: + **kwargs: + Optional. The arguments of the `.custom_method(...)` method. + + Returns: + dict[str, Any]: The response from serving the user request. +""" +_TEST_CUSTOM_ASYNC_METHOD_DEFAULT_DOCSTRING = """ + Runs the Agent Engine to serve the user request. + + This will be based on the `.custom_async_method(...)` of the python object that + was passed in when creating the Agent Engine. The method will invoke the + `async_query` API client of the python object. + + Args: + **kwargs: + Optional. The arguments of the `.custom_async_method(...)` method. + + Returns: + Coroutine[Any]: The response from serving the user request. +""" +_TEST_CUSTOM_STREAM_METHOD_DEFAULT_DOCSTRING = """ + Runs the Agent Engine to serve the user request. + + This will be based on the `.custom_stream_method(...)` of the python object that + was passed in when creating the Agent Engine. The method will invoke the + `stream_query` API client of the python object. + + Args: + **kwargs: + Optional. The arguments of the `.custom_stream_method(...)` method. + + Returns: + Iterable[Any]: The response from serving the user request. +""" +_TEST_CUSTOM_ASYNC_STREAM_METHOD_DEFAULT_DOCSTRING = """ + Runs the Agent Engine to serve the user request. + + This will be based on the `.custom_async_stream_method(...)` of the python object that + was passed in when creating the Agent Engine. The method will invoke the + `async_stream_query` API client of the python object. + + Args: + **kwargs: + Optional. The arguments of the `.custom_async_stream_method(...)` method. + + Returns: + AsyncIterable[Any]: The response from serving the user request. +""" +_TEST_METHOD_TO_BE_UNREGISTERED_NAME = "method_to_be_unregistered" +_TEST_QUERY_PROMPT = "Find the first fibonacci number greater than 999" +_TEST_AGENT_ENGINE_ENV_KEY = "GOOGLE_CLOUD_AGENT_ENGINE_ENV" +_TEST_AGENT_ENGINE_ENV_VALUE = "test_env_value" +_TEST_AGENT_ENGINE_GCS_URI = "{}/{}/{}".format( + _TEST_STAGING_BUCKET, + _TEST_GCS_DIR_NAME, + _TEST_BLOB_FILENAME, +) +_TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI = "{}/{}/{}".format( + _TEST_STAGING_BUCKET, + _TEST_GCS_DIR_NAME, + _TEST_EXTRA_PACKAGES_FILE, +) +_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI = "{}/{}/{}".format( + _TEST_STAGING_BUCKET, + _TEST_GCS_DIR_NAME, + _TEST_REQUIREMENTS_FILE, +) +_TEST_AGENT_ENGINE_REQUIREMENTS = [ + "google-cloud-aiplatform==1.29.0", + "langchain", +] +_TEST_AGENT_ENGINE_INVALID_EXTRA_PACKAGES = [ + "lib", + "main.py", +] +_TEST_AGENT_ENGINE_QUERY_SCHEMA = _agent_engines_utils._generate_schema( + CapitalizeEngine().query, + schema_name=_TEST_DEFAULT_METHOD_NAME, +) +_TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE +_TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}" +_TEST_PYTHON_VERSION_OVERRIDE = "3.11" +_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK +_TEST_AGENT_ENGINE_CLASS_METHOD_1 = { + "description": "Runs the engine.", + "name": "query", + "parameters": { + "type": "object", + "properties": { + "unused_arbitrary_string_name": {"type": "string"}, + }, + "required": ["unused_arbitrary_string_name"], + }, + "api_mode": "", +} +_TEST_AGENT_ENGINE_CLASS_METHOD_ASYNC_QUERY = { + "description": "Runs the engine.", + "name": "async_query", + "parameters": { + "type": "object", + "properties": { + "unused_arbitrary_string_name": {"type": "string"}, + }, + "required": ["unused_arbitrary_string_name"], + }, + "api_mode": "async", +} +_TEST_AGENT_ENGINE_CLASS_METHOD_STREAM_QUERY = { + "description": "Runs the engine.", + "name": "stream_query", + "parameters": { + "type": "object", + "properties": { + "unused_arbitrary_string_name": {"type": "string"}, + }, + "required": ["unused_arbitrary_string_name"], + }, + "api_mode": "stream", +} +_TEST_AGENT_ENGINE_CLASS_METHOD_ASYNC_STREAM_QUERY = { + "description": "Runs the engine.", + "name": "async_stream_query", + "parameters": { + "type": "object", + "properties": { + "unused_arbitrary_string_name": {"type": "string"}, + }, + "required": ["unused_arbitrary_string_name"], + }, + "api_mode": "async_stream", +} +_TEST_AGENT_ENGINE_ENV_VARS_INPUT = { + "TEST_ENV_VAR": "TEST_ENV_VAR_VALUE", + "TEST_ENV_VAR_2": "TEST_ENV_VAR_VALUE_2", + "TEST_SECRET_ENV_VAR": { + "secret": "TEST_SECRET_NAME_1", + "version": "TEST_SECRET_VERSION_1", + }, +} +_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG = { + "network_attachment": "test-network-attachment", + "dns_peering_configs": [ + { + "domain": "test-domain", + "target_project": "test-target-project", + "target_network": "test-target-network", + } + ], +} +_TEST_AGENT_ENGINE_MIN_INSTANCES = 2 +_TEST_AGENT_ENGINE_MAX_INSTANCES = 4 +_TEST_AGENT_ENGINE_RESOURCE_LIMITS = { + "cpu": "2", + "memory": "4Gi", +} +_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY = 4 +_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT = "test-custom-service-account" +_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT = ( + _genai_types.IdentityType.SERVICE_ACCOUNT +) +_TEST_AGENT_ENGINE_ENCRYPTION_SPEC = {"kms_key_name": "test-kms-key"} +_TEST_AGENT_ENGINE_REVISION_SPEC = _genai_types.ReasoningEngineSpecDict( + agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK, + class_methods=[_TEST_AGENT_ENGINE_CLASS_METHOD_1], + deployment_spec={ + "env": [ + {"name": "TEST_ENV_VAR", "value": "TEST_ENV_VAR_VALUE"}, + {"name": "TEST_ENV_VAR_2", "value": "TEST_ENV_VAR_VALUE_2"}, + ], + "secret_env": [ + { + "name": "TEST_SECRET_ENV_VAR", + "secret_ref": { + "secret": "TEST_SECRET_NAME_1", + "version": "TEST_SECRET_VERSION_1", + }, + }, + ], + }, + package_spec=_genai_types.ReasoningEngineSpecPackageSpecDict( + python_version=f"{sys.version_info.major}.{sys.version_info.minor}", + pickle_object_gcs_uri=_TEST_AGENT_ENGINE_GCS_URI, + dependency_files_gcs_uri=_TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI, + requirements_gcs_uri=_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, + ), + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, +) +_TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}] +_TEST_AGENT_ENGINE_OPERATION_SCHEMAS = [] +_TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py" +_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _agent_engines_utils._generate_schema( + AsyncQueryEngine().async_query, + schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME, +) +_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_ASYNC_API_MODE +_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( + OperationRegistrableEngine().custom_method, + schema_name=_TEST_CUSTOM_METHOD_NAME, +) +_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_STANDARD_API_MODE +) +_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( + OperationRegistrableEngine().custom_async_method, + schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME, +) +_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_ASYNC_API_MODE +) +_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( + StreamQueryEngine().stream_query, + schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME, +) +_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STREAM_API_MODE +_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( + OperationRegistrableEngine().custom_stream_method, + schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME, +) +_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_STREAM_API_MODE +) +_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( + AsyncStreamQueryEngine().async_stream_query, + schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, +) +_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_ASYNC_STREAM_API_MODE +) +_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = ( + _agent_engines_utils._generate_schema( + OperationRegistrableEngine().custom_async_stream_method, + schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, + ) +) +_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_ASYNC_STREAM_API_MODE +) +_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( + OperationRegistrableEngine().bidi_stream_query, + schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, +) +_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_BIDI_STREAM_API_MODE +) +_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA = ( + _agent_engines_utils._generate_schema( + OperationRegistrableEngine().custom_bidi_stream_method, + schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, + ) +) +_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_BIDI_STREAM_API_MODE +) +_TEST_OPERATION_REGISTRABLE_SCHEMAS = [ + _TEST_AGENT_ENGINE_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA, + _TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA, + _TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA, + _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA, +] +_TEST_OPERATION_NOT_REGISTERED_SCHEMAS = [ + _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA, +] +_TEST_REGISTERED_OPERATION_NOT_EXIST_SCHEMAS = [ + _TEST_AGENT_ENGINE_QUERY_SCHEMA, + _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA, +] +_TEST_NO_OPERATION_REGISTRABLE_SCHEMAS = [ + _TEST_AGENT_ENGINE_QUERY_SCHEMA, +] +_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _agent_engines_utils._generate_schema( + MethodToBeUnregisteredEngine().method_to_be_unregistered, + schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME, +) +_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( + _TEST_STANDARD_API_MODE +) +_TEST_ASYNC_QUERY_SCHEMAS = [_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA] +_TEST_STREAM_QUERY_SCHEMAS = [ + _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA, +] +_TEST_ASYNC_STREAM_QUERY_SCHEMAS = [ + _TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA, +] +_TEST_PACKAGE_DISTRIBUTIONS = { + "requests": ["requests"], + "cloudpickle": ["cloudpickle"], + "pydantic": ["pydantic"], +} +_TEST_OPERATION_NAME = "test_operation_name" +_TEST_AGENT_ENGINE_ERROR = { + "message": "The following quotas are exceeded", + "code": 8, +} + + +_TEST_AGENT_ENGINE_CLASS_METHODS = [ + { + "name": "query", + "description": "Simple query method", + "parameters": {"type": "object", "properties": {"input": {"type": "string"}}}, + "api_mode": "", + } +] + + +def _create_empty_fake_package(package_name: str) -> str: + """Creates a temporary directory structure representing an empty fake Python package. + + Args: + package_name (str): The name of the fake package. + + Returns: + str: The path to the top-level directory of the fake package. + """ + temp_dir = tempfile.mkdtemp() + package_dir = os.path.join(temp_dir, package_name) + os.makedirs(package_dir) + + # Create an empty __init__.py file to mark it as a package + init_path = os.path.join(package_dir, "__init__.py") + open(init_path, "w").close() + + return temp_dir + + +_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH = _create_empty_fake_package( + _TEST_AGENT_ENGINE_EXTRA_PACKAGE +) + + +@pytest.fixture(scope="module") +def google_auth_mock(): + with mock.patch.object(auth, "default") as google_auth_mock: + google_auth_mock.return_value = ( + auth_credentials.AnonymousCredentials(), + _TEST_PROJECT, + ) + yield google_auth_mock + + +@pytest.fixture(scope="module") +def importlib_metadata_version_mock(): + with mock.patch.object( + importlib.metadata, "version" + ) as importlib_metadata_version_mock: + + def get_version(pkg): + versions = { + "requests": "2.0.0", + "cloudpickle": "3.0.0", + "pydantic": "1.11.1", + } + return versions.get(pkg, "unknown") + + importlib_metadata_version_mock.side_effect = get_version + yield importlib_metadata_version_mock + + +class InvalidCapitalizeEngineWithoutQuerySelf: + """A sample Agent Engine with an invalid query method.""" + + def set_up(self): + pass + + def query() -> str: + """Runs the engine.""" + return "RESPONSE" + + +class InvalidCapitalizeEngineWithoutAsyncQuerySelf: + """A sample Agent Engine with an invalid async_query method.""" + + def set_up(self): + pass + + async def async_query() -> str: + """Runs the engine.""" + return "RESPONSE" + + +class InvalidCapitalizeEngineWithoutStreamQuerySelf: + """A sample Agent Engine with an invalid query_stream_query method.""" + + def set_up(self): + pass + + def stream_query() -> str: + """Runs the engine.""" + return "RESPONSE" + + +class InvalidCapitalizeEngineWithoutAsyncStreamQuerySelf: + """A sample Agent Engine with an invalid async_stream_query method.""" + + def set_up(self): + pass + + async def async_stream_query() -> str: + """Runs the engine.""" + return "RESPONSE" + + +class InvalidCapitalizeEngineWithoutRegisterOperationsSelf: + """A sample Agent Engine with an invalid register_operations method.""" + + def set_up(self): + pass + + def register_operations() -> str: + """Runs the engine.""" + return "RESPONSE" + + +class InvalidCapitalizeEngineWithoutQueryMethod: + """A sample Agent Engine without a query method.""" + + def set_up(self): + pass + + def invoke(self) -> str: + """Runs the engine.""" + return "RESPONSE" + + +class InvalidCapitalizeEngineWithNoncallableQueryStreamQuery: + """A sample Agent Engine with a noncallable query attribute.""" + + def __init__(self): + self.query = "RESPONSE" + + def set_up(self): + pass + + +def _create_fake_object_with_module(module_name): + class FakeObject: + pass + + FakeObject.__module__ = module_name + return FakeObject() + + +@pytest.mark.usefixtures("google_auth_mock") +class TestRuntimeRevisionsHelpers: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + importlib.reload(vertexai) + importlib.reload(os) + os.environ[_TEST_AGENT_ENGINE_ENV_KEY] = _TEST_AGENT_ENGINE_ENV_VALUE + self.client = vertexai.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_get_delete_runtime_revision_operation(self): + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse( + body=json.dumps( + { + "name": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + "done": True, + } + ), + ) + operation = self.client.agent_engines.runtime_revisions._get_delete_runtime_revision_operation( + operation_name=_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + ) + request_mock.assert_called_with( + "get", + _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + {"_url": {"operationName": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME}}, + None, + ) + assert isinstance( + operation, _genai_types.DeleteAgentEngineRuntimeRevisionOperation + ) + assert operation.done + + def test_await_operation(self): + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse( + body=json.dumps( + { + "name": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + "done": True, + } + ), + ) + operation = _agent_engines_utils._await_operation( + operation_name=_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + get_operation_fn=self.client.agent_engines.runtime_revisions._get_delete_runtime_revision_operation, + ) + request_mock.assert_called_with( + "get", + _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + {"_url": {"operationName": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME}}, + None, + ) + assert isinstance( + operation, _genai_types.DeleteAgentEngineRuntimeRevisionOperation + ) + + def test_register_api_methods(self): + agent = self.client.agent_engines.runtime_revisions._register_api_methods( + agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + api_client=self.client.agent_engines.runtime_revisions._api_client, + api_resource=_genai_types.ReasoningEngineRuntimeRevision( + spec=_genai_types.ReasoningEngineSpec( + class_methods=[ + _TEST_AGENT_ENGINE_CLASS_METHOD_1, + ] + ), + ), + ) + ) + assert agent.query.__doc__ == _TEST_AGENT_ENGINE_CLASS_METHOD_1.get( + "description" + ) + + +@pytest.mark.usefixtures("google_auth_mock") +class TestRuntimeRevisions: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + importlib.reload(vertexai) + importlib.reload(os) + os.environ[_TEST_AGENT_ENGINE_ENV_KEY] = _TEST_AGENT_ENGINE_ENV_VALUE + self.client = vertexai.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + self.test_agent = CapitalizeEngine() + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_get_runtime_revision(self): + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.agent_engines.runtime_revisions.get( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME + ) + request_mock.assert_called_with( + "get", + _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + {"_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}}, + None, + ) + + def test_list_runtime_revisions(self): + with mock.patch.object( + self.client.agent_engines.runtime_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.runtime_revisions.list( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config=expected_query_params + ) + ) + request_mock.assert_called_with( + "get", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}/runtimeRevisions?{urlencode(expected_query_params)}", + { + "_query": expected_query_params, + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + }, + None, + ) + + def test_delete_runtime_revision(self): + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, "request" + ) as request_mock: + request_mock.side_effect = [ + # First call: response to delete. + genai_types.HttpResponse(body=""), + # Second call: response to GET operation. + genai_types.HttpResponse( + body=json.dumps( + { + "name": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + "done": True, + } + ), + ), + ] + + self.client.agent_engines.runtime_revisions.delete( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME + ) + request_mock.call_args_list[0].assert_called_with( + "delete", + _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + {"_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}}, + None, + ) + + request_mock.call_args_list[1].assert_called_with( + "get", + _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + {"_url": {"operationName": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME}}, + None, + ) + + def test_query_runtime_revision(self): + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + agent = self.client.agent_engines.runtime_revisions._register_api_methods( + agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + api_client=self.client.agent_engines.runtime_revisions, + api_resource=_genai_types.ReasoningEngineRuntimeRevision( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + spec=_genai_types.ReasoningEngineSpec( + class_methods=[ + _TEST_AGENT_ENGINE_CLASS_METHOD_1, + ] + ), + ), + ) + ) + agent.query(query=_TEST_QUERY_PROMPT) + request_mock.assert_called_with( + "post", + f"{_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}:query", + { + "_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}, + "classMethod": "query", + "input": {"query": _TEST_QUERY_PROMPT}, + }, + None, + ) + + def test_query_agent_engine_async(self): + agent = self.client.agent_engines.runtime_revisions._register_api_methods( + agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + api_async_client=runtime_revisions.AsyncRuntimeRevisions( + api_client_=self.client.agent_engines.runtime_revisions._api_client + ), + api_resource=_genai_types.ReasoningEngineRuntimeRevision( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + spec=_genai_types.ReasoningEngineSpec( + class_methods=[ + _TEST_AGENT_ENGINE_CLASS_METHOD_ASYNC_QUERY, + ] + ), + ), + ) + ) + with mock.patch.object( + self.client.agent_engines.runtime_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)) + request_mock.assert_called_with( + "post", + f"{_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}:query", + { + "_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}, + "classMethod": "async_query", + "input": {"query": _TEST_QUERY_PROMPT}, + }, + None, + ) + + def test_query_agent_engine_stream(self): + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, "request_streamed" + ) as request_mock: + agent = self.client.agent_engines.runtime_revisions._register_api_methods( + agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + api_client=self.client.agent_engines.runtime_revisions, + api_resource=_genai_types.ReasoningEngineRuntimeRevision( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + spec=_genai_types.ReasoningEngineSpec( + class_methods=[ + _TEST_AGENT_ENGINE_CLASS_METHOD_STREAM_QUERY, + ] + ), + ), + ) + ) + list(agent.stream_query(query=_TEST_QUERY_PROMPT)) + request_mock.assert_called_with( + "post", + f"{_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}:streamQuery?alt=sse", + { + "_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}, + "classMethod": "stream_query", + "input": {"query": _TEST_QUERY_PROMPT}, + }, + None, + ) + + def test_query_agent_engine_async_stream(self): + async def mock_async_generator(): + yield genai_types.HttpResponse(body=b"") + + with mock.patch.object( + self.client.agent_engines.runtime_revisions._api_client, + "async_request_streamed", + ) as request_mock: + request_mock.return_value = mock_async_generator() + agent = self.client.agent_engines.runtime_revisions._register_api_methods( + agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + api_client=self.client.agent_engines.runtime_revisions, + api_resource=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + spec=_genai_types.ReasoningEngineSpec( + class_methods=[ + _TEST_AGENT_ENGINE_CLASS_METHOD_ASYNC_STREAM_QUERY, + ] + ), + ), + ) + ) + + async def consume(): + async for response in agent.async_stream_query( + query=_TEST_QUERY_PROMPT + ): + assert isinstance(response, genai_types.HttpResponse) + + asyncio.run(consume()) + request_mock.assert_called_with( + "post", + f"{_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}:streamQuery?alt=sse", + { + "_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}, + "classMethod": "async_stream_query", + "input": {"query": _TEST_QUERY_PROMPT}, + }, + None, + ) + + +@pytest.mark.usefixtures("google_auth_mock") +class TestAsyncRuntimeRevisions: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + importlib.reload(vertexai) + importlib.reload(os) + os.environ[_TEST_AGENT_ENGINE_ENV_KEY] = _TEST_AGENT_ENGINE_ENV_VALUE + self.client = vertexai.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + self.test_agent = CapitalizeEngine() + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_delete_runtime_revision(self): + with mock.patch.object( + self.client.aio.agent_engines.runtime_revisions._api_client, "async_request" + ) as request_mock: + request_mock.side_effect = [ + # First call: response to delete. + genai_types.HttpResponse(body=""), + # Second call: response to GET operation. + genai_types.HttpResponse( + body=json.dumps( + { + "name": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + "done": True, + } + ), + ), + ] + asyncio.run( + self.client.aio.agent_engines.runtime_revisions.delete( + name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME + ) + ) + request_mock.call_args_list[0].assert_called_with( + "delete", + _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, + {"_url": {"name": _TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}}, + None, + ) + + request_mock.call_args_list[1].assert_called_with( + "get", + _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, + {"_url": {"operationName": _TEST_AGENT_ENGINE_REVISION_OPERATION_NAME}}, + None, + ) diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index fc4fa850bc..f9c7bacd14 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -1021,7 +1021,7 @@ def _prepare( def _register_api_methods_or_raise( *, - agent_engine: genai_types.AgentEngine, + agent_engine: genai_types.AgentEngine | genai_types.AgentEngineRuntimeRevision, wrap_operation_fn: Optional[ dict[str, Callable[[str, str], Callable[..., Any]]] ] = None, diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 4272e5b34e..2ac830bf53 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -38,10 +38,12 @@ from . import sessions as sessions_module from . import memories as memories_module from . import a2a_tasks as a2a_tasks_module + from . import runtime_revisions as runtime_revisions_module _ = sessions_module __ = memories_module ___ = a2a_tasks_module + ____ = runtime_revisions_module logger = logging.getLogger("vertexai_genai.agentengines") @@ -377,6 +379,9 @@ def _UpdateAgentEngineConfig_to_vertex( parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) ) + if getv(from_object, ["traffic_config"]) is not None: + setv(parent_object, ["trafficConfig"], getv(from_object, ["traffic_config"])) + return to_object @@ -1037,6 +1042,24 @@ def _update( _memories = None _sandboxes = None _sessions = None + _runtime_revisions = None + + @property + def runtime_revisions(self) -> "runtime_revisions_module.RuntimeRevisions": + if self._runtime_revisions is None: + try: + # We need to lazy load the runtime_revisions module to handle the + # possibility of ImportError when dependencies are not installed. + self._runtime_revisions = importlib.import_module( + ".runtime_revisions", __package__ + ) + except ImportError as e: + raise ImportError( + "The 'agent_engines.runtime_revisions' module requires additional " + "packages. Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._runtime_revisions.RuntimeRevisions(self._api_client) # type: ignore[no-any-return] @property def a2a_tasks(self) -> "a2a_tasks_module.A2aTasks": @@ -1812,6 +1835,7 @@ def _create_config( types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict ] = None, container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, + traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, ) -> types.UpdateAgentEngineConfigDict: import sys @@ -1839,6 +1863,9 @@ def _create_config( if labels is not None: update_masks.append("labels") config["labels"] = labels + if traffic_config is not None: + update_masks.append("traffic_config") + config["traffic_config"] = traffic_config if agent_framework == "google-adk": env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) @@ -2213,6 +2240,9 @@ def update( agent_config_source = config.agent_config_source if agent_config_source is not None: agent_config_source = json.loads(agent_config_source.model_dump_json()) + traffic_config = config.traffic_config + if traffic_config is not None: + traffic_config = json.loads(traffic_config.model_dump_json()) if agent and agent_engine: raise ValueError("Please specify only one of `agent` or `agent_engine`.") elif agent_engine: @@ -2259,6 +2289,7 @@ def update( image_spec=image_spec, agent_config_source=agent_config_source, container_spec=container_spec, + traffic_config=traffic_config, ) operation = self._update(name=name, config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( @@ -3274,6 +3305,7 @@ async def _update( _a2a_tasks = None _memories = None _sessions = None + _runtime_revisions = None async def delete( self, @@ -3303,6 +3335,23 @@ async def delete( logger.info(f"Started AgentEngine delete operation: {operation.name}") return operation + @property + def runtime_revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": + if self._runtime_revisions is None: + try: + # We need to lazy load the runtime_revisions module to handle the + # possibility of ImportError when dependencies are not installed. + self._runtime_revisions = importlib.import_module( + ".runtime_revisions", __package__ + ) + except ImportError as e: + raise ImportError( + "The 'agent_engines.runtime_revisions' module requires additional " + "packages. Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._runtime_revisions.AsyncRuntimeRevisions(self._api_client) # type: ignore[no-any-return] + @property def a2a_tasks(self) -> "a2a_tasks_module.AsyncA2aTasks": if self._a2a_tasks is None: diff --git a/vertexai/_genai/runtime_revisions.py b/vertexai/_genai/runtime_revisions.py new file mode 100644 index 0000000000..47d7363c34 --- /dev/null +++ b/vertexai/_genai/runtime_revisions.py @@ -0,0 +1,1238 @@ +# 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. +# + +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. + +import functools +import json +import logging +from typing import Any, AsyncIterator, Iterator, Optional, Union +from urllib.parse import urlencode + +from google.genai import _api_module +from google.genai import _common +from google.genai._common import get_value_by_path as getv +from google.genai._common import set_value_by_path as setv +from google.genai.pagers import AsyncPager, Pager + +from . import _agent_engines_utils +from . import types + + +logger = logging.getLogger("vertexai_genai.runtimerevisions") + +logger.setLevel(logging.INFO) + + +def _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + +def _GetAgentEngineRuntimeRevisionRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + +def _GetDeleteAgentEngineRuntimeRevisionOperationParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + +def _ListAgentEngineRuntimeRevisionsConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) + + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + + if getv(from_object, ["filter"]) is not None: + setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) + + return to_object + + +def _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _ListAgentEngineRuntimeRevisionsConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) + + return to_object + + +def _QueryAgentEngineRuntimeRevisionConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["class_method"]) is not None: + setv(parent_object, ["classMethod"], getv(from_object, ["class_method"])) + + if getv(from_object, ["input"]) is not None: + setv(parent_object, ["input"], getv(from_object, ["input"])) + + if getv(from_object, ["include_all_fields"]) is not None: + setv(to_object, ["includeAllFields"], getv(from_object, ["include_all_fields"])) + + return to_object + + +def _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _QueryAgentEngineRuntimeRevisionConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) + + return to_object + + +class RuntimeRevisions(_api_module.BaseModule): + + def _get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.ReasoningEngineRuntimeRevision: + """ + Get an agent engine runtime revision instance. + """ + + parameter_model = types._GetAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _GetAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ReasoningEngineRuntimeRevision._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _list( + self, + *, + name: str, + config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, + ) -> types.ListReasoningEnginesRuntimeRevisionsResponse: + """ + Lists reasoning engine runtime revisions. + + Args: + name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. + config (ListAgentEngineRuntimeRevisionsConfig): + Optional. Additional configurations for listing the reasoning engine runtime revisions. + + Returns: + ListReasoningEnginesRuntimeRevisionsResponse: The requested reasoning engine runtime revisions. + + """ + + parameter_model = types._ListAgentEngineRuntimeRevisionsRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/runtimeRevisions".format_map(request_url_dict) + else: + path = "{name}/runtimeRevisions" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = ( + types.ListReasoningEnginesRuntimeRevisionsResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + ) + + self._api_client._verify_response(return_value) + return return_value + + def _delete( + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + """ + Delete an Agent Engine runtime revision. + + Args: + name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. + config (DeleteAgentEngineRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine runtime revision. + + Returns: + DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + + """ + + parameter_model = types._DeleteAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("delete", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get_delete_runtime_revision_operation( + self, + *, + operation_name: str, + config: Optional[ + types.GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict + ] = None, + ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + parameter_model = types._GetDeleteAgentEngineRuntimeRevisionOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = ( + _GetDeleteAgentEngineRuntimeRevisionOperationParameters_to_vertex( + parameter_model + ) + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Engine runtime revision. + """ + + parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.AgentEngineRuntimeRevision: + """Gets an agent engine runtime revision. + + Args: + name (str): Required. The name of the Agent Engine runtime revision to get. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. + config (GetAgentEngineRuntimeRevisionConfigOrDict): + Optional. Additional configurations for getting the Agent Engine runtime revision. + + Returns: + AgentEngineRuntimeRevision: The requested Agent Engine runtime revision instance. + """ + api_resource = self._get(name=name, config=config) + agent_engine_runtime_revision = types.AgentEngineRuntimeRevision( + api_client=self, + api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), + api_resource=api_resource, + ) + if api_resource.spec: + self._register_api_methods( + agent_engine_runtime_revision=agent_engine_runtime_revision + ) + return agent_engine_runtime_revision + + def list( + self, + *, + name: str, + config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, + ) -> Iterator[types.AgentEngineRuntimeRevision]: + """Lists all reasoning engine runtime revision instances matching the given query. + + Args: + name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. + config (ListAgentEngineRuntimeRevisionsConfig): + Optional. Additional configurations for listing the reasoning engine runtime revisions. + + Returns: + Iterable[AgentEngineRuntimeRevision]: An iterable of runtime revisions. + """ + list_pager: Pager[types.ReasoningEngineRuntimeRevision] = Pager( + "reasoning_engine_runtime_revisions", + functools.partial(self._list, name=name), + self._list(name=name, config=config), + config, + ) + + return ( + types.AgentEngineRuntimeRevision( + api_client=self, + api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), + api_resource=rutime_revision, + ) + for rutime_revision in list_pager + ) + + def delete( + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + """Delete an Agent Engine runtime revision. + + Args: + name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. + config (DeleteAgentEngineRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine runtime revision. + + Returns: + DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + """ + if config is None: + config = types.DeleteAgentEngineRuntimeRevisionConfig() + elif isinstance(config, dict): + config = types.DeleteAgentEngineRuntimeRevisionConfig.model_validate(config) + operation = self._delete( + name=name, + config=config, + ) + if config.wait_for_completion and not operation.done: + operation = _agent_engines_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_delete_runtime_revision_operation, + poll_interval_seconds=0.5, + ) + if operation.error: + raise RuntimeError( + f"Failed to delete runtime revision: {operation.error}" + ) + return operation + + def _register_api_methods( + self, + *, + agent_engine_runtime_revision: types.AgentEngineRuntimeRevision, + ) -> types.AgentEngineRuntimeRevision: + """Registers the API methods for the agent engine runtime revision.""" + try: + _agent_engines_utils._register_api_methods_or_raise( + agent_engine=agent_engine_runtime_revision, + wrap_operation_fn={ + "": _agent_engines_utils._wrap_query_operation, # type: ignore[dict-item] + "async": _agent_engines_utils._wrap_async_query_operation, # type: ignore[dict-item] + "stream": _agent_engines_utils._wrap_stream_query_operation, # type: ignore[dict-item] + "async_stream": _agent_engines_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] + "a2a_extension": _agent_engines_utils._wrap_a2a_operation, + }, + ) + except Exception as e: + logger.warning( + _agent_engines_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e + ) + return agent_engine_runtime_revision + + def _stream_query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> Iterator[Any]: + """Streams the response of the agent engine.""" + parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + http_options = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + for response in self._api_client.request_streamed( + "post", path, request_dict, http_options + ): + yield response + + async def _async_stream_query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> AsyncIterator[Any]: + """Streams the response of the agent engine.""" + parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + http_options = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + async_iterator = await self._api_client.async_request_streamed( + "post", path, request_dict, http_options + ) + async for response in async_iterator: + yield response + + +class AsyncRuntimeRevisions(_api_module.BaseModule): + + async def _get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.ReasoningEngineRuntimeRevision: + """ + Get an agent engine runtime revision instance. + """ + + parameter_model = types._GetAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _GetAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ReasoningEngineRuntimeRevision._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _list( + self, + *, + name: str, + config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, + ) -> types.ListReasoningEnginesRuntimeRevisionsResponse: + """ + Lists reasoning engine runtime revisions. + + Args: + name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. + config (ListAgentEngineRuntimeRevisionsConfig): + Optional. Additional configurations for listing the reasoning engine runtime revisions. + + Returns: + ListReasoningEnginesRuntimeRevisionsResponse: The requested reasoning engine runtime revisions. + + """ + + parameter_model = types._ListAgentEngineRuntimeRevisionsRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/runtimeRevisions".format_map(request_url_dict) + else: + path = "{name}/runtimeRevisions" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = ( + types.ListReasoningEnginesRuntimeRevisionsResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete( + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + """ + Delete an Agent Engine runtime revision. + + Args: + name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. + config (DeleteAgentEngineRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine runtime revision. + + Returns: + DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + + """ + + parameter_model = types._DeleteAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get_delete_runtime_revision_operation( + self, + *, + operation_name: str, + config: Optional[ + types.GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict + ] = None, + ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + parameter_model = types._GetDeleteAgentEngineRuntimeRevisionOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = ( + _GetDeleteAgentEngineRuntimeRevisionOperationParameters_to_vertex( + parameter_model + ) + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Engine runtime revision. + """ + + parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.AgentEngineRuntimeRevision: + """Gets an agent engine runtime revision. + + Args: + name (str): Required. The name of the Agent Engine runtime revision to get. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. + config (GetAgentEngineRuntimeRevisionConfigOrDict): + Optional. Additional configurations for getting the Agent Engine runtime revision. + + Returns: + AgentEngineRuntimeRevision: The requested Agent Engine runtime revision instance. + """ + api_resource = await self._get(name=name, config=config) + agent_engine_runtime_revision = types.AgentEngineRuntimeRevision( + api_client=self, + api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), + api_resource=api_resource, + ) + if api_resource.spec: + self._register_api_methods( + agent_engine_runtime_revision=agent_engine_runtime_revision + ) + return agent_engine_runtime_revision + + async def list( + self, + *, + name: str, + config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, + ) -> AsyncIterator[types.AgentEngineRuntimeRevision]: + """Lists reasoning engine runtime revisions. + + Args: + name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. + config (ListAgentEngineRuntimeRevisionsConfig): + Optional. Additional configurations for listing the reasoning engine runtime revisions. + + Returns: + AsyncIterator[AgentEngineRuntimeRevision]: An async iterator of runtime revisions. + """ + list_pager: AsyncPager[types.ReasoningEngineRuntimeRevision] = AsyncPager( + "reasoning_engine_runtime_revisions", + functools.partial(self._list, name=name), + await self._list(name=name, config=config), + config, + ) + + async for rutime_revision in list_pager: + yield types.AgentEngineRuntimeRevision( + api_client=self, + api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), + api_resource=rutime_revision, + ) + + async def delete( + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + """Delete an Agent Engine runtime revision. + + Args: + name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. + config (DeleteAgentEngineRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine runtime revision. + + Returns: + DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + """ + if config is None: + config = types.DeleteAgentEngineRuntimeRevisionConfig() + elif isinstance(config, dict): + config = types.DeleteAgentEngineRuntimeRevisionConfig.model_validate(config) + operation = await self._delete( + name=name, + config=config, + ) + if config.wait_for_completion and not operation.done: + operation = await _agent_engines_utils._await_async_operation( + operation_name=operation.name, + get_operation_fn=self._get_delete_runtime_revision_operation, + poll_interval_seconds=0.5, + ) + if operation.error: + raise RuntimeError( + f"Failed to delete runtime revision: {operation.error}" + ) + return operation + + def _register_api_methods( + self, + *, + agent_engine_runtime_revision: types.AgentEngineRuntimeRevision, + ) -> types.AgentEngineRuntimeRevision: + """Registers the API methods for the agent engine runtime revision.""" + try: + _agent_engines_utils._register_api_methods_or_raise( + agent_engine=agent_engine_runtime_revision, + wrap_operation_fn={ + "": _agent_engines_utils._wrap_query_operation, # type: ignore[dict-item] + "async": _agent_engines_utils._wrap_async_query_operation, # type: ignore[dict-item] + "stream": _agent_engines_utils._wrap_stream_query_operation, # type: ignore[dict-item] + "async_stream": _agent_engines_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] + "a2a_extension": _agent_engines_utils._wrap_a2a_operation, + }, + ) + except Exception as e: + logger.warning( + _agent_engines_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e + ) + return agent_engine_runtime_revision diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index f293f22bc7..190c630a28 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -43,6 +43,7 @@ from .common import _CustomJobParameters from .common import _DeleteAgentEngineMemoryRequestParameters from .common import _DeleteAgentEngineRequestParameters +from .common import _DeleteAgentEngineRuntimeRevisionRequestParameters from .common import _DeleteAgentEngineSandboxRequestParameters from .common import _DeleteAgentEngineSessionRequestParameters from .common import _DeleteAgentEngineTaskRequestParameters @@ -61,6 +62,7 @@ from .common import _GetAgentEngineMemoryRevisionRequestParameters from .common import _GetAgentEngineOperationParameters from .common import _GetAgentEngineRequestParameters +from .common import _GetAgentEngineRuntimeRevisionRequestParameters from .common import _GetAgentEngineSandboxOperationParameters from .common import _GetAgentEngineSandboxRequestParameters from .common import _GetAgentEngineSessionOperationParameters @@ -71,6 +73,7 @@ from .common import _GetDatasetOperationParameters from .common import _GetDatasetParameters from .common import _GetDatasetVersionParameters +from .common import _GetDeleteAgentEngineRuntimeRevisionOperationParameters from .common import _GetEvaluationItemParameters from .common import _GetEvaluationMetricParameters from .common import _GetEvaluationRunParameters @@ -80,6 +83,7 @@ from .common import _ListAgentEngineMemoryRequestParameters from .common import _ListAgentEngineMemoryRevisionsRequestParameters from .common import _ListAgentEngineRequestParameters +from .common import _ListAgentEngineRuntimeRevisionsRequestParameters from .common import _ListAgentEngineSandboxesRequestParameters from .common import _ListAgentEngineSessionEventsRequestParameters from .common import _ListAgentEngineSessionsRequestParameters @@ -93,6 +97,7 @@ from .common import _OptimizeRequestParameters from .common import _PurgeAgentEngineMemoriesRequestParameters from .common import _QueryAgentEngineRequestParameters +from .common import _QueryAgentEngineRuntimeRevisionRequestParameters from .common import _RestoreVersionRequestParameters from .common import _RetrieveAgentEngineMemoriesRequestParameters from .common import _RollbackAgentEngineMemoryRequestParameters @@ -134,6 +139,9 @@ from .common import AgentEngineRollbackMemoryOperation from .common import AgentEngineRollbackMemoryOperationDict from .common import AgentEngineRollbackMemoryOperationOrDict +from .common import AgentEngineRuntimeRevision +from .common import AgentEngineRuntimeRevisionDict +from .common import AgentEngineRuntimeRevisionOrDict from .common import AgentEngineSandboxOperation from .common import AgentEngineSandboxOperationDict from .common import AgentEngineSandboxOperationOrDict @@ -291,6 +299,12 @@ from .common import DeleteAgentEngineOperation from .common import DeleteAgentEngineOperationDict from .common import DeleteAgentEngineOperationOrDict +from .common import DeleteAgentEngineRuntimeRevisionConfig +from .common import DeleteAgentEngineRuntimeRevisionConfigDict +from .common import DeleteAgentEngineRuntimeRevisionConfigOrDict +from .common import DeleteAgentEngineRuntimeRevisionOperation +from .common import DeleteAgentEngineRuntimeRevisionOperationDict +from .common import DeleteAgentEngineRuntimeRevisionOperationOrDict from .common import DeleteAgentEngineSandboxConfig from .common import DeleteAgentEngineSandboxConfigDict from .common import DeleteAgentEngineSandboxConfigOrDict @@ -504,6 +518,9 @@ from .common import GetAgentEngineOperationConfig from .common import GetAgentEngineOperationConfigDict from .common import GetAgentEngineOperationConfigOrDict +from .common import GetAgentEngineRuntimeRevisionConfig +from .common import GetAgentEngineRuntimeRevisionConfigDict +from .common import GetAgentEngineRuntimeRevisionConfigOrDict from .common import GetAgentEngineSandboxConfig from .common import GetAgentEngineSandboxConfigDict from .common import GetAgentEngineSandboxConfigOrDict @@ -516,6 +533,9 @@ from .common import GetDatasetOperationConfig from .common import GetDatasetOperationConfigDict from .common import GetDatasetOperationConfigOrDict +from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfig +from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigDict +from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict from .common import GetEvaluationItemConfig from .common import GetEvaluationItemConfigDict from .common import GetEvaluationItemConfigOrDict @@ -559,6 +579,9 @@ from .common import ListAgentEngineMemoryRevisionsResponse from .common import ListAgentEngineMemoryRevisionsResponseDict from .common import ListAgentEngineMemoryRevisionsResponseOrDict +from .common import ListAgentEngineRuntimeRevisionsConfig +from .common import ListAgentEngineRuntimeRevisionsConfigDict +from .common import ListAgentEngineRuntimeRevisionsConfigOrDict from .common import ListAgentEngineSandboxesConfig from .common import ListAgentEngineSandboxesConfigDict from .common import ListAgentEngineSandboxesConfigOrDict @@ -613,6 +636,9 @@ from .common import ListReasoningEnginesResponse from .common import ListReasoningEnginesResponseDict from .common import ListReasoningEnginesResponseOrDict +from .common import ListReasoningEnginesRuntimeRevisionsResponse +from .common import ListReasoningEnginesRuntimeRevisionsResponseDict +from .common import ListReasoningEnginesRuntimeRevisionsResponseOrDict from .common import ListReasoningEnginesSessionsResponse from .common import ListReasoningEnginesSessionsResponseDict from .common import ListReasoningEnginesSessionsResponseOrDict @@ -800,6 +826,9 @@ from .common import QueryAgentEngineConfig from .common import QueryAgentEngineConfigDict from .common import QueryAgentEngineConfigOrDict +from .common import QueryAgentEngineRuntimeRevisionConfig +from .common import QueryAgentEngineRuntimeRevisionConfigDict +from .common import QueryAgentEngineRuntimeRevisionConfigOrDict from .common import QueryReasoningEngineResponse from .common import QueryReasoningEngineResponseDict from .common import QueryReasoningEngineResponseOrDict @@ -830,6 +859,9 @@ from .common import ReasoningEngineContextSpecOrDict from .common import ReasoningEngineDict from .common import ReasoningEngineOrDict +from .common import ReasoningEngineRuntimeRevision +from .common import ReasoningEngineRuntimeRevisionDict +from .common import ReasoningEngineRuntimeRevisionOrDict from .common import ReasoningEngineSpec from .common import ReasoningEngineSpecContainerSpec from .common import ReasoningEngineSpecContainerSpecDict @@ -1762,15 +1794,15 @@ "ReasoningEngineSpec", "ReasoningEngineSpecDict", "ReasoningEngineSpecOrDict", - "ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest", - "ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict", - "ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict", "ReasoningEngineTrafficConfigTrafficSplitManualTarget", "ReasoningEngineTrafficConfigTrafficSplitManualTargetDict", "ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict", "ReasoningEngineTrafficConfigTrafficSplitManual", "ReasoningEngineTrafficConfigTrafficSplitManualDict", "ReasoningEngineTrafficConfigTrafficSplitManualOrDict", + "ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest", + "ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict", + "ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict", "ReasoningEngineTrafficConfig", "ReasoningEngineTrafficConfigDict", "ReasoningEngineTrafficConfigOrDict", @@ -1918,6 +1950,30 @@ "ListAgentEngineMemoryRevisionsResponse", "ListAgentEngineMemoryRevisionsResponseDict", "ListAgentEngineMemoryRevisionsResponseOrDict", + "GetAgentEngineRuntimeRevisionConfig", + "GetAgentEngineRuntimeRevisionConfigDict", + "GetAgentEngineRuntimeRevisionConfigOrDict", + "ReasoningEngineRuntimeRevision", + "ReasoningEngineRuntimeRevisionDict", + "ReasoningEngineRuntimeRevisionOrDict", + "ListAgentEngineRuntimeRevisionsConfig", + "ListAgentEngineRuntimeRevisionsConfigDict", + "ListAgentEngineRuntimeRevisionsConfigOrDict", + "ListReasoningEnginesRuntimeRevisionsResponse", + "ListReasoningEnginesRuntimeRevisionsResponseDict", + "ListReasoningEnginesRuntimeRevisionsResponseOrDict", + "DeleteAgentEngineRuntimeRevisionConfig", + "DeleteAgentEngineRuntimeRevisionConfigDict", + "DeleteAgentEngineRuntimeRevisionConfigOrDict", + "DeleteAgentEngineRuntimeRevisionOperation", + "DeleteAgentEngineRuntimeRevisionOperationDict", + "DeleteAgentEngineRuntimeRevisionOperationOrDict", + "GetDeleteAgentEngineRuntimeRevisionOperationConfig", + "GetDeleteAgentEngineRuntimeRevisionOperationConfigDict", + "GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict", + "QueryAgentEngineRuntimeRevisionConfig", + "QueryAgentEngineRuntimeRevisionConfigDict", + "QueryAgentEngineRuntimeRevisionConfigOrDict", "SandboxEnvironmentSpecCodeExecutionEnvironment", "SandboxEnvironmentSpecCodeExecutionEnvironmentDict", "SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict", @@ -2263,6 +2319,9 @@ "OptimizeJobConfig", "OptimizeJobConfigDict", "OptimizeJobConfigOrDict", + "AgentEngineRuntimeRevision", + "AgentEngineRuntimeRevisionDict", + "AgentEngineRuntimeRevisionOrDict", "A2aTaskState", "State", "Strategy", @@ -2350,6 +2409,11 @@ "_PurgeAgentEngineMemoriesRequestParameters", "_GetAgentEngineMemoryRevisionRequestParameters", "_ListAgentEngineMemoryRevisionsRequestParameters", + "_GetAgentEngineRuntimeRevisionRequestParameters", + "_ListAgentEngineRuntimeRevisionsRequestParameters", + "_DeleteAgentEngineRuntimeRevisionRequestParameters", + "_GetDeleteAgentEngineRuntimeRevisionOperationParameters", + "_QueryAgentEngineRuntimeRevisionRequestParameters", "_CreateAgentEngineSandboxRequestParameters", "_DeleteAgentEngineSandboxRequestParameters", "_ExecuteCodeAgentEngineSandboxRequestParameters", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index d017fda1d1..368f7b4b3e 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -7596,46 +7596,28 @@ class ReasoningEngineSpecDict(TypedDict, total=False): ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] -class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(_common.BaseModel): - """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" - - pass - - -class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict(TypedDict, total=False): - """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" - - pass - - -ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest, - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict, -] - - class ReasoningEngineTrafficConfigTrafficSplitManualTarget(_common.BaseModel): """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" - percent: Optional[int] = Field( - default=None, - description="""Required. Specifies percent of the traffic to this Runtime Revision.""", - ) runtime_revision_name: Optional[str] = Field( default=None, description="""Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""", ) + percent: Optional[int] = Field( + default=None, + description="""Required. Specifies percent of the traffic to this Runtime Revision.""", + ) class ReasoningEngineTrafficConfigTrafficSplitManualTargetDict(TypedDict, total=False): """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" - percent: Optional[int] - """Required. Specifies percent of the traffic to this Runtime Revision.""" - runtime_revision_name: Optional[str] """Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""" + percent: Optional[int] + """Required. Specifies percent of the traffic to this Runtime Revision.""" + ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict = Union[ ReasoningEngineTrafficConfigTrafficSplitManualTarget, @@ -7667,34 +7649,52 @@ class ReasoningEngineTrafficConfigTrafficSplitManualDict(TypedDict, total=False) ] +class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(_common.BaseModel): + """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + + pass + + +class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict(TypedDict, total=False): + """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + + pass + + +ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest, + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict, +] + + class ReasoningEngineTrafficConfig(_common.BaseModel): - """Traffic distribution configuration.""" + """Traffic distribution configuration for the Reasoning Engine.""" - traffic_split_always_latest: Optional[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest - ] = Field( - default=None, - description="""Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""", - ) traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManual] = ( Field( default=None, description="""Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""", ) ) + traffic_split_always_latest: Optional[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest + ] = Field( + default=None, + description="""Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""", + ) class ReasoningEngineTrafficConfigDict(TypedDict, total=False): - """Traffic distribution configuration.""" + """Traffic distribution configuration for the Reasoning Engine.""" + + traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManualDict] + """Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" traffic_split_always_latest: Optional[ ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict ] """Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" - traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManualDict] - """Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" - ReasoningEngineTrafficConfigOrDict = Union[ ReasoningEngineTrafficConfig, ReasoningEngineTrafficConfigDict @@ -8634,6 +8634,10 @@ class UpdateAgentEngineConfig(_common.BaseModel): description="""The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", ) + traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + default=None, + description="""Traffic distribution configuration for the Reasoning Engine.""", + ) class UpdateAgentEngineConfigDict(TypedDict, total=False): @@ -8769,6 +8773,9 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """Traffic distribution configuration for the Reasoning Engine.""" + UpdateAgentEngineConfigOrDict = Union[ UpdateAgentEngineConfig, UpdateAgentEngineConfigDict @@ -10808,6 +10815,402 @@ class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): ] +class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): + """Config for getting an Agent Engine Runtime Revision.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for getting an Agent Engine Runtime Revision.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +GetAgentEngineRuntimeRevisionConfigOrDict = Union[ + GetAgentEngineRuntimeRevisionConfig, GetAgentEngineRuntimeRevisionConfigDict +] + + +class _GetAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for getting an agent engine runtime revision.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent engine runtime revision.""" + ) + config: Optional[GetAgentEngineRuntimeRevisionConfig] = Field( + default=None, description="""""" + ) + + +class _GetAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent engine runtime revision.""" + + name: Optional[str] + """Name of the agent engine runtime revision.""" + + config: Optional[GetAgentEngineRuntimeRevisionConfigDict] + """""" + + +_GetAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ + _GetAgentEngineRuntimeRevisionRequestParameters, + _GetAgentEngineRuntimeRevisionRequestParametersDict, +] + + +class ReasoningEngineRuntimeRevision(_common.BaseModel): + """A runtime revision.""" + + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, + description="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", + ) + state: Optional[State] = Field( + default=None, description="""Output only. The state of the revision.""" + ) + + +class ReasoningEngineRuntimeRevisionDict(TypedDict, total=False): + """A runtime revision.""" + + name: Optional[str] + """Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""" + + spec: Optional[ReasoningEngineSpecDict] + """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""" + + state: Optional[State] + """Output only. The state of the revision.""" + + +ReasoningEngineRuntimeRevisionOrDict = Union[ + ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict +] + + +class ListAgentEngineRuntimeRevisionsConfig(_common.BaseModel): + """Config for listing reasoning engine runtime revisions.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( + default=None, + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", + ) + + +class ListAgentEngineRuntimeRevisionsConfigDict(TypedDict, total=False): + """Config for listing reasoning engine runtime revisions.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" + + +ListAgentEngineRuntimeRevisionsConfigOrDict = Union[ + ListAgentEngineRuntimeRevisionsConfig, ListAgentEngineRuntimeRevisionsConfigDict +] + + +class _ListAgentEngineRuntimeRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing reasoning engine runtime revisions.""" + + name: Optional[str] = Field( + default=None, description="""Name of the reasoning engine.""" + ) + config: Optional[ListAgentEngineRuntimeRevisionsConfig] = Field( + default=None, description="""""" + ) + + +class _ListAgentEngineRuntimeRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing reasoning engine runtime revisions.""" + + name: Optional[str] + """Name of the reasoning engine.""" + + config: Optional[ListAgentEngineRuntimeRevisionsConfigDict] + """""" + + +_ListAgentEngineRuntimeRevisionsRequestParametersOrDict = Union[ + _ListAgentEngineRuntimeRevisionsRequestParameters, + _ListAgentEngineRuntimeRevisionsRequestParametersDict, +] + + +class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): + """Response for listing agent engine runtime revisions.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevision] + ] = Field( + default=None, description="""List of reasoning engine runtime revisions.""" + ) + + +class ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): + """Response for listing agent engine runtime revisions.""" + + sdk_http_response: Optional[genai_types.HttpResponseDict] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevisionDict] + ] + """List of reasoning engine runtime revisions.""" + + +ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ + ListReasoningEnginesRuntimeRevisionsResponse, + ListReasoningEnginesRuntimeRevisionsResponseDict, +] + + +class DeleteAgentEngineRuntimeRevisionConfig(_common.BaseModel): + """Config for deleting an Agent Engine Runtime Revision.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + + +class DeleteAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for deleting an Agent Engine Runtime Revision.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + + +DeleteAgentEngineRuntimeRevisionConfigOrDict = Union[ + DeleteAgentEngineRuntimeRevisionConfig, DeleteAgentEngineRuntimeRevisionConfigDict +] + + +class _DeleteAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for deleting agent engine runtime revisions.""" + + name: Optional[str] = Field( + default=None, + description="""Name of the agent engine runtime revision to delete.""", + ) + config: Optional[DeleteAgentEngineRuntimeRevisionConfig] = Field( + default=None, description="""""" + ) + + +class _DeleteAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent engine runtime revisions.""" + + name: Optional[str] + """Name of the agent engine runtime revision to delete.""" + + config: Optional[DeleteAgentEngineRuntimeRevisionConfigDict] + """""" + + +_DeleteAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ + _DeleteAgentEngineRuntimeRevisionRequestParameters, + _DeleteAgentEngineRuntimeRevisionRequestParametersDict, +] + + +class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): + """Operation for deleting agent engine runtime revisions.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): + """Operation for deleting agent engine runtime revisions.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteAgentEngineRuntimeRevisionOperationOrDict = Union[ + DeleteAgentEngineRuntimeRevisionOperation, + DeleteAgentEngineRuntimeRevisionOperationDict, +] + + +class GetDeleteAgentEngineRuntimeRevisionOperationConfig(_common.BaseModel): + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetDeleteAgentEngineRuntimeRevisionOperationConfigDict(TypedDict, total=False): + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict = Union[ + GetDeleteAgentEngineRuntimeRevisionOperationConfig, + GetDeleteAgentEngineRuntimeRevisionOperationConfigDict, +] + + +class _GetDeleteAgentEngineRuntimeRevisionOperationParameters(_common.BaseModel): + """Parameters for getting an operation that deletes a agent engine runtime revision.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict( + TypedDict, total=False +): + """Parameters for getting an operation that deletes a agent engine runtime revision.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfigDict] + """Used to override the default configuration.""" + + +_GetDeleteAgentEngineRuntimeRevisionOperationParametersOrDict = Union[ + _GetDeleteAgentEngineRuntimeRevisionOperationParameters, + _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict, +] + + +class QueryAgentEngineRuntimeRevisionConfig(_common.BaseModel): + """Config for querying agent engine runtime revisions.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + class_method: Optional[str] = Field( + default=None, description="""The class method to call.""" + ) + input: Optional[dict[str, Any]] = Field( + default=None, description="""The input to the class method.""" + ) + include_all_fields: Optional[bool] = Field(default=False, description="""""") + + +class QueryAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for querying agent engine runtime revisions.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + class_method: Optional[str] + """The class method to call.""" + + input: Optional[dict[str, Any]] + """The input to the class method.""" + + include_all_fields: Optional[bool] + """""" + + +QueryAgentEngineRuntimeRevisionConfigOrDict = Union[ + QueryAgentEngineRuntimeRevisionConfig, QueryAgentEngineRuntimeRevisionConfigDict +] + + +class _QueryAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for querying agent engine runtime revisions.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent engine runtime revision.""" + ) + config: Optional[QueryAgentEngineRuntimeRevisionConfig] = Field( + default=None, description="""""" + ) + + +class _QueryAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent engine runtime revisions.""" + + name: Optional[str] + """Name of the agent engine runtime revision.""" + + config: Optional[QueryAgentEngineRuntimeRevisionConfigDict] + """""" + + +_QueryAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ + _QueryAgentEngineRuntimeRevisionRequestParameters, + _QueryAgentEngineRuntimeRevisionRequestParametersDict, +] + + class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): """The code execution environment with customized settings.""" @@ -16139,37 +16542,6 @@ def delete( self.api_client.delete(name=self.api_resource.name, force=force, config=config) # type: ignore[union-attr] -RubricContentProperty = evals_types.RubricContentProperty -RubricContentPropertyDict = evals_types.RubricContentPropertyDict -RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict - -RubricContent = evals_types.RubricContent -RubricContentDict = evals_types.RubricContentDict -RubricContentDictOrDict = evals_types.RubricContentOrDict - -Rubric = evals_types.Rubric -RubricDict = evals_types.RubricDict -RubricDictOrDict = evals_types.RubricOrDict - -RubricVerdict = evals_types.RubricVerdict -RubricVerdictDict = evals_types.RubricVerdictDict -RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict - -CandidateResult = evals_types.CandidateResult -CandidateResultDict = evals_types.CandidateResultDict -CandidateResultDictOrDict = evals_types.CandidateResultOrDict - -Event = evals_types.Event -EventDict = evals_types.EventDict -EventDictOrDict = evals_types.EventOrDict - -Message = evals_types.Message -MessageDict = evals_types.MessageDict -MessageDictOrDict = evals_types.MessageOrDict - -Importance = evals_types.Importance - - class AgentEngineDict(TypedDict, total=False): """An agent engine instance.""" @@ -16383,6 +16755,9 @@ class AgentEngineConfig(_common.BaseModel): container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( default=None, description="""The container spec for the Agent Engine.""" ) + traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + default=None, description="""The traffic config for the Agent Engine.""" + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -16557,6 +16932,9 @@ class AgentEngineConfigDict(TypedDict, total=False): container_spec: Optional[ReasoningEngineSpecContainerSpecDict] """The container spec for the Agent Engine.""" + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """The traffic config for the Agent Engine.""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] @@ -17142,3 +17520,104 @@ class OptimizeJobConfigDict(TypedDict, total=False): OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] + + +class AgentEngineRuntimeRevision(_common.BaseModel): + """An agent engine runtime revision instance.""" + + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" + ) + api_async_client: Optional[Any] = Field( + default=None, + description="""The underlying API client for asynchronous operations.""", + ) + api_resource: Optional[ReasoningEngineRuntimeRevision] = Field( + default=None, + description="""The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""", + ) + + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") + + def __repr__(self) -> str: + return ( + f"AgentEngineRuntimeRevision(api_resource.name='{self.api_resource.name}')" + if self.api_resource is not None + else "AgentEngineRuntimeRevision(api_resource.name=None)" + ) + + def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: + """Returns the schemas of all registered operations for the agent.""" + if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): + raise ValueError("api_resource is not initialized.") + if not self.api_resource.spec: + raise ValueError("api_resource.spec is not initialized.") + return self.api_resource.spec.class_methods + + def delete( + self, + config: Optional[DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + ) -> None: + """Deletes the agent engine runtime revision. + + Args: + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Engine has undeleted child resources. Defaults to False. + config (DeleteAgentEngineRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine Runtime Revision. + """ + if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): + raise ValueError("api_resource is not initialized.") + self.api_client.delete(name=self.api_resource.name, config=config) # type: ignore[union-attr] + + +RubricContentProperty = evals_types.RubricContentProperty +RubricContentPropertyDict = evals_types.RubricContentPropertyDict +RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict + +RubricContent = evals_types.RubricContent +RubricContentDict = evals_types.RubricContentDict +RubricContentDictOrDict = evals_types.RubricContentOrDict + +Rubric = evals_types.Rubric +RubricDict = evals_types.RubricDict +RubricDictOrDict = evals_types.RubricOrDict + +RubricVerdict = evals_types.RubricVerdict +RubricVerdictDict = evals_types.RubricVerdictDict +RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict + +CandidateResult = evals_types.CandidateResult +CandidateResultDict = evals_types.CandidateResultDict +CandidateResultDictOrDict = evals_types.CandidateResultOrDict + +Event = evals_types.Event +EventDict = evals_types.EventDict +EventDictOrDict = evals_types.EventOrDict + +Message = evals_types.Message +MessageDict = evals_types.MessageDict +MessageDictOrDict = evals_types.MessageOrDict + +Importance = evals_types.Importance + + +class AgentEngineRuntimeRevisionDict(TypedDict, total=False): + """An agent engine runtime revision instance.""" + + api_client: Optional[Any] + """The underlying API client.""" + + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" + + api_resource: Optional[ReasoningEngineRuntimeRevisionDict] + """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" + + +AgentEngineRuntimeRevisionOrDict = Union[ + AgentEngineRuntimeRevision, AgentEngineRuntimeRevisionDict +]