Update app.py

app.py

@@ -67,29 +67,27 @@ def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
 # File to store the list of uploaded documents
-DOCUMENTS_FILE = "uploaded_documents.json"
+# File to store metadata about uploaded documents
+METADATA_FILE = "document_metadata.json"
 
-def load_uploaded_documents():
-    global uploaded_documents
-    if os.path.exists(DOCUMENTS_FILE):
-        with open(DOCUMENTS_FILE, 'r') as f:
-            uploaded_documents = json.load(f)
-    else:
-        uploaded_documents = []
+def load_document_metadata():
+    if os.path.exists(METADATA_FILE):
+        with open(METADATA_FILE, 'r') as f:
+            return json.load(f)
+    return []
 
-def save_uploaded_documents():
-    with open(DOCUMENTS_FILE, 'w') as f:
-        json.dump(uploaded_documents, f)
+def save_document_metadata(metadata):
+    with open(METADATA_FILE, 'w') as f:
+        json.dump(metadata, f)
 
 def update_vectors(files, parser):
-    global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
     
     if not files:
         logging.warning("No files provided for update_vectors")
         return "Please upload at least one PDF file.", gr.CheckboxGroup(
-            choices=[doc["name"] for doc in uploaded_documents],
-            value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+            choices=[],
+            value=[],
             label="Select documents to query"
         )
     
@@ -97,6 +95,8 @@ def update_vectors(files, parser):
     total_chunks = 0
     
     all_data = []
+    metadata = load_document_metadata()
+    
     for file in files:
         logging.info(f"Processing file: {file.name}")
         try:
@@ -104,12 +104,13 @@ 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}")
+            
+            # Update metadata
+            if not any(doc["name"] == file.name for doc in metadata):
+                metadata.append({"name": file.name, "selected": True})
+                logging.info(f"Added new document to metadata: {file.name}")
             else:
-                logging.info(f"Document already exists in uploaded_documents: {file.name}")
+                logging.info(f"Document already exists in metadata: {file.name}")
         except Exception as e:
             logging.error(f"Error processing file {file.name}: {str(e)}")
     
@@ -126,13 +127,13 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
     
-    # Save the updated list of documents
-    save_uploaded_documents()
-    logging.info("Uploaded documents list saved")
+    # Save the updated metadata
+    save_document_metadata(metadata)
+    logging.info("Document metadata 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"]],
+        choices=[doc["name"] for doc in metadata],
+        value=[doc["name"] for doc in metadata if doc["selected"]],
         label="Select documents to query"
     )
 
@@ -510,6 +511,8 @@ def initial_conversation():
 # Define the checkbox outside the demo block
 document_selector = gr.CheckboxGroup(label="Select documents to query")
 
+document_selector = initialize_document_selector()
+
 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)"
@@ -571,6 +574,10 @@ with demo:
         update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
+    document_selector = gr.CheckboxGroup(label="Select documents to query")
+    
+    # Initialize document selector
+    document_selector = initialize_document_selector()
     
     # Update both the output text and the document selector
     update_button.click(update_vectors,