diff --git a/agentplatform/preview/rag/utils/_gapic_utils.py b/agentplatform/preview/rag/utils/_gapic_utils.py index ef3d166730..d3e5611c46 100644 --- a/agentplatform/preview/rag/utils/_gapic_utils.py +++ b/agentplatform/preview/rag/utils/_gapic_utils.py @@ -1298,7 +1298,7 @@ def set_backend_config( if backend_config.vector_db is not None: vector_config = backend_config.vector_db if vector_config is None: - rag_corpus.vector_db_config.rag_managed_db.CopyFrom( + rag_corpus.vector_db_config.rag_managed_db = ( GapicRagVectorDbConfig.RagManagedDb() ) elif isinstance(vector_config, RagManagedDb): diff --git a/agentplatform/rag/utils/_gapic_utils.py b/agentplatform/rag/utils/_gapic_utils.py index 1077f47592..0a6476b943 100644 --- a/agentplatform/rag/utils/_gapic_utils.py +++ b/agentplatform/rag/utils/_gapic_utils.py @@ -700,7 +700,7 @@ def set_backend_config( if backend_config.vector_db is not None: vector_config = backend_config.vector_db if isinstance(vector_config, RagManagedDb): - rag_corpus.vector_db_config.rag_managed_db.CopyFrom( + rag_corpus.vector_db_config.rag_managed_db = ( GapicRagVectorDbConfig.RagManagedDb() ) elif isinstance(vector_config, VertexVectorSearch): diff --git a/tests/unit/vertex_rag/test_rag_data.py b/tests/unit/vertex_rag/test_rag_data.py index 763d03a92d..41d406f922 100644 --- a/tests/unit/vertex_rag/test_rag_data.py +++ b/tests/unit/vertex_rag/test_rag_data.py @@ -490,6 +490,18 @@ def test_create_corpus_success(self): rag_corpus_eq(rag_corpus, test_rag_constants.TEST_RAG_CORPUS) + @pytest.mark.usefixtures("create_rag_corpus_mock") + def test_create_corpus_rag_managed_db_success(self): + backend_config = rag.RagVectorDbConfig( + vector_db=rag.RagManagedDb() + ) + rag_corpus = rag.create_corpus( + display_name=test_rag_constants.TEST_CORPUS_DISPLAY_NAME, + backend_config=backend_config, + ) + + rag_corpus_eq(rag_corpus, test_rag_constants.TEST_RAG_CORPUS) + @pytest.mark.usefixtures("create_rag_corpus_mock_vertex_vector_search") def test_create_corpus_vertex_vector_search_success(self): rag_corpus = rag.create_corpus( diff --git a/vertexai/preview/rag/utils/_gapic_utils.py b/vertexai/preview/rag/utils/_gapic_utils.py index 273b415432..ea9f003712 100644 --- a/vertexai/preview/rag/utils/_gapic_utils.py +++ b/vertexai/preview/rag/utils/_gapic_utils.py @@ -1304,7 +1304,7 @@ def set_backend_config( if backend_config.vector_db is not None: vector_config = backend_config.vector_db if vector_config is None: - rag_corpus.vector_db_config.rag_managed_db.CopyFrom( + rag_corpus.vector_db_config.rag_managed_db = ( GapicRagVectorDbConfig.RagManagedDb() ) elif isinstance(vector_config, RagManagedDb): diff --git a/vertexai/rag/utils/_gapic_utils.py b/vertexai/rag/utils/_gapic_utils.py index ea140a7c88..83b63cc376 100644 --- a/vertexai/rag/utils/_gapic_utils.py +++ b/vertexai/rag/utils/_gapic_utils.py @@ -702,7 +702,7 @@ def set_backend_config( if backend_config.vector_db is not None: vector_config = backend_config.vector_db if isinstance(vector_config, RagManagedDb): - rag_corpus.vector_db_config.rag_managed_db.CopyFrom( + rag_corpus.vector_db_config.rag_managed_db = ( GapicRagVectorDbConfig.RagManagedDb() ) elif isinstance(vector_config, VertexVectorSearch):