from langchain.vectorstores import FAISS

class FAISSVectorStore:
    def __init__(self, embedding_model):
        self.embedding_model = embedding_model
        self.db = None

    def initialize_from_documents(self, docs):
        self.db = FAISS.from_documents(docs, self.embedding_model.model)

    def initialize_from_file(self, path):
        self.db = FAISS.load_local(path, self.embedding_model.model)

    def save(self, path):
        self.db.save_local(path)

    def add_documents(self, documents):
        return self.db.add_documents(documents)
    
    def query(self, query: str, k: int = 4):
        # TODO adjust fetch_k parameter. It is now set to match the defaults k=4, fetch_k=20 in the original code.
        return self.db.similarity_search_with_score(query, k=k, fetch_k=5*k)