feat: add vertex_rag_source to create_feature_view method

PiperOrigin-RevId: 697000979

Author: holtskinner

samples/model-builder/feature_store/create_feature_view_from_rag_source.py

@@ -0,0 +1,37 @@
+# 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
+#
+#     https://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.
+
+#  [START aiplatform_sdk_create_feature_view_from_rag_source]
+
+from google.cloud import aiplatform
+from vertexai.resources.preview import feature_store
+
+
+def create_feature_view_from_rag_source(
+    project: str,
+    location: str,
+    existing_feature_online_store_id: str,
+    feature_view_id: str,
+    bq_table_uri: str,
+):
+    aiplatform.init(project=project, location=location)
+    fos = feature_store.FeatureOnlineStore(existing_feature_online_store_id)
+    fv = fos.create_feature_view(
+        name=feature_view_id,
+        source=feature_store.utils.FeatureViewVertexRagSource(uri=bq_table_uri),
+    )
+    return fv
+
+
+#  [END aiplatform_sdk_create_feature_view_from_rag_source]

samples/model-builder/feature_store/create_feature_view_from_rag_source_test.py

@@ -0,0 +1,39 @@
+# 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
+#
+#     https://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 feature_store import create_feature_view_from_rag_source
+import test_constants as constants
+
+
+def test_create_feature_view_from_rag_source_sample(
+    mock_sdk_init,
+    mock_get_feature_online_store,
+):
+    create_feature_view_from_rag_source.create_feature_view_from_rag_source(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        existing_feature_online_store_id=constants.FEATURE_ONLINE_STORE_ID,
+        feature_view_id=constants.FEATURE_VIEW_ID,
+        bq_table_uri=constants.FEATURE_VIEW_BQ_URI,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_get_feature_online_store.assert_called_once()
+    mock_get_feature_online_store.return_value.create_feature_view.assert_called_once_with(
+        name=constants.FEATURE_VIEW_ID,
+        source=constants.FEATURE_VIEW_RAG_SOURCE,
+    )

samples/model-builder/test_constants.py

@@ -264,6 +264,11 @@
         entity_id_columns=FEATURE_VIEW_BQ_ENTITY_ID_COLUMNS,
     )
 )
