Merge branch 'main' into patch-1

Author: jaycee-li

google/cloud/aiplatform/compat/__init__.py

@@ -226,6 +226,9 @@
     types.explanation_metadata = types.explanation_metadata_v1
     types.feature = types.feature_v1
     types.feature_group = types.feature_group_v1
+    # TODO(b/293184410): Temporary code. Switch to v1 once v1 is available.
+    types.feature_monitor = types.feature_monitor_v1beta1
+    types.feature_monitor_job = types.feature_monitor_job_v1beta1
     types.feature_monitoring_stats = types.feature_monitoring_stats_v1
     types.feature_online_store = types.feature_online_store_v1
     types.feature_online_store_admin_service = (

google/cloud/aiplatform/utils/__init__.py

@@ -763,6 +763,12 @@ def parse_feature_path(path: str) -> Dict[str, str]:
     )
 
 
+class FeatureRegistryClientV1Beta1WithOverride(FeatureRegistryClientWithOverride):
+    """Adds function override for v1beta1 client classes to support new Feature Store."""
+
+    _default_version = compat.V1BETA1
+
+
 class FeaturestoreClientWithOverride(ClientWithOverride):
     _is_temporary = True
     _default_version = compat.DEFAULT_VERSION

tests/unit/vertexai/conftest.py

@@ -34,6 +34,9 @@
 from google.cloud.aiplatform.compat.services import (
     feature_registry_service_client,
 )
+from google.cloud.aiplatform_v1beta1.services.feature_registry_service import (
+    FeatureRegistryServiceClient,
+)
 from google.cloud.aiplatform.compat.services import job_service_client
 from google.cloud.aiplatform.compat.types import (
     custom_job as gca_custom_job_compat,
@@ -59,6 +62,7 @@
     _TEST_FG1,
     _TEST_FG1_F1,
     _TEST_FG1_F2,
+    _TEST_FG1_FM1,
     _TEST_FV1,
     _TEST_OPTIMIZED_EMBEDDING_FV,
     _TEST_OPTIMIZED_FV1,
@@ -67,6 +71,7 @@
 )
 import pytest
 
+
 _TEST_PROJECT = "test-project"
 _TEST_PROJECT_NUMBER = "12345678"
 _TEST_LOCATION = "us-central1"
@@ -519,3 +524,13 @@ def get_feature_with_version_column_mock():
     ) as get_fg_mock:
         get_fg_mock.return_value = _TEST_FG1_F2
         yield get_fg_mock
+
+
+@pytest.fixture
+def get_feature_monitor_mock():
+    with patch.object(
+        FeatureRegistryServiceClient,
+        "get_feature_monitor",
+    ) as get_fg_mock:
+        get_fg_mock.return_value = _TEST_FG1_FM1
+        yield get_fg_mock

tests/unit/vertexai/feature_store_constants.py

@@ -352,3 +352,15 @@
 )
 
 _TEST_FG1_FEATURE_LIST = [_TEST_FG1_F1, _TEST_FG1_F2]
+
+_TEST_FG1_FM1_ID = "my_fg1_fm1"
+_TEST_FG1_FM1_PATH = (
+    f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}/featureMonitors/{_TEST_FG1_FM1_ID}"
+)
+_TEST_FG1_FM1_DESCRIPTION = "My feature monitor 1 in feature group 1"
+_TEST_FG1_FM1_LABELS = {"my_fg1_feature_monitor": "fm1"}
+_TEST_FG1_FM1 = types.feature_monitor.FeatureMonitor(
+    name=_TEST_FG1_FM1_PATH,
+    description=_TEST_FG1_FM1_DESCRIPTION,
+    labels=_TEST_FG1_FM1_LABELS,
+)

