Content-Length: 862158 | pFad | https://github.com/googleapis/python-aiplatform/commit/216a30ff5a687e270e840a8f564baff3d79d48e1

1E feat: Update v1beta1 sdk for RagVectorDbConfig & RagCorpus protos. Ba… · googleapis/python-aiplatform@216a30f · GitHub
Skip to content

Commit 216a30f

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Update v1beta1 sdk for RagVectorDbConfig & RagCorpus protos. Backward compatible with older fields.
PiperOrigin-RevId: 702069384
1 parent 00d93bd commit 216a30f

File tree

6 files changed

+430
-29
lines changed

6 files changed

+430
-29
lines changed

‎tests/unit/vertex_rag/test_rag_constants_preview.py

Lines changed: 76 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@
4040
VertexAiSearchConfig,
4141
VertexVectorSearch,
4242
VertexFeatureStore,
43+
RagEmbeddingModelConfig,
44+
VertexPredictionEndpoint,
45+
RagVectorDbConfig,
4346
)
4447
from google.cloud.aiplatform_v1beta1 import (
4548
GoogleDriveSource,
@@ -56,7 +59,7 @@
5659
SlackSource as GapicSlackSource,
5760
RagContexts,
5861
RetrieveContextsResponse,
59-
RagVectorDbConfig,
62+
RagVectorDbConfig as GapicRagVectorDbConfig,
6063
VertexAiSearchConfig as GapicVertexAiSearchConfig,
6164
)
6265
from google.cloud.aiplatform_v1beta1.types import api_auth
@@ -112,8 +115,8 @@
112115
name=TEST_RAG_CORPUS_RESOURCE_NAME,
113116
display_name=TEST_CORPUS_DISPLAY_NAME,
114117
description=TEST_CORPUS_DISCRIPTION,
115-
rag_vector_db_config=RagVectorDbConfig(
116-
weaviate=RagVectorDbConfig.Weaviate(
118+
rag_vector_db_config=GapicRagVectorDbConfig(
119+
weaviate=GapicRagVectorDbConfig.Weaviate(
117120
http_endpoint=TEST_WEAVIATE_HTTP_ENDPOINT,
118121
collection_name=TEST_WEAVIATE_COLLECTION_NAME,
119122
),
@@ -128,8 +131,8 @@
128131
name=TEST_RAG_CORPUS_RESOURCE_NAME,
129132
display_name=TEST_CORPUS_DISPLAY_NAME,
130133
description=TEST_CORPUS_DISCRIPTION,
131-
rag_vector_db_config=RagVectorDbConfig(
132-
vertex_feature_store=RagVectorDbConfig.VertexFeatureStore(
134+
rag_vector_db_config=GapicRagVectorDbConfig(
135+
vertex_feature_store=GapicRagVectorDbConfig.VertexFeatureStore(
133136
feature_view_resource_name=TEST_VERTEX_FEATURE_STORE_RESOURCE_NAME
134137
),
135138
),
@@ -138,8 +141,8 @@
138141
name=TEST_RAG_CORPUS_RESOURCE_NAME,
139142
display_name=TEST_CORPUS_DISPLAY_NAME,
140143
description=TEST_CORPUS_DISCRIPTION,
141-
rag_vector_db_config=RagVectorDbConfig(
142-
vertex_vector_search=RagVectorDbConfig.VertexVectorSearch(
144+
rag_vector_db_config=GapicRagVectorDbConfig(
145+
vertex_vector_search=GapicRagVectorDbConfig.VertexVectorSearch(
143146
index_endpoint=TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT,
144147
index=TEST_VERTEX_VECTOR_SEARCH_INDEX,
145148
),
@@ -149,8 +152,8 @@
149152
name=TEST_RAG_CORPUS_RESOURCE_NAME,
150153
display_name=TEST_CORPUS_DISPLAY_NAME,
151154
description=TEST_CORPUS_DISCRIPTION,
152-
rag_vector_db_config=RagVectorDbConfig(
153-
pinecone=RagVectorDbConfig.Pinecone(index_name=TEST_PINECONE_INDEX_NAME),
155+
rag_vector_db_config=GapicRagVectorDbConfig(
156+
pinecone=GapicRagVectorDbConfig.Pinecone(index_name=TEST_PINECONE_INDEX_NAME),
154157
api_auth=api_auth.ApiAuth(
155158
api_key_config=api_auth.ApiAuth.ApiKeyConfig(
156159
api_key_secret_version=TEST_PINECONE_API_KEY_SECRET_VERSION
@@ -161,6 +164,14 @@
161164
TEST_EMBEDDING_MODEL_CONFIG = EmbeddingModelConfig(
162165
publisher_model="publishers/google/models/textembedding-gecko",
163166
)
167+
TEST_RAG_EMBEDDING_MODEL_CONFIG = RagEmbeddingModelConfig(
168+
vertex_prediction_endpoint=VertexPredictionEndpoint(
169+
publisher_model="publishers/google/models/textembedding-gecko",
170+
),
171+
)
172+
TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG = RagVectorDbConfig(
173+
rag_embedding_model_config=TEST_RAG_EMBEDDING_MODEL_CONFIG,
174+
)
164175
TEST_VERTEX_FEATURE_STORE_CONFIG = VertexFeatureStore(
165176
resource_name=TEST_VERTEX_FEATURE_STORE_RESOURCE_NAME,
166177
)
@@ -195,6 +206,62 @@
195206
vector_db=TEST_VERTEX_VECTOR_SEARCH_CONFIG,
196207
)
197208
TEST_PAGE_TOKEN = "test-page-token"
209+
# Backend Config
210+
TEST_GAPIC_RAG_CORPUS_BACKEND_CONFIG = GapicRagCorpus(
211+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
212+
display_name=TEST_CORPUS_DISPLAY_NAME,
213+
description=TEST_CORPUS_DISCRIPTION,
214+
)
215+
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(
216+
TEST_PROJECT, TEST_REGION
217+
)
218+
TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND_CONFIG = GapicRagCorpus(
219+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
220+
display_name=TEST_CORPUS_DISPLAY_NAME,
221+
description=TEST_CORPUS_DISCRIPTION,
222+
vector_db_config=GapicRagVectorDbConfig(
223+
vertex_vector_search=GapicRagVectorDbConfig.VertexVectorSearch(
224+
index_endpoint=TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT,
225+
index=TEST_VERTEX_VECTOR_SEARCH_INDEX,
226+
),
227+
),
228+
)
229+
TEST_GAPIC_RAG_CORPUS_PINECONE_BACKEND_CONFIG = GapicRagCorpus(
230+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
231+
display_name=TEST_CORPUS_DISPLAY_NAME,
232+
description=TEST_CORPUS_DISCRIPTION,
233+
vector_db_config=GapicRagVectorDbConfig(
234+
pinecone=GapicRagVectorDbConfig.Pinecone(index_name=TEST_PINECONE_INDEX_NAME),
235+
api_auth=api_auth.ApiAuth(
236+
api_key_config=api_auth.ApiAuth.ApiKeyConfig(
237+
api_key_secret_version=TEST_PINECONE_API_KEY_SECRET_VERSION
238+
),
239+
),
240+
),
241+
)
242+
TEST_RAG_CORPUS_BACKEND = RagCorpus(
243+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
244+
display_name=TEST_CORPUS_DISPLAY_NAME,
245+
backend_config=TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
246+
)
247+
TEST_BACKEND_CONFIG_PINECONE_CONFIG = RagVectorDbConfig(
248+
vector_db=TEST_PINECONE_CONFIG,
249+
)
250+
TEST_RAG_CORPUS_PINECONE_BACKEND = RagCorpus(
251+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
252+
display_name=TEST_CORPUS_DISPLAY_NAME,
253+
description=TEST_CORPUS_DISCRIPTION,
254+
backend_config=TEST_BACKEND_CONFIG_PINECONE_CONFIG,
255+
)
256+
TEST_BACKEND_CONFIG_VERTEX_VECTOR_SEARCH_CONFIG = RagVectorDbConfig(
257+
vector_db=TEST_VERTEX_VECTOR_SEARCH_CONFIG,
258+
)
259+
TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND = RagCorpus(
260+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
261+
display_name=TEST_CORPUS_DISPLAY_NAME,
262+
description=TEST_CORPUS_DISCRIPTION,
263+
backend_config=TEST_BACKEND_CONFIG_VERTEX_VECTOR_SEARCH_CONFIG,
264+
)
198265
# Vertex AI Search Config
199266
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"
200267
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"

‎tests/unit/vertex_rag/test_rag_data_preview.py

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,21 @@ def create_rag_corpus_mock():
5353
yield create_rag_corpus_mock
5454

5555

56+
@pytest.fixture
57+
def create_rag_corpus_mock_backend():
58+
with mock.patch.object(
59+
VertexRagDataServiceClient,
60+
"create_rag_corpus",
61+
) as create_rag_corpus_mock:
62+
create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
63+
create_rag_corpus_lro_mock.done.return_value = True
64+
create_rag_corpus_lro_mock.result.return_value = (
65+
test_rag_constants_preview.TEST_GAPIC_RAG_CORPUS_BACKEND_CONFIG
66+
)
67+
create_rag_corpus_mock.return_value = create_rag_corpus_lro_mock
68+
yield create_rag_corpus_mock
69+
70+
5671
@pytest.fixture
5772
def create_rag_corpus_mock_weaviate():
5873
with mock.patch.object(
@@ -102,6 +117,23 @@ def create_rag_corpus_mock_vertex_vector_search():
102117
yield create_rag_corpus_mock_vertex_vector_search
103118

104119

120+
@pytest.fixture
121+
def create_rag_corpus_mock_vertex_vector_search_backend():
122+
with mock.patch.object(
123+
VertexRagDataServiceClient,
124+
"create_rag_corpus",
125+
) as create_rag_corpus_mock_vertex_vector_search:
126+
create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
127+
create_rag_corpus_lro_mock.done.return_value = True
128+
create_rag_corpus_lro_mock.result.return_value = (
129+
test_rag_constants_preview.TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND_CONFIG
130+
)
131+
create_rag_corpus_mock_vertex_vector_search.return_value = (
132+
create_rag_corpus_lro_mock
133+
)
134+
yield create_rag_corpus_mock_vertex_vector_search
135+
136+
105137
@pytest.fixture
106138
def create_rag_corpus_mock_pinecone():
107139
with mock.patch.object(
@@ -117,6 +149,21 @@ def create_rag_corpus_mock_pinecone():
117149
yield create_rag_corpus_mock_pinecone
118150

119151

152+
@pytest.fixture
153+
def create_rag_corpus_mock_pinecone_backend():
154+
with mock.patch.object(
155+
VertexRagDataServiceClient,
156+
"create_rag_corpus",
157+
) as create_rag_corpus_mock_pinecone:
158+
create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
159+
create_rag_corpus_lro_mock.done.return_value = True
160+
create_rag_corpus_lro_mock.result.return_value = (
161+
test_rag_constants_preview.TEST_GAPIC_RAG_CORPUS_PINECONE_BACKEND_CONFIG
162+
)
163+
create_rag_corpus_mock_pinecone.return_value = create_rag_corpus_lro_mock
164+
yield create_rag_corpus_mock_pinecone
165+
166+
120167
@pytest.fixture
121168
def create_rag_corpus_mock_vertex_ai_engine_search_config():
122169
with mock.patch.object(
@@ -407,6 +454,15 @@ def test_create_corpus_success(self):
407454

408455
rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS)
409456

457+
@pytest.mark.usefixtures("create_rag_corpus_mock_backend")
458+
def test_create_corpus_backend_success(self):
459+
rag_corpus = rag.create_corpus(
460+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
461+
backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
462+
)
463+
464+
rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_BACKEND)
465+
410466
@pytest.mark.usefixtures("create_rag_corpus_mock_weaviate")
411467
def test_create_corpus_weaviate_success(self):
412468
rag_corpus = rag.create_corpus(
@@ -438,6 +494,18 @@ def test_create_corpus_vertex_vector_search_success(self):
438494
rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH
439495
)
440496

497+
@pytest.mark.usefixtures("create_rag_corpus_mock_vertex_vector_search_backend")
498+
def test_create_corpus_vertex_vector_search_backend_success(self):
499+
rag_corpus = rag.create_corpus(
500+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
501+
backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_VERTEX_VECTOR_SEARCH_CONFIG,
502+
)
503+
504+
rag_corpus_eq(
505+
rag_corpus,
506+
test_rag_constants_preview.TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH_BACKEND,
507+
)
508+
441509
@pytest.mark.usefixtures("create_rag_corpus_mock_pinecone")
442510
def test_create_corpus_pinecone_success(self):
443511
rag_corpus = rag.create_corpus(
@@ -447,6 +515,43 @@ def test_create_corpus_pinecone_success(self):
447515

448516
rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_PINECONE)
449517

518+
@pytest.mark.usefixtures("create_rag_corpus_mock_pinecone_backend")
519+
def test_create_corpus_pinecone_backend_success(self):
520+
rag_corpus = rag.create_corpus(
521+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
522+
backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_PINECONE_CONFIG,
523+
)
524+
525+
rag_corpus_eq(
526+
rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_PINECONE_BACKEND
527+
)
528+
529+
def test_create_corpus_backend_config_with_embedding_model_config_failure(
530+
self,
531+
):
532+
with pytest.raises(ValueError) as e:
533+
rag.create_corpus(
534+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
535+
backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
536+
embedding_model_config=test_rag_constants_preview.TEST_EMBEDDING_MODEL_CONFIG,
537+
)
538+
e.match(
539+
"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."
540+
)
541+
542+
def test_create_corpus_backend_config_with_vector_db_failure(
543+
self,
544+
):
545+
with pytest.raises(ValueError) as e:
546+
rag.create_corpus(
547+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
548+
backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
549+
vector_db=test_rag_constants_preview.TEST_PINECONE_CONFIG,
550+
)
551+
e.match(
552+
"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."
553+
)
554+
450555
@pytest.mark.usefixtures("create_rag_corpus_mock_vertex_ai_engine_search_config")
451556
def test_create_corpus_vais_engine_search_config_success(self):
452557
rag_corpus = rag.create_corpus(
@@ -480,6 +585,17 @@ def test_create_corpus_vais_datastore_search_config_with_vector_db_failure(self)
480585
)
481586
e.match("Only one of vertex_ai_search_config or vector_db can be set.")
482587

588+
def test_create_corpus_vais_datastore_search_config_with_backend_config_failure(
589+
self,
590+
):
591+
with pytest.raises(ValueError) as e:
592+
rag.create_corpus(
593+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
594+
vertex_ai_search_config=test_rag_constants_preview.TEST_VERTEX_AI_SEARCH_CONFIG_DATASTORE,
595+
backend_config=test_rag_constants_preview.TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
596+
)
597+
e.match("Only one of vertex_ai_search_config or backend_config can be set.")
598+
483599
def test_create_corpus_vais_datastore_search_config_with_embedding_model_config_failure(
484600
self,
485601
):

‎vertexai/preview/rag/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,9 @@
5858
VertexFeatureStore,
5959
VertexVectorSearch,
6060
Weaviate,
61+
RagEmbeddingModelConfig,
62+
VertexPredictionEndpoint,
63+
RagVectorDbConfig,
6164
)
6265

6366
__all__ = (
@@ -87,6 +90,9 @@
8790
"VertexRagStore",
8891
"VertexVectorSearch",
8992
"Weaviate",
93+
"RagEmbeddingModelConfig",
94+
"VertexPredictionEndpoint",
95+
"RagVectorDbConfig",
9096
"create_corpus",
9197
"delete_corpus",
9298
"delete_file",

0 commit comments

Comments
 (0)








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: https://github.com/googleapis/python-aiplatform/commit/216a30ff5a687e270e840a8f564baff3d79d48e1

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy