From c81f9124dce841a7cbc310c6c7652ad793c9a58f Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Thu, 30 Oct 2025 14:29:25 -0700 Subject: [PATCH 01/24] chore!: Switch `cloudtrace.googleapis.com` to `telemetry.googleapis.com` for tracing API. PiperOrigin-RevId: 826188567 --- setup.py | 1 + .../test_agent_engine_templates_adk.py | 46 ++---- vertexai/agent_engines/templates/adk.py | 153 +++++++++++------- 3 files changed, 106 insertions(+), 94 deletions(-) diff --git a/setup.py b/setup.py index e95d484ceb..1cef024ab1 100644 --- a/setup.py +++ b/setup.py @@ -162,6 +162,7 @@ "google-cloud-logging < 4", "opentelemetry-sdk < 2", "opentelemetry-exporter-gcp-trace < 2", + "opentelemetry-exporter-otlp-proto-http < 2", "pydantic >= 2.11.1, < 3", "typing_extensions", ] diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 4dfcf5acdd..af89d0ea33 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -18,7 +18,6 @@ import os from unittest import mock from typing import Optional -import dataclasses from google import auth import vertexai @@ -97,27 +96,11 @@ def vertexai_init_mock(): @pytest.fixture -def cloud_trace_exporter_mock(): - import sys - import opentelemetry - - mock_cloud_trace_exporter = mock.Mock() - - opentelemetry.exporter = type(sys)("exporter") - opentelemetry.exporter.cloud_trace = type(sys)("cloud_trace") - opentelemetry.exporter.cloud_trace.CloudTraceSpanExporter = ( - mock_cloud_trace_exporter - ) - - sys.modules["opentelemetry.exporter"] = opentelemetry.exporter - sys.modules["opentelemetry.exporter.cloud_trace"] = ( - opentelemetry.exporter.cloud_trace - ) - - yield mock_cloud_trace_exporter - - del sys.modules["opentelemetry.exporter.cloud_trace"] - del sys.modules["opentelemetry.exporter"] +def otlp_span_exporter_mock(): + with mock.patch( + "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter" + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock @pytest.fixture @@ -609,9 +592,9 @@ def test_custom_instrumentor_enablement( ) def test_tracing_setup( self, - trace_provider_mock: mock.Mock, - cloud_trace_exporter_mock: mock.Mock, monkeypatch, + trace_provider_mock: mock.Mock, + otlp_span_exporter_mock: mock.Mock, ): monkeypatch.setattr( "uuid.uuid4", lambda: uuid.UUID("12345678123456781234567812345678") @@ -633,17 +616,9 @@ def test_tracing_setup( "some-attribute": "some-value", } - @dataclasses.dataclass - class RegexMatchingAll: - keys: set[str] - - def __eq__(self, regex: object) -> bool: - return isinstance(regex, str) and set(regex.split("|")) == self.keys - - cloud_trace_exporter_mock.assert_called_once_with( - project_id=_TEST_PROJECT, - client=mock.ANY, - resource_regex=RegexMatchingAll(keys=set(expected_attributes.keys())), + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", ) assert ( @@ -655,7 +630,6 @@ def __eq__(self, regex: object) -> bool: def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, ): diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index ca283ee1b0..b517f7ce4c 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -231,6 +231,28 @@ def _warn(msg: str): _warn._LOGGER.warning(msg) # pyright: ignore[reportFunctionMemberAccess] +def _force_flush_traces(): + try: + import opentelemetry.trace + except (ImportError, AttributeError): + _warn( + "Could not force flush traces. opentelemetry-api is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + ) + return None + + try: + import opentelemetry.sdk.trace + except (ImportError, AttributeError): + _warn( + "Could not force flush traces. opentelemetry-sdk is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + ) + return None + + provider = opentelemetry.trace.get_tracer_provider() + if isinstance(provider, opentelemetry.sdk.trace.TracerProvider): + _ = provider.force_flush() + + def _default_instrumentor_builder( project_id: str, *, @@ -311,28 +333,23 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: if enable_tracing: try: - import opentelemetry.exporter.cloud_trace + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import google.auth.transport.requests except (ImportError, AttributeError): return _warn_missing_dependency( - "opentelemetry-exporter-gcp-trace", needed_for_tracing=True - ) - - try: - import google.cloud.trace_v2 - except (ImportError, AttributeError): - return _warn_missing_dependency( - "google-cloud-trace", needed_for_tracing=True + "opentelemetry-exporter-otlp-proto-http", needed_for_tracing=True ) import google.auth credentials, _ = google.auth.default() - span_exporter = opentelemetry.exporter.cloud_trace.CloudTraceSpanExporter( - project_id=project_id, - client=google.cloud.trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), - resource_regex="|".join(resource.attributes.keys()), + span_exporter = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials + ), + endpoint="https://telemetry.googleapis.com/v1/traces", + ) ) span_processor = opentelemetry.sdk.trace.export.BatchSpanProcessor( span_exporter=span_exporter, @@ -695,54 +712,17 @@ def set_up(self): else: os.environ["ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS"] = "false" - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( - "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" - ) - - def telemetry_enabled() -> Optional[bool]: - return ( - os.getenv(GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY, "0").lower() - in ("true", "1") - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in os.environ - else None - ) - - # Tracing enablement follows truth table: - def tracing_enabled() -> bool: - """Tracing enablement follows true table: - - | enable_tracing | enable_telemetry(env) | tracing_actually_enabled | - |----------------|-----------------------|--------------------------| - | false | false | false | - | false | true | false | - | false | None | false | - | true | false | false | - | true | true | true | - | true | None | true | - | None(default) | false | false | - | None(default) | true | adk_version >= 1.17 | - | None(default) | None | false | - """ - enable_tracing: Optional[bool] = self._tmpl_attrs.get("enable_tracing") - enable_telemetry: Optional[bool] = telemetry_enabled() - - return (enable_tracing is True and enable_telemetry is not False) or ( - enable_tracing is None - and enable_telemetry is True - and is_version_sufficient("1.17.0") - ) - - enable_logging = bool(telemetry_enabled()) + enable_logging = bool(self._telemetry_enabled()) custom_instrumentor = self._tmpl_attrs.get("instrumentor_builder") - if custom_instrumentor and tracing_enabled(): + if custom_instrumentor and self._tracing_enabled(): self._tmpl_attrs["instrumentor"] = custom_instrumentor(project) if not custom_instrumentor: self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder( project, - enable_tracing=tracing_enabled(), + enable_tracing=self._tracing_enabled(), enable_logging=enable_logging, ) @@ -914,9 +894,14 @@ async def async_stream_query( **kwargs, ) - async for event in events_async: - # Yield the event data as a dictionary - yield _utils.dump_event_for_json(event) + try: + async for event in events_async: + # Yield the event data as a dictionary + yield _utils.dump_event_for_json(event) + finally: + # Avoid trace data loss having to do with CPU throttling on instance turndown + if self._tracing_enabled(): + _ = await asyncio.to_thread(_force_flush_traces) def stream_query( self, @@ -1068,6 +1053,9 @@ async def streaming_agent_run_with_events(self, request_json: str): user_id=request.user_id, session_id=session.id, ) + # Avoid trace data loss having to do with CPU throttling on instance turndown + if self._tracing_enabled(): + _ = await asyncio.to_thread(_force_flush_traces) async def async_get_session( self, @@ -1450,3 +1438,52 @@ def register_operations(self) -> Dict[str, List[str]]: "streaming_agent_run_with_events", ], } + + def _telemetry_enabled(self) -> Optional[bool]: + """Return status of telemetry enablement depending on enablement env variable. + + In detail: + - Logging is always enabled when telemetry is enabled. + - Tracing is enabled depending on the truth table seen in `_tracing_enabled` method, in order to not break existing user enablement. + + Returns: + True if telemetry is enabled, False if telemetry is disabled, or None + if telemetry enablement is not set (i.e. old deployments which don't support this env variable). + """ + import os + + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( + "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" + ) + + return ( + os.getenv(GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY, "0").lower() + in ("true", "1") + if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in os.environ + else None + ) + + # Tracing enablement follows truth table: + def _tracing_enabled(self) -> bool: + """Tracing enablement follows true table: + + | enable_tracing | enable_telemetry(env) | tracing_actually_enabled | + |----------------|-----------------------|--------------------------| + | false | false | false | + | false | true | false | + | false | None | false | + | true | false | false | + | true | true | true | + | true | None | true | + | None(default) | false | false | + | None(default) | true | adk_version >= 1.17 | + | None(default) | None | false | + """ + enable_tracing: Optional[bool] = self._tmpl_attrs.get("enable_tracing") + enable_telemetry: Optional[bool] = self._telemetry_enabled() + + return (enable_tracing is True and enable_telemetry is not False) or ( + enable_tracing is None + and enable_telemetry is True + and is_version_sufficient("1.17.0") + ) From 2f497ee4f8b7f1bfb013628ebe19ef65d2b339a7 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 31 Oct 2025 02:57:51 -0700 Subject: [PATCH 02/24] chore: add cloud.platform attribute to OTel resource This applies to ADK Apps on Agent Engine and follows https://github.com/open-telemetry/semantic-conventions/pull/2957. PiperOrigin-RevId: 826404720 --- tests/unit/vertex_adk/test_agent_engine_templates_adk.py | 1 + tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py | 1 + vertexai/agent_engines/templates/adk.py | 1 + vertexai/preview/reasoning_engines/templates/adk.py | 1 + 4 files changed, 4 insertions(+) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index af89d0ea33..2b5d8fab21 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -609,6 +609,7 @@ def test_tracing_setup( "telemetry.sdk.version": "1.36.0", "gcp.project_id": "test-project", "cloud.account.id": "test-project", + "cloud.platform": "gcp.agent_engine", "service.name": "test_agent_id", "cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project/locations/us-central1/reasoningEngines/test_agent_id", "service.instance.id": "12345678123456781234567812345678-123123123", diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index c6e9c02e7c..6ce5222114 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -636,6 +636,7 @@ def test_tracing_setup( "telemetry.sdk.version": "1.36.0", "gcp.project_id": "test-project", "cloud.account.id": "test-project", + "cloud.platform": "gcp.agent_engine", "service.name": "test_agent_id", "cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project/locations/us-central1/reasoningEngines/test_agent_id", "service.instance.id": "12345678123456781234567812345678-123123123", diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index b517f7ce4c..c16aed3a1a 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -320,6 +320,7 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: attributes={ "gcp.project_id": project_id, "cloud.account.id": project_id, + "cloud.platform": "gcp.agent_engine", "service.name": os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", ""), "service.instance.id": f"{uuid.uuid4().hex}-{os.getpid()}", "cloud.region": os.getenv("GOOGLE_CLOUD_LOCATION", ""), diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 9cc9dc876c..a4d31e5034 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -300,6 +300,7 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: attributes={ "gcp.project_id": project_id, "cloud.account.id": project_id, + "cloud.platform": "gcp.agent_engine", "service.name": os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", ""), "service.instance.id": f"{uuid.uuid4().hex}-{os.getpid()}", "cloud.region": os.getenv("GOOGLE_CLOUD_LOCATION", ""), From 27ef56b8319ba793f6f00e7857fe0c20c2c8b61a Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 31 Oct 2025 03:57:34 -0700 Subject: [PATCH 03/24] chore!: Switch tracing APIs in preview AdkApp. Currently AdkApp uses `cloudtrace.googleapis.com` for GCP tracing. This change switches it to `telemetry.googleapis.com`. It's a breaking change as users might need to enable the new API. PiperOrigin-RevId: 826422072 --- setup.py | 1 + .../test_reasoning_engine_templates_adk.py | 46 +++----------- .../reasoning_engines/templates/adk.py | 61 +++++++++++++------ 3 files changed, 54 insertions(+), 54 deletions(-) diff --git a/setup.py b/setup.py index 1cef024ab1..e199577e8a 100644 --- a/setup.py +++ b/setup.py @@ -151,6 +151,7 @@ "google-cloud-trace < 2", "opentelemetry-sdk < 2", "opentelemetry-exporter-gcp-trace < 2", + "opentelemetry-exporter-otlp-proto-http < 2", "pydantic >= 2.11.1, < 3", "typing_extensions", ] diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index 6ce5222114..32c0319fc4 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -16,7 +16,6 @@ import base64 import importlib import json -import dataclasses import os from unittest import mock from typing import Optional @@ -112,27 +111,11 @@ def simple_span_processor_mock(): @pytest.fixture -def cloud_trace_exporter_mock(): - import sys - import opentelemetry - - mock_cloud_trace_exporter = mock.Mock() - - opentelemetry.exporter = type(sys)("exporter") - opentelemetry.exporter.cloud_trace = type(sys)("cloud_trace") - opentelemetry.exporter.cloud_trace.CloudTraceSpanExporter = ( - mock_cloud_trace_exporter - ) - - sys.modules["opentelemetry.exporter"] = opentelemetry.exporter - sys.modules["opentelemetry.exporter.cloud_trace"] = ( - opentelemetry.exporter.cloud_trace - ) - - yield mock_cloud_trace_exporter - - del sys.modules["opentelemetry.exporter.cloud_trace"] - del sys.modules["opentelemetry.exporter"] +def otlp_span_exporter_mock(): + with mock.patch( + "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter" + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock @pytest.fixture @@ -619,9 +602,9 @@ def test_default_instrumentor_enablement( ) def test_tracing_setup( self, - trace_provider_mock: mock.Mock, - cloud_trace_exporter_mock: mock.Mock, monkeypatch: pytest.MonkeyPatch, + trace_provider_mock: mock.Mock, + otlp_span_exporter_mock: mock.Mock, ): monkeypatch.setattr( "uuid.uuid4", lambda: uuid.UUID("12345678123456781234567812345678") @@ -644,17 +627,9 @@ def test_tracing_setup( "some-attribute": "some-value", } - @dataclasses.dataclass - class RegexMatchingAll: - keys: set[str] - - def __eq__(self, regex: object) -> bool: - return isinstance(regex, str) and set(regex.split("|")) == self.keys - - cloud_trace_exporter_mock.assert_called_once_with( - project_id=_TEST_PROJECT, - client=mock.ANY, - resource_regex=RegexMatchingAll(keys=set(expected_attributes.keys())), + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", ) assert ( @@ -686,7 +661,6 @@ def test_span_content_capture_enabled_with_tracing(self): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, ): diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index a4d31e5034..345ff981f1 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -233,6 +233,28 @@ def _warn(msg: str): _warn._LOGGER.warning(msg) # pyright: ignore[reportFunctionMemberAccess] +def _force_flush_traces(): + try: + import opentelemetry.trace + except (ImportError, AttributeError): + _warn( + "Could not force flush traces. opentelemetry-api is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + ) + return None + + try: + import opentelemetry.sdk.trace + except (ImportError, AttributeError): + _warn( + "Could not force flush traces. opentelemetry-sdk is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + ) + return None + + provider = opentelemetry.trace.get_tracer_provider() + if isinstance(provider, opentelemetry.sdk.trace.TracerProvider): + _ = provider.force_flush() + + def _default_instrumentor_builder( project_id: str, *, @@ -314,28 +336,23 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: if enable_tracing: try: - import opentelemetry.exporter.cloud_trace - except (ImportError, AttributeError): - return _warn_missing_dependency( - "opentelemetry-exporter-gcp-trace", needed_for_tracing=True - ) - - try: - import google.cloud.trace_v2 + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import google.auth.transport.requests except (ImportError, AttributeError): return _warn_missing_dependency( - "google-cloud-trace", needed_for_tracing=True + "opentelemetry-exporter-otlp-proto-http", needed_for_tracing=True ) import google.auth credentials, _ = google.auth.default() - span_exporter = opentelemetry.exporter.cloud_trace.CloudTraceSpanExporter( - project_id=project_id, - client=google.cloud.trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), - resource_regex="|".join(resource.attributes.keys()), + span_exporter = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials + ), + endpoint="https://telemetry.googleapis.com/v1/traces", + ) ) span_processor = opentelemetry.sdk.trace.export.BatchSpanProcessor( span_exporter=span_exporter, @@ -875,9 +892,14 @@ async def async_stream_query( **kwargs, ) - async for event in events_async: - # Yield the event data as a dictionary - yield _utils.dump_event_for_json(event) + try: + async for event in events_async: + # Yield the event data as a dictionary + yield _utils.dump_event_for_json(event) + finally: + # Avoid trace data loss having to do with CPU throttling on instance turndown + if self._tracing_enabled(): + _ = await asyncio.to_thread(_force_flush_traces) def streaming_agent_run_with_events(self, request_json: str): import json @@ -938,6 +960,9 @@ async def _invoke_agent_async(): user_id=request.user_id, session_id=session.id, ) + # Avoid trace data loss having to do with CPU throttling on instance turndown + if self._tracing_enabled(): + _ = await asyncio.to_thread(_force_flush_traces) def _asyncio_thread_main(): try: From f51b813bb5ca436b915b6c4b332272b70c7cd2b6 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 31 Oct 2025 07:08:09 -0700 Subject: [PATCH 04/24] chore: Enable default-on telemetry for ADK agents. PiperOrigin-RevId: 826475712 --- .../test_agent_engine_templates_adk.py | 228 ++++++++++++++++++ .../vertex_langchain/test_agent_engines.py | 2 +- .../unit/vertexai/genai/test_agent_engines.py | 47 ++++ vertexai/_genai/_agent_engines_utils.py | 47 ++++ vertexai/_genai/agent_engines.py | 2 + vertexai/agent_engines/_agent_engines.py | 81 ++++++- 6 files changed, 405 insertions(+), 2 deletions(-) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 2b5d8fab21..1eb91046bb 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -16,14 +16,25 @@ import importlib import json import os +import cloudpickle +import sys from unittest import mock from typing import Optional from google import auth +from google.auth import credentials as auth_credentials +from google.cloud import storage import vertexai +from google.cloud import aiplatform +from google.cloud.aiplatform_v1 import types as aip_types +from google.cloud.aiplatform_v1.services import reasoning_engine_service +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from vertexai.agent_engines import _utils from vertexai import agent_engines +from vertexai.agent_engines.templates import adk as adk_template +from vertexai.agent_engines import _agent_engines +from google.api_core import operation as ga_operation from google.genai import types import pytest import uuid @@ -75,6 +86,52 @@ def __init__(self, name: str, model: str): "streaming_mode": "sse", "max_llm_calls": 500, } +_TEST_STAGING_BUCKET = "gs://test-bucket" +_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) +_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" +_TEST_RESOURCE_ID = "1028944691210842416" +_TEST_AGENT_ENGINE_RESOURCE_NAME = ( + f"{_TEST_PARENT}/reasoningEngines/{_TEST_RESOURCE_ID}" +) +_TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" +_TEST_GCS_DIR_NAME = _agent_engines._DEFAULT_GCS_DIR_NAME +_TEST_BLOB_FILENAME = _agent_engines._BLOB_FILENAME +_TEST_REQUIREMENTS_FILE = _agent_engines._REQUIREMENTS_FILE +_TEST_EXTRA_PACKAGES_FILE = _agent_engines._EXTRA_PACKAGES_FILE +_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_PACKAGE_SPEC = aip_types.ReasoningEngineSpec.PackageSpec( + 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, +) +_ADK_AGENT_FRAMEWORK = adk_template.AdkApp.agent_framework +_TEST_AGENT_ENGINE_OBJ = aip_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + spec=aip_types.ReasoningEngineSpec( + package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, + agent_framework=_ADK_AGENT_FRAMEWORK, + ), +) + +GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( + "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" +) @pytest.fixture(scope="module") @@ -727,3 +784,174 @@ async def test_async_stream_query_invalid_message_type(self): ): async for _ in app.async_stream_query(user_id=_TEST_USER_ID, message=123): pass + + +@pytest.fixture(scope="module") +def create_agent_engine_mock(): + with mock.patch.object( + reasoning_engine_service.ReasoningEngineServiceClient, + "create_reasoning_engine", + ) as create_agent_engine_mock: + create_agent_engine_lro_mock = mock.Mock(ga_operation.Operation) + create_agent_engine_lro_mock.result.return_value = _TEST_AGENT_ENGINE_OBJ + create_agent_engine_mock.return_value = create_agent_engine_lro_mock + yield create_agent_engine_mock + + +@pytest.fixture(scope="module") +def get_agent_engine_mock(): + with mock.patch.object( + reasoning_engine_service.ReasoningEngineServiceClient, + "get_reasoning_engine", + ) as get_agent_engine_mock: + api_client_mock = mock.Mock() + api_client_mock.get_reasoning_engine.return_value = _TEST_AGENT_ENGINE_OBJ + get_agent_engine_mock.return_value = api_client_mock + yield get_agent_engine_mock + + +@pytest.fixture(scope="module") +def cloud_storage_create_bucket_mock(): + with mock.patch.object(storage, "Client") as cloud_storage_mock: + bucket_mock = mock.Mock(spec=storage.Bucket) + bucket_mock.blob.return_value.open.return_value = "blob_file" + bucket_mock.blob.return_value.upload_from_filename.return_value = None + bucket_mock.blob.return_value.upload_from_string.return_value = None + + cloud_storage_mock.get_bucket = mock.Mock( + side_effect=ValueError("bucket not found") + ) + cloud_storage_mock.bucket.return_value = bucket_mock + cloud_storage_mock.create_bucket.return_value = bucket_mock + + yield cloud_storage_mock + + +@pytest.fixture(scope="module") +def cloudpickle_dump_mock(): + with mock.patch.object(cloudpickle, "dump") as cloudpickle_dump_mock: + yield cloudpickle_dump_mock + + +@pytest.fixture(scope="module") +def cloudpickle_load_mock(): + with mock.patch.object(cloudpickle, "load") as cloudpickle_load_mock: + yield cloudpickle_load_mock + + +@pytest.fixture(scope="function") +def get_gca_resource_mock(): + with mock.patch.object( + base.VertexAiResourceNoun, + "_get_gca_resource", + ) as get_gca_resource_mock: + get_gca_resource_mock.return_value = _TEST_AGENT_ENGINE_OBJ + yield get_gca_resource_mock + + +# Function scope is required for the pytest parameterized tests. +@pytest.fixture(scope="function") +def update_agent_engine_mock(): + with mock.patch.object( + reasoning_engine_service.ReasoningEngineServiceClient, + "update_reasoning_engine", + ) as update_agent_engine_mock: + yield update_agent_engine_mock + + +@pytest.mark.usefixtures("google_auth_mock") +class TestAgentEngines: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + staging_bucket=_TEST_STAGING_BUCKET, + ) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "env_vars,expected_env_vars", + [ + ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), + (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), + ( + {"some_env": "some_val"}, + { + "some_env": "some_val", + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true", + }, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + ), + ], + ) + def test_create_default_telemetry_enablement( + self, + create_agent_engine_mock: mock.Mock, + cloud_storage_create_bucket_mock: mock.Mock, + cloudpickle_dump_mock: mock.Mock, + cloudpickle_load_mock: mock.Mock, + get_gca_resource_mock: mock.Mock, + env_vars: dict[str, str], + expected_env_vars: dict[str, str], + ): + agent_engines.create( + agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT), + env_vars=env_vars, + ) + create_agent_engine_mock.assert_called_once() + deployment_spec = create_agent_engine_mock.call_args.kwargs[ + "reasoning_engine" + ].spec.deployment_spec + assert _utils.to_dict(deployment_spec)["env"] == [ + {"name": key, "value": value} for key, value in expected_env_vars.items() + ] + + @pytest.mark.parametrize( + "env_vars,expected_env_vars", + [ + ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), + (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), + ( + {"some_env": "some_val"}, + { + "some_env": "some_val", + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true", + }, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + ), + ], + ) + def test_update_default_telemetry_enablement( + self, + update_agent_engine_mock: mock.Mock, + cloud_storage_create_bucket_mock: mock.Mock, + cloudpickle_dump_mock: mock.Mock, + cloudpickle_load_mock: mock.Mock, + get_gca_resource_mock: mock.Mock, + get_agent_engine_mock: mock.Mock, + env_vars: dict[str, str], + expected_env_vars: dict[str, str], + ): + agent_engines.update( + resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + description="foobar", # avoid "At least one of ... must be specified" errors. + env_vars=env_vars, + ) + update_agent_engine_mock.assert_called_once() + deployment_spec = update_agent_engine_mock.call_args.kwargs[ + "request" + ].reasoning_engine.spec.deployment_spec + assert _utils.to_dict(deployment_spec)["env"] == [ + {"name": key, "value": value} for key, value in expected_env_vars.items() + ] diff --git a/tests/unit/vertex_langchain/test_agent_engines.py b/tests/unit/vertex_langchain/test_agent_engines.py index 133b2d3b19..8c3610577e 100644 --- a/tests/unit/vertex_langchain/test_agent_engines.py +++ b/tests/unit/vertex_langchain/test_agent_engines.py @@ -3260,7 +3260,7 @@ def test_create_agent_engine_with_invalid_type_env_var( "TEST_ENV_VAR": 0.01, # should be a string or dict or SecretRef }, ) - with pytest.raises(TypeError, match="env_vars must be a list or a dict"): + with pytest.raises(TypeError, match="env_vars must be a list, tuple or a dict"): agent_engines.create( self.test_agent, display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index bc204d5f79..37e678f7f6 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -31,6 +31,7 @@ from google.cloud import aiplatform import vertexai from google.cloud.aiplatform import initializer +from vertexai.agent_engines.templates import adk from vertexai._genai import _agent_engines_utils from vertexai._genai import agent_engines from vertexai._genai import types as _genai_types @@ -40,6 +41,9 @@ _TEST_AGENT_FRAMEWORK = "test-agent-framework" +GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( + "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" +) class CapitalizeEngine: @@ -848,6 +852,49 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): "description": _TEST_AGENT_ENGINE_DESCRIPTION, } + @mock.patch.object(_agent_engines_utils, "_prepare") + @pytest.mark.parametrize( + "env_vars,expected_env_vars", + [ + ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), + (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), + ( + {"some_env": "some_val"}, + { + "some_env": "some_val", + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true", + }, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + ), + ], + ) + def test_agent_engine_adk_telemetry_enablement( + self, + mock_prepare: mock.Mock, + env_vars: dict[str, str], + expected_env_vars: dict[str, str], + ): + agent = mock.Mock(spec=adk.AdkApp) + agent.clone = lambda: agent + agent.register_operations = lambda: {} + + config = self.client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + description=_TEST_AGENT_ENGINE_DESCRIPTION, + env_vars=env_vars, + ) + assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME + assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION + assert config["spec"]["deployment_spec"]["env"] == [ + {"name": key, "value": value} for key, value in expected_env_vars.items() + ] + @mock.patch.object(_agent_engines_utils, "_prepare") def test_create_agent_engine_config_full(self, mock_prepare): config = self.client.agent_engines._create_config( diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 52cb009df2..8364212528 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -1845,3 +1845,50 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None: f"Memory size of {memory_str} requires at least {min_cpu} CPUs." f" Got {cpu}" ) + + +def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool: + """Checks if the agent engine is an ADK agent. + + Args: + agent_engine: The agent engine to check. + + Returns: + True if the agent engine is an ADK agent, False otherwise. + """ + + from vertexai.agent_engines.templates import adk + + return isinstance(agent_engine, adk.AdkApp) + + +def _add_telemetry_enablement_env( + env_vars: Optional[Dict[str, Union[str, Any]]] +) -> Optional[Dict[str, Union[str, Any]]]: + """Adds telemetry enablement env var to the env vars. + + This is in order to achieve default-on telemetry. + If the telemetry enablement env var is already set, we do not override it. + + Args: + env_vars: The env vars to add the telemetry enablement env var to. + + Returns: + The env vars with the telemetry enablement env var added. + """ + + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( + "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" + ) + env_to_add = {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"} + + if env_vars is None: + return env_to_add + + if not isinstance(env_vars, dict): + raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") + + if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars: + return env_vars + + return env_vars | env_to_add diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 98367a9421..2b3ad56d8f 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -1032,6 +1032,8 @@ def _create_config( raise ValueError("location must be set using `vertexai.Client`.") gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME agent = _agent_engines_utils._validate_agent_or_raise(agent=agent) + if _agent_engines_utils._is_adk_agent(agent): + env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise( staging_bucket=staging_bucket, ) diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index 9173883d3d..596056b15a 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -514,9 +514,13 @@ def create( _validate_sys_version_or_raise(sys_version) gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME staging_bucket = initializer.global_config.staging_bucket + if agent_engine is not None: agent_engine = _validate_agent_engine_or_raise(agent_engine) staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) + if _is_adk_agent(None, agent_engine): + env_vars = _add_telemetry_enablement_env(env_vars=env_vars) + if agent_engine is None: if requirements is not None: raise ValueError("requirements must be None if agent_engine is None.") @@ -533,6 +537,7 @@ def create( sdk_resource = cls.__new__(cls) base.VertexAiResourceNounWithFutureManager.__init__(sdk_resource) + # Prepares the Agent Engine for creation in Vertex AI. # This involves packaging and uploading the artifacts for # agent_engine, requirements and extra_packages to @@ -798,6 +803,9 @@ def update( if agent_engine is not None: agent_engine = _validate_agent_engine_or_raise(agent_engine) + if _is_adk_agent(self, agent_engine): + env_vars = _add_telemetry_enablement_env(env_vars=env_vars) + # Prepares the Agent Engine for update in Vertex AI. This involves # packaging and uploading the artifacts for agent_engine, requirements # and extra_packages to `staging_bucket/gcs_dir_name`. @@ -1056,6 +1064,77 @@ def _validate_agent_engine_or_raise( return agent_engine +def _is_adk_agent( + agent_engine_to_update: Optional[AgentEngine], + new_agent_engine: Optional[_AgentEngineInterface], +) -> bool: + """Checks if the agent engine is an ADK agent. + + Args: + agent_engine_to_update: Existing agent engine, None if creating new one. + new_agent_engine: The new agent engine to deploy. Can be None during an update, if the Python agent implementation is not provided, and should remain unchanged. + + Returns: + True if the agent after the create/update operation, will be an ADK agent. + """ + + from vertexai.agent_engines.templates import adk + + if new_agent_engine is not None: + return ( + getattr(new_agent_engine, "agent_framework", None) + == adk.AdkApp.agent_framework + ) + if agent_engine_to_update is not None: + return ( + agent_engine_to_update.gca_resource.spec.agent_framework + == adk.AdkApp.agent_framework + ) + return False + + +EnvVars = Optional[Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]] + + +def _add_telemetry_enablement_env(*, env_vars: EnvVars) -> EnvVars: + """Adds telemetry enablement env var to the env vars. + + This is in order to achieve default-on telemetry. + If the telemetry enablement env var is already set, we do not override it. + + Args: + env_vars: The env vars to add the telemetry enablement env var to. + + Returns: + The env vars with the telemetry enablement env var added. + """ + + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( + "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" + ) + + if env_vars is None: + return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"} + if isinstance(env_vars, dict): + return ( + env_vars + if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars + else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"} + ) + if isinstance(env_vars, list) or isinstance(env_vars, tuple): + if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY not in os.environ: + os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "true" + + if isinstance(env_vars, list): + return env_vars + [GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] + else: + return env_vars + (GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY,) + + raise TypeError( + f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." + ) + + def _validate_requirements_or_raise( *, agent_engine: _AgentEngineInterface, @@ -1309,7 +1388,7 @@ def _generate_deployment_spec_or_raise( ) else: raise TypeError( - f"env_vars must be a list or a dict, but got {type(env_vars)}." + f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." ) if deployment_spec.env: update_masks.append("spec.deployment_spec.env") From e60027755f7663b3ce15d22570e1b8ffa8af99e8 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Fri, 31 Oct 2025 11:22:16 -0700 Subject: [PATCH 05/24] chore: move CandidateResult, Event, Message, SessionInput and sub-fields to types/evals.py PiperOrigin-RevId: 826562043 --- .../genai/replays/test_get_evaluation_run.py | 8 +- .../replays/test_public_generate_rubrics.py | 2 +- tests/unit/vertexai/genai/test_evals.py | 22 +- vertexai/_genai/_evals_common.py | 11 +- vertexai/_genai/_evals_data_converters.py | 112 ++--- vertexai/_genai/_evals_metric_handlers.py | 4 +- .../_genai/_observability_data_converter.py | 2 +- vertexai/_genai/evals.py | 2 +- vertexai/_genai/types/__init__.py | 50 +-- vertexai/_genai/types/common.py | 381 +++--------------- vertexai/_genai/types/evals.py | 317 ++++++++++++++- 11 files changed, 473 insertions(+), 438 deletions(-) diff --git a/tests/unit/vertexai/genai/replays/test_get_evaluation_run.py b/tests/unit/vertexai/genai/replays/test_get_evaluation_run.py index 3db3bea517..6d07a52178 100644 --- a/tests/unit/vertexai/genai/replays/test_get_evaluation_run.py +++ b/tests/unit/vertexai/genai/replays/test_get_evaluation_run.py @@ -251,10 +251,10 @@ def check_run_5133048044039700480_evaluation_item_results( assert universal_metric_result.explanation is None # Check the first rubric verdict. rubric_verdict_0 = universal_metric_result.rubric_verdicts[0] - assert isinstance(rubric_verdict_0, types.RubricVerdict) - assert rubric_verdict_0.evaluated_rubric == types.Rubric( - content=types.RubricContent( - property=types.RubricContentProperty( + assert isinstance(rubric_verdict_0, types.evals.RubricVerdict) + assert rubric_verdict_0.evaluated_rubric == types.evals.Rubric( + content=types.evals.RubricContent( + property=types.evals.RubricContentProperty( description="The response is in English." ) ), diff --git a/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py b/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py index c21ca0e312..d3085bec74 100644 --- a/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py +++ b/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py @@ -173,7 +173,7 @@ def test_public_method_generate_rubrics(client): assert "text_quality_rubrics" in first_rubric_group assert isinstance(first_rubric_group["text_quality_rubrics"], list) assert first_rubric_group["text_quality_rubrics"] - assert isinstance(first_rubric_group["text_quality_rubrics"][0], types.Rubric) + assert isinstance(first_rubric_group["text_quality_rubrics"][0], types.evals.Rubric) pytestmark = pytest_helper.setup( diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py index b6da92e69c..dec2bb447d 100644 --- a/tests/unit/vertexai/genai/test_evals.py +++ b/tests/unit/vertexai/genai/test_evals.py @@ -2314,7 +2314,7 @@ def test_convert_with_intermediate_events_as_event_objects(self): "response": ["Hi"], "intermediate_events": [ [ - vertexai_genai_types.Event( + vertexai_genai_types.evals.Event( event_id="event1", content=genai_types.Content( parts=[genai_types.Part(text="intermediate event")] @@ -2577,14 +2577,14 @@ def test_convert_with_conversation_history(self): ) assert len(eval_case.conversation_history) == 2 - assert eval_case.conversation_history[0] == vertexai_genai_types.Message( + assert eval_case.conversation_history[0] == vertexai_genai_types.evals.Message( content=genai_types.Content( parts=[genai_types.Part(text="Hello")], role="user" ), turn_id="0", author="user", ) - assert eval_case.conversation_history[1] == vertexai_genai_types.Message( + assert eval_case.conversation_history[1] == vertexai_genai_types.evals.Message( content=genai_types.Content( parts=[genai_types.Part(text="Hi")], role="system" ), @@ -2786,7 +2786,7 @@ class TestEvent: """Unit tests for the Event class.""" def test_event_creation(self): - event = vertexai_genai_types.Event( + event = vertexai_genai_types.evals.Event( event_id="event1", content=genai_types.Content( parts=[genai_types.Part(text="intermediate event")] @@ -2820,7 +2820,7 @@ def test_eval_case_with_agent_eval_fields(self): tool_declarations=[tool], ) intermediate_events = [ - vertexai_genai_types.Event( + vertexai_genai_types.evals.Event( event_id="event1", content=genai_types.Content( parts=[genai_types.Part(text="intermediate event")] @@ -2846,7 +2846,7 @@ class TestSessionInput: """Unit tests for the SessionInput class.""" def test_session_input_creation(self): - session_input = vertexai_genai_types.SessionInput( + session_input = vertexai_genai_types.evals.SessionInput( user_id="user1", state={"key": "value"}, ) @@ -3692,7 +3692,7 @@ def test_eval_case_to_agent_data(self): tool_declarations=[tool], ) intermediate_events = [ - vertexai_genai_types.Event( + vertexai_genai_types.evals.Event( event_id="event1", content=genai_types.Content( parts=[genai_types.Part(text="intermediate event")] @@ -3722,7 +3722,7 @@ def test_eval_case_to_agent_data(self): def test_eval_case_to_agent_data_events_only(self): intermediate_events = [ - vertexai_genai_types.Event( + vertexai_genai_types.evals.Event( event_id="event1", content=genai_types.Content( parts=[genai_types.Part(text="intermediate event")] @@ -3751,7 +3751,7 @@ def test_eval_case_to_agent_data_events_only(self): def test_eval_case_to_agent_data_empty_event_content(self): intermediate_events = [ - vertexai_genai_types.Event( + vertexai_genai_types.evals.Event( event_id="event1", content=None, ) @@ -3933,12 +3933,12 @@ def test_build_request_payload_various_field_types(self): ) ], conversation_history=[ - vertexai_genai_types.Message( + vertexai_genai_types.evals.Message( content=genai_types.Content( parts=[genai_types.Part(text="Turn 1 user")], role="user" ) ), - vertexai_genai_types.Message( + vertexai_genai_types.evals.Message( content=genai_types.Content( parts=[genai_types.Part(text="Turn 1 model")], role="model" ) diff --git a/vertexai/_genai/_evals_common.py b/vertexai/_genai/_evals_common.py index bd43229bd9..1383822d64 100644 --- a/vertexai/_genai/_evals_common.py +++ b/vertexai/_genai/_evals_common.py @@ -1271,16 +1271,19 @@ def _execute_agent_run_with_retry( """Executes agent run for a single prompt.""" try: if isinstance(row["session_inputs"], str): - session_inputs = types.SessionInput.model_validate( + session_inputs = types.evals.SessionInput.model_validate( json.loads(row["session_inputs"]) ) elif isinstance(row["session_inputs"], dict): - session_inputs = types.SessionInput.model_validate(row["session_inputs"]) - elif isinstance(row["session_inputs"], types.SessionInput): + session_inputs = types.evals.SessionInput.model_validate( + row["session_inputs"] + ) + elif isinstance(row["session_inputs"], types.evals.SessionInput): session_inputs = row["session_inputs"] else: raise TypeError( - f"Unsupported session_inputs type: {type(row['session_inputs'])}. Expecting string or dict in types.SessionInput format." + f"Unsupported session_inputs type: {type(row['session_inputs'])}. " + "Expecting string or dict in types.evals.SessionInput format." ) user_id = session_inputs.user_id session_state = session_inputs.state diff --git a/vertexai/_genai/_evals_data_converters.py b/vertexai/_genai/_evals_data_converters.py index 337abaaae8..459600caff 100644 --- a/vertexai/_genai/_evals_data_converters.py +++ b/vertexai/_genai/_evals_data_converters.py @@ -60,7 +60,7 @@ class _GeminiEvalDataConverter(_evals_utils.EvalDataConverter): def _parse_request(self, request_data: dict[str, Any]) -> tuple[ genai_types.Content, genai_types.Content, - list[types.Message], + list[types.evals.Message], types.ResponseCandidate, ]: """Parses a request from a Gemini dataset.""" @@ -76,16 +76,16 @@ def _parse_request(self, request_data: dict[str, Any]) -> tuple[ for turn_id, content_dict in enumerate(request_data.get("contents", [])): if not isinstance(content_dict, dict): raise TypeError( - f"Expected a dictionary for content at turn {turn_id}, but got" - f" {type(content_dict).__name__}: {content_dict}" + "Expected a dictionary for content at turn %s, but got %s: %s" + % (turn_id, type(content_dict).__name__, content_dict) ) if "parts" not in content_dict: raise ValueError( - f"Missing 'parts' key in content structure at turn {turn_id}:" - f" {content_dict}" + "Missing 'parts' key in content structure at turn %s: %s" + % (turn_id, content_dict) ) conversation_history.append( - types.Message( + types.evals.Message( turn_id=str(turn_id), content=genai_types.Content.model_validate(content_dict), ) @@ -121,7 +121,7 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: eval_cases = [] for i, item in enumerate(raw_data): - eval_case_id = f"gemini_eval_case_{i}" + eval_case_id = "gemini_eval_case_%s" % i request_data = item.get("request", {}) response_data = item.get("response", {}) @@ -187,11 +187,11 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: for i, item_dict in enumerate(raw_data): if not isinstance(item_dict, dict): raise TypeError( - f"Expected a dictionary for item at index {i}, but got" - f" {type(item_dict).__name__}: {item_dict}" + "Expected a dictionary for item at index %s, but got %s: %s" + % (i, type(item_dict).__name__, item_dict) ) item = copy.deepcopy(item_dict) - eval_case_id = f"eval_case_{i}" + eval_case_id = "eval_case_%s" % i prompt_data = item.pop("prompt", None) if not prompt_data: prompt_data = item.pop("source", None) @@ -205,10 +205,12 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: if not response_data: raise ValueError( - f"Response is required but missing for {eval_case_id}." + "Response is required but missing for %s." % eval_case_id ) if not prompt_data: - raise ValueError(f"Prompt is required but missing for {eval_case_id}.") + raise ValueError( + "Prompt is required but missing for %s." % eval_case_id + ) prompt: genai_types.Content if isinstance(prompt_data, str): @@ -219,16 +221,16 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: prompt = prompt_data else: raise ValueError( - f"Invalid prompt type for case {i}: {type(prompt_data)}" + "Invalid prompt type for case %s: %s" % (i, type(prompt_data)) ) - conversation_history: Optional[list[types.Message]] = None + conversation_history: Optional[list[types.evals.Message]] = None if isinstance(conversation_history_data, list): conversation_history = [] for turn_id, content in enumerate(conversation_history_data): if isinstance(content, genai_types.Content): conversation_history.append( - types.Message( + types.evals.Message( turn_id=str(turn_id), content=content, ) @@ -239,7 +241,7 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: content ) conversation_history.append( - types.Message( + types.evals.Message( turn_id=str(turn_id), content=validated_content, ) @@ -282,7 +284,7 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: responses = [types.ResponseCandidate(response=response_data)] else: raise ValueError( - f"Invalid response type for case {i}: {type(response_data)}" + "Invalid response type for case %s: %s" % (i, type(response_data)) ) reference: Optional[types.ResponseCandidate] = None @@ -322,14 +324,14 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: try: validated_rubrics = [ ( - types.Rubric.model_validate(r) + types.evals.Rubric.model_validate(r) if isinstance(r, dict) else r ) for r in value ] if all( - isinstance(r, types.Rubric) + isinstance(r, types.evals.Rubric) for r in validated_rubrics ): rubric_groups[key] = types.RubricGroup( @@ -337,11 +339,16 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: ) else: logger.warning( - f"Invalid item type in rubric list for group '{key}' in case {i}." + "Invalid item type in rubric list for group '%s' in case %s.", + key, + i, ) except Exception as e: logger.warning( - f"Failed to validate rubrics for group '{key}' in case {i}: {e}" + "Failed to validate rubrics for group '%s' in case %s: %s", + key, + i, + e, ) elif isinstance(value, types.RubricGroup): rubric_groups[key] = value @@ -352,44 +359,56 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: ) except Exception as e: logger.warning( - f"Failed to validate RubricGroup dict for group '{key}' in case {i}: {e}" + "Failed to validate RubricGroup dict for group '%s' in case %s: %s", + key, + i, + e, ) else: logger.warning( - f"Invalid type for rubric group '{key}' in case {i}." - " Expected list of rubrics, dict, or RubricGroup." + "Invalid type for rubric group '%s' in case %s." + " Expected list of rubrics, dict, or RubricGroup.", + key, + i, ) else: logger.warning( - f"Invalid type for rubric_groups in case {i}. Expected dict." + "Invalid type for rubric_groups in case %s. Expected dict.", + i, ) - intermediate_events: Optional[list[types.Event]] = None + intermediate_events: Optional[list[types.evals.Event]] = None if intermediate_events_data: if isinstance(intermediate_events_data, list): intermediate_events = [] for event in intermediate_events_data: if isinstance(event, dict): try: - validated_event = types.Event.model_validate(event) + validated_event = types.evals.Event.model_validate( + event + ) intermediate_events.append(validated_event) except Exception as e: logger.warning( "Failed to validate intermediate event dict for" - f" case {i}: {e}" + " case %s: %s", + i, + e, ) - elif isinstance(event, types.Event): + elif isinstance(event, types.evals.Event): intermediate_events.append(event) else: logger.warning( "Invalid type for intermediate_event in case" - f" {i}. Expected list of dicts or list of" - " types.Event objects." + " %s. Expected list of dicts or list of" + " types.evals.Event objects.", + i, ) else: logger.warning( - f"Invalid type for intermediate_events in case {i}. Expected" - " list of types.Event objects." + "Invalid type for intermediate_events in case %s. Expected" + " list of types.evals.Event objects.", + i, ) eval_case = types.EvalCase( @@ -414,7 +433,7 @@ class _OpenAIDataConverter(_evals_utils.EvalDataConverter): def _parse_messages(self, messages: list[dict[str, Any]]) -> tuple[ Optional[genai_types.Content], - list[types.Message], + list[types.evals.Message], Optional[genai_types.Content], Optional[types.ResponseCandidate], ]: @@ -434,7 +453,7 @@ def _parse_messages(self, messages: list[dict[str, Any]]) -> tuple[ role = msg.get("role", "user") content = msg.get("content", "") conversation_history.append( - types.Message( + types.evals.Message( turn_id=str(turn_id), content=genai_types.Content( parts=[genai_types.Part(text=content)], role=role @@ -460,11 +479,11 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset: """Converts a list of OpenAI ChatCompletion data into an EvaluationDataset.""" eval_cases = [] for i, item in enumerate(raw_data): - eval_case_id = f"openai_eval_case_{i}" + eval_case_id = "openai_eval_case_%s" % i if "request" not in item or "response" not in item: logger.warning( - f"Skipping case {i} due to missing 'request' or 'response' key." + "Skipping case %s due to missing 'request' or 'response' key.", i ) continue @@ -610,7 +629,7 @@ def get_dataset_converter( if dataset_schema in _CONVERTER_REGISTRY: return _CONVERTER_REGISTRY[dataset_schema]() # type: ignore[abstract] else: - raise ValueError(f"Unsupported dataset schema: {dataset_schema}") + raise ValueError("Unsupported dataset schema: %s" % dataset_schema) def _get_first_part_text(content: genai_types.Content) -> str: @@ -695,7 +714,7 @@ def merge_response_datasets_into_canonical_format( """ if not isinstance(raw_datasets, list): raise TypeError( - f"Input 'raw_datasets' must be a list, got {type(raw_datasets)}." + "Input 'raw_datasets' must be a list, got %s." % type(raw_datasets) ) if not raw_datasets or not all(isinstance(ds, list) for ds in raw_datasets): raise ValueError( @@ -704,7 +723,7 @@ def merge_response_datasets_into_canonical_format( if not schemas or len(schemas) != len(raw_datasets): raise ValueError( "A list of schemas must be provided, one for each raw dataset. " - f"Got {len(schemas)} schemas for {len(raw_datasets)} datasets." + "Got %s schemas for %s datasets." % (len(schemas), len(raw_datasets)) ) num_expected_cases = len(raw_datasets[0]) @@ -719,8 +738,8 @@ def merge_response_datasets_into_canonical_format( if len(raw_ds_entry) != num_expected_cases: raise ValueError( "All datasets must have the same number of evaluation cases. " - f"Base dataset (0) has {num_expected_cases}, but dataset {i} " - f"(schema: {schema}) has {len(raw_ds_entry)}." + "Base dataset (0) has %s, but dataset %s (schema: %s) has %s." + % (num_expected_cases, i, schema, len(raw_ds_entry)) ) converter = get_dataset_converter(schema) parsed_evaluation_datasets.append(converter.convert(raw_ds_entry)) @@ -746,7 +765,7 @@ def merge_response_datasets_into_canonical_format( ) candidate_responses.append( _create_placeholder_response_candidate( - f"Missing response from base dataset (0) for case {case_idx}" + "Missing response from base dataset (0) for case %s" % case_idx ) ) @@ -799,13 +818,14 @@ def merge_response_datasets_into_canonical_format( ) candidate_responses.append( _create_placeholder_response_candidate( - f"Missing response from dataset {dataset_idx_offset} " - f"for case {case_idx}" + "Missing response from dataset %s for case %s" + % (dataset_idx_offset, case_idx) ) ) merged_case = types.EvalCase( - eval_case_id=base_eval_case.eval_case_id or f"merged_eval_case_{case_idx}", + eval_case_id=base_eval_case.eval_case_id + or "merged_eval_case_%s" % case_idx, prompt=base_eval_case.prompt, responses=candidate_responses, reference=base_eval_case.reference, diff --git a/vertexai/_genai/_evals_metric_handlers.py b/vertexai/_genai/_evals_metric_handlers.py index 322d3aff71..9f68bc353d 100644 --- a/vertexai/_genai/_evals_metric_handlers.py +++ b/vertexai/_genai/_evals_metric_handlers.py @@ -480,7 +480,7 @@ def _build_rubric_based_input( ) rubrics_list = [] - parsed_rubrics = [types.Rubric(**r) for r in rubrics_list] + parsed_rubrics = [types.evals.Rubric(**r) for r in rubrics_list] rubric_enhanced_contents = { "prompt": ( [eval_case.prompt.model_dump(mode="json", exclude_none=True)] @@ -535,7 +535,7 @@ def _build_pointwise_input( elif isinstance(value, list) and value: if isinstance(value[0], genai_types.Content): content_list_to_serialize = value - elif isinstance(value[0], types.Message): + elif isinstance(value[0], types.evals.Message): history_texts = [] for msg_obj in value: msg_text = _extract_text_from_content(msg_obj.content) diff --git a/vertexai/_genai/_observability_data_converter.py b/vertexai/_genai/_observability_data_converter.py index f7e7f11a08..c52f80b8ea 100644 --- a/vertexai/_genai/_observability_data_converter.py +++ b/vertexai/_genai/_observability_data_converter.py @@ -129,7 +129,7 @@ def _parse_messages( if len(request_msgs) > 1: for i, msg in enumerate(request_msgs[:-1]): conversation_history.append( - types.Message( + types.evals.Message( turn_id=str(i), content=self._message_to_content(msg), author=msg.get("role", ""), diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py index bc60977a82..7b7ad2c023 100644 --- a/vertexai/_genai/evals.py +++ b/vertexai/_genai/evals.py @@ -361,7 +361,7 @@ def _RubricBasedMetricSpec_to_vertex( setv( to_object, ["inline_rubrics", "rubrics"], - [item for item in getv(from_object, ["inline_rubrics"])], + getv(from_object, ["inline_rubrics"]), ) if getv(from_object, ["rubric_group_key"]) is not None: diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index 735e567dc1..ed2eff77c4 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -146,7 +146,6 @@ from .common import CandidateResponseOrDict from .common import CandidateResult from .common import CandidateResultDict -from .common import CandidateResultOrDict from .common import Chunk from .common import ChunkDict from .common import ChunkOrDict @@ -343,7 +342,6 @@ from .common import EventMetadata from .common import EventMetadataDict from .common import EventMetadataOrDict -from .common import EventOrDict from .common import ExactMatchInput from .common import ExactMatchInputDict from .common import ExactMatchInputOrDict @@ -552,7 +550,6 @@ from .common import MemoryTopicIdOrDict from .common import Message from .common import MessageDict -from .common import MessageOrDict from .common import Metadata from .common import MetadataDict from .common import MetadataOrDict @@ -755,10 +752,8 @@ from .common import RubricBasedMetricSpecOrDict from .common import RubricContent from .common import RubricContentDict -from .common import RubricContentOrDict from .common import RubricContentProperty from .common import RubricContentPropertyDict -from .common import RubricContentPropertyOrDict from .common import RubricContentType from .common import RubricDict from .common import RubricEnhancedContents @@ -773,10 +768,8 @@ from .common import RubricGroup from .common import RubricGroupDict from .common import RubricGroupOrDict -from .common import RubricOrDict from .common import RubricVerdict from .common import RubricVerdictDict -from .common import RubricVerdictOrDict from .common import SamplingConfig from .common import SamplingConfigDict from .common import SamplingConfigOrDict @@ -876,9 +869,6 @@ from .common import SessionEvent from .common import SessionEventDict from .common import SessionEventOrDict -from .common import SessionInput -from .common import SessionInputDict -from .common import SessionInputOrDict from .common import SessionOrDict from .common import State from .common import Strategy @@ -990,21 +980,6 @@ "EvaluationItemRequest", "EvaluationItemRequestDict", "EvaluationItemRequestOrDict", - "RubricContentProperty", - "RubricContentPropertyDict", - "RubricContentPropertyOrDict", - "RubricContent", - "RubricContentDict", - "RubricContentOrDict", - "Rubric", - "RubricDict", - "RubricOrDict", - "RubricVerdict", - "RubricVerdictDict", - "RubricVerdictOrDict", - "CandidateResult", - "CandidateResultDict", - "CandidateResultOrDict", "EvaluationItemResult", "EvaluationItemResultDict", "EvaluationItemResultOrDict", @@ -1059,12 +1034,6 @@ "ResponseCandidate", "ResponseCandidateDict", "ResponseCandidateOrDict", - "Event", - "EventDict", - "EventOrDict", - "Message", - "MessageDict", - "MessageOrDict", "EvalCase", "EvalCaseDict", "EvalCaseOrDict", @@ -1770,9 +1739,6 @@ "EvaluationRunInferenceConfig", "EvaluationRunInferenceConfigDict", "EvaluationRunInferenceConfigOrDict", - "SessionInput", - "SessionInputDict", - "SessionInputOrDict", "WinRateStats", "WinRateStatsDict", "WinRateStatsOrDict", @@ -1835,7 +1801,6 @@ "SamplingMethod", "RubricContentType", "EvaluationRunState", - "Importance", "OptimizeTarget", "GenerateMemoriesResponseGeneratedMemoryAction", "PromptOptimizerMethod", @@ -1844,6 +1809,21 @@ "PromptDataOrDict", "LLMMetric", "MetricPromptBuilder", + "RubricContentProperty", + "RubricContentPropertyDict", + "RubricContent", + "RubricContentDict", + "Rubric", + "RubricDict", + "RubricVerdict", + "RubricVerdictDict", + "CandidateResult", + "CandidateResultDict", + "Event", + "EventDict", + "Message", + "MessageDict", + "Importance", "_CreateEvaluationItemParameters", "_CreateEvaluationRunParameters", "_CreateEvaluationSetParameters", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 3e6622acf4..d193042586 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -318,19 +318,6 @@ class EvaluationRunState(_common.CaseInSensitiveEnum): """Evaluation run is performing rubric generation.""" -class Importance(_common.CaseInSensitiveEnum): - """Importance level of the rubric.""" - - IMPORTANCE_UNSPECIFIED = "IMPORTANCE_UNSPECIFIED" - """Importance is not specified.""" - HIGH = "HIGH" - """High importance.""" - MEDIUM = "MEDIUM" - """Medium importance.""" - LOW = "LOW" - """Low importance.""" - - class OptimizeTarget(_common.CaseInSensitiveEnum): """None""" @@ -537,196 +524,6 @@ class EvaluationItemRequestDict(TypedDict, total=False): EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] -class RubricContentProperty(_common.BaseModel): - """Defines criteria based on a specific property.""" - - description: Optional[str] = Field( - default=None, - description="""Description of the property being evaluated. - Example: "The model's response is grammatically correct." """, - ) - - -class RubricContentPropertyDict(TypedDict, total=False): - """Defines criteria based on a specific property.""" - - description: Optional[str] - """Description of the property being evaluated. - Example: "The model's response is grammatically correct." """ - - -RubricContentPropertyOrDict = Union[RubricContentProperty, RubricContentPropertyDict] - - -class RubricContent(_common.BaseModel): - """Content of the rubric, defining the testable criteria.""" - - property: Optional[RubricContentProperty] = Field( - default=None, - description="""Evaluation criteria based on a specific property.""", - ) - - -class RubricContentDict(TypedDict, total=False): - """Content of the rubric, defining the testable criteria.""" - - property: Optional[RubricContentPropertyDict] - """Evaluation criteria based on a specific property.""" - - -RubricContentOrDict = Union[RubricContent, RubricContentDict] - - -class Rubric(_common.BaseModel): - """Message representing a single testable criterion for evaluation. - - One input prompt could have multiple rubrics. - """ - - rubric_id: Optional[str] = Field( - default=None, - description="""Required. Unique identifier for the rubric. - This ID is used to refer to this rubric, e.g., in RubricVerdict.""", - ) - content: Optional[RubricContent] = Field( - default=None, - description="""Required. The actual testable criteria for the rubric.""", - ) - type: Optional[str] = Field( - default=None, - description="""Optional. A type designator for the rubric, which can inform how it's - evaluated or interpreted by systems or users. - It's recommended to use consistent, well-defined, upper snake_case strings. - Examples: "SUMMARIZATION_QUALITY", "SAFETY_HARMFUL_CONTENT", - "INSTRUCTION_ADHERENCE".""", - ) - importance: Optional[Importance] = Field( - default=None, - description="""Optional. The relative importance of this rubric.""", - ) - - -class RubricDict(TypedDict, total=False): - """Message representing a single testable criterion for evaluation. - - One input prompt could have multiple rubrics. - """ - - rubric_id: Optional[str] - """Required. Unique identifier for the rubric. - This ID is used to refer to this rubric, e.g., in RubricVerdict.""" - - content: Optional[RubricContentDict] - """Required. The actual testable criteria for the rubric.""" - - type: Optional[str] - """Optional. A type designator for the rubric, which can inform how it's - evaluated or interpreted by systems or users. - It's recommended to use consistent, well-defined, upper snake_case strings. - Examples: "SUMMARIZATION_QUALITY", "SAFETY_HARMFUL_CONTENT", - "INSTRUCTION_ADHERENCE".""" - - importance: Optional[Importance] - """Optional. The relative importance of this rubric.""" - - -RubricOrDict = Union[Rubric, RubricDict] - - -class RubricVerdict(_common.BaseModel): - """Represents the verdict of an evaluation against a single rubric.""" - - evaluated_rubric: Optional[Rubric] = Field( - default=None, - description="""Required. The full rubric definition that was evaluated. - Storing this ensures the verdict is self-contained and understandable, - especially if the original rubric definition changes or was dynamically - generated.""", - ) - verdict: Optional[bool] = Field( - default=None, - description="""Required. Outcome of the evaluation against the rubric, represented as a - boolean. `true` indicates a "Pass", `false` indicates a "Fail".""", - ) - reasoning: Optional[str] = Field( - default=None, - description="""Optional. Human-readable reasoning or explanation for the verdict. - This can include specific examples or details from the evaluated content - that justify the given verdict.""", - ) - - -class RubricVerdictDict(TypedDict, total=False): - """Represents the verdict of an evaluation against a single rubric.""" - - evaluated_rubric: Optional[RubricDict] - """Required. The full rubric definition that was evaluated. - Storing this ensures the verdict is self-contained and understandable, - especially if the original rubric definition changes or was dynamically - generated.""" - - verdict: Optional[bool] - """Required. Outcome of the evaluation against the rubric, represented as a - boolean. `true` indicates a "Pass", `false` indicates a "Fail".""" - - reasoning: Optional[str] - """Optional. Human-readable reasoning or explanation for the verdict. - This can include specific examples or details from the evaluated content - that justify the given verdict.""" - - -RubricVerdictOrDict = Union[RubricVerdict, RubricVerdictDict] - - -class CandidateResult(_common.BaseModel): - """Result for a single candidate.""" - - candidate: Optional[str] = Field( - default=None, - description="""The candidate that is being evaluated. The value is the same as the candidate name in the EvaluationRequest.""", - ) - metric: Optional[str] = Field( - default=None, description="""The metric that was evaluated.""" - ) - score: Optional[float] = Field( - default=None, description="""The score of the metric.""" - ) - explanation: Optional[str] = Field( - default=None, description="""The explanation for the metric.""" - ) - rubric_verdicts: Optional[list[RubricVerdict]] = Field( - default=None, description="""The rubric verdicts for the metric.""" - ) - additional_results: Optional[dict[str, Any]] = Field( - default=None, description="""Additional results for the metric.""" - ) - - -class CandidateResultDict(TypedDict, total=False): - """Result for a single candidate.""" - - candidate: Optional[str] - """The candidate that is being evaluated. The value is the same as the candidate name in the EvaluationRequest.""" - - metric: Optional[str] - """The metric that was evaluated.""" - - score: Optional[float] - """The score of the metric.""" - - explanation: Optional[str] - """The explanation for the metric.""" - - rubric_verdicts: Optional[list[RubricVerdictDict]] - """The rubric verdicts for the metric.""" - - additional_results: Optional[dict[str, Any]] - """Additional results for the metric.""" - - -CandidateResultOrDict = Union[CandidateResult, CandidateResultDict] - - class EvaluationItemResult(_common.BaseModel): """Represents the result of an evaluation item.""" @@ -743,7 +540,7 @@ class EvaluationItemResult(_common.BaseModel): metric: Optional[str] = Field( default=None, description="""The metric that was evaluated.""" ) - candidate_results: Optional[list[CandidateResult]] = Field( + candidate_results: Optional[list[evals_types.CandidateResult]] = Field( default=None, description="""TThe results for the metric.""" ) metadata: Optional[dict[str, Any]] = Field( @@ -766,7 +563,7 @@ class EvaluationItemResultDict(TypedDict, total=False): metric: Optional[str] """The metric that was evaluated.""" - candidate_results: Optional[list[CandidateResultDict]] + candidate_results: Optional[list[evals_types.CandidateResult]] """TThe results for the metric.""" metadata: Optional[dict[str, Any]] @@ -1440,89 +1237,6 @@ class ResponseCandidateDict(TypedDict, total=False): ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] -class Event(_common.BaseModel): - """Represents an event in a conversation between agents and users. - - It is used to store the content of the conversation, as well as the actions - taken by the agents like function calls, function responses, intermediate NL - responses etc. - """ - - event_id: Optional[str] = Field( - default=None, description="""Unique identifier for the agent event.""" - ) - content: Optional[genai_types.Content] = Field( - default=None, description="""Content of the event.""" - ) - creation_timestamp: Optional[datetime.datetime] = Field( - default=None, description="""The creation timestamp of the event.""" - ) - author: Optional[str] = Field( - default=None, description="""Name of the entity that produced the event.""" - ) - - -class EventDict(TypedDict, total=False): - """Represents an event in a conversation between agents and users. - - It is used to store the content of the conversation, as well as the actions - taken by the agents like function calls, function responses, intermediate NL - responses etc. - """ - - event_id: Optional[str] - """Unique identifier for the agent event.""" - - content: Optional[genai_types.ContentDict] - """Content of the event.""" - - creation_timestamp: Optional[datetime.datetime] - """The creation timestamp of the event.""" - - author: Optional[str] - """Name of the entity that produced the event.""" - - -EventOrDict = Union[Event, EventDict] - - -class Message(_common.BaseModel): - """Represents a single message turn in a conversation.""" - - turn_id: Optional[str] = Field( - default=None, description="""Unique identifier for the message turn.""" - ) - content: Optional[genai_types.Content] = Field( - default=None, description="""Content of the message, including function call.""" - ) - creation_timestamp: Optional[datetime.datetime] = Field( - default=None, - description="""Timestamp indicating when the message was created.""", - ) - author: Optional[str] = Field( - default=None, description="""Name of the entity that produced the message.""" - ) - - -class MessageDict(TypedDict, total=False): - """Represents a single message turn in a conversation.""" - - turn_id: Optional[str] - """Unique identifier for the message turn.""" - - content: Optional[genai_types.ContentDict] - """Content of the message, including function call.""" - - creation_timestamp: Optional[datetime.datetime] - """Timestamp indicating when the message was created.""" - - author: Optional[str] - """Name of the entity that produced the message.""" - - -MessageOrDict = Union[Message, MessageDict] - - class EvalCase(_common.BaseModel): """A comprehensive representation of a GenAI interaction for evaluation.""" @@ -1540,7 +1254,7 @@ class EvalCase(_common.BaseModel): system_instruction: Optional[genai_types.Content] = Field( default=None, description="""System instruction for the model.""" ) - conversation_history: Optional[list[Message]] = Field( + conversation_history: Optional[list[evals_types.Message]] = Field( default=None, description="""List of all prior messages in the conversation (chat history).""", ) @@ -1551,7 +1265,7 @@ class EvalCase(_common.BaseModel): eval_case_id: Optional[str] = Field( default=None, description="""Unique identifier for the evaluation case.""" ) - intermediate_events: Optional[list[Event]] = Field( + intermediate_events: Optional[list[evals_types.Event]] = Field( default=None, description="""This field is experimental and may change in future versions. Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""", ) @@ -1578,7 +1292,7 @@ class EvalCaseDict(TypedDict, total=False): system_instruction: Optional[genai_types.ContentDict] """System instruction for the model.""" - conversation_history: Optional[list[MessageDict]] + conversation_history: Optional[list[evals_types.Message]] """List of all prior messages in the conversation (chat history).""" rubric_groups: Optional[dict[str, "RubricGroupDict"]] @@ -1587,7 +1301,7 @@ class EvalCaseDict(TypedDict, total=False): eval_case_id: Optional[str] """Unique identifier for the evaluation case.""" - intermediate_events: Optional[list[EventDict]] + intermediate_events: Optional[list[evals_types.Event]] """This field is experimental and may change in future versions. Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""" agent_info: Optional[evals_types.AgentInfo] @@ -2727,7 +2441,7 @@ class RubricBasedMetricSpec(_common.BaseModel): default=None, description="""Optional configuration for the judge LLM (Autorater).""", ) - inline_rubrics: Optional[list[Rubric]] = Field( + inline_rubrics: Optional[list[evals_types.Rubric]] = Field( default=None, description="""Use rubrics provided directly in the spec.""" ) rubric_group_key: Optional[str] = Field( @@ -2752,7 +2466,7 @@ class RubricBasedMetricSpecDict(TypedDict, total=False): judge_autorater_config: Optional[genai_types.AutoraterConfigDict] """Optional configuration for the judge LLM (Autorater).""" - inline_rubrics: Optional[list[RubricDict]] + inline_rubrics: Optional[list[evals_types.Rubric]] """Use rubrics provided directly in the spec.""" rubric_group_key: Optional[str] @@ -3220,7 +2934,7 @@ class MetricResult(_common.BaseModel): default=None, description="""The score for the metric. Please refer to each metric's documentation for the meaning of the score.""", ) - rubric_verdicts: Optional[list[RubricVerdict]] = Field( + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( default=None, description="""For rubric-based metrics, the verdicts for each rubric.""", ) @@ -3238,7 +2952,7 @@ class MetricResultDict(TypedDict, total=False): score: Optional[float] """The score for the metric. Please refer to each metric's documentation for the meaning of the score.""" - rubric_verdicts: Optional[list[RubricVerdictDict]] + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] """For rubric-based metrics, the verdicts for each rubric.""" explanation: Optional[str] @@ -3257,7 +2971,7 @@ class RubricBasedMetricResult(_common.BaseModel): score: Optional[float] = Field( default=None, description="""Passing rate of all the rubrics.""" ) - rubric_verdicts: Optional[list[RubricVerdict]] = Field( + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( default=None, description="""The details of all the rubrics and their verdicts.""", ) @@ -3269,7 +2983,7 @@ class RubricBasedMetricResultDict(TypedDict, total=False): score: Optional[float] """Passing rate of all the rubrics.""" - rubric_verdicts: Optional[list[RubricVerdictDict]] + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] """The details of all the rubrics and their verdicts.""" @@ -3855,7 +3569,7 @@ class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): class GenerateInstanceRubricsResponse(_common.BaseModel): """Response for generating rubrics.""" - generated_rubrics: Optional[list[Rubric]] = Field( + generated_rubrics: Optional[list[evals_types.Rubric]] = Field( default=None, description="""A list of generated rubrics.""" ) @@ -3863,7 +3577,7 @@ class GenerateInstanceRubricsResponse(_common.BaseModel): class GenerateInstanceRubricsResponseDict(TypedDict, total=False): """Response for generating rubrics.""" - generated_rubrics: Optional[list[RubricDict]] + generated_rubrics: Optional[list[evals_types.Rubric]] """A list of generated rubrics.""" @@ -12611,7 +12325,7 @@ class EvalCaseMetricResult(_common.BaseModel): explanation: Optional[str] = Field( default=None, description="""Explanation of the metric.""" ) - rubric_verdicts: Optional[list[RubricVerdict]] = Field( + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( default=None, description="""The details of all the rubrics and their verdicts for rubric-based metrics.""", ) @@ -12635,7 +12349,7 @@ class EvalCaseMetricResultDict(TypedDict, total=False): explanation: Optional[str] """Explanation of the metric.""" - rubric_verdicts: Optional[list[RubricVerdictDict]] + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] """The details of all the rubrics and their verdicts for rubric-based metrics.""" raw_output: Optional[list[str]] @@ -12713,34 +12427,6 @@ class EvaluationRunInferenceConfigDict(TypedDict, total=False): ] -class SessionInput(_common.BaseModel): - """This field is experimental and may change in future versions. - - Input to initialize a session and run an agent, used for agent evaluation. - """ - - user_id: Optional[str] = Field(default=None, description="""The user id.""") - state: Optional[dict[str, str]] = Field( - default=None, description="""The state of the session.""" - ) - - -class SessionInputDict(TypedDict, total=False): - """This field is experimental and may change in future versions. - - Input to initialize a session and run an agent, used for agent evaluation. - """ - - user_id: Optional[str] - """The user id.""" - - state: Optional[dict[str, str]] - """The state of the session.""" - - -SessionInputOrDict = Union[SessionInput, SessionInputDict] - - class WinRateStats(_common.BaseModel): """Statistics for win rates for a single metric.""" @@ -12953,7 +12639,7 @@ class RubricGroup(_common.BaseModel): Example: "Instruction Following V1", "Content Quality - Summarization Task".""", ) - rubrics: Optional[list[Rubric]] = Field( + rubrics: Optional[list[evals_types.Rubric]] = Field( default=None, description="""Rubrics that are part of this group.""" ) @@ -12970,7 +12656,7 @@ class RubricGroupDict(TypedDict, total=False): Example: "Instruction Following V1", "Content Quality - Summarization Task".""" - rubrics: Optional[list[RubricDict]] + rubrics: Optional[list[evals_types.Rubric]] """Rubrics that are part of this group.""" @@ -13030,6 +12716,37 @@ 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.""" diff --git a/vertexai/_genai/types/evals.py b/vertexai/_genai/types/evals.py index 9ecade7ee3..2c8aba057a 100644 --- a/vertexai/_genai/types/evals.py +++ b/vertexai/_genai/types/evals.py @@ -15,13 +15,27 @@ # Code generated by the Google Gen AI SDK generator DO NOT EDIT. -from typing import Optional, Union +import datetime +from typing import Any, Optional, Union from google.genai import _common from google.genai import types as genai_types from pydantic import Field from typing_extensions import TypedDict +class Importance(_common.CaseInSensitiveEnum): + """Importance level of the rubric.""" + + IMPORTANCE_UNSPECIFIED = "IMPORTANCE_UNSPECIFIED" + """Importance is not specified.""" + HIGH = "HIGH" + """High importance.""" + MEDIUM = "MEDIUM" + """Medium importance.""" + LOW = "LOW" + """Low importance.""" + + class AgentInfo(_common.BaseModel): """The agent info of an agent, used for agent eval.""" @@ -67,6 +81,307 @@ class AgentInfoDict(TypedDict, total=False): AgentInfoOrDict = Union[AgentInfo, AgentInfoDict] +class RubricContentProperty(_common.BaseModel): + """Defines criteria based on a specific property.""" + + description: Optional[str] = Field( + default=None, + description="""Description of the property being evaluated. + Example: "The model's response is grammatically correct." """, + ) + + +class RubricContentPropertyDict(TypedDict, total=False): + """Defines criteria based on a specific property.""" + + description: Optional[str] + """Description of the property being evaluated. + Example: "The model's response is grammatically correct." """ + + +RubricContentPropertyOrDict = Union[RubricContentProperty, RubricContentPropertyDict] + + +class RubricContent(_common.BaseModel): + """Content of the rubric, defining the testable criteria.""" + + property: Optional[RubricContentProperty] = Field( + default=None, + description="""Evaluation criteria based on a specific property.""", + ) + + +class RubricContentDict(TypedDict, total=False): + """Content of the rubric, defining the testable criteria.""" + + property: Optional[RubricContentPropertyDict] + """Evaluation criteria based on a specific property.""" + + +RubricContentOrDict = Union[RubricContent, RubricContentDict] + + +class Rubric(_common.BaseModel): + """Message representing a single testable criterion for evaluation. + + One input prompt could have multiple rubrics. + """ + + rubric_id: Optional[str] = Field( + default=None, + description="""Required. Unique identifier for the rubric. + This ID is used to refer to this rubric, e.g., in RubricVerdict.""", + ) + content: Optional[RubricContent] = Field( + default=None, + description="""Required. The actual testable criteria for the rubric.""", + ) + type: Optional[str] = Field( + default=None, + description="""Optional. A type designator for the rubric, which can inform how it's + evaluated or interpreted by systems or users. + It's recommended to use consistent, well-defined, upper snake_case strings. + Examples: "SUMMARIZATION_QUALITY", "SAFETY_HARMFUL_CONTENT", + "INSTRUCTION_ADHERENCE".""", + ) + importance: Optional[Importance] = Field( + default=None, + description="""Optional. The relative importance of this rubric.""", + ) + + +class RubricDict(TypedDict, total=False): + """Message representing a single testable criterion for evaluation. + + One input prompt could have multiple rubrics. + """ + + rubric_id: Optional[str] + """Required. Unique identifier for the rubric. + This ID is used to refer to this rubric, e.g., in RubricVerdict.""" + + content: Optional[RubricContentDict] + """Required. The actual testable criteria for the rubric.""" + + type: Optional[str] + """Optional. A type designator for the rubric, which can inform how it's + evaluated or interpreted by systems or users. + It's recommended to use consistent, well-defined, upper snake_case strings. + Examples: "SUMMARIZATION_QUALITY", "SAFETY_HARMFUL_CONTENT", + "INSTRUCTION_ADHERENCE".""" + + importance: Optional[Importance] + """Optional. The relative importance of this rubric.""" + + +RubricOrDict = Union[Rubric, RubricDict] + + +class RubricVerdict(_common.BaseModel): + """Represents the verdict of an evaluation against a single rubric.""" + + evaluated_rubric: Optional[Rubric] = Field( + default=None, + description="""Required. The full rubric definition that was evaluated. + Storing this ensures the verdict is self-contained and understandable, + especially if the original rubric definition changes or was dynamically + generated.""", + ) + verdict: Optional[bool] = Field( + default=None, + description="""Required. Outcome of the evaluation against the rubric, represented as a + boolean. `true` indicates a "Pass", `false` indicates a "Fail".""", + ) + reasoning: Optional[str] = Field( + default=None, + description="""Optional. Human-readable reasoning or explanation for the verdict. + This can include specific examples or details from the evaluated content + that justify the given verdict.""", + ) + + +class RubricVerdictDict(TypedDict, total=False): + """Represents the verdict of an evaluation against a single rubric.""" + + evaluated_rubric: Optional[RubricDict] + """Required. The full rubric definition that was evaluated. + Storing this ensures the verdict is self-contained and understandable, + especially if the original rubric definition changes or was dynamically + generated.""" + + verdict: Optional[bool] + """Required. Outcome of the evaluation against the rubric, represented as a + boolean. `true` indicates a "Pass", `false` indicates a "Fail".""" + + reasoning: Optional[str] + """Optional. Human-readable reasoning or explanation for the verdict. + This can include specific examples or details from the evaluated content + that justify the given verdict.""" + + +RubricVerdictOrDict = Union[RubricVerdict, RubricVerdictDict] + + +class CandidateResult(_common.BaseModel): + """Result for a single candidate.""" + + candidate: Optional[str] = Field( + default=None, + description="""The candidate that is being evaluated. The value is the same as the candidate name in the EvaluationRequest.""", + ) + metric: Optional[str] = Field( + default=None, description="""The metric that was evaluated.""" + ) + score: Optional[float] = Field( + default=None, description="""The score of the metric.""" + ) + explanation: Optional[str] = Field( + default=None, description="""The explanation for the metric.""" + ) + rubric_verdicts: Optional[list[RubricVerdict]] = Field( + default=None, description="""The rubric verdicts for the metric.""" + ) + additional_results: Optional[dict[str, Any]] = Field( + default=None, description="""Additional results for the metric.""" + ) + + +class CandidateResultDict(TypedDict, total=False): + """Result for a single candidate.""" + + candidate: Optional[str] + """The candidate that is being evaluated. The value is the same as the candidate name in the EvaluationRequest.""" + + metric: Optional[str] + """The metric that was evaluated.""" + + score: Optional[float] + """The score of the metric.""" + + explanation: Optional[str] + """The explanation for the metric.""" + + rubric_verdicts: Optional[list[RubricVerdictDict]] + """The rubric verdicts for the metric.""" + + additional_results: Optional[dict[str, Any]] + """Additional results for the metric.""" + + +CandidateResultOrDict = Union[CandidateResult, CandidateResultDict] + + +class Event(_common.BaseModel): + """Represents an event in a conversation between agents and users. + + It is used to store the content of the conversation, as well as the actions + taken by the agents like function calls, function responses, intermediate NL + responses etc. + """ + + event_id: Optional[str] = Field( + default=None, description="""Unique identifier for the agent event.""" + ) + content: Optional[genai_types.Content] = Field( + default=None, description="""Content of the event.""" + ) + creation_timestamp: Optional[datetime.datetime] = Field( + default=None, description="""The creation timestamp of the event.""" + ) + author: Optional[str] = Field( + default=None, description="""Name of the entity that produced the event.""" + ) + + +class EventDict(TypedDict, total=False): + """Represents an event in a conversation between agents and users. + + It is used to store the content of the conversation, as well as the actions + taken by the agents like function calls, function responses, intermediate NL + responses etc. + """ + + event_id: Optional[str] + """Unique identifier for the agent event.""" + + content: Optional[genai_types.ContentDict] + """Content of the event.""" + + creation_timestamp: Optional[datetime.datetime] + """The creation timestamp of the event.""" + + author: Optional[str] + """Name of the entity that produced the event.""" + + +EventOrDict = Union[Event, EventDict] + + +class Message(_common.BaseModel): + """Represents a single message turn in a conversation.""" + + turn_id: Optional[str] = Field( + default=None, description="""Unique identifier for the message turn.""" + ) + content: Optional[genai_types.Content] = Field( + default=None, description="""Content of the message, including function call.""" + ) + creation_timestamp: Optional[datetime.datetime] = Field( + default=None, + description="""Timestamp indicating when the message was created.""", + ) + author: Optional[str] = Field( + default=None, description="""Name of the entity that produced the message.""" + ) + + +class MessageDict(TypedDict, total=False): + """Represents a single message turn in a conversation.""" + + turn_id: Optional[str] + """Unique identifier for the message turn.""" + + content: Optional[genai_types.ContentDict] + """Content of the message, including function call.""" + + creation_timestamp: Optional[datetime.datetime] + """Timestamp indicating when the message was created.""" + + author: Optional[str] + """Name of the entity that produced the message.""" + + +MessageOrDict = Union[Message, MessageDict] + + +class SessionInput(_common.BaseModel): + """This field is experimental and may change in future versions. + + Input to initialize a session and run an agent, used for agent evaluation. + """ + + user_id: Optional[str] = Field(default=None, description="""The user id.""") + state: Optional[dict[str, str]] = Field( + default=None, description="""The state of the session.""" + ) + + +class SessionInputDict(TypedDict, total=False): + """This field is experimental and may change in future versions. + + Input to initialize a session and run an agent, used for agent evaluation. + """ + + user_id: Optional[str] + """The user id.""" + + state: Optional[dict[str, str]] + """The state of the session.""" + + +SessionInputOrDict = Union[SessionInput, SessionInputDict] + + class Tools(_common.BaseModel): """Represents a list of tools for an agent.""" From 13faa27376814f7b0a223ff9455c289a9af75288 Mon Sep 17 00:00:00 2001 From: Tongzhou Jiang Date: Fri, 31 Oct 2025 16:14:36 -0700 Subject: [PATCH 06/24] feat: Alow VertexAiSession for streaming_agent_run_with_events PiperOrigin-RevId: 826664694 --- vertexai/agent_engines/templates/adk.py | 50 ++++++++++++------- .../reasoning_engines/templates/adk.py | 50 +++++++++++-------- 2 files changed, 61 insertions(+), 39 deletions(-) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index c16aed3a1a..ca7c715075 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -568,7 +568,6 @@ async def _init_session( ): """Initializes the session, and returns the session id.""" from google.adk.events.event import Event - import random session_state = None if request.authorizations: @@ -577,14 +576,9 @@ async def _init_session( auth = _Authorization(**auth) session_state[f"temp:{auth_id}"] = auth.access_token - if request.session_id: - session_id = request.session_id - else: - session_id = f"temp_session_{random.randbytes(8).hex()}" session = await session_service.create_session( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, - session_id=session_id, state=session_state, ) if not session: @@ -602,7 +596,7 @@ async def _init_session( saved_version = await artifact_service.save_artifact( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, - session_id=session_id, + session_id=session.id, filename=artifact.file_name, artifact=version_data.data, ) @@ -998,35 +992,53 @@ async def streaming_agent_run_with_events(self, request_json: str): import json from google.genai import types + from google.genai.errors import ClientError request = _StreamRunRequest(**json.loads(request_json)) if not self._tmpl_attrs.get("in_memory_runner"): self.set_up() + if not self._tmpl_attrs.get("runner"): + self.set_up() # Prepare the in-memory session. if not self._tmpl_attrs.get("in_memory_artifact_service"): self.set_up() + if not self._tmpl_attrs.get("artifact_service"): + self.set_up() if not self._tmpl_attrs.get("in_memory_session_service"): self.set_up() - session_service = self._tmpl_attrs.get("in_memory_session_service") - artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") + if not self._tmpl_attrs.get("session_service"): + self.set_up() app = self._tmpl_attrs.get("app") + # Try to get the session, if it doesn't exist, create a new one. - session = None if request.session_id: + session_service = self._tmpl_attrs.get("session_service") + artifact_service = self._tmpl_attrs.get("artifact_service") + runner = self._tmpl_attrs.get("runner") try: session = await session_service.get_session( app_name=app.name if app else self._tmpl_attrs.get("app_name"), user_id=request.user_id, session_id=request.session_id, ) - except RuntimeError: - pass - if not session: - # Fall back to create session if the session is not found. - session = await self._init_session( - session_service=session_service, - artifact_service=artifact_service, - request=request, + except ClientError: + # Fall back to create session if the session is not found. + # Specifying session_id on creation is not supported, + # so session id will be regenerated. + session = await self._init_session( + session_service=session_service, + artifact_service=artifact_service, + request=request, + ) + else: + # Not providing a session ID will create a new in-memory session. + session_service = self._tmpl_attrs.get("in_memory_session_service") + artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") + runner = self._tmpl_attrs.get("in_memory_runner") + session = await session_service.create_session( + app_name=self._tmpl_attrs.get("app_name"), + user_id=request.user_id, + session_id=request.session_id, ) if not session: raise RuntimeError("Session initialization failed.") @@ -1034,7 +1046,7 @@ async def streaming_agent_run_with_events(self, request_json: str): # Run the agent message_for_agent = types.Content(**request.message) try: - async for event in self._tmpl_attrs.get("in_memory_runner").run_async( + async for event in runner.run_async( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 345ff981f1..b07b1ddc98 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -520,7 +520,6 @@ async def _init_session( ): """Initializes the session, and returns the session id.""" from google.adk.events.event import Event - import random session_state = None if request.authorizations: @@ -529,14 +528,9 @@ async def _init_session( auth = _Authorization(**auth) session_state[f"temp:{auth_id}"] = auth.access_token - if request.session_id: - session_id = request.session_id - else: - session_id = f"temp_session_{random.randbytes(8).hex()}" session = await session_service.create_session( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, - session_id=session_id, state=session_state, ) if not session: @@ -554,7 +548,7 @@ async def _init_session( saved_version = await artifact_service.save_artifact( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, - session_id=session_id, + session_id=session.id, filename=artifact.file_name, artifact=version_data.data, ) @@ -904,6 +898,7 @@ async def async_stream_query( def streaming_agent_run_with_events(self, request_json: str): import json from google.genai import types + from google.genai.errors import ClientError event_queue = queue.Queue(maxsize=1) @@ -911,37 +906,52 @@ async def _invoke_agent_async(): request = _StreamRunRequest(**json.loads(request_json)) if not self._tmpl_attrs.get("in_memory_runner"): self.set_up() + if not self._tmpl_attrs.get("runner"): + self.set_up() # Prepare the in-memory session. if not self._tmpl_attrs.get("in_memory_artifact_service"): self.set_up() + if not self._tmpl_attrs.get("artifact_service"): + self.set_up() if not self._tmpl_attrs.get("in_memory_session_service"): self.set_up() - session_service = self._tmpl_attrs.get("in_memory_session_service") - artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") - # Try to get the session, if it doesn't exist, create a new one. - session = None + if not self._tmpl_attrs.get("session_service"): + self.set_up() if request.session_id: + session_service = self._tmpl_attrs.get("session_service") + artifact_service = self._tmpl_attrs.get("artifact_service") + runner = self._tmpl_attrs.get("runner") try: session = await session_service.get_session( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, session_id=request.session_id, ) - except RuntimeError: - pass - if not session: - # Fall back to create session if the session is not found. - session = await self._init_session( - session_service=session_service, - artifact_service=artifact_service, - request=request, + except ClientError: + # Fall back to create session if the session is not found. + # Specifying session_id on creation is not supported, + # so session id will be regenerated. + session = await self._init_session( + session_service=session_service, + artifact_service=artifact_service, + request=request, + ) + else: + # Not providing a session ID will create a new in-memory session. + session_service = self._tmpl_attrs.get("in_memory_session_service") + artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") + runner = self._tmpl_attrs.get("in_memory_runner") + session = await session_service.create_session( + app_name=self._tmpl_attrs.get("app_name"), + user_id=request.user_id, + session_id=request.session_id, ) if not session: raise RuntimeError("Session initialization failed.") # Run the agent. message_for_agent = types.Content(**request.message) try: - for event in self._tmpl_attrs.get("in_memory_runner").run( + for event in runner.run( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, From 9a46e67a8c341673b14bece88bc635b455314711 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Sat, 1 Nov 2025 18:35:27 -0700 Subject: [PATCH 07/24] feat: GenAI Client(evals) - Add retry to predefine metric PiperOrigin-RevId: 826984457 --- tests/unit/vertexai/genai/test_evals.py | 105 ++++++++++++++++++++++ vertexai/_genai/_evals_metric_handlers.py | 30 ++++++- 2 files changed, 132 insertions(+), 3 deletions(-) diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py index dec2bb447d..bcfd8cc888 100644 --- a/tests/unit/vertexai/genai/test_evals.py +++ b/tests/unit/vertexai/genai/test_evals.py @@ -33,6 +33,7 @@ from vertexai._genai import evals from vertexai._genai import types as vertexai_genai_types from google.genai import client +from google.genai import errors as genai_errors from google.genai import types as genai_types import pandas as pd import pytest @@ -4861,6 +4862,110 @@ def test_execute_evaluation_adds_creation_timestamp( assert result.metadata is not None assert result.metadata.creation_timestamp == mock_now + @mock.patch( + "vertexai._genai._evals_metric_handlers._evals_constant.SUPPORTED_PREDEFINED_METRICS", + frozenset(["summarization_quality"]), + ) + @mock.patch("time.sleep", return_value=None) + @mock.patch("vertexai._genai.evals.Evals._evaluate_instances") + def test_predefined_metric_retry_on_resource_exhausted( + self, + mock_private_evaluate_instances, + mock_sleep, + mock_api_client_fixture, + ): + dataset_df = pd.DataFrame( + [{"prompt": "Test prompt", "response": "Test response"}] + ) + input_dataset = vertexai_genai_types.EvaluationDataset( + eval_dataset_df=dataset_df + ) + metric = vertexai_genai_types.Metric(name="summarization_quality") + metric_result = vertexai_genai_types.MetricResult( + score=0.9, + explanation="Mocked predefined explanation", + rubric_verdicts=[], + error=None, + ) + error_response_json = { + "error": { + "code": 429, + "message": ("Judge model resource exhausted. Please try again later."), + "status": "RESOURCE_EXHAUSTED", + } + } + mock_private_evaluate_instances.side_effect = [ + genai_errors.ClientError(code=429, response_json=error_response_json), + genai_errors.ClientError(code=429, response_json=error_response_json), + vertexai_genai_types.EvaluateInstancesResponse( + metric_results=[metric_result] + ), + ] + + result = _evals_common._execute_evaluation( + api_client=mock_api_client_fixture, + dataset=input_dataset, + metrics=[metric], + ) + + assert mock_private_evaluate_instances.call_count == 3 + assert mock_sleep.call_count == 2 + assert len(result.summary_metrics) == 1 + summary_metric = result.summary_metrics[0] + assert summary_metric.metric_name == "summarization_quality" + assert summary_metric.mean_score == 0.9 + + @mock.patch( + "vertexai._genai._evals_metric_handlers._evals_constant.SUPPORTED_PREDEFINED_METRICS", + frozenset(["summarization_quality"]), + ) + @mock.patch("time.sleep", return_value=None) + @mock.patch("vertexai._genai.evals.Evals._evaluate_instances") + def test_predefined_metric_retry_fail_on_resource_exhausted( + self, + mock_private_evaluate_instances, + mock_sleep, + mock_api_client_fixture, + ): + dataset_df = pd.DataFrame( + [{"prompt": "Test prompt", "response": "Test response"}] + ) + input_dataset = vertexai_genai_types.EvaluationDataset( + eval_dataset_df=dataset_df + ) + error_response_json = { + "error": { + "code": 429, + "message": ("Judge model resource exhausted. Please try again later."), + "status": "RESOURCE_EXHAUSTED", + } + } + metric = vertexai_genai_types.Metric(name="summarization_quality") + mock_private_evaluate_instances.side_effect = [ + genai_errors.ClientError(code=429, response_json=error_response_json), + genai_errors.ClientError(code=429, response_json=error_response_json), + genai_errors.ClientError(code=429, response_json=error_response_json), + ] + + result = _evals_common._execute_evaluation( + api_client=mock_api_client_fixture, + dataset=input_dataset, + metrics=[metric], + ) + + assert mock_private_evaluate_instances.call_count == 3 + assert mock_sleep.call_count == 2 + assert len(result.summary_metrics) == 1 + summary_metric = result.summary_metrics[0] + assert summary_metric.metric_name == "summarization_quality" + assert summary_metric.mean_score is None + assert summary_metric.num_cases_error == 1 + assert ( + "Judge model resource exhausted after 3 retries" + ) in result.eval_case_results[0].response_candidate_results[0].metric_results[ + "summarization_quality" + ].error_message + class TestEvaluationDataset: """Contains set of tests for the EvaluationDataset class methods.""" diff --git a/vertexai/_genai/_evals_metric_handlers.py b/vertexai/_genai/_evals_metric_handlers.py index 9f68bc353d..eec98cfaa0 100644 --- a/vertexai/_genai/_evals_metric_handlers.py +++ b/vertexai/_genai/_evals_metric_handlers.py @@ -20,8 +20,10 @@ import json import logging import statistics +import time from typing import Any, Callable, Optional, TypeVar, Union +from google.genai import errors as genai_errors from google.genai import _common from google.genai import types as genai_types from tqdm import tqdm @@ -34,6 +36,7 @@ logger = logging.getLogger(__name__) +_MAX_RETRIES = 3 def _extract_text_from_content( @@ -964,9 +967,30 @@ def get_metric_result( metric_name = self.metric.name try: payload = self._build_request_payload(eval_case, response_index) - api_response = self.module._evaluate_instances( - metrics=[self.metric], instance=payload.get("instance") - ) + for attempt in range(_MAX_RETRIES): + try: + api_response = self.module._evaluate_instances( + metrics=[self.metric], instance=payload.get("instance") + ) + break + except genai_errors.ClientError as e: + if e.code == 429: + logger.warning( + "Resource Exhausted error on attempt %d/%d: %s. Retrying in %s" + " seconds...", + attempt + 1, + _MAX_RETRIES, + e, + 2**attempt, + ) + if attempt == _MAX_RETRIES - 1: + return types.EvalCaseMetricResult( + metric_name=metric_name, + error_message=f"Judge model resource exhausted after {_MAX_RETRIES} retries: {e}", + ) + time.sleep(2**attempt) + else: + raise e if ( api_response From 59e3004338a0a5c810a237ce821bd97f7f77d65d Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Nov 2025 10:35:08 -0800 Subject: [PATCH 08/24] chore: force flush OTel logs at the end of each request Similar to what is already done for spans. Logs and spans are flushed in concurrently. PiperOrigin-RevId: 827559067 --- .../test_agent_engine_templates_adk.py | 91 ++++++++++++++++++ .../test_reasoning_engine_templates_adk.py | 93 +++++++++++++++++++ vertexai/agent_engines/templates/adk.py | 41 +++++--- .../reasoning_engines/templates/adk.py | 41 +++++--- 4 files changed, 242 insertions(+), 24 deletions(-) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 1eb91046bb..980e985433 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -170,6 +170,34 @@ def trace_provider_mock(): yield tracer_provider_mock +@pytest.fixture +def trace_provider_force_flush_mock(): + import opentelemetry.trace + import opentelemetry.sdk.trace + + with mock.patch.object( + opentelemetry.trace, "get_tracer_provider" + ) as get_tracer_provider_mock: + get_tracer_provider_mock.return_value = mock.Mock( + spec=opentelemetry.sdk.trace.TracerProvider() + ) + yield get_tracer_provider_mock.return_value.force_flush + + +@pytest.fixture +def logger_provider_force_flush_mock(): + import opentelemetry._logs + import opentelemetry.sdk._logs + + with mock.patch.object( + opentelemetry._logs, "get_logger_provider" + ) as get_logger_provider_mock: + get_logger_provider_mock.return_value = mock.Mock( + spec=opentelemetry.sdk._logs.LoggerProvider() + ) + yield get_logger_provider_mock.return_value.force_flush + + @pytest.fixture def default_instrumentor_builder_mock(): with mock.patch( @@ -351,6 +379,29 @@ async def test_async_stream_query(self): events.append(event) assert len(events) == 1 + @pytest.mark.asyncio + @mock.patch.dict( + os.environ, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + ) + async def test_async_stream_query_force_flush_otel( + self, + trace_provider_force_flush_mock: mock.Mock, + logger_provider_force_flush_mock: mock.Mock, + ): + app = agent_engines.AdkApp(agent=_TEST_AGENT) + assert app._tmpl_attrs.get("runner") is None + app.set_up() + app._tmpl_attrs["runner"] = _MockRunner() + async for _ in app.async_stream_query( + user_id=_TEST_USER_ID, + message="test message", + ): + pass + + trace_provider_force_flush_mock.assert_called_once() + logger_provider_force_flush_mock.assert_called_once() + @pytest.mark.asyncio async def test_async_stream_query_with_content(self): app = agent_engines.AdkApp(agent=_TEST_AGENT) @@ -403,6 +454,46 @@ async def test_streaming_agent_run_with_events(self): events.append(event) assert len(events) == 1 + @pytest.mark.asyncio + @mock.patch.dict( + os.environ, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + ) + async def test_streaming_agent_run_with_events_force_flush_otel( + self, + trace_provider_force_flush_mock: mock.Mock, + logger_provider_force_flush_mock: mock.Mock, + ): + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app.set_up() + app._tmpl_attrs["in_memory_runner"] = _MockRunner() + request_json = json.dumps( + { + "artifacts": [ + { + "file_name": "test_file_name", + "versions": [{"version": "v1", "data": "v1data"}], + } + ], + "authorizations": { + "test_user_id1": {"access_token": "test_access_token"}, + "test_user_id2": {"accessToken": "test-access-token"}, + }, + "user_id": _TEST_USER_ID, + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + } + ) + async for _ in app.streaming_agent_run_with_events( + request_json=request_json, + ): + pass + + trace_provider_force_flush_mock.assert_called_once() + logger_provider_force_flush_mock.assert_called_once() + @pytest.mark.asyncio async def test_async_create_session(self): app = agent_engines.AdkApp(agent=_TEST_AGENT) diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index 32c0319fc4..917b9430d1 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -128,6 +128,34 @@ def trace_provider_mock(): yield tracer_provider_mock +@pytest.fixture +def trace_provider_force_flush_mock(): + import opentelemetry.trace + import opentelemetry.sdk.trace + + with mock.patch.object( + opentelemetry.trace, "get_tracer_provider" + ) as get_tracer_provider_mock: + get_tracer_provider_mock.return_value = mock.Mock( + spec=opentelemetry.sdk.trace.TracerProvider() + ) + yield get_tracer_provider_mock.return_value.force_flush + + +@pytest.fixture +def logger_provider_force_flush_mock(): + import opentelemetry._logs + import opentelemetry.sdk._logs + + with mock.patch.object( + opentelemetry._logs, "get_logger_provider" + ) as get_logger_provider_mock: + get_logger_provider_mock.return_value = mock.Mock( + spec=opentelemetry.sdk._logs.LoggerProvider() + ) + yield get_logger_provider_mock.return_value.force_flush + + @pytest.fixture def default_instrumentor_builder_mock(): with mock.patch( @@ -353,6 +381,31 @@ async def test_async_stream_query(self): events.append(event) assert len(events) == 1 + @pytest.mark.asyncio + @mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY": "true"}, + ) + async def test_async_stream_query_force_flush_otel( + self, + trace_provider_force_flush_mock: mock.Mock, + logger_provider_force_flush_mock: mock.Mock, + ): + app = reasoning_engines.AdkApp( + agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL), enable_tracing=True + ) + assert app._tmpl_attrs.get("runner") is None + app.set_up() + app._tmpl_attrs["runner"] = _MockRunner() + async for _ in app.async_stream_query( + user_id=_TEST_USER_ID, + message="test message", + ): + pass + + trace_provider_force_flush_mock.assert_called_once() + logger_provider_force_flush_mock.assert_called_once() + @pytest.mark.asyncio async def test_async_stream_query_with_content(self): app = reasoning_engines.AdkApp( @@ -404,6 +457,46 @@ def test_streaming_agent_run_with_events(self): events = list(app.streaming_agent_run_with_events(request_json=request_json)) assert len(events) == 1 + @pytest.mark.asyncio + @mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY": "true"}, + ) + async def test_streaming_agent_run_with_events_force_flush_otel( + self, + trace_provider_force_flush_mock: mock.Mock, + logger_provider_force_flush_mock: mock.Mock, + ): + app = reasoning_engines.AdkApp( + agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL), + enable_tracing=True, + ) + app.set_up() + app._tmpl_attrs["in_memory_runner"] = _MockRunner() + request_json = json.dumps( + { + "artifacts": [ + { + "file_name": "test_file_name", + "versions": [{"version": "v1", "data": "v1data"}], + } + ], + "authorizations": { + "test_user_id1": {"access_token": "test_access_token"}, + "test_user_id2": {"accessToken": "test-access-token"}, + }, + "user_id": _TEST_USER_ID, + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + } + ) + list(app.streaming_agent_run_with_events(request_json=request_json)) + + trace_provider_force_flush_mock.assert_called_once() + logger_provider_force_flush_mock.assert_called_once() + @pytest.mark.asyncio async def test_async_bidi_stream_query(self): app = reasoning_engines.AdkApp( diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index ca7c715075..db1b1ff423 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -25,6 +25,7 @@ ) import asyncio +from collections.abc import Awaitable import queue import threading import warnings @@ -231,26 +232,38 @@ def _warn(msg: str): _warn._LOGGER.warning(msg) # pyright: ignore[reportFunctionMemberAccess] -def _force_flush_traces(): +async def _force_flush_otel(tracing_enabled: bool, logging_enabled: bool): try: import opentelemetry.trace + import opentelemetry._logs except (ImportError, AttributeError): _warn( - "Could not force flush traces. opentelemetry-api is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + "Could not force flush telemetry data. opentelemetry-api is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." ) return None try: import opentelemetry.sdk.trace + import opentelemetry.sdk._logs except (ImportError, AttributeError): _warn( - "Could not force flush traces. opentelemetry-sdk is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + "Could not force flush telemetry data. opentelemetry-sdk is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." ) return None - provider = opentelemetry.trace.get_tracer_provider() - if isinstance(provider, opentelemetry.sdk.trace.TracerProvider): - _ = provider.force_flush() + coros: List[Awaitable[bool]] = [] + + if tracing_enabled: + tracer_provider = opentelemetry.trace.get_tracer_provider() + if isinstance(tracer_provider, opentelemetry.sdk.trace.TracerProvider): + coros.append(asyncio.to_thread(tracer_provider.force_flush)) + + if logging_enabled: + logger_provider = opentelemetry._logs.get_logger_provider() + if isinstance(logger_provider, opentelemetry.sdk._logs.LoggerProvider): + coros.append(asyncio.to_thread(logger_provider.force_flush)) + + await asyncio.gather(*coros, return_exceptions=True) def _default_instrumentor_builder( @@ -894,9 +907,11 @@ async def async_stream_query( # Yield the event data as a dictionary yield _utils.dump_event_for_json(event) finally: - # Avoid trace data loss having to do with CPU throttling on instance turndown - if self._tracing_enabled(): - _ = await asyncio.to_thread(_force_flush_traces) + # Avoid telemetry data loss having to do with CPU throttling on instance turndown + _ = await _force_flush_otel( + tracing_enabled=self._tracing_enabled(), + logging_enabled=bool(self._telemetry_enabled()), + ) def stream_query( self, @@ -1066,9 +1081,11 @@ async def streaming_agent_run_with_events(self, request_json: str): user_id=request.user_id, session_id=session.id, ) - # Avoid trace data loss having to do with CPU throttling on instance turndown - if self._tracing_enabled(): - _ = await asyncio.to_thread(_force_flush_traces) + # Avoid telemetry data loss having to do with CPU throttling on instance turndown + _ = await _force_flush_otel( + tracing_enabled=self._tracing_enabled(), + logging_enabled=bool(self._telemetry_enabled()), + ) async def async_get_session( self, diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index b07b1ddc98..a74e1bb6e6 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -25,6 +25,7 @@ ) import asyncio +from collections.abc import Awaitable import queue import threading @@ -233,26 +234,38 @@ def _warn(msg: str): _warn._LOGGER.warning(msg) # pyright: ignore[reportFunctionMemberAccess] -def _force_flush_traces(): +async def _force_flush_otel(tracing_enabled: bool, logging_enabled: bool): try: import opentelemetry.trace + import opentelemetry._logs except (ImportError, AttributeError): _warn( - "Could not force flush traces. opentelemetry-api is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + "Could not force flush telemetry data. opentelemetry-api is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." ) return None try: import opentelemetry.sdk.trace + import opentelemetry.sdk._logs except (ImportError, AttributeError): _warn( - "Could not force flush traces. opentelemetry-sdk is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." + "Could not force flush telemetry data. opentelemetry-sdk is not installed. Please call 'pip install google-cloud-aiplatform[agent_engines]'." ) return None - provider = opentelemetry.trace.get_tracer_provider() - if isinstance(provider, opentelemetry.sdk.trace.TracerProvider): - _ = provider.force_flush() + coros: List[Awaitable[bool]] = [] + + if tracing_enabled: + tracer_provider = opentelemetry.trace.get_tracer_provider() + if isinstance(tracer_provider, opentelemetry.sdk.trace.TracerProvider): + coros.append(asyncio.to_thread(tracer_provider.force_flush)) + + if logging_enabled: + logger_provider = opentelemetry._logs.get_logger_provider() + if isinstance(logger_provider, opentelemetry.sdk._logs.LoggerProvider): + coros.append(asyncio.to_thread(logger_provider.force_flush)) + + await asyncio.gather(*coros, return_exceptions=True) def _default_instrumentor_builder( @@ -891,9 +904,11 @@ async def async_stream_query( # Yield the event data as a dictionary yield _utils.dump_event_for_json(event) finally: - # Avoid trace data loss having to do with CPU throttling on instance turndown - if self._tracing_enabled(): - _ = await asyncio.to_thread(_force_flush_traces) + # Avoid telemetry data loss having to do with CPU throttling on instance turndown + _ = await _force_flush_otel( + tracing_enabled=self._tracing_enabled(), + logging_enabled=bool(self._telemetry_enabled()), + ) def streaming_agent_run_with_events(self, request_json: str): import json @@ -970,9 +985,11 @@ async def _invoke_agent_async(): user_id=request.user_id, session_id=session.id, ) - # Avoid trace data loss having to do with CPU throttling on instance turndown - if self._tracing_enabled(): - _ = await asyncio.to_thread(_force_flush_traces) + # Avoid telemetry data loss having to do with CPU throttling on instance turndown + _ = await _force_flush_otel( + tracing_enabled=self._tracing_enabled(), + logging_enabled=bool(self._telemetry_enabled()), + ) def _asyncio_thread_main(): try: From 6fd298460026e76fcb290469e15ac2d8977b402a Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Nov 2025 11:36:06 -0800 Subject: [PATCH 09/24] chore: make default value of enable_tracing in ADK preview template to match the GA one This will allow for default-on telemetry with preview template if enable_tracing was not specified. See https://github.com/googleapis/python-aiplatform/blob/9a46e67a8c341673b14bece88bc635b455314711/vertexai/agent_engines/templates/adk.py#L481. PiperOrigin-RevId: 827584607 --- vertexai/preview/reasoning_engines/templates/adk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index a74e1bb6e6..eefe01ed53 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -491,7 +491,7 @@ def __init__( *, agent: "BaseAgent", plugins: Optional[List["BasePlugin"]] = None, - enable_tracing: bool = False, + enable_tracing: Optional[bool] = None, session_service_builder: Optional[Callable[..., "BaseSessionService"]] = None, artifact_service_builder: Optional[Callable[..., "BaseArtifactService"]] = None, memory_service_builder: Optional[Callable[..., "BaseMemoryService"]] = None, From a9171221e3bafecdc75580c3f25347f1c3d18851 Mon Sep 17 00:00:00 2001 From: Jason Dai Date: Mon, 3 Nov 2025 12:04:38 -0800 Subject: [PATCH 10/24] fix: GenAI Client(evals) - Support direct pandas DataFrame dataset in evaluate() PiperOrigin-RevId: 827595694 --- .../vertexai/genai/replays/test_evaluate.py | 41 +++++++++++++++++++ vertexai/_genai/evals.py | 17 ++++++-- 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/tests/unit/vertexai/genai/replays/test_evaluate.py b/tests/unit/vertexai/genai/replays/test_evaluate.py index 7b497ccd11..d934b7f5d4 100644 --- a/tests/unit/vertexai/genai/replays/test_evaluate.py +++ b/tests/unit/vertexai/genai/replays/test_evaluate.py @@ -54,6 +54,47 @@ def test_evaluation_result(client): assert case_result.response_candidate_results is not None +def test_evaluation_byor(client): + """Tests that evaluate() with BYOR (Bring-Your-Own Response) produces a correctly structured EvaluationResult.""" + byor_df = pd.DataFrame( + { + "prompt": [ + "Write a simple story about a dinosaur", + "Generate a poem about Vertex AI", + ], + "response": [ + "Once upon a time, there was a T-Rex named Rexy.", + "In clouds of code, a mind of silicon born...", + ], + } + ) + + metrics_to_run = [ + types.RubricMetric.GENERAL_QUALITY, + ] + + evaluation_result = client.evals.evaluate( + dataset=byor_df, + metrics=metrics_to_run, + ) + + assert isinstance(evaluation_result, types.EvaluationResult) + + assert evaluation_result.summary_metrics is not None + assert len(evaluation_result.summary_metrics) > 0 + for summary in evaluation_result.summary_metrics: + assert isinstance(summary, types.AggregatedMetricResult) + assert summary.metric_name is not None + assert summary.mean_score is not None + + assert evaluation_result.eval_case_results is not None + assert len(evaluation_result.eval_case_results) > 0 + for case_result in evaluation_result.eval_case_results: + assert isinstance(case_result, types.EvalCaseResult) + assert case_result.eval_case_index is not None + assert case_result.response_candidate_results is not None + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py index 7b7ad2c023..c21b186ecd 100644 --- a/vertexai/_genai/evals.py +++ b/vertexai/_genai/evals.py @@ -970,7 +970,9 @@ def evaluate( self, *, dataset: Union[ - types.EvaluationDatasetOrDict, list[types.EvaluationDatasetOrDict] + pd.DataFrame, + types.EvaluationDatasetOrDict, + list[types.EvaluationDatasetOrDict], ], metrics: list[types.MetricOrDict] = None, config: Optional[types.EvaluateMethodConfigOrDict] = None, @@ -979,10 +981,13 @@ def evaluate( """Evaluates candidate responses in the provided dataset(s) using the specified metrics. Args: - dataset: The dataset(s) to evaluate. Can be a single `types.EvaluationDataset` or a list of `types.EvaluationDataset`. + dataset: The dataset(s) to evaluate. Can be a pandas DataFrame, a single + `types.EvaluationDataset` or a list of `types.EvaluationDataset`. metrics: The list of metrics to use for evaluation. - config: Optional configuration for the evaluation. Can be a dictionary or a `types.EvaluateMethodConfig` object. - - dataset_schema: Schema to use for the dataset. If not specified, the dataset schema will be inferred from the dataset automatically. + config: Optional configuration for the evaluation. Can be a dictionary or a + `types.EvaluateMethodConfig` object. + - dataset_schema: Schema to use for the dataset. If not specified, the + dataset schema will be inferred from the dataset automatically. - dest: Destination path for storing evaluation results. **kwargs: Extra arguments to pass to evaluation, such as `agent_info`. @@ -993,6 +998,10 @@ def evaluate( config = types.EvaluateMethodConfig() if isinstance(config, dict): config = types.EvaluateMethodConfig.model_validate(config) + + if isinstance(dataset, pd.DataFrame): + dataset = types.EvaluationDataset(eval_dataset_df=dataset) + if isinstance(dataset, list): dataset = [ ( From 7d18e72e3a868152ceae6c6dcdc06062aa0e9ba0 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Nov 2025 12:28:31 -0800 Subject: [PATCH 11/24] chore: add warning to ADK deployment if tracing is disabled. PiperOrigin-RevId: 827604495 --- .../test_agent_engine_templates_adk.py | 16 ++++++++++++ .../test_reasoning_engine_templates_adk.py | 19 ++++++++++++++ vertexai/agent_engines/templates/adk.py | 26 +++++++++++++++++++ .../reasoning_engines/templates/adk.py | 26 +++++++++++++++++++ 4 files changed, 87 insertions(+) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 980e985433..19fb6c79d4 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -800,6 +800,22 @@ def test_enable_tracing_warning(self, caplog): # app.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + # TODO(b/384730642): Re-enable this test once the parent issue is fixed. + # @pytest.mark.parametrize( + # "enable_tracing,want_warning", + # [ + # (True, False), + # (False, True), + # (None, False), + # ], + # ) + # @pytest.mark.usefixtures("caplog") + # def test_tracing_disabled_warning(self, enable_tracing, want_warning, caplog): + # _ = agent_engines.AdkApp(agent=_TEST_AGENT, enable_tracing=enable_tracing) + # assert ( + # "[WARNING] Your 'enable_tracing=False' setting" in caplog.text + # ) == want_warning + @mock.patch.dict(os.environ) def test_span_content_capture_disabled_by_default(self): app = agent_engines.AdkApp(agent=_TEST_AGENT) diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index 917b9430d1..add706859a 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -781,6 +781,25 @@ def test_enable_tracing_warning(self, caplog): # app.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + # TODO(b/384730642): Re-enable this test once the parent issue is fixed. + # @pytest.mark.parametrize( + # "enable_tracing,want_warning", + # [ + # (True, False), + # (False, True), + # (None, False), + # ], + # ) + # @pytest.mark.usefixtures("caplog") + # def test_tracing_disabled_warning(self, enable_tracing, want_warning, caplog): + # app = reasoning_engines.AdkApp( + # agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL), + # enable_tracing=enable_tracing + # ) + # assert ( + # "[WARNING] Your 'enable_tracing=False' setting" in caplog.text + # ) == want_warning + def test_dump_event_for_json(): from google.adk.events import event diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index db1b1ff423..4176b58e4a 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -724,6 +724,32 @@ def set_up(self): custom_instrumentor = self._tmpl_attrs.get("instrumentor_builder") + if self._tmpl_attrs.get("enable_tracing") is False: + _warn( + ( + "Your 'enable_tracing=False' setting is being deprecated " + "and will be removed in a future release.\n" + "This legacy setting overrides the new Cloud Console " + "toggle and environment variable controls.\n" + "Impact: The Cloud Console may incorrectly show telemetry " + "as 'On' when it is actually 'Off', and the UI toggle will " + "not work.\n" + "Action: To fix this and control telemetry, please remove " + "the 'enable_tracing' parameter from your deployment " + "code.\n" + "You can then use the " + "'GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY' " + "environment variable:\n" + "agent_engines.create(\n" + " env_vars={\n" + ' "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY": true|false\n' + " }\n" + ")\n" + "or the toggle in the Cloud Console: " + "https://console.cloud.google.com/vertex-ai/agents." + ), + ) + if custom_instrumentor and self._tracing_enabled(): self._tmpl_attrs["instrumentor"] = custom_instrumentor(project) diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index eefe01ed53..8b49b992c6 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -657,6 +657,32 @@ def set_up(self): else: os.environ["ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS"] = "false" + if self._tmpl_attrs.get("enable_tracing") is False: + _warn( + ( + "Your 'enable_tracing=False' setting is being deprecated " + "and will be removed in a future release.\n" + "This legacy setting overrides the new Cloud Console " + "toggle and environment variable controls.\n" + "Impact: The Cloud Console may incorrectly show telemetry " + "as 'On' when it is actually 'Off', and the UI toggle will " + "not work.\n" + "Action: To fix this and control telemetry, please remove " + "the 'enable_tracing' parameter from your deployment " + "code.\n" + "You can then use the " + "'GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY' " + "environment variable:\n" + "agent_engines.create(\n" + " env_vars={\n" + ' "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY": true|false\n' + " }\n" + ")\n" + "or the toggle in the Cloud Console: " + "https://console.cloud.google.com/vertex-ai/agents." + ), + ) + enable_logging = bool(self._telemetry_enabled()) self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder( From 05834cb43302cf2bfc34f25f400d075e44aa9df3 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Nov 2025 13:30:58 -0800 Subject: [PATCH 12/24] feat: Add support for Vertex Express Mode API key in AdkApp PiperOrigin-RevId: 827628603 --- .../test_agent_engine_templates_adk.py | 36 +++++++++++++++++++ vertexai/agent_engines/templates/adk.py | 23 ++++++++++-- 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 19fb6c79d4..2b0f06cece 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -54,6 +54,7 @@ def __init__(self, name: str, model: str): _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" +_TEST_API_KEY = "test-api-key" _TEST_MODEL = "gemini-2.0-flash" _TEST_USER_ID = "test_user_id" _TEST_AGENT_NAME = "test_agent" @@ -868,6 +869,41 @@ def test_dump_event_for_json(): assert base64.b64decode(part["thought_signature"]) == raw_signature +def test_adk_app_initialization_with_api_key(): + importlib.reload(initializer) + importlib.reload(vertexai) + try: + vertexai.init(api_key=_TEST_API_KEY) + app = agent_engines.AdkApp(agent=_TEST_AGENT) + assert app._tmpl_attrs.get("project") is None + assert app._tmpl_attrs.get("location") is None + assert app._tmpl_attrs.get("express_mode_api_key") == _TEST_API_KEY + assert app._tmpl_attrs.get("runner") is None + app.set_up() + assert app._tmpl_attrs.get("runner") is not None + assert os.environ.get("GOOGLE_API_KEY") == _TEST_API_KEY + assert "GOOGLE_CLOUD_LOCATION" not in os.environ + assert "GOOGLE_CLOUD_PROJECT" not in os.environ + finally: + initializer.global_pool.shutdown(wait=True) + + +def test_adk_app_initialization_with_env_api_key(): + try: + os.environ["GOOGLE_API_KEY"] == _TEST_API_KEY + app = agent_engines.AdkApp(agent=_TEST_AGENT) + assert app._tmpl_attrs.get("project") is None + assert app._tmpl_attrs.get("location") is None + assert app._tmpl_attrs.get("express_mode_api_key") == _TEST_API_KEY + assert app._tmpl_attrs.get("runner") is None + app.set_up() + assert app._tmpl_attrs.get("runner") is not None + assert "GOOGLE_CLOUD_LOCATION" not in os.environ + assert "GOOGLE_CLOUD_PROJECT" not in os.environ + finally: + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.usefixtures("mock_adk_version") class TestAdkAppErrors: @pytest.mark.asyncio diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 4176b58e4a..f3c186e903 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -534,6 +534,7 @@ def __init__( If not provided, a default instrumentor builder will be used. This parameter is ignored if `enable_tracing` is False. """ + import os from google.cloud.aiplatform import initializer adk_version = get_adk_version() @@ -571,6 +572,9 @@ def __init__( "artifact_service_builder": artifact_service_builder, "memory_service_builder": memory_service_builder, "instrumentor_builder": instrumentor_builder, + "express_mode_api_key": ( + initializer.global_config.api_key or os.environ.get("GOOGLE_API_KEY") + ), } async def _init_session( @@ -708,9 +712,18 @@ def set_up(self): os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" project = self._tmpl_attrs.get("project") - os.environ["GOOGLE_CLOUD_PROJECT"] = project + if project: + os.environ["GOOGLE_CLOUD_PROJECT"] = project location = self._tmpl_attrs.get("location") - os.environ["GOOGLE_CLOUD_LOCATION"] = location + if location: + os.environ["GOOGLE_CLOUD_LOCATION"] = location + express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key") + if express_mode_api_key and not project: + os.environ["GOOGLE_API_KEY"] = express_mode_api_key + # Clear location and project env vars if express mode api key is provided. + os.environ.pop("GOOGLE_CLOUD_LOCATION", None) + os.environ.pop("GOOGLE_CLOUD_PROJECT", None) + location = None # Disable content capture in custom ADK spans unless user enabled # tracing explicitly with the old flag @@ -783,6 +796,8 @@ def set_up(self): VertexAiSessionService, ) + # If the express mode api key is set, it will be read from the + # environment variable when initializing the session service. self._tmpl_attrs["session_service"] = VertexAiSessionService( project=project, location=location, @@ -793,6 +808,8 @@ def set_up(self): VertexAiSessionService, ) + # If the express mode api key is set, it will be read from the + # environment variable when initializing the session service. self._tmpl_attrs["session_service"] = VertexAiSessionService( project=project, location=location, @@ -813,6 +830,8 @@ def set_up(self): VertexAiMemoryBankService, ) + # If the express mode api key is set, it will be read from the + # environment variable when initializing the memory service. self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService( project=project, location=location, From c8f38a0a51c318a5065438067f85f31be5088af1 Mon Sep 17 00:00:00 2001 From: Ray Xu Date: Mon, 3 Nov 2025 14:54:15 -0800 Subject: [PATCH 13/24] feat: Add reservation affinity support to preview BatchPredictionJob PiperOrigin-RevId: 827661239 --- google/cloud/aiplatform/jobs.py | 27 +- google/cloud/aiplatform/preview/jobs.py | 555 +++++++++++++++++- .../test_batch_prediction_job_preview.py | 301 ++++++++++ 3 files changed, 880 insertions(+), 3 deletions(-) create mode 100644 tests/unit/aiplatform/test_batch_prediction_job_preview.py diff --git a/google/cloud/aiplatform/jobs.py b/google/cloud/aiplatform/jobs.py index e854faa3e6..472c1661d5 100644 --- a/google/cloud/aiplatform/jobs.py +++ b/google/cloud/aiplatform/jobs.py @@ -319,7 +319,7 @@ def cancel(self) -> None: getattr(self.api_client, self._cancel_method)(name=self.resource_name) -class BatchPredictionJob(_Job): +class BatchPredictionJob(_Job, base.PreviewMixin): _resource_noun = "batchPredictionJobs" _getter_method = "get_batch_prediction_job" @@ -329,6 +329,9 @@ class BatchPredictionJob(_Job): _job_type = "batch-predictions" _parse_resource_name_method = "parse_batch_prediction_job_path" _format_resource_name_method = "batch_prediction_job_path" + _preview_class = ( + "google.cloud.aiplatform.aiplatform.preview.jobs.BatchPredictionJob" + ) def __init__( self, @@ -949,6 +952,9 @@ def _submit_impl( ] = None, analysis_instance_schema_uri: Optional[str] = None, service_account: Optional[str] = None, + reservation_affinity_type: Optional[str] = None, + reservation_affinity_key: Optional[str] = None, + reservation_affinity_values: Optional[Sequence[str]] = None, wait_for_completion: bool = False, ) -> "BatchPredictionJob": """Create a batch prediction job. @@ -1136,6 +1142,18 @@ def _submit_impl( service_account (str): Optional. Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. + reservation_affinity_type (str): + Optional. The type of reservation affinity. + One of NO_RESERVATION, ANY_RESERVATION, SPECIFIC_RESERVATION, + SPECIFIC_THEN_ANY_RESERVATION, SPECIFIC_THEN_NO_RESERVATION + reservation_affinity_key (str): + Optional. Corresponds to the label key of a reservation resource. + To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key + and specify the name of your reservation as its value. + reservation_affinity_values (List[str]): + Optional. Corresponds to the label values of a reservation resource. + This must be the full resource name of the reservation. + Format: 'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}' wait_for_completion (bool): Whether to wait for the job completion. Returns: @@ -1268,6 +1286,13 @@ def _submit_impl( machine_spec.accelerator_type = accelerator_type machine_spec.accelerator_count = accelerator_count + if reservation_affinity_type: + machine_spec.reservation_affinity = utils.get_reservation_affinity( + reservation_affinity_type, + reservation_affinity_key, + reservation_affinity_values, + ) + dedicated_resources = gca_machine_resources_compat.BatchDedicatedResources() dedicated_resources.machine_spec = machine_spec diff --git a/google/cloud/aiplatform/preview/jobs.py b/google/cloud/aiplatform/preview/jobs.py index b8ed5519e7..645800e408 100644 --- a/google/cloud/aiplatform/preview/jobs.py +++ b/google/cloud/aiplatform/preview/jobs.py @@ -15,9 +15,8 @@ # limitations under the License. # -from typing import Dict, List, Optional, Union - import copy +from typing import Dict, List, Optional, Sequence, Union import uuid from google.api_core import retry @@ -68,6 +67,12 @@ gca_job_state_v1beta1.JobState.JOB_STATE_CANCELLED, ) +# _block_until_complete wait times +_JOB_WAIT_TIME = 5 # start at five seconds +_LOG_WAIT_TIME = 5 +_MAX_WAIT_TIME = 60 * 5 # 5 minute wait +_WAIT_TIME_MULTIPLIER = 2 # scale wait by 2 every iteration + class CustomJob(jobs.CustomJob): """Deprecated. Vertex AI Custom Job (preview).""" @@ -867,3 +872,549 @@ def _run( ) self._block_until_complete() + + +class BatchPredictionJob(jobs.BatchPredictionJob): + """Vertex AI Batch Prediction Job.""" + + @classmethod + def create( + cls, + # TODO(b/223262536): Make the job_display_name parameter optional in the next major release + job_display_name: str, + model_name: Union[str, "aiplatform.Model"], + instances_format: str = "jsonl", + predictions_format: str = "jsonl", + gcs_source: Optional[Union[str, Sequence[str]]] = None, + bigquery_source: Optional[str] = None, + gcs_destination_prefix: Optional[str] = None, + bigquery_destination_prefix: Optional[str] = None, + model_parameters: Optional[Dict] = None, + machine_type: Optional[str] = None, + accelerator_type: Optional[str] = None, + accelerator_count: Optional[int] = None, + starting_replica_count: Optional[int] = None, + max_replica_count: Optional[int] = None, + generate_explanation: Optional[bool] = False, + explanation_metadata: Optional["aiplatform.explain.ExplanationMetadata"] = None, + explanation_parameters: Optional[ + "aiplatform.explain.ExplanationParameters" + ] = None, + labels: Optional[Dict[str, str]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + encryption_spec_key_name: Optional[str] = None, + sync: bool = True, + create_request_timeout: Optional[float] = None, + batch_size: Optional[int] = None, + model_monitoring_objective_config: Optional[ + "aiplatform.model_monitoring.ObjectiveConfig" + ] = None, + model_monitoring_alert_config: Optional[ + "aiplatform.model_monitoring.AlertConfig" + ] = None, + analysis_instance_schema_uri: Optional[str] = None, + service_account: Optional[str] = None, + reservation_affinity_type: Optional[str] = None, + reservation_affinity_key: Optional[str] = None, + reservation_affinity_values: Optional[List[str]] = None, + ) -> "BatchPredictionJob": + """Create a batch prediction job. + + Args: + job_display_name (str): + Required. The user-defined name of the BatchPredictionJob. + The name can be up to 128 characters long and can be consist + of any UTF-8 characters. + model_name (Union[str, aiplatform.Model]): + Required. A fully-qualified model resource name or model ID. + Example: "projects/123/locations/us-central1/models/456" or + "456" when project and location are initialized or passed. + May optionally contain a version ID or alias in + {model_name}@{version} form. + + Or an instance of aiplatform.Model. + instances_format (str): + Required. The format in which instances are provided. Must be one + of the formats listed in `Model.supported_input_storage_formats`. + Default is "jsonl" when using `gcs_source`. If a `bigquery_source` + is provided, this is overridden to "bigquery". + predictions_format (str): + Required. The format in which Vertex AI outputs the + predictions, must be one of the formats specified in + `Model.supported_output_storage_formats`. + Default is "jsonl" when using `gcs_destination_prefix`. If a + `bigquery_destination_prefix` is provided, this is overridden to + "bigquery". + gcs_source (Optional[Sequence[str]]): + Google Cloud Storage URI(-s) to your instances to run + batch prediction on. They must match `instances_format`. + + bigquery_source (Optional[str]): + BigQuery URI to a table, up to 2000 characters long. For example: + `bq://projectId.bqDatasetId.bqTableId` + gcs_destination_prefix (Optional[str]): + The Google Cloud Storage location of the directory where the + output is to be written to. In the given directory a new + directory is created. Its name is + ``prediction--``, where + timestamp is in YYYY-MM-DDThh:mm:ss.sssZ ISO-8601 format. + Inside of it files ``predictions_0001.``, + ``predictions_0002.``, ..., + ``predictions_N.`` are created where + ```` depends on chosen ``predictions_format``, + and N may equal 0001 and depends on the total number of + successfully predicted instances. If the Model has both + ``instance`` and ``prediction`` schemata defined then each such + file contains predictions as per the ``predictions_format``. + If prediction for any instance failed (partially or + completely), then an additional ``errors_0001.``, + ``errors_0002.``,..., ``errors_N.`` + files are created (N depends on total number of failed + predictions). These files contain the failed instances, as + per their schema, followed by an additional ``error`` field + which as value has ```google.rpc.Status`` `__ + containing only ``code`` and ``message`` fields. + bigquery_destination_prefix (Optional[str]): + The BigQuery project or dataset location where the output is + to be written to. If project is provided, a new dataset is + created with name + ``prediction__`` where + is made BigQuery-dataset-name compatible (for example, most + special characters become underscores), and timestamp is in + YYYY_MM_DDThh_mm_ss_sssZ "based on ISO-8601" format. In the + dataset two tables will be created, ``predictions``, and + ``errors``. If the Model has both + [instance][google.cloud.aiplatform.v1.PredictSchemata.instance_schema_uri] + and + [prediction][google.cloud.aiplatform.v1.PredictSchemata.parameters_schema_uri] + schemata defined then the tables have columns as follows: + The ``predictions`` table contains instances for which the + prediction succeeded, it has columns as per a concatenation + of the Model's instance and prediction schemata. The + ``errors`` table contains rows for which the prediction has + failed, it has instance columns, as per the instance schema, + followed by a single "errors" column, which as values has + [google.rpc.Status][google.rpc.Status] represented as a + STRUCT, and containing only ``code`` and ``message``. + model_parameters (Optional[Dict]): + The parameters that govern the predictions. The schema of + the parameters may be specified via the Model's `parameters_schema_uri`. + machine_type (Optional[str]): + The type of machine for running batch prediction on + dedicated resources. Not specifying machine type will result in + batch prediction job being run with automatic resources. + accelerator_type (Optional[str]): + The type of accelerator(s) that may be attached + to the machine as per `accelerator_count`. Only used if + `machine_type` is set. + accelerator_count (Optional[int]): + The number of accelerators to attach to the + `machine_type`. Only used if `machine_type` is set. + starting_replica_count (Optional[int]): + The number of machine replicas used at the start of the batch + operation. If not set, Vertex AI decides starting number, not + greater than `max_replica_count`. Only used if `machine_type` is + set. + max_replica_count (Optional[int]): + The maximum number of machine replicas the batch operation may + be scaled to. Only used if `machine_type` is set. + Default is 10. + generate_explanation (bool): + Optional. Generate explanation along with the batch prediction + results. This will cause the batch prediction output to include + explanations based on the `prediction_format`: + - `bigquery`: output includes a column named `explanation`. The value + is a struct that conforms to the [aiplatform.gapic.Explanation] object. + - `jsonl`: The JSON objects on each line include an additional entry + keyed `explanation`. The value of the entry is a JSON object that + conforms to the [aiplatform.gapic.Explanation] object. + - `csv`: Generating explanations for CSV format is not supported. + explanation_metadata (aiplatform.explain.ExplanationMetadata): + Optional. Explanation metadata configuration for this BatchPredictionJob. + Can be specified only if `generate_explanation` is set to `True`. + + This value overrides the value of `Model.explanation_metadata`. + All fields of `explanation_metadata` are optional in the request. If + a field of the `explanation_metadata` object is not populated, the + corresponding field of the `Model.explanation_metadata` object is inherited. + For more details, see `Ref docs ` + explanation_parameters (aiplatform.explain.ExplanationParameters): + Optional. Parameters to configure explaining for Model's predictions. + Can be specified only if `generate_explanation` is set to `True`. + + This value overrides the value of `Model.explanation_parameters`. + All fields of `explanation_parameters` are optional in the request. If + a field of the `explanation_parameters` object is not populated, the + corresponding field of the `Model.explanation_parameters` object is inherited. + For more details, see `Ref docs ` + labels (Dict[str, str]): + Optional. The labels with user-defined metadata to organize your + BatchPredictionJobs. Label keys and values can be no longer than + 64 characters (Unicode codepoints), can only contain lowercase + letters, numeric characters, underscores and dashes. + International characters are allowed. See https://goo.gl/xmQnxf + for more information and examples of labels. + credentials (Optional[auth_credentials.Credentials]): + Custom credentials to use to create this batch prediction + job. Overrides credentials set in aiplatform.init. + encryption_spec_key_name (Optional[str]): + Optional. The Cloud KMS resource identifier of the customer + managed encryption key used to protect the job. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. + + If this is set, then all + resources created by the BatchPredictionJob will + be encrypted with the provided encryption key. + + Overrides encryption_spec_key_name set in aiplatform.init. + sync (bool): + Whether to execute this method synchronously. If False, this method + will be executed in concurrent Future and any downstream object will + be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. + batch_size (int): + Optional. The number of the records (e.g. instances) of the operation given in each batch + to a machine replica. Machine type, and size of a single record should be considered + when setting this parameter, higher value speeds up the batch operation's execution, + but too high value will result in a whole batch not fitting in a machine's memory, + and the whole operation will fail. + The default value is 64. + model_monitoring_objective_config (aiplatform.model_monitoring.ObjectiveConfig): + Optional. The objective config for model monitoring. Passing this parameter enables + monitoring on the model associated with this batch prediction job. + model_monitoring_alert_config (aiplatform.model_monitoring.EmailAlertConfig): + Optional. Configures how model monitoring alerts are sent to the user. Right now + only email alert is supported. + analysis_instance_schema_uri (str): + Optional. Only applicable if model_monitoring_objective_config is also passed. + This parameter specifies the YAML schema file uri describing the format of a single + instance that you want Tensorflow Data Validation (TFDV) to + analyze. If this field is empty, all the feature data types are + inferred from predict_instance_schema_uri, meaning that TFDV + will use the data in the exact format as prediction request/response. + If there are any data type differences between predict instance + and TFDV instance, this field can be used to override the schema. + For models trained with Vertex AI, this field must be set as all the + fields in predict instance formatted as string. + service_account (str): + Optional. Specifies the service account for workload run-as account. + Users submitting jobs must have act-as permission on this run-as account. + reservation_affinity_type (str): + Optional. The type of reservation affinity. + One of NO_RESERVATION, ANY_RESERVATION, SPECIFIC_RESERVATION, + SPECIFIC_THEN_ANY_RESERVATION, SPECIFIC_THEN_NO_RESERVATION + reservation_affinity_key (str): + Optional. Corresponds to the label key of a reservation resource. + To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key + and specify the name of your reservation as its value. + reservation_affinity_values (List[str]): + Optional. Corresponds to the label values of a reservation resource. + This must be the full resource name of the reservation. + Format: 'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}' + Returns: + (jobs.BatchPredictionJob): + Instantiated representation of the created batch prediction job. + """ + return cls._submit_impl( + job_display_name=job_display_name, + model_name=model_name, + instances_format=instances_format, + predictions_format=predictions_format, + gcs_source=gcs_source, + bigquery_source=bigquery_source, + gcs_destination_prefix=gcs_destination_prefix, + bigquery_destination_prefix=bigquery_destination_prefix, + model_parameters=model_parameters, + machine_type=machine_type, + accelerator_type=accelerator_type, + accelerator_count=accelerator_count, + starting_replica_count=starting_replica_count, + max_replica_count=max_replica_count, + generate_explanation=generate_explanation, + explanation_metadata=explanation_metadata, + explanation_parameters=explanation_parameters, + labels=labels, + project=project, + location=location, + credentials=credentials, + encryption_spec_key_name=encryption_spec_key_name, + sync=sync, + create_request_timeout=create_request_timeout, + batch_size=batch_size, + model_monitoring_objective_config=model_monitoring_objective_config, + model_monitoring_alert_config=model_monitoring_alert_config, + analysis_instance_schema_uri=analysis_instance_schema_uri, + service_account=service_account, + reservation_affinity_type=reservation_affinity_type, + reservation_affinity_key=reservation_affinity_key, + reservation_affinity_values=reservation_affinity_values, + # Main distinction of `create` vs `submit`: + wait_for_completion=True, + ) + + @classmethod + def submit( + cls, + *, + job_display_name: Optional[str] = None, + model_name: Union[str, "aiplatform.Model"], + instances_format: str = "jsonl", + predictions_format: str = "jsonl", + gcs_source: Optional[Union[str, Sequence[str]]] = None, + bigquery_source: Optional[str] = None, + gcs_destination_prefix: Optional[str] = None, + bigquery_destination_prefix: Optional[str] = None, + model_parameters: Optional[Dict] = None, + machine_type: Optional[str] = None, + accelerator_type: Optional[str] = None, + accelerator_count: Optional[int] = None, + starting_replica_count: Optional[int] = None, + max_replica_count: Optional[int] = None, + generate_explanation: Optional[bool] = False, + explanation_metadata: Optional["aiplatform.explain.ExplanationMetadata"] = None, + explanation_parameters: Optional[ + "aiplatform.explain.ExplanationParameters" + ] = None, + labels: Optional[Dict[str, str]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + encryption_spec_key_name: Optional[str] = None, + create_request_timeout: Optional[float] = None, + batch_size: Optional[int] = None, + model_monitoring_objective_config: Optional[ + "aiplatform.model_monitoring.ObjectiveConfig" + ] = None, + model_monitoring_alert_config: Optional[ + "aiplatform.model_monitoring.AlertConfig" + ] = None, + analysis_instance_schema_uri: Optional[str] = None, + service_account: Optional[str] = None, + reservation_affinity_type: Optional[str] = None, + reservation_affinity_key: Optional[str] = None, + reservation_affinity_values: Optional[List[str]] = None, + ) -> "BatchPredictionJob": + """Sumbit a batch prediction job (not waiting for completion). + + Args: + job_display_name (str): Required. The user-defined name of the + BatchPredictionJob. The name can be up to 128 characters long and + can be consist of any UTF-8 characters. + model_name (Union[str, aiplatform.Model]): Required. A fully-qualified + model resource name or model ID. + Example: "projects/123/locations/us-central1/models/456" or "456" + when project and location are initialized or passed. May + optionally contain a version ID or alias in + {model_name}@{version} form. Or an instance of + aiplatform.Model. + instances_format (str): Required. The format in which instances are + provided. Must be one of the formats listed in + `Model.supported_input_storage_formats`. Default is "jsonl" when + using `gcs_source`. If a `bigquery_source` is provided, this is + overridden to "bigquery". + predictions_format (str): Required. The format in which Vertex AI + outputs the predictions, must be one of the formats specified in + `Model.supported_output_storage_formats`. Default is "jsonl" when + using `gcs_destination_prefix`. If a `bigquery_destination_prefix` + is provided, this is overridden to "bigquery". + gcs_source (Optional[Sequence[str]]): Google Cloud Storage URI(-s) to + your instances to run batch prediction on. They must match + `instances_format`. + bigquery_source (Optional[str]): BigQuery URI to a table, up to 2000 + characters long. For example: `bq://projectId.bqDatasetId.bqTableId` + gcs_destination_prefix (Optional[str]): The Google Cloud Storage + location of the directory where the output is to be written to. In + the given directory a new directory is created. Its name is + ``prediction--``, where + timestamp is in YYYY-MM-DDThh:mm:ss.sssZ ISO-8601 format. Inside of + it files ``predictions_0001.``, + ``predictions_0002.``, ..., ``predictions_N.`` + are created where ```` depends on chosen + ``predictions_format``, and N may equal 0001 and depends on the + total number of successfully predicted instances. If the Model has + both ``instance`` and ``prediction`` schemata defined then each such + file contains predictions as per the ``predictions_format``. If + prediction for any instance failed (partially or completely), then + an additional ``errors_0001.``, + ``errors_0002.``,..., ``errors_N.`` files are + created (N depends on total number of failed predictions). These + files contain the failed instances, as per their schema, followed by + an additional ``error`` field which as value has + ```google.rpc.Status`` `__ containing only ``code`` and + ``message`` fields. + bigquery_destination_prefix (Optional[str]): The BigQuery project or + dataset location where the output is to be written to. If project is + provided, a new dataset is created with name + ``prediction__`` where is made + BigQuery-dataset-name compatible (for example, most special + characters become underscores), and timestamp is in + YYYY_MM_DDThh_mm_ss_sssZ "based on ISO-8601" format. In the dataset + two tables will be created, ``predictions``, and ``errors``. If the + Model has both + [instance][google.cloud.aiplatform.v1.PredictSchemata.instance_schema_uri] + and + [prediction][google.cloud.aiplatform.v1.PredictSchemata.parameters_schema_uri] + schemata defined then the tables have columns as follows: The + ``predictions`` table contains instances for which the prediction + succeeded, it has columns as per a concatenation of the Model's + instance and prediction schemata. The ``errors`` table contains rows + for which the prediction has failed, it has instance columns, as per + the instance schema, followed by a single "errors" column, which as + values has [google.rpc.Status][google.rpc.Status] represented as a + STRUCT, and containing only ``code`` and ``message``. + model_parameters (Optional[Dict]): The parameters that govern the + predictions. The schema of the parameters may be specified via the + Model's `parameters_schema_uri`. + machine_type (Optional[str]): The type of machine for running batch + prediction on dedicated resources. Not specifying machine type will + result in batch prediction job being run with automatic resources. + accelerator_type (Optional[str]): The type of accelerator(s) that may + be attached to the machine as per `accelerator_count`. Only used if + `machine_type` is set. + accelerator_count (Optional[int]): The number of accelerators to + attach to the `machine_type`. Only used if `machine_type` is set. + starting_replica_count (Optional[int]): The number of machine replicas + used at the start of the batch operation. If not set, Vertex AI + decides starting number, not greater than `max_replica_count`. Only + used if `machine_type` is set. + max_replica_count (Optional[int]): The maximum number of machine + replicas the batch operation may be scaled to. Only used if + `machine_type` is set. Default is 10. + generate_explanation (bool): Optional. Generate explanation along with + the batch prediction results. This will cause the batch prediction + output to include explanations based on the `prediction_format`: - + `bigquery`: output includes a column named `explanation`. The value + is a struct that conforms to the [aiplatform.gapic.Explanation] + object. - `jsonl`: The JSON objects on each line include an + additional entry keyed `explanation`. The value of the entry is a + JSON object that conforms to the [aiplatform.gapic.Explanation] + object. - `csv`: Generating explanations for CSV format is not + supported. + explanation_metadata (aiplatform.explain.ExplanationMetadata): + Optional. Explanation metadata configuration for this + BatchPredictionJob. Can be specified only if `generate_explanation` + is set to `True`. This value overrides the value of + `Model.explanation_metadata`. All fields of `explanation_metadata` + are optional in the request. If a field of the + `explanation_metadata` object is not populated, the corresponding + field of the `Model.explanation_metadata` object is inherited. For + more details, see `Ref docs ` + explanation_parameters (aiplatform.explain.ExplanationParameters): + Optional. Parameters to configure explaining for Model's + predictions. Can be specified only if `generate_explanation` is set + to `True`. This value overrides the value of + `Model.explanation_parameters`. All fields of + `explanation_parameters` are optional in the request. If a field of + the `explanation_parameters` object is not populated, the + corresponding field of the `Model.explanation_parameters` object is + inherited. For more details, see `Ref docs + ` + labels (Dict[str, str]): Optional. The labels with user-defined + metadata to organize your BatchPredictionJobs. Label keys and values + can be no longer than 64 characters (Unicode codepoints), can only + contain lowercase letters, numeric characters, underscores and + dashes. International characters are allowed. See + https://goo.gl/xmQnxf for more information and examples of labels. + credentials (Optional[auth_credentials.Credentials]): Custom + credentials to use to create this batch prediction job. Overrides + credentials set in aiplatform.init. + encryption_spec_key_name (Optional[str]): Optional. The Cloud KMS + resource identifier of the customer managed encryption key used to + protect the job. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. If this is set, then all resources created + by the BatchPredictionJob will be encrypted with the provided + encryption key. Overrides encryption_spec_key_name set in + aiplatform.init. + create_request_timeout (float): Optional. The timeout for the create + request in seconds. + batch_size (int): Optional. The number of the records (e.g. instances) + of the operation given in each batch to a machine replica. Machine + type, and size of a single record should be considered when setting + this parameter, higher value speeds up the batch operation's + execution, but too high value will result in a whole batch not + fitting in a machine's memory, and the whole operation will fail. + The default value is 64. model_monitoring_objective_config + (aiplatform.model_monitoring.ObjectiveConfig): Optional. The + objective config for model monitoring. Passing this parameter + enables monitoring on the model associated with this batch + prediction job. model_monitoring_alert_config + (aiplatform.model_monitoring.EmailAlertConfig): Optional. Configures + how model monitoring alerts are sent to the user. Right now only + email alert is supported. + analysis_instance_schema_uri (str): Optional. Only applicable if + model_monitoring_objective_config is also passed. This parameter + specifies the YAML schema file uri describing the format of a single + instance that you want Tensorflow Data Validation (TFDV) to analyze. + If this field is empty, all the feature data types are inferred from + predict_instance_schema_uri, meaning that TFDV will use the data in + the exact format as prediction request/response. If there are any + data type differences between predict instance and TFDV instance, + this field can be used to override the schema. For models trained + with Vertex AI, this field must be set as all the fields in predict + instance formatted as string. + service_account (str): Optional. Specifies the service account for + workload run-as account. Users submitting jobs must have act-as + permission on this run-as account. + reservation_affinity_type (str): Optional. The type of reservation + affinity. One of NO_RESERVATION, ANY_RESERVATION, + SPECIFIC_RESERVATION, SPECIFIC_THEN_ANY_RESERVATION, + SPECIFIC_THEN_NO_RESERVATION + reservation_affinity_key (str): Optional. Corresponds to the label key + of a reservation resource. To target a SPECIFIC_RESERVATION by name, + use `compute.googleapis.com/reservation-name` as the key and specify + the name of your reservation as its value. + reservation_affinity_values (List[str]): Optional. Corresponds to the + label values of a reservation resource. This must be the full + resource name of the reservation. + Format: + 'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}' + + Returns: + (jobs.BatchPredictionJob): + Instantiated representation of the created batch prediction job. + """ + return cls._submit_impl( + job_display_name=job_display_name, + model_name=model_name, + instances_format=instances_format, + predictions_format=predictions_format, + gcs_source=gcs_source, + bigquery_source=bigquery_source, + gcs_destination_prefix=gcs_destination_prefix, + bigquery_destination_prefix=bigquery_destination_prefix, + model_parameters=model_parameters, + machine_type=machine_type, + accelerator_type=accelerator_type, + accelerator_count=accelerator_count, + starting_replica_count=starting_replica_count, + max_replica_count=max_replica_count, + generate_explanation=generate_explanation, + explanation_metadata=explanation_metadata, + explanation_parameters=explanation_parameters, + labels=labels, + project=project, + location=location, + credentials=credentials, + encryption_spec_key_name=encryption_spec_key_name, + create_request_timeout=create_request_timeout, + batch_size=batch_size, + model_monitoring_objective_config=model_monitoring_objective_config, + model_monitoring_alert_config=model_monitoring_alert_config, + analysis_instance_schema_uri=analysis_instance_schema_uri, + service_account=service_account, + reservation_affinity_type=reservation_affinity_type, + reservation_affinity_key=reservation_affinity_key, + reservation_affinity_values=reservation_affinity_values, + # Main distinction of `create` vs `submit`: + wait_for_completion=False, + sync=True, + ) diff --git a/tests/unit/aiplatform/test_batch_prediction_job_preview.py b/tests/unit/aiplatform/test_batch_prediction_job_preview.py new file mode 100644 index 0000000000..9e8c28902f --- /dev/null +++ b/tests/unit/aiplatform/test_batch_prediction_job_preview.py @@ -0,0 +1,301 @@ +# -*- coding: utf-8 -*- + +# 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. +# + +from importlib import reload +from unittest import mock +from unittest.mock import patch + +from google.cloud import aiplatform +from google.cloud.aiplatform import base +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.compat.services import ( + job_service_client, +) +from google.cloud.aiplatform.compat.types import ( + batch_prediction_job as gca_batch_prediction_job_compat, + io as gca_io_compat, + job_state as gca_job_state_compat, + machine_resources as gca_machine_resources_compat, + manual_batch_tuning_parameters as gca_manual_batch_tuning_parameters_compat, + reservation_affinity_v1 as gca_reservation_affinity_compat, +) +from google.cloud.aiplatform.preview import jobs as preview_jobs +import constants as test_constants +import pytest + +# TODO(b/242108750): remove temporary logic once model monitoring for batch prediction is GA +_TEST_API_CLIENT = job_service_client.JobServiceClient + +_TEST_PROJECT = test_constants.ProjectConstants._TEST_PROJECT +_TEST_LOCATION = test_constants.ProjectConstants._TEST_LOCATION +_TEST_ID = test_constants.TrainingJobConstants._TEST_ID +_TEST_ALT_ID = "8834795523125638878" +_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME +_TEST_SERVICE_ACCOUNT = test_constants.ProjectConstants._TEST_SERVICE_ACCOUNT + +_TEST_JOB_STATE_SUCCESS = gca_job_state_compat.JobState(4) +_TEST_JOB_STATE_RUNNING = gca_job_state_compat.JobState(3) +_TEST_JOB_STATE_PENDING = gca_job_state_compat.JobState(2) + +_TEST_PARENT = test_constants.ProjectConstants._TEST_PARENT + +_TEST_MODEL_NAME = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_ALT_ID}" +) + +_TEST_BATCH_PREDICTION_JOB_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/batchPredictionJobs/{_TEST_ID}" +_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME = "test-batch-prediction-job" + +_TEST_BATCH_PREDICTION_GCS_SOURCE = "gs://example-bucket/folder/instance.jsonl" + +_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX = "gs://example-bucket/folder/output" + +_TEST_MACHINE_TYPE = "n1-standard-4" +_TEST_ACCELERATOR_TYPE = "NVIDIA_TESLA_P100" +_TEST_ACCELERATOR_COUNT = 2 +_TEST_RESERVATION_AFFINITY_TYPE = "SPECIFIC_RESERVATION" +_TEST_RESERVATION_AFFINITY_KEY = "compute.googleapis.com/reservation-name" +_TEST_RESERVATION_AFFINITY_VALUES = [ + "projects/fake-project-id/zones/fake-zone/reservations/fake-reservation-name" +] + + +@pytest.fixture +def get_batch_prediction_job_mock(): + with patch.object( + _TEST_API_CLIENT, "get_batch_prediction_job" + ) as get_batch_prediction_job_mock: + get_batch_prediction_job_mock.side_effect = [ + gca_batch_prediction_job_compat.BatchPredictionJob( + name=_TEST_BATCH_PREDICTION_JOB_NAME, + display_name=_TEST_DISPLAY_NAME, + state=_TEST_JOB_STATE_PENDING, + ), + gca_batch_prediction_job_compat.BatchPredictionJob( + name=_TEST_BATCH_PREDICTION_JOB_NAME, + display_name=_TEST_DISPLAY_NAME, + state=_TEST_JOB_STATE_RUNNING, + ), + gca_batch_prediction_job_compat.BatchPredictionJob( + name=_TEST_BATCH_PREDICTION_JOB_NAME, + display_name=_TEST_DISPLAY_NAME, + state=_TEST_JOB_STATE_SUCCESS, + ), + gca_batch_prediction_job_compat.BatchPredictionJob( + name=_TEST_BATCH_PREDICTION_JOB_NAME, + display_name=_TEST_DISPLAY_NAME, + state=_TEST_JOB_STATE_SUCCESS, + ), + ] + yield get_batch_prediction_job_mock + + +@pytest.fixture +def create_batch_prediction_job_mock(): + with mock.patch.object( + _TEST_API_CLIENT, "create_batch_prediction_job" + ) as create_batch_prediction_job_mock: + create_batch_prediction_job_mock.return_value = ( + gca_batch_prediction_job_compat.BatchPredictionJob( + name=_TEST_BATCH_PREDICTION_JOB_NAME, + display_name=_TEST_DISPLAY_NAME, + state=_TEST_JOB_STATE_SUCCESS, + ) + ) + yield create_batch_prediction_job_mock + + +@pytest.mark.usefixtures("google_auth_mock") +class TestBatchPredictionJobPreview: + + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_init_batch_prediction_job(self, get_batch_prediction_job_mock): + preview_jobs.BatchPredictionJob( + batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME + ) + get_batch_prediction_job_mock.assert_called_once_with( + name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=base._DEFAULT_RETRY + ) + + def test_batch_prediction_job_status(self, get_batch_prediction_job_mock): + bp = preview_jobs.BatchPredictionJob( + batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME + ) + + # get_batch_prediction() is called again here + bp_job_state = bp.state + + assert get_batch_prediction_job_mock.call_count == 2 + assert bp_job_state == _TEST_JOB_STATE_RUNNING + + get_batch_prediction_job_mock.assert_called_with( + name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=base._DEFAULT_RETRY + ) + + def test_batch_prediction_job_done_get(self, get_batch_prediction_job_mock): + bp = preview_jobs.BatchPredictionJob( + batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME + ) + + assert bp.done() is False + assert get_batch_prediction_job_mock.call_count == 2 + + @mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 1) + @mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 1) + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_batch_prediction_job_mock") + def test_batch_predict_create_with_reservation( + self, create_batch_prediction_job_mock, sync + ): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + # Make SDK batch_predict method call + batch_prediction_job = preview_jobs.BatchPredictionJob.create( + model_name=_TEST_MODEL_NAME, + job_display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, + gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, + sync=sync, + create_request_timeout=None, + service_account=_TEST_SERVICE_ACCOUNT, + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + reservation_affinity_type=_TEST_RESERVATION_AFFINITY_TYPE, + reservation_affinity_key=_TEST_RESERVATION_AFFINITY_KEY, + reservation_affinity_values=_TEST_RESERVATION_AFFINITY_VALUES, + ) + + batch_prediction_job.wait_for_resource_creation() + + batch_prediction_job.wait() + + # Construct expected request + expected_gapic_batch_prediction_job = gca_batch_prediction_job_compat.BatchPredictionJob( + display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + model=_TEST_MODEL_NAME, + input_config=gca_batch_prediction_job_compat.BatchPredictionJob.InputConfig( + instances_format="jsonl", + gcs_source=gca_io_compat.GcsSource( + uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE] + ), + ), + output_config=gca_batch_prediction_job_compat.BatchPredictionJob.OutputConfig( + gcs_destination=gca_io_compat.GcsDestination( + output_uri_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX + ), + predictions_format="jsonl", + ), + service_account=_TEST_SERVICE_ACCOUNT, + dedicated_resources=gca_machine_resources_compat.BatchDedicatedResources( + machine_spec=gca_machine_resources_compat.MachineSpec( + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + reservation_affinity=gca_reservation_affinity_compat.ReservationAffinity( + reservation_affinity_type=_TEST_RESERVATION_AFFINITY_TYPE, + key=_TEST_RESERVATION_AFFINITY_KEY, + values=_TEST_RESERVATION_AFFINITY_VALUES, + ), + ), + ), + manual_batch_tuning_parameters=gca_manual_batch_tuning_parameters_compat.ManualBatchTuningParameters(), + ) + + create_batch_prediction_job_mock.assert_called_once_with( + parent=_TEST_PARENT, + batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, + ) + + @mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 1) + @mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 1) + @pytest.mark.usefixtures("get_batch_prediction_job_mock") + def test_batch_predict_job_submit(self, create_batch_prediction_job_mock): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + # Make SDK batch_predict method call + batch_prediction_job = preview_jobs.BatchPredictionJob.submit( + model_name=_TEST_MODEL_NAME, + job_display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, + gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, + service_account=_TEST_SERVICE_ACCOUNT, + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + reservation_affinity_type=_TEST_RESERVATION_AFFINITY_TYPE, + reservation_affinity_key=_TEST_RESERVATION_AFFINITY_KEY, + reservation_affinity_values=_TEST_RESERVATION_AFFINITY_VALUES, + ) + + batch_prediction_job.wait_for_resource_creation() + assert batch_prediction_job.done() is False + assert ( + batch_prediction_job.state + != preview_jobs.gca_job_state.JobState.JOB_STATE_SUCCEEDED + ) + + batch_prediction_job.wait_for_completion() + assert ( + batch_prediction_job.state + == preview_jobs.gca_job_state.JobState.JOB_STATE_SUCCEEDED + ) + + # Construct expected request + expected_gapic_batch_prediction_job = gca_batch_prediction_job_compat.BatchPredictionJob( + display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + model=_TEST_MODEL_NAME, + input_config=gca_batch_prediction_job_compat.BatchPredictionJob.InputConfig( + instances_format="jsonl", + gcs_source=gca_io_compat.GcsSource( + uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE] + ), + ), + output_config=gca_batch_prediction_job_compat.BatchPredictionJob.OutputConfig( + gcs_destination=gca_io_compat.GcsDestination( + output_uri_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX + ), + predictions_format="jsonl", + ), + service_account=_TEST_SERVICE_ACCOUNT, + dedicated_resources=gca_machine_resources_compat.BatchDedicatedResources( + machine_spec=gca_machine_resources_compat.MachineSpec( + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + reservation_affinity=gca_reservation_affinity_compat.ReservationAffinity( + reservation_affinity_type=_TEST_RESERVATION_AFFINITY_TYPE, + key=_TEST_RESERVATION_AFFINITY_KEY, + values=_TEST_RESERVATION_AFFINITY_VALUES, + ), + ), + ), + manual_batch_tuning_parameters=gca_manual_batch_tuning_parameters_compat.ManualBatchTuningParameters(), + ) + + create_batch_prediction_job_mock.assert_called_once_with( + parent=_TEST_PARENT, + batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, + ) From 7a5a0667b915bb0ed1009ffc8034a8d29fb2d20e Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Nov 2025 15:03:57 -0800 Subject: [PATCH 14/24] fix: GenAI Client(evals) - Change dataset visualization table to fixed to prevent horizontal expansion. PiperOrigin-RevId: 827664884 --- vertexai/_genai/_evals_visualization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vertexai/_genai/_evals_visualization.py b/vertexai/_genai/_evals_visualization.py index 7b7e6174c7..c3e3dc69e7 100644 --- a/vertexai/_genai/_evals_visualization.py +++ b/vertexai/_genai/_evals_visualization.py @@ -708,7 +708,7 @@ def _get_inference_html(dataframe_json: str) -> str: body {{ font-family: 'Roboto', sans-serif; margin: 2em; background-color: #f8f9fa; color: #202124;}} .container {{ max-width: 95%; margin: 20px auto; padding: 20px; background: #fff; border-radius: 8px; box-shadow: 0 1px 3px rgba(0,0,0,0.12); }} h1 {{ color: #3c4043; border-bottom: 2px solid #4285F4; padding-bottom: 8px; }} - table {{ border-collapse: collapse; width: 100%; }} + table {{ border-collapse: collapse; width: 100%; table-layout: fixed; }} th, td {{ border: 1px solid #dadce0; padding: 12px; text-align: left; vertical-align: top; }} th {{ background-color: #f2f2f2; font-weight: 500;}} td > div {{ white-space: pre-wrap; word-wrap: break-word; max-height: 400px; overflow-y: auto; overflow-wrap: break-word; }} From 0143c07047b27cffc1e1f3603bb571c4cfae3bea Mon Sep 17 00:00:00 2001 From: Shawn Yang Date: Mon, 3 Nov 2025 15:26:28 -0800 Subject: [PATCH 15/24] chore: Support specifying `agent_framework` in Agent Engine creation and update. PiperOrigin-RevId: 827673214 --- .../unit/vertexai/genai/test_agent_engines.py | 98 ++++++++++++++++++- vertexai/_genai/_agent_engines_utils.py | 54 ++++++++-- vertexai/_genai/agent_engines.py | 14 ++- vertexai/_genai/types/common.py | 75 ++++++++++++++ 4 files changed, 231 insertions(+), 10 deletions(-) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index 37e678f7f6..508bf0377b 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -40,7 +40,7 @@ import pytest -_TEST_AGENT_FRAMEWORK = "test-agent-framework" +_TEST_AGENT_FRAMEWORK = "google-adk" GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" ) @@ -976,9 +976,11 @@ def test_create_agent_engine_config_with_source_packages( entrypoint_object="app", requirements_file=requirements_file_path, class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + agent_framework=_TEST_AGENT_FRAMEWORK, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION + assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK assert config["spec"]["source_code_spec"] == { "inline_source": {"source_archive": "test_tarball"}, "python_spec": { @@ -1500,6 +1502,7 @@ def test_create_agent_engine_with_env_vars_dict( entrypoint_module=None, entrypoint_object=None, requirements_file=None, + agent_framework=None, ) request_mock.assert_called_with( "post", @@ -1513,7 +1516,9 @@ def test_create_agent_engine_with_env_vars_dict( "package_spec": { "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, "python_version": _TEST_PYTHON_VERSION, - "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, + "requirements_gcs_uri": ( + _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI + ), }, }, }, @@ -1586,6 +1591,7 @@ def test_create_agent_engine_with_custom_service_account( entrypoint_module=None, entrypoint_object=None, requirements_file=None, + agent_framework=None, ) request_mock.assert_called_with( "post", @@ -1674,6 +1680,7 @@ def test_create_agent_engine_with_experimental_mode( entrypoint_module=None, entrypoint_object=None, requirements_file=None, + agent_framework=None, ) request_mock.assert_called_with( "post", @@ -1826,6 +1833,7 @@ def test_create_agent_engine_with_class_methods( entrypoint_module=None, entrypoint_object=None, requirements_file=None, + agent_framework=None, ) request_mock.assert_called_with( "post", @@ -1845,6 +1853,92 @@ def test_create_agent_engine_with_class_methods( None, ) + @mock.patch.object(agent_engines.AgentEngines, "_create_config") + @mock.patch.object(_agent_engines_utils, "_await_operation") + def test_create_agent_engine_with_agent_framework( + self, + mock_await_operation, + mock_create_config, + ): + mock_create_config.return_value = { + "display_name": _TEST_AGENT_ENGINE_DISPLAY_NAME, + "description": _TEST_AGENT_ENGINE_DESCRIPTION, + "spec": { + "package_spec": { + "python_version": _TEST_PYTHON_VERSION, + "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, + "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, + }, + "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], + "agent_framework": _TEST_AGENT_FRAMEWORK, + }, + } + mock_await_operation.return_value = _genai_types.AgentEngineOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.agent_engines._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.agent_engines.create( + agent=self.test_agent, + config=_genai_types.AgentEngineConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + staging_bucket=_TEST_STAGING_BUCKET, + agent_framework=_TEST_AGENT_FRAMEWORK, + ), + ) + mock_create_config.assert_called_with( + mode="create", + agent=self.test_agent, + staging_bucket=_TEST_STAGING_BUCKET, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + description=None, + gcs_dir_name=None, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + env_vars=None, + service_account=None, + context_spec=None, + psc_interface_config=None, + min_instances=None, + max_instances=None, + resource_limits=None, + container_concurrency=None, + encryption_spec=None, + labels=None, + agent_server_mode=None, + class_methods=None, + source_packages=None, + entrypoint_module=None, + entrypoint_object=None, + requirements_file=None, + agent_framework=_TEST_AGENT_FRAMEWORK, + ) + request_mock.assert_called_with( + "post", + "reasoningEngines", + { + "displayName": _TEST_AGENT_ENGINE_DISPLAY_NAME, + "description": _TEST_AGENT_ENGINE_DESCRIPTION, + "spec": { + "agent_framework": _TEST_AGENT_FRAMEWORK, + "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], + "package_spec": { + "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, + "python_version": _TEST_PYTHON_VERSION, + "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, + }, + }, + }, + None, + ) + @pytest.mark.usefixtures("caplog") @mock.patch.object(_agent_engines_utils, "_prepare") @mock.patch.object(_agent_engines_utils, "_await_operation") diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 8364212528..0e063a2c1b 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -128,6 +128,16 @@ _BASE_MODULES = set(_BUILTIN_MODULE_NAMES + tuple(_STDLIB_MODULE_NAMES)) _BLOB_FILENAME = "agent_engine.pkl" _DEFAULT_AGENT_FRAMEWORK = "custom" +_SUPPORTED_AGENT_FRAMEWORKS = frozenset( + [ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + ] +) _DEFAULT_ASYNC_METHOD_NAME = "async_query" _DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any]" _DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query" @@ -705,13 +715,43 @@ def _generate_schema( return schema -def _get_agent_framework(*, agent: _AgentEngineInterface) -> str: - if ( - hasattr(agent, _AGENT_FRAMEWORK_ATTR) - and getattr(agent, _AGENT_FRAMEWORK_ATTR) is not None - and isinstance(getattr(agent, _AGENT_FRAMEWORK_ATTR), str) - ): - return getattr(agent, _AGENT_FRAMEWORK_ATTR) +def _get_agent_framework( + *, + agent_framework: Optional[str], + agent: _AgentEngineInterface, +) -> str: + """Gets the agent framework to use. + + The agent framework is determined in the following order of priority: + 1. The `agent_framework` passed to this function. + 2. The `agent_framework` attribute on the `agent` object. + 3. The default framework, "custom". + + Args: + agent_framework (str): + The agent framework provided by the user. + agent (_AgentEngineInterface): + The agent engine instance. + + Returns: + str: The name of the agent framework to use. + """ + if agent_framework is not None and agent_framework in _SUPPORTED_AGENT_FRAMEWORKS: + logger.info(f"Using agent framework: {agent_framework}") + return agent_framework + if hasattr(agent, _AGENT_FRAMEWORK_ATTR): + agent_framework_attr = getattr(agent, _AGENT_FRAMEWORK_ATTR) + if ( + agent_framework_attr is not None + and isinstance(agent_framework_attr, str) + and agent_framework_attr in _SUPPORTED_AGENT_FRAMEWORKS + ): + logger.info(f"Using agent framework: {agent_framework_attr}") + return agent_framework_attr + logger.info( + f"The provided agent framework {agent_framework} is not supported." + f" Defaulting to {_DEFAULT_AGENT_FRAMEWORK}." + ) return _DEFAULT_AGENT_FRAMEWORK diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 2b3ad56d8f..238c5d8d69 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -94,6 +94,9 @@ def _CreateAgentEngineConfig_to_vertex( getv(from_object, ["requirements_file"]), ) + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + return to_object @@ -285,6 +288,9 @@ def _UpdateAgentEngineConfig_to_vertex( getv(from_object, ["requirements_file"]), ) + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + if getv(from_object, ["update_mask"]) is not None: setv( parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) @@ -923,6 +929,7 @@ def create( entrypoint_module=config.entrypoint_module, entrypoint_object=config.entrypoint_object, requirements_file=config.requirements_file, + agent_framework=config.agent_framework, ) operation = self._create(config=api_config) # TODO: Use a more specific link. @@ -986,6 +993,7 @@ def _create_config( entrypoint_module: Optional[str] = None, entrypoint_object: Optional[str] = None, requirements_file: Optional[str] = None, + agent_framework: Optional[str] = None, ) -> types.UpdateAgentEngineConfigDict: import sys @@ -1195,7 +1203,10 @@ def _create_config( ] = agent_server_mode agent_engine_spec["agent_framework"] = ( - _agent_engines_utils._get_agent_framework(agent=agent) + _agent_engines_utils._get_agent_framework( + agent_framework=agent_framework, + agent=agent, + ) ) update_masks.append("spec.agent_framework") config["spec"] = agent_engine_spec @@ -1423,6 +1434,7 @@ def update( entrypoint_module=config.entrypoint_module, entrypoint_object=config.entrypoint_object, requirements_file=config.requirements_file, + agent_framework=config.agent_framework, ) operation = self._update(name=name, config=api_config) logger.info( diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index d193042586..d82a6903e3 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -5366,6 +5366,19 @@ class CreateAgentEngineConfig(_common.BaseModel): the source package. """, ) + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] = Field( + default=None, + description="""The agent framework to be used for the Agent Engine. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""", + ) class CreateAgentEngineConfigDict(TypedDict, total=False): @@ -5464,6 +5477,18 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): the source package. """ + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Engine. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" + CreateAgentEngineConfigOrDict = Union[ CreateAgentEngineConfig, CreateAgentEngineConfigDict @@ -6067,6 +6092,19 @@ class UpdateAgentEngineConfig(_common.BaseModel): the source package. """, ) + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] = Field( + default=None, + description="""The agent framework to be used for the Agent Engine. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""", + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -6170,6 +6208,18 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): the source package. """ + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Engine. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -12907,6 +12957,19 @@ class AgentEngineConfig(_common.BaseModel): the source package. """, ) + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] = Field( + default=None, + description="""The agent framework to be used for the Agent Engine. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""", + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -13034,6 +13097,18 @@ class AgentEngineConfigDict(TypedDict, total=False): the source package. """ + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Engine. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] From 7c8c218f47a6fd919cae6869736d35e4206a123e Mon Sep 17 00:00:00 2001 From: Tongzhou Jiang Date: Mon, 3 Nov 2025 17:13:43 -0800 Subject: [PATCH 16/24] fix: revert: Alow VertexAiSession for streaming_agent_run_with_events PiperOrigin-RevId: 827710498 --- vertexai/agent_engines/templates/adk.py | 50 +++++++------------ .../reasoning_engines/templates/adk.py | 50 ++++++++----------- 2 files changed, 39 insertions(+), 61 deletions(-) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index f3c186e903..97be3505c6 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -585,6 +585,7 @@ async def _init_session( ): """Initializes the session, and returns the session id.""" from google.adk.events.event import Event + import random session_state = None if request.authorizations: @@ -593,9 +594,14 @@ async def _init_session( auth = _Authorization(**auth) session_state[f"temp:{auth_id}"] = auth.access_token + if request.session_id: + session_id = request.session_id + else: + session_id = f"temp_session_{random.randbytes(8).hex()}" session = await session_service.create_session( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, + session_id=session_id, state=session_state, ) if not session: @@ -613,7 +619,7 @@ async def _init_session( saved_version = await artifact_service.save_artifact( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, - session_id=session.id, + session_id=session_id, filename=artifact.file_name, artifact=version_data.data, ) @@ -1052,53 +1058,35 @@ async def streaming_agent_run_with_events(self, request_json: str): import json from google.genai import types - from google.genai.errors import ClientError request = _StreamRunRequest(**json.loads(request_json)) if not self._tmpl_attrs.get("in_memory_runner"): self.set_up() - if not self._tmpl_attrs.get("runner"): - self.set_up() # Prepare the in-memory session. if not self._tmpl_attrs.get("in_memory_artifact_service"): self.set_up() - if not self._tmpl_attrs.get("artifact_service"): - self.set_up() if not self._tmpl_attrs.get("in_memory_session_service"): self.set_up() - if not self._tmpl_attrs.get("session_service"): - self.set_up() + session_service = self._tmpl_attrs.get("in_memory_session_service") + artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") app = self._tmpl_attrs.get("app") - # Try to get the session, if it doesn't exist, create a new one. + session = None if request.session_id: - session_service = self._tmpl_attrs.get("session_service") - artifact_service = self._tmpl_attrs.get("artifact_service") - runner = self._tmpl_attrs.get("runner") try: session = await session_service.get_session( app_name=app.name if app else self._tmpl_attrs.get("app_name"), user_id=request.user_id, session_id=request.session_id, ) - except ClientError: - # Fall back to create session if the session is not found. - # Specifying session_id on creation is not supported, - # so session id will be regenerated. - session = await self._init_session( - session_service=session_service, - artifact_service=artifact_service, - request=request, - ) - else: - # Not providing a session ID will create a new in-memory session. - session_service = self._tmpl_attrs.get("in_memory_session_service") - artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") - runner = self._tmpl_attrs.get("in_memory_runner") - session = await session_service.create_session( - app_name=self._tmpl_attrs.get("app_name"), - user_id=request.user_id, - session_id=request.session_id, + except RuntimeError: + pass + if not session: + # Fall back to create session if the session is not found. + session = await self._init_session( + session_service=session_service, + artifact_service=artifact_service, + request=request, ) if not session: raise RuntimeError("Session initialization failed.") @@ -1106,7 +1094,7 @@ async def streaming_agent_run_with_events(self, request_json: str): # Run the agent message_for_agent = types.Content(**request.message) try: - async for event in runner.run_async( + async for event in self._tmpl_attrs.get("in_memory_runner").run_async( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 8b49b992c6..1bf5bb64a0 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -533,6 +533,7 @@ async def _init_session( ): """Initializes the session, and returns the session id.""" from google.adk.events.event import Event + import random session_state = None if request.authorizations: @@ -541,9 +542,14 @@ async def _init_session( auth = _Authorization(**auth) session_state[f"temp:{auth_id}"] = auth.access_token + if request.session_id: + session_id = request.session_id + else: + session_id = f"temp_session_{random.randbytes(8).hex()}" session = await session_service.create_session( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, + session_id=session_id, state=session_state, ) if not session: @@ -561,7 +567,7 @@ async def _init_session( saved_version = await artifact_service.save_artifact( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, - session_id=session.id, + session_id=session_id, filename=artifact.file_name, artifact=version_data.data, ) @@ -939,7 +945,6 @@ async def async_stream_query( def streaming_agent_run_with_events(self, request_json: str): import json from google.genai import types - from google.genai.errors import ClientError event_queue = queue.Queue(maxsize=1) @@ -947,52 +952,37 @@ async def _invoke_agent_async(): request = _StreamRunRequest(**json.loads(request_json)) if not self._tmpl_attrs.get("in_memory_runner"): self.set_up() - if not self._tmpl_attrs.get("runner"): - self.set_up() # Prepare the in-memory session. if not self._tmpl_attrs.get("in_memory_artifact_service"): self.set_up() - if not self._tmpl_attrs.get("artifact_service"): - self.set_up() if not self._tmpl_attrs.get("in_memory_session_service"): self.set_up() - if not self._tmpl_attrs.get("session_service"): - self.set_up() + session_service = self._tmpl_attrs.get("in_memory_session_service") + artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") + # Try to get the session, if it doesn't exist, create a new one. + session = None if request.session_id: - session_service = self._tmpl_attrs.get("session_service") - artifact_service = self._tmpl_attrs.get("artifact_service") - runner = self._tmpl_attrs.get("runner") try: session = await session_service.get_session( app_name=self._tmpl_attrs.get("app_name"), user_id=request.user_id, session_id=request.session_id, ) - except ClientError: - # Fall back to create session if the session is not found. - # Specifying session_id on creation is not supported, - # so session id will be regenerated. - session = await self._init_session( - session_service=session_service, - artifact_service=artifact_service, - request=request, - ) - else: - # Not providing a session ID will create a new in-memory session. - session_service = self._tmpl_attrs.get("in_memory_session_service") - artifact_service = self._tmpl_attrs.get("in_memory_artifact_service") - runner = self._tmpl_attrs.get("in_memory_runner") - session = await session_service.create_session( - app_name=self._tmpl_attrs.get("app_name"), - user_id=request.user_id, - session_id=request.session_id, + except RuntimeError: + pass + if not session: + # Fall back to create session if the session is not found. + session = await self._init_session( + session_service=session_service, + artifact_service=artifact_service, + request=request, ) if not session: raise RuntimeError("Session initialization failed.") # Run the agent. message_for_agent = types.Content(**request.message) try: - for event in runner.run( + for event in self._tmpl_attrs.get("in_memory_runner").run( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, From 16318e2782001a9c09538501d9677da4df65b562 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 4 Nov 2025 09:18:47 -0800 Subject: [PATCH 17/24] chore: Add warning if tracing is enabled but telemetry API is disabled. PiperOrigin-RevId: 828004811 --- vertexai/agent_engines/templates/adk.py | 28 +++++++++++++++++++ .../reasoning_engines/templates/adk.py | 28 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 97be3505c6..0276b4f434 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -94,6 +94,16 @@ _DEFAULT_APP_NAME = "default-app-name" _DEFAULT_USER_ID = "default-user-id" +_TELEMETRY_API_DISABLED_WARNING = ( + "Tracing integration for Agent Engine has migrated to a new API.\n" + "The 'telemetry.googleapis.com' has not been enabled in project %s. \n" + "**Impact:** Until this API is enabled, telemetry data will not be stored." + "\n" + "**Action:** Please enable the API by visiting " + "https://console.developers.google.com/apis/api/telemetry.googleapis.com/overview?project=%s." + "\n" + "(If you enabled this API recently, you can safely ignore this warning.)" +) def get_adk_version() -> Optional[str]: @@ -743,6 +753,9 @@ def set_up(self): custom_instrumentor = self._tmpl_attrs.get("instrumentor_builder") + if self._tmpl_attrs.get("enable_tracing"): + self._warn_if_telemetry_api_disabled() + if self._tmpl_attrs.get("enable_tracing") is False: _warn( ( @@ -1550,3 +1563,18 @@ def _tracing_enabled(self) -> bool: and enable_telemetry is True and is_version_sufficient("1.17.0") ) + + def _warn_if_telemetry_api_disabled(self): + """Warn if telemetry API is disabled.""" + try: + import google.auth.transport.requests + import google.auth + except (ImportError, AttributeError): + return + credentials, project = google.auth.default() + session = google.auth.transport.requests.AuthorizedSession( + credentials=credentials + ) + r = session.post("https://telemetry.googleapis.com/v1/traces", data=None) + if "Telemetry API has not been used in project" in r.text: + _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project)) diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 1bf5bb64a0..9d29b71cea 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -96,6 +96,16 @@ _DEFAULT_APP_NAME = "default-app-name" _DEFAULT_USER_ID = "default-user-id" +_TELEMETRY_API_DISABLED_WARNING = ( + "Tracing integration for Agent Engine has migrated to a new API.\n" + "The 'telemetry.googleapis.com' has not been enabled in project %s. \n" + "**Impact:** Until this API is enabled, telemetry data will not be stored." + "\n" + "**Action:** Please enable the API by visiting " + "https://console.developers.google.com/apis/api/telemetry.googleapis.com/overview?project=%s." + "\n" + "(If you enabled this API recently, you can safely ignore this warning.)" +) def get_adk_version() -> Optional[str]: @@ -663,6 +673,9 @@ def set_up(self): else: os.environ["ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS"] = "false" + if self._tmpl_attrs.get("enable_tracing"): + self._warn_if_telemetry_api_disabled() + if self._tmpl_attrs.get("enable_tracing") is False: _warn( ( @@ -1486,3 +1499,18 @@ def _tracing_enabled(self) -> bool: and enable_telemetry is True and is_version_sufficient("1.17.0") ) + + def _warn_if_telemetry_api_disabled(self): + """Warn if telemetry API is disabled.""" + try: + import google.auth.transport.requests + import google.auth + except (ImportError, AttributeError): + return + credentials, project = google.auth.default() + session = google.auth.transport.requests.AuthorizedSession( + credentials=credentials + ) + r = session.post("https://telemetry.googleapis.com/v1/traces", data=None) + if "Telemetry API has not been used in project" in r.text: + _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project)) From d34485880bdefd21dccac4548656af1eaedb3727 Mon Sep 17 00:00:00 2001 From: Yvonne Yu Date: Tue, 4 Nov 2025 09:24:47 -0800 Subject: [PATCH 18/24] chore: release 1.125.0 Release-As: 1.125.0 PiperOrigin-RevId: 828007059 --- google/cloud/aiplatform/releases.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/releases.txt b/google/cloud/aiplatform/releases.txt index 52d033f99a..7f4834f510 100644 --- a/google/cloud/aiplatform/releases.txt +++ b/google/cloud/aiplatform/releases.txt @@ -1,4 +1,4 @@ Use this file when you need to force a patch release with release-please. Edit line 4 below with the version for the release. -1.124.0 \ No newline at end of file +1.125.0 \ No newline at end of file From 08397c36a9d3a418cbd9c385c03fc94dc18a46d0 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 4 Nov 2025 10:15:52 -0800 Subject: [PATCH 19/24] chore: set user agent for OTLP http exporter in adk templates PiperOrigin-RevId: 828027851 --- .../unit/vertex_adk/test_agent_engine_templates_adk.py | 10 ++++++++++ .../vertex_adk/test_reasoning_engine_templates_adk.py | 10 ++++++++++ vertexai/agent_engines/templates/adk.py | 7 +++++++ vertexai/preview/reasoning_engines/templates/adk.py | 7 +++++++ 4 files changed, 34 insertions(+) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 2b0f06cece..b4f41aeb3d 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -18,6 +18,7 @@ import os import cloudpickle import sys +import re from unittest import mock from typing import Optional @@ -769,8 +770,17 @@ def test_tracing_setup( otlp_span_exporter_mock.assert_called_once_with( session=mock.ANY, endpoint="https://telemetry.googleapis.com/v1/traces", + headers=mock.ANY, ) + user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"] + assert ( + re.fullmatch( + r"Vertex-Agent-Engine\/[\d\.]+ OTel-OTLP-Exporter-Python\/[\d\.]+", + user_agent, + ) + is not None + ) assert ( trace_provider_mock.call_args.kwargs["resource"].attributes == expected_attributes diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index add706859a..e06be088a9 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -17,6 +17,7 @@ import importlib import json import os +import re from unittest import mock from typing import Optional @@ -723,8 +724,17 @@ def test_tracing_setup( otlp_span_exporter_mock.assert_called_once_with( session=mock.ANY, endpoint="https://telemetry.googleapis.com/v1/traces", + headers=mock.ANY, ) + user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"] + assert ( + re.fullmatch( + r"Vertex-Agent-Engine\/[\d\.]+ OTel-OTLP-Exporter-Python\/[\d\.]+", + user_agent, + ) + is not None + ) assert ( trace_provider_mock.call_args.kwargs["resource"].attributes == expected_attributes diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 0276b4f434..7e21e6c680 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -357,8 +357,10 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: if enable_tracing: try: + import opentelemetry.exporter.otlp.proto.http.version import opentelemetry.exporter.otlp.proto.http.trace_exporter import google.auth.transport.requests + from google.cloud.aiplatform import version as aip_version except (ImportError, AttributeError): return _warn_missing_dependency( "opentelemetry-exporter-otlp-proto-http", needed_for_tracing=True @@ -367,12 +369,17 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: import google.auth credentials, _ = google.auth.default() + vertex_sdk_version = aip_version.__version__ + otlp_http_version = opentelemetry.exporter.otlp.proto.http.version.__version__ + user_agent = f"Vertex-Agent-Engine/{vertex_sdk_version} OTel-OTLP-Exporter-Python/{otlp_http_version}" + span_exporter = ( opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter( session=google.auth.transport.requests.AuthorizedSession( credentials=credentials ), endpoint="https://telemetry.googleapis.com/v1/traces", + headers={"User-Agent": user_agent}, ) ) span_processor = opentelemetry.sdk.trace.export.BatchSpanProcessor( diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 9d29b71cea..a3220228ac 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -359,8 +359,10 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: if enable_tracing: try: + import opentelemetry.exporter.otlp.proto.http.version import opentelemetry.exporter.otlp.proto.http.trace_exporter import google.auth.transport.requests + from google.cloud.aiplatform import version as aip_version except (ImportError, AttributeError): return _warn_missing_dependency( "opentelemetry-exporter-otlp-proto-http", needed_for_tracing=True @@ -369,12 +371,17 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: import google.auth credentials, _ = google.auth.default() + vertex_sdk_version = aip_version.__version__ + otlp_http_version = opentelemetry.exporter.otlp.proto.http.version.__version__ + user_agent = f"Vertex-Agent-Engine/{vertex_sdk_version} OTel-OTLP-Exporter-Python/{otlp_http_version}" + span_exporter = ( opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter( session=google.auth.transport.requests.AuthorizedSession( credentials=credentials ), endpoint="https://telemetry.googleapis.com/v1/traces", + headers={"User-Agent": user_agent}, ) ) span_processor = opentelemetry.sdk.trace.export.BatchSpanProcessor( From 1513d8fa15c38d3b82f4f1c1bf7cc0d542e49eb9 Mon Sep 17 00:00:00 2001 From: Amy Wu Date: Tue, 4 Nov 2025 11:43:19 -0800 Subject: [PATCH 20/24] chore: internal change PiperOrigin-RevId: 828068374 --- vertexai/_genai/types/__init__.py | 18 ++- vertexai/_genai/types/common.py | 217 ++++++++++++++++++------------ 2 files changed, 140 insertions(+), 95 deletions(-) diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index ed2eff77c4..c1a01ad8a3 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -491,6 +491,9 @@ from .common import LLMBasedMetricSpecDict from .common import LLMBasedMetricSpecOrDict from .common import LLMMetric +from .common import LustreMount +from .common import LustreMountDict +from .common import LustreMountOrDict from .common import MachineConfig from .common import MachineSpec from .common import MachineSpecDict @@ -1265,6 +1268,9 @@ "DiskSpec", "DiskSpecDict", "DiskSpecOrDict", + "LustreMount", + "LustreMountDict", + "LustreMountOrDict", "ReservationAffinity", "ReservationAffinityDict", "ReservationAffinityOrDict", @@ -1313,12 +1319,6 @@ "ReasoningEngineSpec", "ReasoningEngineSpecDict", "ReasoningEngineSpecOrDict", - "ReasoningEngineContextSpecMemoryBankConfigGenerationConfig", - "ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict", - "ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict", - "ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig", - "ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict", - "ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict", "MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic", "MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict", "MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict", @@ -1346,6 +1346,12 @@ "MemoryBankCustomizationConfig", "MemoryBankCustomizationConfigDict", "MemoryBankCustomizationConfigOrDict", + "ReasoningEngineContextSpecMemoryBankConfigGenerationConfig", + "ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict", + "ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict", + "ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig", + "ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict", + "ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict", "ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig", "ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict", "ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index d82a6903e3..33791bdd92 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -151,6 +151,8 @@ class AcceleratorType(_common.CaseInSensitiveEnum): """Nvidia B200 GPU.""" NVIDIA_GB200 = "NVIDIA_GB200" """Nvidia GB200 GPU.""" + NVIDIA_RTX_PRO_6000 = "NVIDIA_RTX_PRO_6000" + """Nvidia RTX Pro 6000 GPU.""" TPU_V2 = "TPU_V2" """TPU v2.""" TPU_V3 = "TPU_V3" @@ -3969,6 +3971,44 @@ class DiskSpecDict(TypedDict, total=False): DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] +class LustreMount(_common.BaseModel): + """Represents a mount configuration for Lustre file system.""" + + filesystem: Optional[str] = Field( + default=None, description="""Required. The name of the Lustre filesystem.""" + ) + instance_ip: Optional[str] = Field( + default=None, description="""Required. IP address of the Lustre instance.""" + ) + mount_point: Optional[str] = Field( + default=None, + description="""Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""", + ) + volume_handle: Optional[str] = Field( + default=None, + description="""Required. The unique identifier of the Lustre volume.""", + ) + + +class LustreMountDict(TypedDict, total=False): + """Represents a mount configuration for Lustre file system.""" + + filesystem: Optional[str] + """Required. The name of the Lustre filesystem.""" + + instance_ip: Optional[str] + """Required. IP address of the Lustre instance.""" + + mount_point: Optional[str] + """Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""" + + volume_handle: Optional[str] + """Required. The unique identifier of the Lustre volume.""" + + +LustreMountOrDict = Union[LustreMount, LustreMountDict] + + class ReservationAffinity(_common.BaseModel): """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity.""" @@ -4013,6 +4053,10 @@ class MachineSpec(_common.BaseModel): default=None, description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""", ) + gpu_partition_size: Optional[str] = Field( + default=None, + description="""Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""", + ) machine_type: Optional[str] = Field( default=None, description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""", @@ -4040,6 +4084,9 @@ class MachineSpecDict(TypedDict, total=False): accelerator_type: Optional[AcceleratorType] """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" + gpu_partition_size: Optional[str] + """Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""" + machine_type: Optional[str] """Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""" @@ -4142,6 +4189,9 @@ class WorkerPoolSpec(_common.BaseModel): default=None, description="""The custom container task.""" ) disk_spec: Optional[DiskSpec] = Field(default=None, description="""Disk spec.""") + lustre_mounts: Optional[list[LustreMount]] = Field( + default=None, description="""Optional. List of Lustre mounts.""" + ) machine_spec: Optional[MachineSpec] = Field( default=None, description="""Optional. Immutable. The specification of a single machine.""", @@ -4167,6 +4217,9 @@ class WorkerPoolSpecDict(TypedDict, total=False): disk_spec: Optional[DiskSpecDict] """Disk spec.""" + lustre_mounts: Optional[list[LustreMountDict]] + """Optional. List of Lustre mounts.""" + machine_spec: Optional[MachineSpecDict] """Optional. Immutable. The specification of a single machine.""" @@ -4531,7 +4584,7 @@ class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): ) max_instances: Optional[int] = Field( default=None, - description="""Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100.""", + description="""Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100].""", ) min_instances: Optional[int] = Field( default=None, @@ -4563,7 +4616,7 @@ class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): """Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""" max_instances: Optional[int] - """Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100.""" + """Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100].""" min_instances: Optional[int] """Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1.""" @@ -4770,56 +4823,6 @@ class ReasoningEngineSpecDict(TypedDict, total=False): ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): - """Configuration for how to generate memories.""" - - model: Optional[str] = Field( - default=None, - description="""Required. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", - ) - - -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( - TypedDict, total=False -): - """Configuration for how to generate memories.""" - - model: Optional[str] - """Required. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" - - -ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, - ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, -] - - -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( - _common.BaseModel -): - """Configuration for how to perform similarity search on memories.""" - - embedding_model: Optional[str] = Field( - default=None, - description="""Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", - ) - - -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( - TypedDict, total=False -): - """Configuration for how to perform similarity search on memories.""" - - embedding_model: Optional[str] - """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" - - -ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, -] - - class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): """A custom memory topic defined by the developer.""" @@ -5098,6 +5101,56 @@ class MemoryBankCustomizationConfigDict(TypedDict, total=False): ] +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): + """Configuration for how to generate memories.""" + + model: Optional[str] = Field( + default=None, + description="""Required. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + ) + + +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( + TypedDict, total=False +): + """Configuration for how to generate memories.""" + + model: Optional[str] + """Required. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + + +ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, +] + + +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( + _common.BaseModel +): + """Configuration for how to perform similarity search on memories.""" + + embedding_model: Optional[str] = Field( + default=None, + description="""Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + ) + + +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( + TypedDict, total=False +): + """Configuration for how to perform similarity search on memories.""" + + embedding_model: Optional[str] + """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + + +ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, +] + + class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig( _common.BaseModel ): @@ -5181,6 +5234,10 @@ class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=F class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): """Specification for a Memory Bank.""" + customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + default=None, + description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", + ) generation_config: Optional[ ReasoningEngineContextSpecMemoryBankConfigGenerationConfig ] = Field( @@ -5193,10 +5250,6 @@ class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): default=None, description="""Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""", ) - customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( - default=None, - description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", - ) ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( default=None, description="""Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""", @@ -5210,6 +5263,9 @@ class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): """Specification for a Memory Bank.""" + customization_configs: Optional[list[MemoryBankCustomizationConfigDict]] + """Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""" + generation_config: Optional[ ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict ] @@ -5220,9 +5276,6 @@ class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): ] """Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""" - customization_configs: Optional[list[MemoryBankCustomizationConfigDict]] - """Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""" - ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict] """Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""" @@ -6395,7 +6448,7 @@ class Memory(_common.BaseModel): expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", + description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", ) ttl: Optional[str] = Field( default=None, @@ -6448,7 +6501,7 @@ class MemoryDict(TypedDict, total=False): """A memory.""" expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" + """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" ttl: Optional[str] """Optional. Input only. The TTL for this resource. The expiration time is computed: now + TTL.""" @@ -7913,14 +7966,6 @@ class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): default=None, description="""The coding language supported in this environment.""", ) - dependencies: Optional[list[str]] = Field( - default=None, - description="""Optional. The additional dependencies to install in the code execution environment. For example, "pandas==2.2.3".""", - ) - env: Optional[list[EnvVar]] = Field( - default=None, - description="""Optional. The environment variables to set in the code execution environment.""", - ) machine_config: Optional[MachineConfig] = Field( default=None, description="""The machine config of the code execution environment.""", @@ -7933,12 +7978,6 @@ class SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False) code_language: Optional[Language] """The coding language supported in this environment.""" - dependencies: Optional[list[str]] - """Optional. The additional dependencies to install in the code execution environment. For example, "pandas==2.2.3".""" - - env: Optional[list[EnvVarDict]] - """Optional. The environment variables to set in the code execution environment.""" - machine_config: Optional[MachineConfig] """The machine config of the code execution environment.""" @@ -8127,10 +8166,6 @@ class SandboxEnvironment(_common.BaseModel): default=None, description="""Required. The display name of the SandboxEnvironment.""", ) - metadata: Optional[Any] = Field( - default=None, - description="""Output only. Additional information about the SandboxEnvironment.""", - ) name: Optional[str] = Field( default=None, description="""Identifier. The name of the SandboxEnvironment.""" ) @@ -8142,6 +8177,10 @@ class SandboxEnvironment(_common.BaseModel): default=None, description="""Output only. The runtime state of the SandboxEnvironment.""", ) + ttl: Optional[str] = Field( + default=None, + description="""Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""", + ) update_time: Optional[datetime.datetime] = Field( default=None, description="""Output only. The timestamp when this SandboxEnvironment was most recently updated.""", @@ -8164,9 +8203,6 @@ class SandboxEnvironmentDict(TypedDict, total=False): display_name: Optional[str] """Required. The display name of the SandboxEnvironment.""" - metadata: Optional[Any] - """Output only. Additional information about the SandboxEnvironment.""" - name: Optional[str] """Identifier. The name of the SandboxEnvironment.""" @@ -8176,6 +8212,9 @@ class SandboxEnvironmentDict(TypedDict, total=False): state: Optional[State] """Output only. The runtime state of the SandboxEnvironment.""" + ttl: Optional[str] + """Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""" + update_time: Optional[datetime.datetime] """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" @@ -8342,10 +8381,6 @@ class MetadataDict(TypedDict, total=False): class Chunk(_common.BaseModel): """A chunk of data.""" - mime_type: Optional[str] = Field( - default=None, - description="""Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""", - ) data: Optional[bytes] = Field( default=None, description="""Required. The data in the chunk.""" ) @@ -8353,20 +8388,24 @@ class Chunk(_common.BaseModel): default=None, description="""Optional. Metadata that is associated with the data in the payload.""", ) + mime_type: Optional[str] = Field( + default=None, + description="""Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""", + ) class ChunkDict(TypedDict, total=False): """A chunk of data.""" - mime_type: Optional[str] - """Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""" - data: Optional[bytes] """Required. The data in the chunk.""" metadata: Optional[MetadataDict] """Optional. Metadata that is associated with the data in the payload.""" + mime_type: Optional[str] + """Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""" + ChunkOrDict = Union[Chunk, ChunkDict] From d02a7da393a6bb8b1a9c462243b399b62911378c Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 4 Nov 2025 14:06:49 -0800 Subject: [PATCH 21/24] fix: GenAI Client(evals) - Remove requirement for `agent_info.agent` in `create_evaluation_run` in Vertex AI GenAI SDK evals. PiperOrigin-RevId: 828127361 --- vertexai/_genai/evals.py | 36 ++++++++++-------------------------- 1 file changed, 10 insertions(+), 26 deletions(-) diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py index c21b186ecd..afc7868627 100644 --- a/vertexai/_genai/evals.py +++ b/vertexai/_genai/evals.py @@ -1397,15 +1397,6 @@ def create_evaluation_run( ) inference_configs = {} if agent_info: - if isinstance(agent_info, dict): - agent_info = types.evals.AgentInfo.model_validate(agent_info) - if ( - not agent_info.agent - or len(agent_info.agent.split("reasoningEngines/")) != 2 - ): - raise ValueError( - "agent_info.agent cannot be empty. Please provide a valid reasoning engine resource name in the format of projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}." - ) inference_configs[agent_info.name] = types.EvaluationRunInferenceConfig( agent_config=types.EvaluationRunAgentConfig( developer_instruction=genai_types.Content( @@ -1414,10 +1405,11 @@ def create_evaluation_run( tools=agent_info.tool_declarations, ) ) - labels = labels or {} - labels["vertex-ai-evaluation-agent-engine-id"] = agent_info.agent.split( - "reasoningEngines/" - )[-1] + if agent_info.agent: + labels = labels or {} + labels["vertex-ai-evaluation-agent-engine-id"] = agent_info.agent.split( + "reasoningEngines/" + )[-1] if not name: name = f"evaluation_run_{uuid.uuid4()}" @@ -2252,15 +2244,6 @@ async def create_evaluation_run( ) inference_configs = {} if agent_info: - if isinstance(agent_info, dict): - agent_info = types.evals.AgentInfo.model_validate(agent_info) - if ( - not agent_info.agent - or len(agent_info.agent.split("reasoningEngines/")) != 2 - ): - raise ValueError( - "agent_info.agent cannot be empty. Please provide a valid reasoning engine resource name in the format of projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}." - ) inference_configs[agent_info.name] = types.EvaluationRunInferenceConfig( agent_config=types.EvaluationRunAgentConfig( developer_instruction=genai_types.Content( @@ -2269,10 +2252,11 @@ async def create_evaluation_run( tools=agent_info.tool_declarations, ) ) - labels = labels or {} - labels["vertex-ai-evaluation-agent-engine-id"] = agent_info.agent.split( - "reasoningEngines/" - )[-1] + if agent_info.agent: + labels = labels or {} + labels["vertex-ai-evaluation-agent-engine-id"] = agent_info.agent.split( + "reasoningEngines/" + )[-1] if not name: name = f"evaluation_run_{uuid.uuid4()}" From bf1851e59cb34e63b509a2a610e72691e1c4ca28 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 4 Nov 2025 14:11:16 -0800 Subject: [PATCH 22/24] feat: Add the identity type option for the agent engine and add effective identity to the resource PiperOrigin-RevId: 828129186 --- .../unit/vertexai/genai/test_agent_engines.py | 49 ++++++++++++++++++- vertexai/_genai/agent_engines.py | 21 ++++++-- vertexai/_genai/types/__init__.py | 2 + vertexai/_genai/types/common.py | 31 ++++++++++++ 4 files changed, 99 insertions(+), 4 deletions(-) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index 508bf0377b..ade08e7b2f 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -533,6 +533,9 @@ def register_operations(self) -> Dict[str, List[str]]: } _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_SPEC = _genai_types.ReasoningEngineSpecDict( agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK, @@ -559,6 +562,7 @@ def register_operations(self) -> Dict[str, List[str]]: requirements_gcs_uri=_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, ), service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, ) _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}] _TEST_AGENT_ENGINE_OPERATION_SCHEMAS = [] @@ -908,6 +912,7 @@ def test_create_agent_engine_config_full(self, mock_prepare): extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG, min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES, max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES, @@ -950,6 +955,10 @@ def test_create_agent_engine_config_full(self, mock_prepare): config["spec"]["service_account"] == _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT ) + assert ( + config["spec"]["identity_type"] + == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT + ) @mock.patch.object( _agent_engines_utils, @@ -977,6 +986,7 @@ def test_create_agent_engine_config_with_source_packages( requirements_file=requirements_file_path, class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, agent_framework=_TEST_AGENT_FRAMEWORK, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION @@ -994,6 +1004,10 @@ def test_create_agent_engine_config_with_source_packages( mock_create_base64_encoded_tarball.assert_called_once_with( source_packages=[test_file_path] ) + assert ( + config["spec"]["identity_type"] + == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT + ) @mock.patch.object(_agent_engines_utils, "_prepare") def test_update_agent_engine_config_full(self, mock_prepare): @@ -1008,6 +1022,7 @@ def test_update_agent_engine_config_full(self, mock_prepare): extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION @@ -1038,6 +1053,10 @@ def test_update_agent_engine_config_full(self, mock_prepare): config["spec"]["service_account"] == _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT ) + assert ( + config["spec"]["identity_type"] + == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT + ) assert config["update_mask"] == ",".join( [ "display_name", @@ -1048,8 +1067,28 @@ def test_update_agent_engine_config_full(self, mock_prepare): "spec.class_methods", "spec.deployment_spec.env", "spec.deployment_spec.secret_env", - "spec.service_account", "spec.agent_framework", + "spec.identity_type", + "spec.service_account", + ] + ) + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_update_agent_engine_clear_service_account(self, mock_prepare): + config = self.client.agent_engines._create_config( + mode="update", + service_account="", + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, + ) + assert "service_account" not in config["spec"].keys() + assert ( + config["spec"]["identity_type"] + == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT + ) + assert config["update_mask"] == ",".join( + [ + "spec.identity_type", + "spec.service_account", ] ) @@ -1488,6 +1527,7 @@ def test_create_agent_engine_with_env_vars_dict( extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, service_account=None, + identity_type=None, context_spec=None, psc_interface_config=None, min_instances=None, @@ -1543,6 +1583,7 @@ def test_create_agent_engine_with_custom_service_account( }, "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], "service_account": _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + "identity_type": _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, }, } @@ -1564,6 +1605,7 @@ def test_create_agent_engine_with_custom_service_account( extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], staging_bucket=_TEST_STAGING_BUCKET, service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, ), ) mock_create_config.assert_called_with( @@ -1577,6 +1619,7 @@ def test_create_agent_engine_with_custom_service_account( extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], env_vars=None, service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, context_spec=None, psc_interface_config=None, min_instances=None, @@ -1608,6 +1651,7 @@ def test_create_agent_engine_with_custom_service_account( "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, }, "service_account": _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, + "identity_type": _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, }, }, None, @@ -1666,6 +1710,7 @@ def test_create_agent_engine_with_experimental_mode( extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], env_vars=None, service_account=None, + identity_type=None, context_spec=None, psc_interface_config=None, min_instances=None, @@ -1819,6 +1864,7 @@ def test_create_agent_engine_with_class_methods( extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], env_vars=None, service_account=None, + identity_type=None, context_spec=None, psc_interface_config=None, min_instances=None, @@ -1919,6 +1965,7 @@ def test_create_agent_engine_with_agent_framework( entrypoint_object=None, requirements_file=None, agent_framework=_TEST_AGENT_FRAMEWORK, + identity_type=None, ) request_mock.assert_called_with( "post", diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 238c5d8d69..9b0a882de2 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -907,6 +907,7 @@ def create( api_config = self._create_config( mode="create", agent=agent, + identity_type=config.identity_type, staging_bucket=config.staging_bucket, requirements=config.requirements, display_name=config.display_name, @@ -971,6 +972,7 @@ def _create_config( *, mode: str, agent: Any = None, + identity_type: Optional[types.IdentityType] = None, staging_bucket: Optional[str] = None, requirements: Optional[Union[str, Sequence[str]]] = None, display_name: Optional[str] = None, @@ -1189,9 +1191,6 @@ def _create_config( ) update_masks.extend(deployment_update_masks) agent_engine_spec["deployment_spec"] = deployment_spec - if service_account is not None: - agent_engine_spec["service_account"] = service_account - update_masks.append("spec.service_account") if agent_server_mode: if not agent_engine_spec.get("deployment_spec"): @@ -1209,6 +1208,21 @@ def _create_config( ) ) update_masks.append("spec.agent_framework") + + if identity_type is not None or service_account is not None: + if agent_engine_spec is None: + agent_engine_spec = {} + + if identity_type is not None: + agent_engine_spec["identity_type"] = identity_type + update_masks.append("spec.identity_type") + if service_account is not None: + # Clear the field in case of empty service_account. + if service_account: + agent_engine_spec["service_account"] = service_account + update_masks.append("spec.service_account") + + if agent_engine_spec is not None: config["spec"] = agent_engine_spec if update_masks and mode == "update": @@ -1414,6 +1428,7 @@ def update( api_config = self._create_config( mode="update", agent=agent, + identity_type=config.identity_type, staging_bucket=config.staging_bucket, requirements=config.requirements, display_name=config.display_name, diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index c1a01ad8a3..4b9ad32eb7 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -430,6 +430,7 @@ from .common import GetPromptConfig from .common import GetPromptConfigDict from .common import GetPromptConfigOrDict +from .common import IdentityType from .common import Importance from .common import IntermediateExtractedMemory from .common import IntermediateExtractedMemoryDict @@ -1798,6 +1799,7 @@ "AcceleratorType", "Type", "JobState", + "IdentityType", "AgentServerMode", "ManagedTopicEnum", "Language", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 33791bdd92..34619b654e 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -205,6 +205,17 @@ class JobState(_common.CaseInSensitiveEnum): """The job is partially succeeded, some results may be missing due to errors.""" +class IdentityType(_common.CaseInSensitiveEnum): + """The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""" + + IDENTITY_TYPE_UNSPECIFIED = "IDENTITY_TYPE_UNSPECIFIED" + """Default value. Use a custom service account if the `service_account` field is set, otherwise use the default Vertex AI Reasoning Engine Service Agent in the project. Same behavior as SERVICE_ACCOUNT.""" + SERVICE_ACCOUNT = "SERVICE_ACCOUNT" + """Use a custom service account if the `service_account` field is set, otherwise use the default Vertex AI Reasoning Engine Service Agent in the project.""" + AGENT_IDENTITY = "AGENT_IDENTITY" + """Use Agent Identity. The `service_account` field must not be set.""" + + class AgentServerMode(_common.CaseInSensitiveEnum): """The agent server mode.""" @@ -4784,6 +4795,14 @@ class ReasoningEngineSpec(_common.BaseModel): default=None, description="""Optional. The specification of a Reasoning Engine deployment.""", ) + effective_identity: Optional[str] = Field( + default=None, + description="""Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""", + ) + identity_type: Optional[IdentityType] = Field( + default=None, + description="""Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""", + ) package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field( default=None, description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""", @@ -4810,6 +4829,12 @@ class ReasoningEngineSpecDict(TypedDict, total=False): deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict] """Optional. The specification of a Reasoning Engine deployment.""" + effective_identity: Optional[str] + """Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""" + + identity_type: Optional[IdentityType] + """Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""" + package_spec: Optional[ReasoningEngineSpecPackageSpecDict] """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""" @@ -12905,6 +12930,9 @@ class AgentEngineConfig(_common.BaseModel): If not specified, the default Reasoning Engine P6SA service agent will be used.""", ) + identity_type: Optional[IdentityType] = Field( + default=None, description="""The identity type to use for the Agent Engine.""" + ) context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, description="""The context spec to be used for the Agent Engine.""", @@ -13057,6 +13085,9 @@ class AgentEngineConfigDict(TypedDict, total=False): If not specified, the default Reasoning Engine P6SA service agent will be used.""" + identity_type: Optional[IdentityType] + """The identity type to use for the Agent Engine.""" + context_spec: Optional[ReasoningEngineContextSpecDict] """The context spec to be used for the Agent Engine.""" From 65f8bba96896fbe13313069f259781009fe8c42b Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 4 Nov 2025 15:06:54 -0800 Subject: [PATCH 23/24] chore: disable default-on telemetry enablement in ADK on AE PiperOrigin-RevId: 828151286 --- .../test_agent_engine_templates_adk.py | 194 +++++++++--------- .../unit/vertexai/genai/test_agent_engines.py | 86 ++++---- vertexai/_genai/agent_engines.py | 5 +- vertexai/agent_engines/_agent_engines.py | 5 +- 4 files changed, 146 insertions(+), 144 deletions(-) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index b4f41aeb3d..ee3a0677f5 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -26,7 +26,6 @@ from google.auth import credentials as auth_credentials from google.cloud import storage import vertexai -from google.cloud import aiplatform from google.cloud.aiplatform_v1 import types as aip_types from google.cloud.aiplatform_v1.services import reasoning_engine_service from google.cloud.aiplatform import base @@ -1012,99 +1011,100 @@ def update_agent_engine_mock(): yield update_agent_engine_mock -@pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngines: - def setup_method(self): - importlib.reload(initializer) - importlib.reload(aiplatform) - aiplatform.init( - project=_TEST_PROJECT, - location=_TEST_LOCATION, - credentials=_TEST_CREDENTIALS, - staging_bucket=_TEST_STAGING_BUCKET, - ) - - def teardown_method(self): - initializer.global_pool.shutdown(wait=True) - - @pytest.mark.parametrize( - "env_vars,expected_env_vars", - [ - ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - ( - {"some_env": "some_val"}, - { - "some_env": "some_val", - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true", - }, - ), - ( - {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, - {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, - ), - ], - ) - def test_create_default_telemetry_enablement( - self, - create_agent_engine_mock: mock.Mock, - cloud_storage_create_bucket_mock: mock.Mock, - cloudpickle_dump_mock: mock.Mock, - cloudpickle_load_mock: mock.Mock, - get_gca_resource_mock: mock.Mock, - env_vars: dict[str, str], - expected_env_vars: dict[str, str], - ): - agent_engines.create( - agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT), - env_vars=env_vars, - ) - create_agent_engine_mock.assert_called_once() - deployment_spec = create_agent_engine_mock.call_args.kwargs[ - "reasoning_engine" - ].spec.deployment_spec - assert _utils.to_dict(deployment_spec)["env"] == [ - {"name": key, "value": value} for key, value in expected_env_vars.items() - ] - - @pytest.mark.parametrize( - "env_vars,expected_env_vars", - [ - ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - ( - {"some_env": "some_val"}, - { - "some_env": "some_val", - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true", - }, - ), - ( - {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, - {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, - ), - ], - ) - def test_update_default_telemetry_enablement( - self, - update_agent_engine_mock: mock.Mock, - cloud_storage_create_bucket_mock: mock.Mock, - cloudpickle_dump_mock: mock.Mock, - cloudpickle_load_mock: mock.Mock, - get_gca_resource_mock: mock.Mock, - get_agent_engine_mock: mock.Mock, - env_vars: dict[str, str], - expected_env_vars: dict[str, str], - ): - agent_engines.update( - resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - description="foobar", # avoid "At least one of ... must be specified" errors. - env_vars=env_vars, - ) - update_agent_engine_mock.assert_called_once() - deployment_spec = update_agent_engine_mock.call_args.kwargs[ - "request" - ].reasoning_engine.spec.deployment_spec - assert _utils.to_dict(deployment_spec)["env"] == [ - {"name": key, "value": value} for key, value in expected_env_vars.items() - ] +# TODO(jawoszek): Uncomment once we're ready for default-on. +# @pytest.mark.usefixtures("google_auth_mock") +# class TestAgentEngines: +# def setup_method(self): +# importlib.reload(initializer) +# importlib.reload(aiplatform) +# aiplatform.init( +# project=_TEST_PROJECT, +# location=_TEST_LOCATION, +# credentials=_TEST_CREDENTIALS, +# staging_bucket=_TEST_STAGING_BUCKET, +# ) + +# def teardown_method(self): +# initializer.global_pool.shutdown(wait=True) + +# @pytest.mark.parametrize( +# "env_vars,expected_env_vars", +# [ +# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), +# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), +# ( +# {"some_env": "some_val"}, +# { +# "some_env": "some_val", +# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false", +# }, +# ), +# ( +# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, +# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, +# ), +# ], +# ) +# def test_create_default_telemetry_enablement( +# self, +# create_agent_engine_mock: mock.Mock, +# cloud_storage_create_bucket_mock: mock.Mock, +# cloudpickle_dump_mock: mock.Mock, +# cloudpickle_load_mock: mock.Mock, +# get_gca_resource_mock: mock.Mock, +# env_vars: dict[str, str], +# expected_env_vars: dict[str, str], +# ): +# agent_engines.create( +# agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT), +# env_vars=env_vars, +# ) +# create_agent_engine_mock.assert_called_once() +# deployment_spec = create_agent_engine_mock.call_args.kwargs[ +# "reasoning_engine" +# ].spec.deployment_spec +# assert _utils.to_dict(deployment_spec)["env"] == [ +# {"name": key, "value": value} for key, value in expected_env_vars.items() +# ] + +# @pytest.mark.parametrize( +# "env_vars,expected_env_vars", +# [ +# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), +# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), +# ( +# {"some_env": "some_val"}, +# { +# "some_env": "some_val", +# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false", +# }, +# ), +# ( +# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, +# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, +# ), +# ], +# ) +# def test_update_default_telemetry_enablement( +# self, +# update_agent_engine_mock: mock.Mock, +# cloud_storage_create_bucket_mock: mock.Mock, +# cloudpickle_dump_mock: mock.Mock, +# cloudpickle_load_mock: mock.Mock, +# get_gca_resource_mock: mock.Mock, +# get_agent_engine_mock: mock.Mock, +# env_vars: dict[str, str], +# expected_env_vars: dict[str, str], +# ): +# agent_engines.update( +# resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME, +# description="foobar", # avoid "At least one of ... must be specified" errors. +# env_vars=env_vars, +# ) +# update_agent_engine_mock.assert_called_once() +# deployment_spec = update_agent_engine_mock.call_args.kwargs[ +# "request" +# ].reasoning_engine.spec.deployment_spec +# assert _utils.to_dict(deployment_spec)["env"] == [ +# {"name": key, "value": value} for key, value in expected_env_vars.items() +# ] diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index ade08e7b2f..b8971b8f92 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -31,7 +31,6 @@ from google.cloud import aiplatform import vertexai from google.cloud.aiplatform import initializer -from vertexai.agent_engines.templates import adk from vertexai._genai import _agent_engines_utils from vertexai._genai import agent_engines from vertexai._genai import types as _genai_types @@ -856,48 +855,49 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): "description": _TEST_AGENT_ENGINE_DESCRIPTION, } - @mock.patch.object(_agent_engines_utils, "_prepare") - @pytest.mark.parametrize( - "env_vars,expected_env_vars", - [ - ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - ( - {"some_env": "some_val"}, - { - "some_env": "some_val", - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true", - }, - ), - ( - {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, - {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, - ), - ], - ) - def test_agent_engine_adk_telemetry_enablement( - self, - mock_prepare: mock.Mock, - env_vars: dict[str, str], - expected_env_vars: dict[str, str], - ): - agent = mock.Mock(spec=adk.AdkApp) - agent.clone = lambda: agent - agent.register_operations = lambda: {} - - config = self.client.agent_engines._create_config( - mode="create", - agent=agent, - staging_bucket=_TEST_STAGING_BUCKET, - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - description=_TEST_AGENT_ENGINE_DESCRIPTION, - env_vars=env_vars, - ) - assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME - assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION - assert config["spec"]["deployment_spec"]["env"] == [ - {"name": key, "value": value} for key, value in expected_env_vars.items() - ] + # TODO(jawoszek): Uncomment once we're ready for default-on. + # @mock.patch.object(_agent_engines_utils, "_prepare") + # @pytest.mark.parametrize( + # "env_vars,expected_env_vars", + # [ + # ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), + # (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), + # ( + # {"some_env": "some_val"}, + # { + # "some_env": "some_val", + # GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false", + # }, + # ), + # ( + # {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + # {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + # ), + # ], + # ) + # def test_agent_engine_adk_telemetry_enablement( + # self, + # mock_prepare: mock.Mock, + # env_vars: dict[str, str], + # expected_env_vars: dict[str, str], + # ): + # agent = mock.Mock(spec=adk.AdkApp) + # agent.clone = lambda: agent + # agent.register_operations = lambda: {} + + # config = self.client.agent_engines._create_config( + # mode="create", + # agent=agent, + # staging_bucket=_TEST_STAGING_BUCKET, + # display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + # description=_TEST_AGENT_ENGINE_DESCRIPTION, + # env_vars=env_vars, + # ) + # assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME + # assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION + # assert config["spec"]["deployment_spec"]["env"] == [ + # {"name": key, "value": value} for key, value in expected_env_vars.items() + # ] @mock.patch.object(_agent_engines_utils, "_prepare") def test_create_agent_engine_config_full(self, mock_prepare): diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 9b0a882de2..b0a5fa737f 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -1042,8 +1042,9 @@ def _create_config( raise ValueError("location must be set using `vertexai.Client`.") gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME agent = _agent_engines_utils._validate_agent_or_raise(agent=agent) - if _agent_engines_utils._is_adk_agent(agent): - env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) + # TODO(jawoszek): Uncomment once we're ready for default-on. + # if _agent_engines_utils._is_adk_agent(agent): + # env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise( staging_bucket=staging_bucket, ) diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index 596056b15a..3696fd4a33 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -518,8 +518,9 @@ def create( if agent_engine is not None: agent_engine = _validate_agent_engine_or_raise(agent_engine) staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - if _is_adk_agent(None, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) + # TODO(jawoszek): Uncomment once we're ready for default-on. + # if _is_adk_agent(None, agent_engine): + # env_vars = _add_telemetry_enablement_env(env_vars=env_vars) if agent_engine is None: if requirements is not None: From 366d4a1932d90aa3bbd5fd471b8cd8d23b1d747c Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Tue, 4 Nov 2025 16:17:49 -0800 Subject: [PATCH 24/24] chore(main): release 1.125.0 (#6070) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> Co-authored-by: Yvonne Yu <150068659+yyyu-google@users.noreply.github.com> --- .release-please-manifest.json | 2 +- CHANGELOG.md | 31 +++++++++++++++++++ google/cloud/aiplatform/gapic_version.py | 2 +- .../schema/predict/instance/gapic_version.py | 2 +- .../predict/instance_v1/gapic_version.py | 2 +- .../v1/schema/predict/params/gapic_version.py | 2 +- .../schema/predict/params_v1/gapic_version.py | 2 +- .../predict/prediction/gapic_version.py | 2 +- .../predict/prediction_v1/gapic_version.py | 2 +- .../trainingjob/definition/gapic_version.py | 2 +- .../definition_v1/gapic_version.py | 2 +- .../schema/predict/instance/gapic_version.py | 2 +- .../predict/instance_v1beta1/gapic_version.py | 2 +- .../schema/predict/params/gapic_version.py | 2 +- .../predict/params_v1beta1/gapic_version.py | 2 +- .../predict/prediction/gapic_version.py | 2 +- .../prediction_v1beta1/gapic_version.py | 2 +- .../trainingjob/definition/gapic_version.py | 2 +- .../definition_v1beta1/gapic_version.py | 2 +- google/cloud/aiplatform/version.py | 2 +- google/cloud/aiplatform_v1/gapic_version.py | 2 +- .../cloud/aiplatform_v1beta1/gapic_version.py | 2 +- pypi/_vertex_ai_placeholder/version.py | 2 +- ...t_metadata_google.cloud.aiplatform.v1.json | 2 +- ...adata_google.cloud.aiplatform.v1beta1.json | 2 +- 25 files changed, 55 insertions(+), 24 deletions(-) diff --git a/.release-please-manifest.json b/.release-please-manifest.json index 768ea8b693..71e4e764a1 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "1.124.0" + ".": "1.125.0" } diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b00dd0dc4..5119d90488 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,36 @@ # Changelog +## [1.125.0](https://github.com/googleapis/python-aiplatform/compare/v1.124.0...v1.125.0) (2025-11-04) + + +### ⚠ BREAKING CHANGES + +* Switch tracing APIs in preview AdkApp. +* Switch `cloudtrace.googleapis.com` to `telemetry.googleapis.com` for tracing API. + +### Features + +* Add reservation affinity support to preview BatchPredictionJob ([c8f38a0](https://github.com/googleapis/python-aiplatform/commit/c8f38a0a51c318a5065438067f85f31be5088af1)) +* Add support for Vertex Express Mode API key in AdkApp ([05834cb](https://github.com/googleapis/python-aiplatform/commit/05834cb43302cf2bfc34f25f400d075e44aa9df3)) +* Add the identity type option for the agent engine and add effective identity to the resource ([bf1851e](https://github.com/googleapis/python-aiplatform/commit/bf1851e59cb34e63b509a2a610e72691e1c4ca28)) +* Alow VertexAiSession for streaming_agent_run_with_events ([13faa27](https://github.com/googleapis/python-aiplatform/commit/13faa27376814f7b0a223ff9455c289a9af75288)) +* GenAI Client(evals) - Add retry to predefine metric ([9a46e67](https://github.com/googleapis/python-aiplatform/commit/9a46e67a8c341673b14bece88bc635b455314711)) + + +### Bug Fixes + +* GenAI Client(evals) - Change dataset visualization table to fixed to prevent horizontal expansion. ([7a5a066](https://github.com/googleapis/python-aiplatform/commit/7a5a0667b915bb0ed1009ffc8034a8d29fb2d20e)) +* GenAI Client(evals) - Remove requirement for `agent_info.agent` in `create_evaluation_run` in Vertex AI GenAI SDK evals. ([d02a7da](https://github.com/googleapis/python-aiplatform/commit/d02a7da393a6bb8b1a9c462243b399b62911378c)) +* GenAI Client(evals) - Support direct pandas DataFrame dataset in evaluate() ([a917122](https://github.com/googleapis/python-aiplatform/commit/a9171221e3bafecdc75580c3f25347f1c3d18851)) +* Revert: Alow VertexAiSession for streaming_agent_run_with_events ([7c8c218](https://github.com/googleapis/python-aiplatform/commit/7c8c218f47a6fd919cae6869736d35e4206a123e)) + + +### Miscellaneous Chores + +* Release 1.125.0 ([d344858](https://github.com/googleapis/python-aiplatform/commit/d34485880bdefd21dccac4548656af1eaedb3727)) +* Switch `cloudtrace.googleapis.com` to `telemetry.googleapis.com` for tracing API. ([c81f912](https://github.com/googleapis/python-aiplatform/commit/c81f9124dce841a7cbc310c6c7652ad793c9a58f)) +* Switch tracing APIs in preview AdkApp. ([27ef56b](https://github.com/googleapis/python-aiplatform/commit/27ef56b8319ba793f6f00e7857fe0c20c2c8b61a)) + ## [1.124.0](https://github.com/googleapis/python-aiplatform/compare/v1.123.0...v1.124.0) (2025-10-30) diff --git a/google/cloud/aiplatform/gapic_version.py b/google/cloud/aiplatform/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/gapic_version.py +++ b/google/cloud/aiplatform/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/version.py b/google/cloud/aiplatform/version.py index cfd51ec79a..d10e1bf741 100644 --- a/google/cloud/aiplatform/version.py +++ b/google/cloud/aiplatform/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.124.0" +__version__ = "1.125.0" diff --git a/google/cloud/aiplatform_v1/gapic_version.py b/google/cloud/aiplatform_v1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform_v1/gapic_version.py +++ b/google/cloud/aiplatform_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform_v1beta1/gapic_version.py b/google/cloud/aiplatform_v1beta1/gapic_version.py index 91c21609ea..a2635f5568 100644 --- a/google/cloud/aiplatform_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.124.0" # {x-release-please-version} +__version__ = "1.125.0" # {x-release-please-version} diff --git a/pypi/_vertex_ai_placeholder/version.py b/pypi/_vertex_ai_placeholder/version.py index be23177662..b5c782d347 100644 --- a/pypi/_vertex_ai_placeholder/version.py +++ b/pypi/_vertex_ai_placeholder/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.124.0" +__version__ = "1.125.0" diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json index 8f2c7f65c0..5923b5022d 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.124.0" + "version": "1.125.0" }, "snippets": [ { diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json index 6a065b5413..e778517b37 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.124.0" + "version": "1.125.0" }, "snippets": [ {