feat: Update v1beta1 sdk for RagVectorDbConfig & RagCorpus protos. Ba…

…ckward compatible with older fields. PiperOrigin-RevId: 702069384
googleapis · Dec 2, 2024 · 216a30f · 216a30f
1 parent 00d93bd
commit 216a30f
Show file tree

Hide file tree

Showing 6 changed files with 430 additions and 29 deletions.
diff --git a/tests/unit/vertex_rag/test_rag_constants_preview.py b/tests/unit/vertex_rag/test_rag_constants_preview.py
@@ -40,6 +40,9 @@
     VertexAiSearchConfig,
     VertexVectorSearch,
     VertexFeatureStore,
+    RagEmbeddingModelConfig,
+    VertexPredictionEndpoint,
+    RagVectorDbConfig,
 )
 from google.cloud.aiplatform_v1beta1 import (
     GoogleDriveSource,
@@ -56,7 +59,7 @@
     SlackSource as GapicSlackSource,
     RagContexts,
     RetrieveContextsResponse,
-    RagVectorDbConfig,
+    RagVectorDbConfig as GapicRagVectorDbConfig,
     VertexAiSearchConfig as GapicVertexAiSearchConfig,
 )
 from google.cloud.aiplatform_v1beta1.types import api_auth
@@ -112,8 +115,8 @@
     name=TEST_RAG_CORPUS_RESOURCE_NAME,
     display_name=TEST_CORPUS_DISPLAY_NAME,
     description=TEST_CORPUS_DISCRIPTION,
