No public description

PiperOrigin-RevId: 764358461

Author: vertex-sdk-bot

owlbot.py

@@ -83,6 +83,7 @@
             "noxfile.py",
             "testing",
             "docs/conf.py",
+            "*.tar.gz"
         ],
     )
     has_generator_updates = True

setup.py

@@ -49,11 +49,8 @@
 profiler_extra_require = [
     "tensorboard-plugin-profile >= 2.4.0, <2.18.0",  # <3.0.0",
     "werkzeug >= 2.0.0, <4.0.0",
-    "tensorflow >=2.4.0, <3.0.0",
 ]
-tensorboard_extra_require = [
-    "tensorflow >=2.3.0, <3.0.0; python_version<='3.11'"
-] + profiler_extra_require
+tensorboard_extra_require = profiler_extra_require
 
 metadata_extra_require = ["pandas >= 1.0.0", "numpy>=1.15.0"]
 xai_extra_require = ["tensorflow >=2.3.0, <3.0.0"]

tests/unit/vertexai/test_generative_models.py

@@ -21,13 +21,15 @@
 from typing import Dict, Iterable, List, MutableSequence, Optional
 from unittest import mock
 
+from google.api_core import operation as ga_operation
 import vertexai
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform_v1 import types as types_v1
 from google.cloud.aiplatform_v1.services import (
     prediction_service as prediction_service_v1,
 )
 from google.cloud.aiplatform_v1beta1 import types as types_v1beta1
+from google.cloud.aiplatform_v1beta1.services import endpoint_service
 from vertexai import generative_models
 from vertexai.preview import (
     generative_models as preview_generative_models,
@@ -48,6 +50,7 @@
 )
 from vertexai.generative_models import _function_calling_utils
 from vertexai.caching import CachedContent
+from google.protobuf import field_mask_pb2
 
 
 _TEST_PROJECT = "test-project"
@@ -1710,6 +1713,115 @@ def test_defs_ref_renaming(self):
             _fix_schema_dict_for_gapic_in_place(actual)
             assert actual == expected
 
+    @pytest.mark.parametrize(
+        "generative_models",
+        [preview_generative_models],  # Only preview supports set_logging_config
+    )
+    @mock.patch.object(endpoint_service.EndpointServiceClient, "update_endpoint")
+    def test_set_logging_config_for_endpoint(
+        self, mock_update_endpoint, generative_models: generative_models
+    ):
+        endpoint_name = (
+            f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/12345"
+        )
+        model = generative_models.GenerativeModel(endpoint_name)
+
+        mock_update_endpoint.return_value = types_v1beta1.Endpoint(name=endpoint_name)
+
+        enabled = True
+        sampling_rate = 0.5
+        bigquery_destination = f"bq://{_TEST_PROJECT}.my_dataset.my_table"
+        enable_otel_logging = True
+
+        model.set_request_response_logging_config(
+            enabled=enabled,
+            sampling_rate=sampling_rate,
+            bigquery_destination=bigquery_destination,
+            enable_otel_logging=enable_otel_logging,
+        )
+
+        expected_logging_config = types_v1beta1.PredictRequestResponseLoggingConfig(
+            enabled=enabled,
+            sampling_rate=sampling_rate,
+            bigquery_destination=types_v1beta1.BigQueryDestination(
+                output_uri=bigquery_destination
+            ),
+            enable_otel_logging=enable_otel_logging,
+        )
+        expected_endpoint = types_v1beta1.Endpoint(
+            name=endpoint_name,
+            predict_request_response_logging_config=expected_logging_config,
+        )
+        expected_update_mask = field_mask_pb2.FieldMask(
+            paths=["predict_request_response_logging_config"]
+        )
+
+        mock_update_endpoint.assert_called_once_with(
+            types_v1beta1.UpdateEndpointRequest(
+                endpoint=expected_endpoint,
+                update_mask=expected_update_mask,
+            )
+        )
+
+    @pytest.mark.parametrize(
+        "generative_models",
+        [preview_generative_models],  # Only preview supports set_logging_config
+    )
+    @mock.patch.object(
+        endpoint_service.EndpointServiceClient, "set_publisher_model_config"
+    )
+    def test_set_logging_config_for_publisher_model(
+        self, mock_set_publisher_model_config, generative_models: generative_models
+    ):
+        model_name = "gemini-pro"
+        model = generative_models.GenerativeModel(model_name)
+        full_model_name = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/{model_name}"
+
+        enabled = False
+        sampling_rate = 1.0
+        bigquery_destination = f"bq://{_TEST_PROJECT}.another_dataset"
+        enable_otel_logging = False
+
+        mock_operation = mock.Mock(spec=ga_operation.Operation)
+        mock_set_publisher_model_config.return_value = mock_operation
+        mock_operation.result.return_value = types_v1beta1.PublisherModelConfig(
+            logging_config=types_v1beta1.PredictRequestResponseLoggingConfig(
+                enabled=enabled,
+                sampling_rate=sampling_rate,
+                bigquery_destination=types_v1beta1.BigQueryDestination(
+                    output_uri=bigquery_destination
+                ),
+                enable_otel_logging=enable_otel_logging,
+            )
+        )
+
+        model.set_request_response_logging_config(
+            enabled=enabled,
+            sampling_rate=sampling_rate,
+            bigquery_destination=bigquery_destination,
+            enable_otel_logging=enable_otel_logging,
+        )
+
+        expected_logging_config = types_v1beta1.PredictRequestResponseLoggingConfig(
+            enabled=enabled,
+            sampling_rate=sampling_rate,
+            bigquery_destination=types_v1beta1.BigQueryDestination(
+                output_uri=bigquery_destination
+            ),
+            enable_otel_logging=enable_otel_logging,
+        )
+        expected_publisher_model_config = types_v1beta1.PublisherModelConfig(
+            logging_config=expected_logging_config
+        )
+
+        mock_set_publisher_model_config.assert_called_once_with(
+            types_v1beta1.SetPublisherModelConfigRequest(
+                name=full_model_name,
+                publisher_model_config=expected_publisher_model_config,
+            )
+        )
+        mock_operation.result.assert_called_once()
+
 
 EXPECTED_SCHEMA_FOR_GET_CURRENT_WEATHER = {
     "title": "get_current_weather",

vertexai/generative_models/_generative_models.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# 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.
@@ -48,6 +48,7 @@
     llm_utility_service as llm_utility_service_v1,
 )
 from google.cloud.aiplatform_v1beta1 import types as aiplatform_types
