Add GRPC pinecone integration (dgarnitz#51)

Update requirements.txt to solve dependencies
sabakilai · Sep 14, 2023 · 74b1f50 · 74b1f50
1 parent 78dba8e
commit 74b1f50
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/worker/requirements.txt b/src/worker/requirements.txt
@@ -14,6 +14,8 @@ dnspython==2.4.1
 environs==9.5.0
 exceptiongroup==1.1.2
 frozenlist==1.4.0
+googleapis-common-protos==1.60.0
+grpc-gateway-protoc-gen-openapiv2==0.1.0
 grpcio
 grpcio-tools
 h11==0.14.0
@@ -24,6 +26,7 @@ httpx==0.24.1
 hyperframe==6.0.1
 idna==3.4
 loguru==0.7.0
+lz4==4.3.2
 Mako==1.2.4
 MarkupSafe==2.1.3
 marshmallow==3.20.1
@@ -35,7 +38,7 @@ pandas==2.0.3
 pika==1.3.2
 pinecone-client==2.2.2
 portalocker==2.7.0
-protobuf==4.24.0
+protobuf==3.20.2
 psycopg2-binary==2.9.6
 pycparser==2.21
 pydantic==2.1.1

diff --git a/src/worker/vdb_upload_worker.py b/src/worker/vdb_upload_worker.py
@@ -75,7 +75,7 @@ def create_pinecone_source_chunk_dict(text_embeddings_list, batch_id, job_id):
 def write_embeddings_to_pinecone(upsert_list, vector_db_metadata):
     pinecone_api_key = os.getenv('VECTOR_DB_KEY')
     pinecone.init(api_key=pinecone_api_key, environment=vector_db_metadata.environment)
-    index = pinecone.Index(vector_db_metadata.index_name)
+    index = pinecone.GRPCIndex(vector_db_metadata.index_name)
     if not index:
         logging.error(f"Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}")
         return None
@@ -88,7 +88,7 @@ def write_embeddings_to_pinecone(upsert_list, vector_db_metadata):
     for i in range(0,len(upsert_list), batch_size):
         try:
             upsert_response = index.upsert(vectors=upsert_list[i:i+batch_size])
-            vectors_uploaded += upsert_response["upserted_count"]
+            vectors_uploaded += upsert_response.upserted_count
         except Exception as e:
             logging.error('Error writing embeddings to pinecone:', e)
             return None