Content-Length: 693073 | pFad | https://github.com/googleapis/python-aiplatform/commit/f882657183e34c8e07baa4b8dc9f45ed8bca9db7

10 feat: Adding Vertex Vector Search Vector DB option for RAG corpuses t… · googleapis/python-aiplatform@f882657 · GitHub
Skip to content

Commit f882657

Browse files
speedstorm1copybara-github
speedstorm1
authored and
copybara-github
committed
feat: Adding Vertex Vector Search Vector DB option for RAG corpuses to SDK
PiperOrigin-RevId: 675710933
1 parent 07e471e commit f882657

File tree

6 files changed

+96
-6
lines changed

6 files changed

+96
-6
lines changed

‎tests/unit/vertex_rag/test_rag_constants.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
JiraSource,
3030
JiraQuery,
3131
Weaviate,
32+
VertexVectorSearch,
3233
VertexFeatureStore,
3334
)
3435
from google.cloud.aiplatform_v1beta1 import (
@@ -78,6 +79,12 @@
7879
index_name=TEST_PINECONE_INDEX_NAME,
7980
api_key=TEST_PINECONE_API_KEY_SECRET_VERSION,
8081
)
82+
TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT = "test-vector-search-index-endpoint"
83+
TEST_VERTEX_VECTOR_SEARCH_INDEX = "test-vector-search-index"
84+
TEST_VERTEX_VECTOR_SEARCH_CONFIG = VertexVectorSearch(
85+
index_endpoint=TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT,
86+
index=TEST_VERTEX_VECTOR_SEARCH_INDEX,
87+
)
8188
TEST_VERTEX_FEATURE_STORE_RESOURCE_NAME = "test-feature-view-resource-name"
8289
TEST_GAPIC_RAG_CORPUS = GapicRagCorpus(
8390
name=TEST_RAG_CORPUS_RESOURCE_NAME,
@@ -115,6 +122,17 @@
115122
),
116123
),
117124
)
125+
TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH = GapicRagCorpus(
126+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
127+
display_name=TEST_CORPUS_DISPLAY_NAME,
128+
description=TEST_CORPUS_DISCRIPTION,
129+
rag_vector_db_config=RagVectorDbConfig(
130+
vertex_vector_search=RagVectorDbConfig.VertexVectorSearch(
131+
index_endpoint=TEST_VERTEX_VECTOR_SEARCH_INDEX_ENDPOINT,
132+
index=TEST_VERTEX_VECTOR_SEARCH_INDEX,
133+
),
134+
),
135+
)
118136
TEST_GAPIC_RAG_CORPUS_PINECONE = GapicRagCorpus(
119137
name=TEST_RAG_CORPUS_RESOURCE_NAME,
120138
display_name=TEST_CORPUS_DISPLAY_NAME,
@@ -158,6 +176,12 @@
158176
description=TEST_CORPUS_DISCRIPTION,
159177
vector_db=TEST_PINECONE_CONFIG,
160178
)
179+
TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH = RagCorpus(
180+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
181+
display_name=TEST_CORPUS_DISPLAY_NAME,
182+
description=TEST_CORPUS_DISCRIPTION,
183+
vector_db=TEST_VERTEX_VECTOR_SEARCH_CONFIG,
184+
)
161185
TEST_PAGE_TOKEN = "test-page-token"
162186

163187
# RagFiles

‎tests/unit/vertex_rag/test_rag_data.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,23 @@ def create_rag_corpus_mock_vertex_feature_store():
7979
yield create_rag_corpus_mock_vertex_feature_store
8080

8181

82+
@pytest.fixture
83+
def create_rag_corpus_mock_vertex_vector_search():
84+
with mock.patch.object(
85+
VertexRagDataServiceClient,
86+
"create_rag_corpus",
87+
) as create_rag_corpus_mock_vertex_vector_search:
88+
create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
89+
create_rag_corpus_lro_mock.done.return_value = True
90+
create_rag_corpus_lro_mock.result.return_value = (
91+
tc.TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH
92+
)
93+
create_rag_corpus_mock_vertex_vector_search.return_value = (
94+
create_rag_corpus_lro_mock
95+
)
96+
yield create_rag_corpus_mock_vertex_vector_search
97+
98+
8299
@pytest.fixture
83100
def create_rag_corpus_mock_pinecone():
84101
with mock.patch.object(
@@ -257,6 +274,15 @@ def test_create_corpus_vertex_feature_store_success(self):
257274

258275
rag_corpus_eq(rag_corpus, tc.TEST_RAG_CORPUS_VERTEX_FEATURE_STORE)
259276

277+
@pytest.mark.usefixtures("create_rag_corpus_mock_vertex_vector_search")
278+
def test_create_corpus_vertex_vector_search_success(self):
279+
rag_corpus = rag.create_corpus(
280+
display_name=tc.TEST_CORPUS_DISPLAY_NAME,
281+
vector_db=tc.TEST_VERTEX_VECTOR_SEARCH_CONFIG,
282+
)
283+
284+
rag_corpus_eq(rag_corpus, tc.TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH)
285+
260286
@pytest.mark.usefixtures("create_rag_corpus_mock_pinecone")
261287
def test_create_corpus_pinecone_success(self):
262288
rag_corpus = rag.create_corpus(

‎vertexai/preview/rag/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
SlackChannel,
4848
SlackChannelsSource,
4949
VertexFeatureStore,
50+
VertexVectorSearch,
5051
Weaviate,
5152
)
5253

@@ -64,6 +65,7 @@
6465
"SlackChannelsSource",
6566
"VertexFeatureStore",
6667
"VertexRagStore",
68+
"VertexVectorSearch",
6769
"Weaviate",
6870
"create_corpus",
6971
"delete_corpus",

‎vertexai/preview/rag/rag_data.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
RagFile,
5151
SlackChannelsSource,
5252
VertexFeatureStore,
53+
VertexVectorSearch,
5354
Weaviate,
5455
)
5556