+from google.cloud.aiplatform_v1beta1.services import endpoint_service
 from google.cloud.aiplatform_v1beta1.services import prediction_service
 from google.cloud.aiplatform_v1beta1.services import llm_utility_service
 from google.cloud.aiplatform_v1beta1.types import (
@@ -59,6 +60,7 @@
 )
 from google.cloud.aiplatform_v1beta1.types import tool as gapic_tool_types
 from google.protobuf import json_format
+from google.protobuf import field_mask_pb2
 import warnings
 
 if TYPE_CHECKING:
@@ -502,6 +504,19 @@ def _llm_utility_client(self) -> llm_utility_service.LlmUtilityServiceClient:
             api_key=api_key,
         )
 
+    @functools.cached_property
+    def _endpoint_client(self) -> endpoint_service.EndpointServiceClient:
+        # Note this doesn't work with GCP Express but it's better to set the
+        # client correctly and allow the service to throw
+        api_key = aiplatform_initializer.global_config.api_key
+        if api_key and aiplatform_initializer.global_config.project:
+            api_key = None
+        return aiplatform_initializer.global_config.create_client(
+            client_class=endpoint_service.EndpointServiceClient,
+            location_override=self._location,
+            api_key=api_key,
+        )
+
     @functools.cached_property
     def _llm_utility_async_client(
         self,
@@ -3612,3 +3627,68 @@ def start_chat(
             response_validation=response_validation,
             responder=responder,
         )
+
+    def set_request_response_logging_config(
+        self,
+        *,
+        enabled: bool,
+        sampling_rate: float,
+        bigquery_destination: str,
+        enable_otel_logging: Optional[bool] = None,
+    ) -> Union[aiplatform_types.PublisherModelConfig, aiplatform_types.Endpoint]:
+        """
+        Sets the request/response logging config.
+
+        Args:
+            enabled: If logging is enabled or not.
+            sampling_rate: Percentage of requests to be logged, expressed as a
+                fraction in range(0,1].
+        bigquery_destination: BigQuery table for logging. If only given a project,
+            a new dataset will be created with name
+            ``logging_<endpoint-display-name>_<endpoint-id>`` where will
+            be made BigQuery-dataset-name compatible (e.g. most special
+            characters will become underscores). If no table name is
+            given, a new table will be created with name
+            ``request_response_logging``
+        enable_otel_logging: This field is used for large models. If true, in
+            addition to the original large model logs, logs will be converted in
+            OTel schema format, and saved in otel_log column. Default
+            value is false.
+        Returns:
+            The updated PublisherModelConfig or Endpoint.
+        """
+
+        logging_config = aiplatform_types.PredictRequestResponseLoggingConfig(
+            enabled=enabled,
+            sampling_rate=sampling_rate,
+            bigquery_destination=aiplatform_types.BigQueryDestination(
+                output_uri=bigquery_destination
+            ),
+            enable_otel_logging=enable_otel_logging,
+        )
+
+        if self._endpoint_client.parse_endpoint_path(self._prediction_resource_name):
+            return self._endpoint_client.update_endpoint(
+                aiplatform_types.UpdateEndpointRequest(
+                    endpoint=aiplatform_types.Endpoint(
+                        name=self._prediction_resource_name,
+                        predict_request_response_logging_config=logging_config,
+                    ),
+                    update_mask=field_mask_pb2.FieldMask(
+                        paths=["predict_request_response_logging_config"]
+                    ),
+                )
+            )
+
+        else:
+
+            operation = self._endpoint_client.set_publisher_model_config(
+                aiplatform_types.SetPublisherModelConfigRequest(
+                    name=self._prediction_resource_name,
+                    publisher_model_config=aiplatform_types.PublisherModelConfig(
+                        logging_config=logging_config
+                    ),
+                )
+            )
+
+            return operation.result()