From ca61c5a4ed5138f01496ee1a80c4f7a221f671d3 Mon Sep 17 00:00:00 2001 From: Holt Skinner Date: Fri, 15 Nov 2024 14:25:24 -0800 Subject: [PATCH] feat: add vertex_rag_source to create_feature_view method PiperOrigin-RevId: 697000979 --- .../create_feature_view_from_rag_source.py | 37 ++++++++ ...reate_feature_view_from_rag_source_test.py | 39 +++++++++ samples/model-builder/test_constants.py | 5 ++ tests/unit/vertexai/conftest.py | 23 +++++ .../unit/vertexai/feature_store_constants.py | 16 +++- .../vertexai/test_feature_online_store.py | 87 ++++++++++++++++++- tests/unit/vertexai/test_feature_view.py | 11 +++ vertexai/resources/preview/__init__.py | 2 + .../preview/feature_store/__init__.py | 2 + .../feature_store/feature_online_store.py | 49 +++++++---- .../resources/preview/feature_store/utils.py | 6 ++ 11 files changed, 260 insertions(+), 17 deletions(-) create mode 100644 samples/model-builder/feature_store/create_feature_view_from_rag_source.py create mode 100644 samples/model-builder/feature_store/create_feature_view_from_rag_source_test.py diff --git a/samples/model-builder/feature_store/create_feature_view_from_rag_source.py b/samples/model-builder/feature_store/create_feature_view_from_rag_source.py new file mode 100644 index 0000000000..f4340969b5 --- /dev/null +++ b/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] diff --git a/samples/model-builder/feature_store/create_feature_view_from_rag_source_test.py b/samples/model-builder/feature_store/create_feature_view_from_rag_source_test.py new file mode 100644 index 0000000000..dbcc8322b7 --- /dev/null +++ b/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, + ) diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py index 6af770216c..a85c7677f9 100644 --- a/samples/model-builder/test_constants.py +++ b/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 = ( diff --git a/tests/unit/vertexai/conftest.py b/tests/unit/vertexai/conftest.py index 1aa84bf751..7d5bb4af45 100644 --- a/tests/unit/vertexai/conftest.py +++ b/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( diff --git a/tests/unit/vertexai/feature_store_constants.py b/tests/unit/vertexai/feature_store_constants.py index eacc06b2f9..1525544753 100644 --- a/tests/unit/vertexai/feature_store_constants.py +++ b/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" diff --git a/tests/unit/vertexai/test_feature_online_store.py b/tests/unit/vertexai/test_feature_online_store.py index a131041d22..dd79f245e3 100644 --- a/tests/unit/vertexai/test_feature_online_store.py +++ b/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')" + ), + ] + ) diff --git a/tests/unit/vertexai/test_feature_view.py b/tests/unit/vertexai/test_feature_view.py index 8972374370..c310bbd1e3 100644 --- a/tests/unit/vertexai/test_feature_view.py +++ b/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): diff --git a/vertexai/resources/preview/__init__.py b/vertexai/resources/preview/__init__.py index 6a719dbfb0..b7bbc36920 100644 --- a/vertexai/resources/preview/__init__.py +++ b/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", diff --git a/vertexai/resources/preview/feature_store/__init__.py b/vertexai/resources/preview/feature_store/__init__.py index 2592328827..a7c0d7023d 100644 --- a/vertexai/resources/preview/feature_store/__init__.py +++ b/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, diff --git a/vertexai/resources/preview/feature_store/feature_online_store.py b/vertexai/resources/preview/feature_store/feature_online_store.py index 205706b049..2a0ecd949c 100644 --- a/vertexai/resources/preview/feature_store/feature_online_store.py +++ b/vertexai/resources/preview/feature_store/feature_online_store.py @@ -21,6 +21,7 @@ Optional, Sequence, Tuple, + Union, ) from google.auth import credentials as auth_credentials @@ -40,6 +41,7 @@ from vertexai.resources.preview.feature_store.utils import ( IndexConfig, FeatureViewBigQuerySource, + FeatureViewVertexRagSource, ) @@ -404,7 +406,10 @@ def labels(self) -> Dict[str, str]: def create_feature_view( self, name: str, - source: FeatureViewBigQuerySource, + source: Union[ + FeatureViewBigQuerySource, + FeatureViewVertexRagSource, + ], labels: Optional[Dict[str, str]] = None, sync_config: Optional[str] = None, index_config: Optional[IndexConfig] = None, @@ -420,7 +425,7 @@ def create_feature_view( Example Usage: ``` existing_fos = FeatureOnlineStore('my_fos') - new_fv = existing_fos.create_feature_view_from_bigquery( + new_fv = existing_fos.create_feature_view( 'my_fos', BigQuerySource( uri='bq://my-proj/dataset/table', @@ -428,7 +433,7 @@ def create_feature_view( ) ) # Example for how to create an embedding FeatureView. - embedding_fv = existing_fos.create_feature_view_from_bigquery( + embedding_fv = existing_fos.create_feature_view( 'my_fos', BigQuerySource( uri='bq://my-proj/dataset/table', @@ -448,7 +453,7 @@ def create_feature_view( name: The name of the feature view. source: The source to load data from when a feature view sync runs. - Currently supports a BigQuery source. + Currently supports a BigQuery source or a Vertex RAG source. labels: The labels with user-defined metadata to organize your FeatureViews. @@ -499,22 +504,36 @@ def create_feature_view( if not source: raise ValueError("Please specify a valid source.") - # Only BigQuery source is supported right now. - if not isinstance(source, FeatureViewBigQuerySource): - raise ValueError("Only FeatureViewBigQuerySource is a supported source.") + big_query_source = None + vertex_rag_source = None - # BigQuery source validation. - if not source.uri: - raise ValueError("Please specify URI in BigQuery source.") + if isinstance(source, FeatureViewBigQuerySource): + if not source.uri: + raise ValueError("Please specify URI in BigQuery source.") - if not source.entity_id_columns: - raise ValueError("Please specify entity ID columns in BigQuery source.") + if not source.entity_id_columns: + raise ValueError("Please specify entity ID columns in BigQuery source.") - gapic_feature_view = gca_feature_view.FeatureView( - big_query_source=gca_feature_view.FeatureView.BigQuerySource( + big_query_source = gca_feature_view.FeatureView.BigQuerySource( uri=source.uri, entity_id_columns=source.entity_id_columns, - ), + ) + elif isinstance(source, FeatureViewVertexRagSource): + if not source.uri: + raise ValueError("Please specify URI in Vertex RAG source.") + + vertex_rag_source = gca_feature_view.FeatureView.VertexRagSource( + uri=source.uri, + rag_corpus_id=source.rag_corpus_id or None, + ) + else: + raise ValueError( + "Only FeatureViewBigQuerySource and FeatureViewVertexRagSource are supported sources." + ) + + gapic_feature_view = gca_feature_view.FeatureView( + big_query_source=big_query_source, + vertex_rag_source=vertex_rag_source, sync_config=gca_feature_view.FeatureView.SyncConfig(cron=sync_config) if sync_config else None, diff --git a/vertexai/resources/preview/feature_store/utils.py b/vertexai/resources/preview/feature_store/utils.py index 0896a863b2..497d0dc670 100644 --- a/vertexai/resources/preview/feature_store/utils.py +++ b/vertexai/resources/preview/feature_store/utils.py @@ -51,6 +51,12 @@ class FeatureViewBigQuerySource: entity_id_columns: List[str] +@dataclass +class FeatureViewVertexRagSource: + uri: str + rag_corpus_id: Optional[str] = None + + @dataclass(frozen=True) class ConnectionOptions: """Represents connection options used for sending RPCs to the online store."""