Create vector_store.py

app/vector_store.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+import faiss
+import numpy as np
+import os
+import logging
+import sys
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stdout)
+logger = logging.getLogger(__name__)
+
+class BaseVectorStore(ABC):
+    @abstractmethod
+    def add_documents(self, documents, embeddings):
+        pass
+
+    @abstractmethod
+    def search(self, query_vector, num_results=5):
+        pass
+
+class FaissVectorStore(BaseVectorStore):
+    def __init__(self, dimension):
+        self.dimension = dimension
+        self.index = faiss.IndexFlatL2(dimension)
+        self.documents = []
+        self.index_path = "data/faiss_index"
+        os.makedirs("data", exist_ok=True)
+        self.load_index()
+
+    def load_index(self):
+        if os.path.exists(self.index_path):
+            try:
+                self.index = faiss.read_index(self.index_path)
+            except Exception as e:
+                logger.error(f"Error loading FAISS index: {e}")
+
+    def save_index(self):
+        try:
+            faiss.write_index(self.index, self.index_path)
+        except Exception as e:
+            logger.error(f"Error saving FAISS index: {e}")
+
+    def add_documents(self, documents, embeddings):
+        self.index.add(np.array(embeddings))
+        self.documents.extend(documents)
+        self.save_index()
+
+    def search(self, query_vector, num_results=5):
+        if len(self.documents) == 0:
+            return []
+        D, I = self.index.search(np.array([query_vector]), num_results)
+        return [self.documents[i] for i in I[0]]
+
+def get_vector_store(config):
+    """Factory function to get the appropriate vector store"""
+    return FaissVectorStore  # Always return FAISS vector store