From f1f0c8a21b44672eb8d27559553f0aae7fd678ba Mon Sep 17 00:00:00 2001
From: henri123lemoine <henri123lemoine@gmail.com>
Date: Sun, 2 Jul 2023 06:01:03 -0400
Subject: [PATCH] Implemented upsert_entries in the Pinecone handler

---
 src/dataset/pinecone_db_handler.py | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py
index 39e695a..b0090d3 100644
--- a/src/dataset/pinecone_db_handler.py
+++ b/src/dataset/pinecone_db_handler.py
@@ -50,11 +50,38 @@ class PineconeDB:
             ),
             batch_size=upsert_size
         )
-        
+    
+    def upsert_entries(self, entries_batch, chunks_batch, chunks_ids_batch, embeddings, upsert_size=100):
+        self.index.upsert(
+            vectors=list(
+                zip(
+                    chunks_ids_batch,
+                    embeddings.tolist(),
+                    [
+                        {
+                            'entry_id': entry['id'],
+                            'source': entry['source'],
+                            'title': entry['title'],
+                            'authors': entry['authors'],
+                            'text': chunk,
+                        }
+                        for entry in entries_batch
+                        for chunk in chunks_batch
+                    ]
+                )
+            ),
+            batch_size=upsert_size
+        )
+
     def delete_entry(self, id):
         self.index.delete(
             filter={"entry_id": {"$eq": id}}
         )
+    
+    def delete_entries(self, ids):
+        self.index.delete(
+            filter={"entry_id": {"$in": ids}}
+        )
 
     def create_index(self, replace_current_index: bool = True):
         if replace_current_index: