From da79e218cacfa39e00d57ada7f0ec2b12fa35a84 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 9 Dec 2025 00:31:20 -0800 Subject: [PATCH] feat: Adding RagManagedVertexVectorSearch Vector DB option for RAG corpuses to SDK PiperOrigin-RevId: 842103706 --- .../vertex_rag/test_rag_constants_preview.py | 25 +++++++++++++++ .../unit/vertex_rag/test_rag_data_preview.py | 29 +++++++++++++++++ vertexai/preview/rag/__init__.py | 2 ++ vertexai/preview/rag/rag_data.py | 2 ++ vertexai/preview/rag/utils/_gapic_utils.py | 31 +++++++++++++++++-- vertexai/preview/rag/utils/resources.py | 25 +++++++++++++-- 6 files changed, 110 insertions(+), 4 deletions(-) diff --git a/tests/unit/vertex_rag/test_rag_constants_preview.py b/tests/unit/vertex_rag/test_rag_constants_preview.py index 9fd79fe457..594a4c3cad 100644 --- a/tests/unit/vertex_rag/test_rag_constants_preview.py +++ b/tests/unit/vertex_rag/test_rag_constants_preview.py @@ -62,6 +62,7 @@ RagFile, RagManagedDb, RagManagedDbConfig, + RagManagedVertexVectorSearch, RagResource, RagRetrievalConfig, RagVectorDbConfig, @@ -137,6 +138,12 @@ display_name=TEST_CORPUS_DISPLAY_NAME, description=TEST_CORPUS_DISCRIPTION, ) +TEST_RAG_MANAGED_VERTEX_VECTOR_SEARCH_COLLECTION_NAME = ( + "test-rag-managed-vertex-vector-search-collection" +) +TEST_RAG_MANAGED_VERTEX_VECTOR_SEARCH_CONFIG = RagManagedVertexVectorSearch( + collection_name=TEST_RAG_MANAGED_VERTEX_VECTOR_SEARCH_COLLECTION_NAME +) TEST_GAPIC_RAG_CORPUS.rag_embedding_model_config.vertex_prediction_endpoint.endpoint = ( "projects/{}/locations/{}/publishers/google/models/textembedding-gecko".format( TEST_PROJECT, TEST_REGION @@ -200,6 +207,17 @@ ), ), ) +TEST_GAPIC_RAG_CORPUS_RAG_MANAGED_VERTEX_VECTOR_SEARCH = GapicRagCorpus( + name=TEST_RAG_CORPUS_RESOURCE_NAME, + display_name=TEST_CORPUS_DISPLAY_NAME, + description=TEST_CORPUS_DISCRIPTION, + rag_vector_db_config=GapicRagVectorDbConfig( + rag_managed_vertex_vector_search=GapicRagVectorDbConfig.RagManagedVertexVectorSearch( + collection_name=TEST_RAG_MANAGED_VERTEX_VECTOR_SEARCH_COLLECTION_NAME, + ), + ), +) + TEST_GAPIC_RAG_CORPUS_RAG_MANAGED_DB = GapicRagCorpus( name=TEST_RAG_CORPUS_RESOURCE_NAME, display_name=TEST_CORPUS_DISPLAY_NAME, @@ -301,6 +319,13 @@ description=TEST_CORPUS_DISCRIPTION, vector_db=TEST_VERTEX_VECTOR_SEARCH_CONFIG, ) +TEST_RAG_CORPUS_RAG_MANAGED_VERTEX_VECTOR_SEARCH = RagCorpus( + name=TEST_RAG_CORPUS_RESOURCE_NAME, + display_name=TEST_CORPUS_DISPLAY_NAME, + description=TEST_CORPUS_DISCRIPTION, + vector_db=TEST_RAG_MANAGED_VERTEX_VECTOR_SEARCH_CONFIG, +) + TEST_PAGE_TOKEN = "test-page-token" # Backend Config TEST_GAPIC_RAG_CORPUS_BACKEND_CONFIG = GapicRagCorpus( diff --git a/tests/unit/vertex_rag/test_rag_data_preview.py b/tests/unit/vertex_rag/test_rag_data_preview.py index cbd6bdd30c..b887d40f52 100644 --- a/tests/unit/vertex_rag/test_rag_data_preview.py +++ b/tests/unit/vertex_rag/test_rag_data_preview.py @@ -164,6 +164,23 @@ def create_rag_corpus_mock_pinecone(): yield create_rag_corpus_mock_pinecone +@pytest.fixture +def create_rag_corpus_mock_rag_managed_vertex_vector_search(): + with mock.patch.object( + VertexRagDataServiceClient, + "create_rag_corpus", + ) as create_rag_corpus_mock_rag_managed_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_RAG_MANAGED_VERTEX_VECTOR_SEARCH + ) + create_rag_corpus_mock_rag_managed_vertex_vector_search.return_value = ( + create_rag_corpus_lro_mock + ) + yield create_rag_corpus_mock_rag_managed_vertex_vector_search + + @pytest.fixture def create_rag_corpus_mock_rag_managed_db(): with mock.patch.object( @@ -805,6 +822,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_rag_managed_vertex_vector_search") + def test_create_corpus_rag_managed_vertex_vector_search_success(self): + rag_corpus = rag.create_corpus( + display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME, + vector_db=test_rag_constants_preview.TEST_RAG_MANAGED_VERTEX_VECTOR_SEARCH_CONFIG, + ) + + rag_corpus_eq( + rag_corpus, + test_rag_constants_preview.TEST_RAG_CORPUS_RAG_MANAGED_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( diff --git a/vertexai/preview/rag/__init__.py b/vertexai/preview/rag/__init__.py index 9140c49fe4..44d9644901 100644 --- a/vertexai/preview/rag/__init__.py +++ b/vertexai/preview/rag/__init__.py @@ -59,6 +59,7 @@ RagFile, RagManagedDb, RagManagedDbConfig, + RagManagedVertexVectorSearch, RagResource, RagRetrievalConfig, RagVectorDbConfig, @@ -102,6 +103,7 @@ "RagFile", "RagManagedDb", "RagManagedDbConfig", + "RagManagedVertexVectorSearch", "RagResource", "RagRetrievalConfig", "RagVectorDbConfig", diff --git a/vertexai/preview/rag/rag_data.py b/vertexai/preview/rag/rag_data.py index 731d6e3e84..3b1d9553c6 100644 --- a/vertexai/preview/rag/rag_data.py +++ b/vertexai/preview/rag/rag_data.py @@ -56,6 +56,7 @@ RagEngineConfig, RagFile, RagManagedDb, + RagManagedVertexVectorSearch, RagVectorDbConfig, SharePointSources, SlackChannelsSource, @@ -79,6 +80,7 @@ def create_corpus( VertexVectorSearch, Pinecone, RagManagedDb, + RagManagedVertexVectorSearch, ] ] = None, vertex_ai_search_config: Optional[VertexAiSearchConfig] = None, diff --git a/vertexai/preview/rag/utils/_gapic_utils.py b/vertexai/preview/rag/utils/_gapic_utils.py index 4730c37b95..d003e866d3 100644 --- a/vertexai/preview/rag/utils/_gapic_utils.py +++ b/vertexai/preview/rag/utils/_gapic_utils.py @@ -59,6 +59,7 @@ RagFile, RagManagedDb, RagManagedDbConfig, + RagManagedVertexVectorSearch, RagVectorDbConfig, Basic, Enterprise, @@ -176,6 +177,15 @@ def _check_vertex_vector_search(gapic_vector_db: GapicRagVectorDbConfig) -> bool return gapic_vector_db.vertex_vector_search.ByteSize() > 0 +def _check_rag_managed_vertex_vector_search( + gapic_vector_db: GapicRagVectorDbConfig, +) -> bool: + try: + return gapic_vector_db.__contains__("rag_managed_vertex_vector_search") + except AttributeError: + return gapic_vector_db.rag_managed_vertex_vector_search.ByteSize() > 0 + + def _check_rag_embedding_model_config( gapic_vector_db: GapicRagVectorDbConfig, ) -> bool: @@ -240,6 +250,10 @@ def convert_gapic_to_vector_db( index_name=gapic_vector_db.pinecone.index_name, api_key=gapic_vector_db.api_auth.api_key_config.api_key_secret_version, ) + elif _check_rag_managed_vertex_vector_search(gapic_vector_db): + return RagManagedVertexVectorSearch( + collection_name=gapic_vector_db.rag_managed_vertex_vector_search.collection_name, + ) elif _check_vertex_vector_search(gapic_vector_db): return VertexVectorSearch( index_endpoint=gapic_vector_db.vertex_vector_search.index_endpoint, @@ -299,6 +313,10 @@ def convert_gapic_to_backend_config( index_name=gapic_vector_db.pinecone.index_name, api_key=gapic_vector_db.api_auth.api_key_config.api_key_secret_version, ) + elif _check_rag_managed_vertex_vector_search(gapic_vector_db): + vector_config.vector_db = RagManagedVertexVectorSearch( + collection_name=gapic_vector_db.rag_managed_vertex_vector_search.collection_name, + ) elif _check_vertex_vector_search(gapic_vector_db): vector_config.vector_db = VertexVectorSearch( index_endpoint=gapic_vector_db.vertex_vector_search.index_endpoint, @@ -904,9 +922,14 @@ def set_vector_db( ), ), ) + elif isinstance(vector_db, RagManagedVertexVectorSearch): + rag_corpus.rag_vector_db_config = GapicRagVectorDbConfig( + rag_managed_vertex_vector_search=GapicRagVectorDbConfig.RagManagedVertexVectorSearch(), + ) + else: raise TypeError( - "vector_db must be a Weaviate, VertexFeatureStore, VertexVectorSearch, RagManagedDb, or Pinecone." + "vector_db must be a Weaviate, VertexFeatureStore, VertexVectorSearch, RagManagedDb, Pinecone, or RagManagedVertexVectorSearch." ) @@ -973,10 +996,14 @@ def set_backend_config( rag_corpus.vector_db_config.api_auth.api_key_config.api_key_secret_version = ( api_key ) + elif isinstance(vector_config, RagManagedVertexVectorSearch): + rag_corpus.vector_db_config.rag_managed_vertex_vector_search.CopyFrom( + GapicRagVectorDbConfig.RagManagedVertexVectorSearch() + ) else: raise TypeError( "backend_config must be a VertexFeatureStore," - "RagManagedDb, or Pinecone." + "RagManagedDb, Pinecone, or RagManagedVertexVectorSearch." ) if backend_config.rag_embedding_model_config: set_embedding_model_config( diff --git a/vertexai/preview/rag/utils/resources.py b/vertexai/preview/rag/utils/resources.py index 459c924e13..613db44eec 100644 --- a/vertexai/preview/rag/utils/resources.py +++ b/vertexai/preview/rag/utils/resources.py @@ -225,6 +225,20 @@ class Pinecone: api_key: Optional[str] = None +@dataclasses.dataclass +class RagManagedVertexVectorSearch: + """RagManagedVertexVectorSearch. + + Attributes: + collection_name: The resource name of the Vector Search 2.0 Collection that + RAG Created for the corpus. Only populated after the corpus is successfully + created. Format: + ``projects/{project}/locations/{location}/collections/{collection_id}`` + """ + + collection_name: Optional[str] = None + + @dataclasses.dataclass class VertexAiSearchConfig: """VertexAiSearchConfig. @@ -246,12 +260,19 @@ class RagVectorDbConfig: Attributes: vector_db: Can be one of the following: Weaviate, VertexFeatureStore, - VertexVectorSearch, Pinecone, RagManagedDb. + VertexVectorSearch, Pinecone, RagManagedDb, RagManagedVertexVectorSearch. rag_embedding_model_config: The embedding model config of the Vector DB. """ vector_db: Optional[ - Union[Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone, RagManagedDb] + Union[ + Weaviate, + VertexFeatureStore, + VertexVectorSearch, + Pinecone, + RagManagedDb, + RagManagedVertexVectorSearch, + ] ] = None rag_embedding_model_config: Optional[RagEmbeddingModelConfig] = None