persistent database fix (#124)

* persistent database fix

* changed to abs path

* added faiss method and updated hybrid

* abs path fix 2

* fixes subsequent similarity chains saving

---------

Signed-off-by: Kevin Guan <[email protected]>

Author: kevinwguan

backend/src/chains/hybrid_retriever_chain.py

@@ -1,3 +1,4 @@
+import os
 from typing import Optional, Union, Any
 
 from langchain.retrievers import EnsembleRetriever
@@ -75,9 +76,21 @@ def create_hybrid_retriever(self) -> None:
             use_cuda=self.use_cuda,
         )
         if self.vector_db is None:
-            similarity_retriever_chain.embed_docs(return_docs=True)
+            cur_path = os.path.abspath(__file__)
+            path = os.path.join(cur_path, "../../../", "faiss_db")
+            path = os.path.abspath(path)  # Ensure proper parent directory
+            path_flag = os.path.isdir(path)  # Checks if database already exists
+            database_name = similarity_retriever_chain.name
+            if path_flag and database_name in os.listdir(path):
+                if database_name in os.listdir(path):
+                    similarity_retriever_chain.create_vector_db()
+                    similarity_retriever_chain.vector_db.load_db(database_name)
+                    self.vector_db = similarity_retriever_chain.vector_db
+                    self.vector_db.processed_docs = similarity_retriever_chain.vector_db.get_documents()
+            else:
+                similarity_retriever_chain.embed_docs(return_docs=True)
+                self.vector_db = similarity_retriever_chain.vector_db
 
-        self.vector_db = similarity_retriever_chain.vector_db
         similarity_retriever_chain.create_similarity_retriever(search_k=self.search_k)
         similarity_retriever = similarity_retriever_chain.retriever
 

backend/src/chains/similarity_retriever_chain.py

@@ -12,6 +12,7 @@
 
 
 class SimilarityRetrieverChain(BaseChain):
+    count = 0
     def __init__(
         self,
         llm_model: Optional[
@@ -33,6 +34,9 @@ def __init__(
             vector_db=vector_db,
         )
 
+        SimilarityRetrieverChain.count += 1
+        self.name = f"similarity_INST{SimilarityRetrieverChain.count}"
+
         self.embeddings_config: Optional[dict[str, str]] = embeddings_config
         self.use_cuda: bool = use_cuda
 
@@ -96,6 +100,8 @@ def embed_docs(
                 folder_paths=self.html_docs_path,
                 return_docs=return_docs,
             )
+        if self.vector_db is not None:
+            self.vector_db.save_db(self.name)
 
         return (
             self.processed_docs,

backend/src/vectorstores/faiss.py

@@ -190,16 +190,25 @@ def add_documents(
 
         return None
 
-    def save_db(self) -> None:
+    def get_db_path(self) -> str:
+        cur_path = os.path.abspath(__file__)
+        path = os.path.join(cur_path, "../../../", "faiss_db")
+        path = os.path.abspath(path)  # Ensure proper parent directory
+        return path
+
+    def save_db(self, name) -> None:
         if self._faiss_db is None:
             raise ValueError("No documents in FAISS database")
+        else:
+            save_path = f"{self.get_db_path()}/{name}"
+            self._faiss_db.save_local(save_path)
 
-        self._faiss_db.save_local(os.getenv("FAISS_DB_PATH", "faiss_db"))
+    def load_db(self, name) -> None:
+        load_path = f"{self.get_db_path()}/{name}"
+        self._faiss_db = FAISS.load_local(load_path, self.embedding_model, allow_dangerous_deserialization=True)
 
-    def load_db(self) -> None:
-        self._faiss_db = FAISS.load_local(
-            os.getenv("FAISS_DB_PATH", "faiss db"), self.embedding_model
-        )
+    def get_documents(self) -> list[Document]:
+        return self._faiss_db.docstore._dict.values()
 
     def process_json(self, folder_paths: list[str]) -> FAISS:
         logging.info("Processing json files...")