@@ -58,7 +59,9 @@ def create_corpus(
5859
display_name: Optional[str] = None,
5960
description: Optional[str] = None,
6061
embedding_model_config: Optional[EmbeddingModelConfig] = None,
61-
vector_db: Optional[Union[Weaviate, VertexFeatureStore, Pinecone]] = None,
62+
vector_db: Optional[
63+
Union[Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone]
64+
] = None,
6265
) -> RagCorpus:
6366
"""Creates a new RagCorpus resource.
6467

‎vertexai/preview/rag/utils/_gapic_utils.py

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
SlackChannelsSource,
4545
JiraSource,
4646
VertexFeatureStore,
47+
VertexVectorSearch,
4748
Weaviate,
4849
)
4950

@@ -99,8 +100,8 @@ def convert_gapic_to_embedding_model_config(
99100

100101
def convert_gapic_to_vector_db(
101102
gapic_vector_db: RagVectorDbConfig,
102-
) -> Union[Weaviate, VertexFeatureStore, Pinecone]:
103-
"""Convert Gapic RagVectorDbConfig to Weaviate, VertexFeatureStore, or Pinecone."""
103+
) -> Union[Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone]:
104+
"""Convert Gapic RagVectorDbConfig to Weaviate, VertexFeatureStore, VertexVectorSearch, or Pinecone."""
104105
if gapic_vector_db.__contains__("weaviate"):
105106
return Weaviate(
106107
weaviate_http_endpoint=gapic_vector_db.weaviate.http_endpoint,
@@ -116,6 +117,11 @@ def convert_gapic_to_vector_db(
116117
index_name=gapic_vector_db.pinecone.index_name,
117118
api_key=gapic_vector_db.api_auth.api_key_config.api_key_secret_version,
118119
)
120+
elif gapic_vector_db.__contains__("vertex_vector_search"):
121+
return VertexVectorSearch(
122+
index_endpoint=gapic_vector_db.vertex_vector_search.index_endpoint,
123+
index=gapic_vector_db.vertex_vector_search.index,
124+
)
119125
else:
120126
return None
121127

@@ -418,7 +424,7 @@ def set_embedding_model_config(
418424

419425

420426
def set_vector_db(
421-
vector_db: Union[Weaviate, VertexFeatureStore, Pinecone],
427+
vector_db: Union[Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone],
422428
rag_corpus: GapicRagCorpus,
423429
) -> None:
424430
"""Sets the vector db configuration for the rag corpus."""
@@ -446,6 +452,16 @@ def set_vector_db(
446452
feature_view_resource_name=resource_name,
447453
),
448454
)
455+
elif isinstance(vector_db, VertexVectorSearch):
456+
index_endpoint = vector_db.index_endpoint
457+
index = vector_db.index
458+
459+
rag_corpus.rag_vector_db_config = RagVectorDbConfig(
460+
vertex_vector_search=RagVectorDbConfig.VertexVectorSearch(
461+
index_endpoint=index_endpoint,
462+
index=index,
463+
),
464+
)
449465
elif isinstance(vector_db, Pinecone):
450466
index_name = vector_db.index_name
451467
api_key = vector_db.api_key
@@ -462,5 +478,5 @@ def set_vector_db(
462478
)
463479
else:
464480
raise TypeError(
465-
"vector_db must be a Weaviate, VertexFeatureStore, or Pinecone."
481+
"vector_db must be a Weaviate, VertexFeatureStore, VertexVectorSearch, or Pinecone."
466482
)

‎vertexai/preview/rag/utils/resources.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,23 @@ class VertexFeatureStore:
9898
resource_name: str
9999

100100

101+
@dataclasses.dataclass
102+
class VertexVectorSearch:
103+
"""VertexVectorSearch.
104+
105+
Attributes:
106+
index_endpoint (str):
107+
The resource name of the Index Endpoint. Format:
108+
``projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}``
109+
index (str):
110+
The resource name of the Index. Format:
111+
``projects/{project}/locations/{location}/indexes/{index}``
112+
"""
113+
114+
index_endpoint: str
115+
index: str
116+
117+
101118
@dataclasses.dataclass
102119
class Pinecone:
103120
"""Pinecone.
@@ -129,7 +146,9 @@ class RagCorpus:
129146
display_name: Optional[str] = None
130147
description: Optional[str] = None
131148
embedding_model_config: Optional[EmbeddingModelConfig] = None
132-
vector_db: Optional[Union[Weaviate, VertexFeatureStore, Pinecone]] = None
149+
vector_db: Optional[
150+
Union[Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone]
151+
] = None
133152

134153

135154
@dataclasses.dataclass

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/f882657183e34c8e07baa4b8dc9f45ed8bca9db7

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy