Update app.py

app.py

@@ -66,6 +66,26 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
+def scan_faiss_database():
+    global uploaded_documents
+    uploaded_documents = []
+    
+    if os.path.exists("faiss_database"):
+        embed = HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        
+        # Assuming each document in FAISS has a 'source' metadata field with the file name
+        for doc in database.docstore._dict.values():
+            file_name = os.path.basename(doc.metadata['source'])
+            if not any(d['name'] == file_name for d in uploaded_documents):
+                uploaded_documents.append({"name": file_name, "selected": True})
+    
+    return uploaded_documents
+
+# Call this function when the application starts
+uploaded_documents = scan_faiss_database()
+
+
 def update_vectors(files, parser):
     global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
@@ -111,6 +131,9 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
     
+    # After processing new files, rescan the database
+    uploaded_documents = scan_faiss_database()
+    
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", gr.CheckboxGroup(
         choices=[doc["name"] for doc in uploaded_documents],
         value=[doc["name"] for doc in uploaded_documents if doc["selected"]],