app/vector_store.py

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