+FEATURE_VIEW_RAG_SOURCE = (
+    vertexai.resources.preview.feature_store.utils.FeatureViewVertexRagSource(
+        uri=FEATURE_VIEW_BQ_URI,
+    )
+)
 FEATURE_VIEW_BQ_EMBEDDING_COLUMN = "embedding"
 FEATURE_VIEW_BQ_EMBEDDING_DIMENSIONS = 10
 FEATURE_VIEW_BQ_INDEX_CONFIG = (

tests/unit/vertexai/conftest.py

@@ -64,6 +64,7 @@
     _TEST_FG1_F2,
     _TEST_FG1_FM1,
     _TEST_FV1,
+    _TEST_FV3,
     _TEST_OPTIMIZED_EMBEDDING_FV,
     _TEST_OPTIMIZED_FV1,
     _TEST_OPTIMIZED_FV2,
@@ -432,6 +433,16 @@ def get_fv_mock():
         yield get_fv_mock
 
 
+@pytest.fixture
+def get_rag_fv_mock():
+    with patch.object(
+        feature_online_store_admin_service_client.FeatureOnlineStoreAdminServiceClient,
+        "get_feature_view",
+    ) as get_rag_fv_mock:
+        get_rag_fv_mock.return_value = _TEST_FV3
+        yield get_rag_fv_mock
+
+
 @pytest.fixture
 def create_bq_fv_mock():
     with patch.object(
@@ -444,6 +455,18 @@ def create_bq_fv_mock():
         yield create_bq_fv_mock
 
 
+@pytest.fixture
+def create_rag_fv_mock():
+    with patch.object(
+        feature_online_store_admin_service_client.FeatureOnlineStoreAdminServiceClient,
+        "create_feature_view",
+    ) as create_rag_fv_mock:
+        create_rag_fv_lro_mock = mock.Mock(ga_operation.Operation)
+        create_rag_fv_lro_mock.result.return_value = _TEST_FV3
+        create_rag_fv_mock.return_value = create_rag_fv_lro_mock
+        yield create_rag_fv_mock
+
+
 @pytest.fixture
 def create_embedding_fv_from_bq_mock():
     with patch.object(

tests/unit/vertexai/feature_store_constants.py

@@ -155,7 +155,21 @@
     labels=_TEST_FV2_LABELS,
 )
 
-_TEST_FV_LIST = [_TEST_FV1, _TEST_FV2]
+# Test feature view 3
+_TEST_FV3_ID = "my_fv3"
+_TEST_FV3_PATH = f"{_TEST_BIGTABLE_FOS1_PATH}/featureViews/my_fv3"
+_TEST_FV3_LABELS = {"my_key": "my_fv3"}
+_TEST_FV3_BQ_URI = f"bq://{_TEST_PROJECT}.my_dataset.my_table"
+_TEST_FV3 = types.feature_view.FeatureView(
+    name=_TEST_FV3_PATH,
+    vertex_rag_source=types.feature_view.FeatureView.VertexRagSource(
+        uri=_TEST_FV3_BQ_URI,
+    ),
+    labels=_TEST_FV3_LABELS,
+)
+
+
+_TEST_FV_LIST = [_TEST_FV1, _TEST_FV2, _TEST_FV3]
 
 # Test feature view sync 1
 _TEST_FV_SYNC1_ID = "my_fv_sync1"

tests/unit/vertexai/test_feature_online_store.py

@@ -47,6 +47,10 @@
     _TEST_FV1_ID,
     _TEST_FV1_LABELS,
     _TEST_FV1_PATH,
+    _TEST_FV3_BQ_URI,
+    _TEST_FV3_ID,
+    _TEST_FV3_LABELS,
+    _TEST_FV3_PATH,
     _TEST_LOCATION,
     _TEST_OPTIMIZED_EMBEDDING_FV_ID,
     _TEST_OPTIMIZED_EMBEDDING_FV_PATH,
@@ -63,6 +67,7 @@
     FeatureOnlineStore,
     FeatureOnlineStoreType,
     FeatureViewBigQuerySource,
+    FeatureViewVertexRagSource,
     IndexConfig,
     TreeAhConfig,
 )
@@ -463,7 +468,9 @@ def test_create_fv_wrong_object_type_raises_error(get_fos_mock):
 
     with pytest.raises(
         ValueError,
-        match=re.escape("Only FeatureViewBigQuerySource is a supported source."),
+        match=re.escape(
+            "Only FeatureViewBigQuerySource and FeatureViewVertexRagSource are supported sources."
+        ),
     ):
         fos.create_feature_view("bq_fv", fos)
 
@@ -594,3 +601,81 @@ def test_create_embedding_fv(
         location=_TEST_LOCATION,
         labels=_TEST_FV1_LABELS,
     )
+
+
+def test_create_rag_fv_bad_uri_raises_error(get_fos_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape("Please specify URI in Vertex RAG source."),
+    ):
+        fos.create_feature_view(
+            "rag_fv",
+            FeatureViewVertexRagSource(uri=None),
+        )
+
+
+@pytest.mark.parametrize("create_request_timeout", [None, 1.0])
+@pytest.mark.parametrize("sync", [True, False])
+def test_create_rag_fv(
+    create_request_timeout,
+    sync,
+    get_fos_mock,
+    create_rag_fv_mock,
+    get_rag_fv_mock,
+    fos_logger_mock,
+):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+
+    rag_fv = fos.create_feature_view(
+        _TEST_FV3_ID,
+        FeatureViewVertexRagSource(uri=_TEST_FV3_BQ_URI),
+        labels=_TEST_FV3_LABELS,
+        create_request_timeout=create_request_timeout,
+    )
+
+    if not sync:
+        fos.wait()
+
+    # When creating, the FeatureView object doesn't have the path set.
+    expected_fv = types.feature_view.FeatureView(
+        vertex_rag_source=types.feature_view.FeatureView.VertexRagSource(
+            uri=_TEST_FV3_BQ_URI,
+        ),
+        labels=_TEST_FV3_LABELS,
+    )
+    create_rag_fv_mock.assert_called_with(
+        parent=_TEST_BIGTABLE_FOS1_PATH,
+        feature_view=expected_fv,
+        feature_view_id=_TEST_FV3_ID,
+        metadata=(),
+        timeout=create_request_timeout,
+    )
+
+    fv_eq(
+        fv_to_check=rag_fv,
+        name=_TEST_FV3_ID,
+        resource_name=_TEST_FV3_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        labels=_TEST_FV3_LABELS,
+    )
+
+    fos_logger_mock.assert_has_calls(
+        [
+            call("Creating FeatureView"),
+            call(
+                f"Create FeatureView backing LRO: {create_rag_fv_mock.return_value.operation.name}"
+            ),
+            call(
+                "FeatureView created. Resource name: projects/test-project/locations/us-central1/featureOnlineStores/my_fos1/featureViews/my_fv3"
+            ),
+            call("To use this FeatureView in another session:"),
+            call(
+                "feature_view = aiplatform.FeatureView('projects/test-project/locations/us-central1/featureOnlineStores/my_fos1/featureViews/my_fv3')"
+            ),
+        ]
+    )

tests/unit/vertexai/test_feature_view.py

@@ -47,6 +47,9 @@
     _TEST_FV2_ID,
     _TEST_FV2_LABELS,
     _TEST_FV2_PATH,
+    _TEST_FV3_ID,
+    _TEST_FV3_LABELS,
+    _TEST_FV3_PATH,
     _TEST_FV_FETCH1,
     _TEST_FV_LIST,
     _TEST_FV_SEARCH1,
@@ -289,6 +292,14 @@ def test_list(list_fv_mock, get_fos_mock):
         location=_TEST_LOCATION,
         labels=_TEST_FV2_LABELS,
     )
+    fv_eq(
+        feature_views[2],
+        name=_TEST_FV3_ID,
+        resource_name=_TEST_FV3_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        labels=_TEST_FV3_LABELS,
+    )
 
 
 def test_delete(delete_fv_mock, fv_logger_mock, get_fos_mock, get_fv_mock, sync=True):

vertexai/resources/preview/__init__.py

@@ -46,6 +46,7 @@
     FeatureView,
     FeatureViewBigQuerySource,
     FeatureViewReadResponse,
+    FeatureViewVertexRagSource,
     IndexConfig,
     TreeAhConfig,
     BruteForceConfig,
@@ -77,6 +78,7 @@
     "FeatureView",
     "FeatureViewBigQuerySource",
     "FeatureViewReadResponse",
+    "FeatureViewVertexRagSource",
     "IndexConfig",
     "TreeAhConfig",
     "BruteForceConfig",

vertexai/resources/preview/feature_store/__init__.py

@@ -41,6 +41,7 @@
     FeatureGroupBigQuerySource,
     FeatureViewBigQuerySource,
     FeatureViewReadResponse,
+    FeatureViewVertexRagSource,
     IndexConfig,
     TreeAhConfig,
     BruteForceConfig,
@@ -58,6 +59,7 @@
     FeatureView,
     FeatureViewBigQuerySource,
     FeatureViewReadResponse,
+    FeatureViewVertexRagSource,
     IndexConfig,
     IndexConfig,
     TreeAhConfig,