tests/unit/vertexai/test_feature_monitor.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import re
+from typing import Dict
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform import base
+from feature_store_constants import _TEST_FG1_FM1_DESCRIPTION
+from feature_store_constants import _TEST_FG1_FM1_ID
+from feature_store_constants import _TEST_FG1_FM1_LABELS
+from feature_store_constants import _TEST_FG1_FM1_PATH
+from feature_store_constants import _TEST_FG1_ID
+from feature_store_constants import _TEST_LOCATION
+from feature_store_constants import _TEST_PROJECT
+from vertexai.resources.preview import FeatureMonitor
+import pytest
+
+
+pytestmark = pytest.mark.usefixtures("google_auth_mock")
+
+
+def feature_monitor_eq(
+    feature_monitor_to_check: FeatureMonitor,
+    name: str,
+    resource_name: str,
+    project: str,
+    location: str,
+    description: str,
+    labels: Dict[str, str],
+):
+    """Check if a Feature Monitor has the appropriate values set."""
+    assert feature_monitor_to_check.name == name
+    assert feature_monitor_to_check.resource_name == resource_name
+    assert feature_monitor_to_check.project == project
+    assert feature_monitor_to_check.location == location
+    assert feature_monitor_to_check.description == description
+    assert feature_monitor_to_check.labels == labels
+
+
+def test_init_with_feature_monitor_id_and_no_fg_id_raises_error():
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Since feature monitor 'my_fg1_fm1' is not provided as a path, please"
+            " specify feature_group_id."
+        ),
+    ):
+        FeatureMonitor(_TEST_FG1_FM1_ID)
+
+
+def test_init_with_feature_monitor_path_and_fg_id_raises_error():
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Since feature monitor 'projects/test-project/locations/us-central1/"
+            "featureGroups/my_fg1/featureMonitors/my_fg1_fm1' is provided as a "
+            "path, feature_group_id should not be specified."
+        ),
+    ):
+        FeatureMonitor(_TEST_FG1_FM1_PATH, feature_group_id=_TEST_FG1_ID)
+
+
+def test_init_with_feature_monitor_id(get_feature_monitor_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature_monitor = FeatureMonitor(_TEST_FG1_FM1_ID, feature_group_id=_TEST_FG1_ID)
+
+    get_feature_monitor_mock.assert_called_once_with(
+        name=_TEST_FG1_FM1_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_monitor_eq(
+        feature_monitor,
+        name=_TEST_FG1_FM1_ID,
+        resource_name=_TEST_FG1_FM1_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_FM1_DESCRIPTION,
+        labels=_TEST_FG1_FM1_LABELS,
+    )
+
+
+def test_init_with_feature_path(get_feature_monitor_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature_monitor = FeatureMonitor(_TEST_FG1_FM1_PATH)
+
+    get_feature_monitor_mock.assert_called_once_with(
+        name=_TEST_FG1_FM1_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_monitor_eq(
+        feature_monitor,
+        name=_TEST_FG1_FM1_ID,
+        resource_name=_TEST_FG1_FM1_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_FM1_DESCRIPTION,
+        labels=_TEST_FG1_FM1_LABELS,
+    )

vertexai/generative_models/_generative_models.py

@@ -484,6 +484,7 @@ def _prepare_request(
         self,
         contents: ContentsType,
         *,
+        model: Optional[str] = None,
         generation_config: Optional[GenerationConfigType] = None,
         safety_settings: Optional[SafetySettingsType] = None,
         tools: Optional[List["Tool"]] = None,
@@ -495,6 +496,7 @@ def _prepare_request(
         if not contents:
             raise TypeError("contents must not be empty")
 
+        model = model or self._prediction_resource_name
         generation_config = generation_config or self._generation_config
         safety_settings = safety_settings or self._safety_settings
         tools = tools or self._tools
@@ -563,7 +565,7 @@ def _prepare_request(
             # The `model` parameter now needs to be set for the vision models.
             # Always need to pass the resource via the `model` parameter.
             # Even when resource is an endpoint.
-            model=self._prediction_resource_name,
+            model=model,
             contents=contents,
             generation_config=gapic_generation_config,
             safety_settings=gapic_safety_settings,
@@ -2072,11 +2074,26 @@ def __init__(
                 )
             )
 
+        def __repr__(self) -> str:
+            return self._gapic_function_calling_config.__repr__()
+
     def __init__(self, function_calling_config: "ToolConfig.FunctionCallingConfig"):
         self._gapic_tool_config = gapic_tool_types.ToolConfig(
             function_calling_config=function_calling_config._gapic_function_calling_config
         )
 
+    @classmethod
+    def _from_gapic(
+        cls,
+        gapic_tool_config: gapic_tool_types.ToolConfig,
+    ) -> "ToolConfig":
+        response = cls.__new__(cls)
+        response._gapic_tool_config = gapic_tool_config
+        return response
+
+    def __repr__(self) -> str:
+        return self._gapic_tool_config.__repr__()
+
 
 class FunctionDeclaration:
     r"""A representation of a function declaration.
@@ -3249,6 +3266,7 @@ def _prepare_request(
         self,
         contents: ContentsType,
         *,
+        model: Optional[str] = None,
         generation_config: Optional[GenerationConfigType] = None,
         safety_settings: Optional[SafetySettingsType] = None,
         tools: Optional[List["Tool"]] = None,
@@ -3259,6 +3277,7 @@ def _prepare_request(
         """Prepares a GAPIC GenerateContentRequest."""
         request_v1beta1 = super()._prepare_request(
             contents=contents,
+            model=model,
             generation_config=generation_config,
             safety_settings=safety_settings,
             tools=tools,

vertexai/resources/preview/__init__.py

@@ -40,6 +40,7 @@
     Feature,
     FeatureGroup,
     FeatureGroupBigQuerySource,
+    FeatureMonitor,
     FeatureOnlineStore,
     FeatureOnlineStoreType,
     FeatureView,
@@ -70,6 +71,7 @@
     "Feature",
     "FeatureGroup",
     "FeatureGroupBigQuerySource",
+    "FeatureMonitor",
     "FeatureOnlineStoreType",
     "FeatureOnlineStore",
     "FeatureView",

vertexai/resources/preview/feature_store/__init__.py

@@ -24,6 +24,10 @@
     FeatureGroup,
 )
 
+from vertexai.resources.preview.feature_store.feature_monitor import (
+    FeatureMonitor,
+)
+
 from vertexai.resources.preview.feature_store.feature_online_store import (
     FeatureOnlineStore,
     FeatureOnlineStoreType,
@@ -48,6 +52,7 @@
     Feature,
     FeatureGroup,
     FeatureGroupBigQuerySource,
+    FeatureMonitor,
     FeatureOnlineStoreType,
     FeatureOnlineStore,
     FeatureView,

vertexai/resources/preview/feature_store/feature_group.py

@@ -36,6 +36,9 @@
 from vertexai.resources.preview.feature_store import (
     Feature,
 )
+from vertexai.resources.preview.feature_store.feature_monitor import (
+    FeatureMonitor,
+)
 
 
 _LOGGER = base.Logger(__name__)
@@ -398,6 +401,33 @@ def list_features(
             credentials=credentials,
         )
 
+    def get_feature_monitor(
+        self,
+        feature_monitor_id: str,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> FeatureMonitor:
+        """Retrieves an existing feature monitor.
+
+        Args:
+            feature_monitor_id: The ID of the feature monitor.
+            credentials:
+                Custom credentials to use to retrieve the feature monitor under this
+                feature group. The order of which credentials are used is as
+                follows: (1) this parameter (2) credentials passed to FeatureGroup
+                constructor (3) credentials set in aiplatform.init.
+
+        Returns:
+            FeatureMonitor - the Feature Monitor resource object under this
+            feature group.
+        """
+        credentials = (
+            credentials or self.credentials or initializer.global_config.credentials
+        )
+        return FeatureMonitor(
+            f"{self.resource_name}/featureMonitors/{feature_monitor_id}",
+            credentials=credentials,
+        )
+
     @property
     def source(self) -> FeatureGroupBigQuerySource:
         return FeatureGroupBigQuerySource(