From c76ad2df2807fd56a582efbb7271ad53c6ba00d0 Mon Sep 17 00:00:00 2001
From: henri123lemoine <henri123lemoine@gmail.com>
Date: Fri, 23 Jun 2023 18:23:10 -0400
Subject: [PATCH] Added a Pinecone database handler

---
 src/dataset/pinecone_handler.py | 79 +++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 src/dataset/pinecone_handler.py

diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py
new file mode 100644
index 0000000..786567f
--- /dev/null
+++ b/src/dataset/pinecone_handler.py
@@ -0,0 +1,79 @@
+# dataset/pinecone_handler.py
+
+import pinecone
+import os
+from typing import List
+
+
+class PineconeHandler:
+    def __init__(
+        self, 
+        index_name: str,
+        dimensions: int = 1536,
+        metric: str = "cosine",
+        location: str = "us-central1-gcp"
+    ):
+        self.index_name = index_name
+        self.dimensions = dimensions
+        self.metric = metric
+        
+        PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
+        assert PINECONE_API_KEY, "PINECONE_API_KEY environment variable not set."
+        
+        pinecone.init(
+            api_key = PINECONE_API_KEY,
+            environment = location,
+        )
+        
+        self.index = pinecone.Index(index_name=self.index_name)
+        index_stats_response = self.index.describe_index_stats()
+        
+        print(f"Index info:\n\t{index_stats_response}\n\n")
+    
+    def insert_entry(self, entry, chunks, embeddings, upsert_size=100):
+        assert len(chunks) == len(embeddings), f"len(chunks) != len(embeddings) for {entry['title']} of {entry['source']}"
+        
+        chunk_len = len(chunks)
+        
+        vectors = [
+            {
+                'id': f"{entry['id']}_{str(i).zfill(6)}",
+                'values': embeddings[i],
+                'metadata': {
+                    'entry_id': entry['id'],
+                    'source': entry['source'],
+                    'title': entry['title'],
+                    'authors': entry['authors']
+                }
+            } for i in range(chunk_len)
+        ]
+        
+        self.index.upsert(
+            vectors=vectors,
+            batch_size=upsert_size
+        )
+        # print(f"Successfully inserted {chunk_len} chunks from {entry['source']} article \'{entry['title']}\'.")
+        
+    def delete_entry(self, id):
+        self.index.delete(
+            filter={"entry_id": {"$eq": id}}
+        )
+        # print(f"Successfully deleted elements from id {id}.")
+
+    def info(self):
+        info = pinecone.describe_index(self.index_name)
+        return info        
+    
+    def create_index(self):
+        pinecone.create_index(
+            name=self.index_name,
+            dimension=self.dimensions,
+            metric=self.metric,
+            metadata_config = {
+                "indexed": ["title", "author", "date", "url", "source"]
+            }
+        )
+
+    def delete_index(self):
+        if self.index_name in pinecone.list_indexes():
+            pinecone.delete_index(self.index_name)
\ No newline at end of file