Update app.py

app.py

@@ -66,7 +66,19 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
-def update_vectors(files, parser):
+def save_uploaded_documents():
+    with open("uploaded_documents.json", "w") as f:
+        json.dump(uploaded_documents, f)
+
+def load_uploaded_documents():
+    global uploaded_documents
+    if os.path.exists("uploaded_documents.json"):
+        with open("uploaded_documents.json", "r") as f:
+            uploaded_documents = json.load(f)
+    else:
+        uploaded_documents = []
+
+def update_vectors(files: List[NamedTemporaryFile], parser: str):
     global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
     
@@ -82,6 +94,7 @@ def update_vectors(files, parser):
     total_chunks = 0
     
     all_data = []
+    new_documents = []
     for file in files:
         logging.info(f"Processing file: {file.name}")
         try:
@@ -89,15 +102,22 @@ def update_vectors(files, parser):
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
-            # Append new documents instead of replacing
-            if not any(doc["name"] == file.name for doc in uploaded_documents):
-                uploaded_documents.append({"name": file.name, "selected": True})
-                logging.info(f"Added new document to uploaded_documents: {file.name}")
-            else:
+            
+            # Check if the document is already in the list
+            existing_doc = next((doc for doc in uploaded_documents if doc["name"] == file.name), None)
+            if existing_doc:
                 logging.info(f"Document already exists in uploaded_documents: {file.name}")
+                existing_doc["selected"] = True  # Ensure it's selected
+            else:
+                new_doc = {"name": file.name, "selected": True}
+                new_documents.append(new_doc)
+                logging.info(f"Added new document to uploaded_documents: {file.name}")
         except Exception as e:
             logging.error(f"Error processing file {file.name}: {str(e)}")
     
+    # Add new documents to the list
+    uploaded_documents.extend(new_documents)
+    
     logging.info(f"Total chunks processed: {total_chunks}")
     
     if os.path.exists("faiss_database"):
@@ -111,6 +131,10 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
     
+    # Save the updated uploaded_documents list
+    save_uploaded_documents()
+    logging.info("Uploaded documents list saved")
+    
     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"]],
@@ -492,6 +516,8 @@ use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
+load_uploaded_documents()
+
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[