File size: 801 Bytes
7f7b773
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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)