Skip to content

Commit 6083b93

Browse files
speedstorm1copybara-github
authored andcommitted
fix: Fix rag corpus creation error
PiperOrigin-RevId: 676157340
1 parent 47c0416 commit 6083b93

File tree

1 file changed

+32
-4
lines changed

1 file changed

+32
-4
lines changed

vertexai/preview/rag/utils/_gapic_utils.py

+32-4
Original file line numberDiff line numberDiff line change
@@ -98,26 +98,54 @@ def convert_gapic_to_embedding_model_config(
9898
return embedding_model_config
9999

100100

101+
def _check_weaviate(gapic_vector_db: RagVectorDbConfig) -> bool:
102+
try:
103+
return gapic_vector_db.__contains__("weaviate")
104+
except AttributeError:
105+
return gapic_vector_db.weaviate.ByteSize() > 0
106+
107+
108+
def _check_vertex_feature_store(gapic_vector_db: RagVectorDbConfig) -> bool:
109+
try:
110+
return gapic_vector_db.__contains__("vertex_feature_store")
111+
except AttributeError:
112+
return gapic_vector_db.vertex_feature_store.ByteSize() > 0
113+
114+
115+
def _check_pinecone(gapic_vector_db: RagVectorDbConfig) -> bool:
116+
try:
117+
return gapic_vector_db.__contains__("pinecone")
118+
except AttributeError:
119+
return gapic_vector_db.pinecone.ByteSize() > 0
120+
121+
122+
def _check_vertex_vector_search(gapic_vector_db: RagVectorDbConfig) -> bool:
123+
try:
124+
return gapic_vector_db.__contains__("vertex_vector_search")
125+
except AttributeError:
126+
return gapic_vector_db.vertex_vector_search.ByteSize() > 0
127+
128+
101129
def convert_gapic_to_vector_db(
102130
gapic_vector_db: RagVectorDbConfig,
103131
) -> Union[Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone]:
104132
"""Convert Gapic RagVectorDbConfig to Weaviate, VertexFeatureStore, VertexVectorSearch, or Pinecone."""
105-
if gapic_vector_db.__contains__("weaviate"):
133+
if _check_weaviate(gapic_vector_db):
106134
return Weaviate(
107135
weaviate_http_endpoint=gapic_vector_db.weaviate.http_endpoint,
108136
collection_name=gapic_vector_db.weaviate.collection_name,
109137
api_key=gapic_vector_db.api_auth.api_key_config.api_key_secret_version,
110138
)
111-
elif gapic_vector_db.__contains__("vertex_feature_store"):
139+
elif _check_vertex_feature_store(gapic_vector_db):
112140
return VertexFeatureStore(
113141
resource_name=gapic_vector_db.vertex_feature_store.feature_view_resource_name,
114142
)
115-
elif gapic_vector_db.__contains__("pinecone"):
143+
elif _check_pinecone(gapic_vector_db):
116144
return Pinecone(
117145
index_name=gapic_vector_db.pinecone.index_name,
118146
api_key=gapic_vector_db.api_auth.api_key_config.api_key_secret_version,
119147
)
120-
elif gapic_vector_db.__contains__("vertex_vector_search"):
148+
elif _check_vertex_vector_search(gapic_vector_db):
121149
return VertexVectorSearch(
122150
index_endpoint=gapic_vector_db.vertex_vector_search.index_endpoint,
123151
index=gapic_vector_db.vertex_vector_search.index,

0 commit comments

Comments
 (0)