-    rag_vector_db_config=RagVectorDbConfig(
-        weaviate=RagVectorDbConfig.Weaviate(
+    rag_vector_db_config=GapicRagVectorDbConfig(
+        weaviate=GapicRagVectorDbConfig.Weaviate(
             http_endpoint=TEST_WEAVIATE_HTTP_ENDPOINT,
             collection_name=TEST_WEAVIATE_COLLECTION_NAME,
         ),
@@ -128,8 +131,8 @@
     name=TEST_RAG_CORPUS_RESOURCE_NAME,
     display_name=TEST_CORPUS_DISPLAY_NAME,
     description=TEST_CORPUS_DISCRIPTION,
-    rag_vector_db_config=RagVectorDbConfig(
-        vertex_feature_store=RagVectorDbConfig.VertexFeatureStore(
+    rag_vector_db_config=GapicRagVectorDbConfig(
+        vertex_feature_store=GapicRagVectorDbConfig.VertexFeatureStore(
             feature_view_resource_name=TEST_VERTEX_FEATURE_STORE_RESOURCE_NAME
         ),
     ),
@@ -138,8 +141,8 @@
     name=TEST_RAG_CORPUS_RESOURCE_NAME,
     display_name=TEST_CORPUS_DISPLAY_NAME,
     description=TEST_CORPUS_DISCRIPTION,
-    rag_vector_db_config=RagVectorDbConfig(
-        vertex_vector_search=RagVectorDbConfig.VertexVectorSearch(
+    rag_vector_db_config=GapicRagVectorDbConfig(
+        vertex_vector_search=GapicRagVectorDbConfig.VertexVectorSearch(
             index_endpoint=TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT,
             index=TEST_VERTEX_VECTOR_SEARCH_INDEX,
         ),
@@ -149,8 +152,8 @@
     name=TEST_RAG_CORPUS_RESOURCE_NAME,
     display_name=TEST_CORPUS_DISPLAY_NAME,
     description=TEST_CORPUS_DISCRIPTION,
-    rag_vector_db_config=RagVectorDbConfig(
-        pinecone=RagVectorDbConfig.Pinecone(index_name=TEST_PINECONE_INDEX_NAME),
+    rag_vector_db_config=GapicRagVectorDbConfig(
+        pinecone=GapicRagVectorDbConfig.Pinecone(index_name=TEST_PINECONE_INDEX_NAME),
         api_auth=api_auth.ApiAuth(
             api_key_config=api_auth.ApiAuth.ApiKeyConfig(
                 api_key_secret_version=TEST_PINECONE_API_KEY_SECRET_VERSION
@@ -161,6 +164,14 @@
 TEST_EMBEDDING_MODEL_CONFIG = EmbeddingModelConfig(
     publisher_model="publishers/google/models/textembedding-gecko",
 )
+TEST_RAG_EMBEDDING_MODEL_CONFIG = RagEmbeddingModelConfig(
+    vertex_prediction_endpoint=VertexPredictionEndpoint(
+        publisher_model="publishers/google/models/textembedding-gecko",
+    ),
+)
+TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG = RagVectorDbConfig(
+    rag_embedding_model_config=TEST_RAG_EMBEDDING_MODEL_CONFIG,
+)
 TEST_VERTEX_FEATURE_STORE_CONFIG = VertexFeatureStore(
     resource_name=TEST_VERTEX_FEATURE_STORE_RESOURCE_NAME,
 )
@@ -195,6 +206,62 @@
     vector_db=TEST_VERTEX_VECTOR_SEARCH_CONFIG,
 )
 TEST_PAGE_TOKEN = "test-page-token"
+# Backend Config
+TEST_GAPIC_RAG_CORPUS_BACKEND_CONFIG = GapicRagCorpus(
+    name=TEST_RAG_CORPUS_RESOURCE_NAME,
+    display_name=TEST_CORPUS_DISPLAY_NAME,
+    description=TEST_CORPUS_DISCRIPTION,
+)
+TEST_GAPIC_RAG_CORPUS_BACKEND_CONFIG.vector_db_config.rag_embedding_model_config.vertex_prediction_endpoint.endpoint = "projects/{}/locations/{}/publishers/google/models/textembedding-gecko".format(
+    TEST_PROJECT, TEST_REGION
+)
+TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND_CONFIG = GapicRagCorpus(
+    name=TEST_RAG_CORPUS_RESOURCE_NAME,
+    display_name=TEST_CORPUS_DISPLAY_NAME,
+    description=TEST_CORPUS_DISCRIPTION,
+    vector_db_config=GapicRagVectorDbConfig(
+        vertex_vector_search=GapicRagVectorDbConfig.VertexVectorSearch(
+            index_endpoint=TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT,
+            index=TEST_VERTEX_VECTOR_SEARCH_INDEX,
+        ),
+    ),
+)
+TEST_GAPIC_RAG_CORPUS_PINECONE_BACKEND_CONFIG = GapicRagCorpus(
+    name=TEST_RAG_CORPUS_RESOURCE_NAME,
+    display_name=TEST_CORPUS_DISPLAY_NAME,
+    description=TEST_CORPUS_DISCRIPTION,
+    vector_db_config=GapicRagVectorDbConfig(
+        pinecone=GapicRagVectorDbConfig.Pinecone(index_name=TEST_PINECONE_INDEX_NAME),
+        api_auth=api_auth.ApiAuth(
+            api_key_config=api_auth.ApiAuth.ApiKeyConfig(
+                api_key_secret_version=TEST_PINECONE_API_KEY_SECRET_VERSION
+            ),
+        ),
+    ),
+)
+TEST_RAG_CORPUS_BACKEND = RagCorpus(
+    name=TEST_RAG_CORPUS_RESOURCE_NAME,
+    display_name=TEST_CORPUS_DISPLAY_NAME,
+    backend_config=TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
+)
+TEST_BACKEND_CONFIG_PINECONE_CONFIG = RagVectorDbConfig(
+    vector_db=TEST_PINECONE_CONFIG,
+)
+TEST_RAG_CORPUS_PINECONE_BACKEND = RagCorpus(
+    name=TEST_RAG_CORPUS_RESOURCE_NAME,
+    display_name=TEST_CORPUS_DISPLAY_NAME,
+    description=TEST_CORPUS_DISCRIPTION,
+    backend_config=TEST_BACKEND_CONFIG_PINECONE_CONFIG,
+)
+TEST_BACKEND_CONFIG_VERTEX_VECTOR_SEARCH_CONFIG = RagVectorDbConfig(
+    vector_db=TEST_VERTEX_VECTOR_SEARCH_CONFIG,
+)
+TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND = RagCorpus(
+    name=TEST_RAG_CORPUS_RESOURCE_NAME,
+    display_name=TEST_CORPUS_DISPLAY_NAME,
+    description=TEST_CORPUS_DISCRIPTION,
+    backend_config=TEST_BACKEND_CONFIG_VERTEX_VECTOR_SEARCH_CONFIG,
+)
 # Vertex AI Search Config
 TEST_VERTEX_AI_SEARCH_ENGINE_SERVING_CONFIG = f"projects/{TEST_PROJECT_NUMBER}/locations/{TEST_REGION}/collections/test-collection/engines/test-engine/servingConfigs/test-serving-config"
 TEST_VERTEX_AI_SEARCH_DATASTORE_SERVING_CONFIG = f"projects/{TEST_PROJECT_NUMBER}/locations/{TEST_REGION}/collections/test-collection/dataStores/test-datastore/servingConfigs/test-serving-config"

diff --git a/tests/unit/vertex_rag/test_rag_data_preview.py b/tests/unit/vertex_rag/test_rag_data_preview.py
@@ -53,6 +53,21 @@ def create_rag_corpus_mock():
         yield create_rag_corpus_mock
 
 
+@pytest.fixture
+def create_rag_corpus_mock_backend():
+    with mock.patch.object(
+        VertexRagDataServiceClient,
+        "create_rag_corpus",
+    ) as create_rag_corpus_mock:
+        create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
+        create_rag_corpus_lro_mock.done.return_value = True
+        create_rag_corpus_lro_mock.result.return_value = (
+            test_rag_constants_preview.TEST_GAPIC_RAG_CORPUS_BACKEND_CONFIG
+        )
+        create_rag_corpus_mock.return_value = create_rag_corpus_lro_mock
+        yield create_rag_corpus_mock
+
+
 @pytest.fixture
 def create_rag_corpus_mock_weaviate():
     with mock.patch.object(
@@ -102,6 +117,23 @@ def create_rag_corpus_mock_vertex_vector_search():
         yield create_rag_corpus_mock_vertex_vector_search
 
 
+@pytest.fixture
+def create_rag_corpus_mock_vertex_vector_search_backend():
+    with mock.patch.object(
+        VertexRagDataServiceClient,
+        "create_rag_corpus",
+    ) as create_rag_corpus_mock_vertex_vector_search:
+        create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
+        create_rag_corpus_lro_mock.done.return_value = True
+        create_rag_corpus_lro_mock.result.return_value = (
+            test_rag_constants_preview.TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND_CONFIG
+        )
+        create_rag_corpus_mock_vertex_vector_search.return_value = (
+            create_rag_corpus_lro_mock
+        )
+        yield create_rag_corpus_mock_vertex_vector_search
+
+
 @pytest.fixture
 def create_rag_corpus_mock_pinecone():
     with mock.patch.object(
@@ -117,6 +149,21 @@ def create_rag_corpus_mock_pinecone():
         yield create_rag_corpus_mock_pinecone
 
 
+@pytest.fixture
+def create_rag_corpus_mock_pinecone_backend():
+    with mock.patch.object(
+        VertexRagDataServiceClient,
+        "create_rag_corpus",
+    ) as create_rag_corpus_mock_pinecone:
+        create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
+        create_rag_corpus_lro_mock.done.return_value = True
+        create_rag_corpus_lro_mock.result.return_value = (
+            test_rag_constants_preview.TEST_GAPIC_RAG_CORPUS_PINECONE_BACKEND_CONFIG
+        )
+        create_rag_corpus_mock_pinecone.return_value = create_rag_corpus_lro_mock
+        yield create_rag_corpus_mock_pinecone
+
+
 @pytest.fixture
 def create_rag_corpus_mock_vertex_ai_engine_search_config():
     with mock.patch.object(
@@ -407,6 +454,15 @@ def test_create_corpus_success(self):
 
         rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS)
 
+    @pytest.mark.usefixtures("create_rag_corpus_mock_backend")
+    def test_create_corpus_backend_success(self):
+        rag_corpus = rag.create_corpus(
+            display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+            backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
+        )
+
+        rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_BACKEND)
+
     @pytest.mark.usefixtures("create_rag_corpus_mock_weaviate")
     def test_create_corpus_weaviate_success(self):
         rag_corpus = rag.create_corpus(
@@ -438,6 +494,18 @@ def test_create_corpus_vertex_vector_search_success(self):
             rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH
         )
 
+    @pytest.mark.usefixtures("create_rag_corpus_mock_vertex_vector_search_backend")
+    def test_create_corpus_vertex_vector_search_backend_success(self):
+        rag_corpus = rag.create_corpus(
+            display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+            backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_VERTEX_VECTOR_SEARCH_CONFIG,
+        )
+
+        rag_corpus_eq(
+            rag_corpus,
+            test_rag_constants_preview.TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND,
+        )
+
     @pytest.mark.usefixtures("create_rag_corpus_mock_pinecone")
     def test_create_corpus_pinecone_success(self):
         rag_corpus = rag.create_corpus(
@@ -447,6 +515,43 @@ def test_create_corpus_pinecone_success(self):
 
         rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_PINECONE)
 
+    @pytest.mark.usefixtures("create_rag_corpus_mock_pinecone_backend")
+    def test_create_corpus_pinecone_backend_success(self):
+        rag_corpus = rag.create_corpus(
+            display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+            backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_PINECONE_CONFIG,
+        )
+
+        rag_corpus_eq(
+            rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_PINECONE_BACKEND
+        )
+
+    def test_create_corpus_backend_config_with_embedding_model_config_failure(
+        self,
+    ):
+        with pytest.raises(ValueError) as e:
+            rag.create_corpus(
+                display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+                backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
+                embedding_model_config=test_rag_constants_preview.TEST_EMBEDDING_MODEL_CONFIG,
+            )
+        e.match(
+            "Only one of backend_config or embedding_model_config and vector_db can be set. embedding_model_config and vector_db are deprecated, use backend_config instead."
+        )
+
+    def test_create_corpus_backend_config_with_vector_db_failure(
+        self,
+    ):
+        with pytest.raises(ValueError) as e:
+            rag.create_corpus(
+                display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+                backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
+                vector_db=test_rag_constants_preview.TEST_PINECONE_CONFIG,
+            )
+        e.match(
+            "Only one of backend_config or embedding_model_config and vector_db can be set. embedding_model_config and vector_db are deprecated, use backend_config instead."
+        )
+
     @pytest.mark.usefixtures("create_rag_corpus_mock_vertex_ai_engine_search_config")
     def test_create_corpus_vais_engine_search_config_success(self):
         rag_corpus = rag.create_corpus(
@@ -480,6 +585,17 @@ def test_create_corpus_vais_datastore_search_config_with_vector_db_failure(self)
             )
         e.match("Only one of vertex_ai_search_config or vector_db can be set.")
 
+    def test_create_corpus_vais_datastore_search_config_with_backend_config_failure(
+        self,
+    ):
+        with pytest.raises(ValueError) as e:
+            rag.create_corpus(
+                display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+                vertex_ai_search_config=test_rag_constants_preview.TEST_VERTEX_AI_SEARCH_CONFIG_DATASTORE,
+                backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
+            )
+        e.match("Only one of vertex_ai_search_config or backend_config can be set.")
+
     def test_create_corpus_vais_datastore_search_config_with_embedding_model_config_failure(
         self,
     ):

diff --git a/vertexai/preview/rag/__init__.py b/vertexai/preview/rag/__init__.py
@@ -58,6 +58,9 @@
     VertexFeatureStore,
     VertexVectorSearch,
     Weaviate,
+    RagEmbeddingModelConfig,
+    VertexPredictionEndpoint,
+    RagVectorDbConfig,
 )
 
 __all__ = (
@@ -87,6 +90,9 @@
     "VertexRagStore",
     "VertexVectorSearch",
     "Weaviate",
+    "RagEmbeddingModelConfig",
+    "VertexPredictionEndpoint",
+    "RagVectorDbConfig",
     "create_corpus",
     "delete_corpus",
     "delete_file",