Update app.py

app.py

@@ -72,11 +72,7 @@ def update_vectors(files, 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"]],
-            label="Select documents to query"
-        )
+        return "Please upload at least one PDF file.", gr.update(choices=[], value=[])
     
     embed = get_embeddings()
     total_chunks = 0
@@ -89,7 +85,6 @@ 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}")
@@ -111,11 +106,49 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database 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"]],
-        label="Select documents to query"
-    )
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", gr.update(choices=[doc["name"] for doc in uploaded_documents], value=[doc["name"] for doc in uploaded_documents if doc["selected"]])
+
+# Add this new function to handle file deletion
+def delete_file(file_name):
+    global uploaded_documents
+    logging.info(f"Attempting to delete file: {file_name}")
+    
+    # Remove the file from uploaded_documents
+    uploaded_documents = [doc for doc in uploaded_documents if doc["name"] != file_name]
+    
+    # Remove the file from the file system if it exists
+    if os.path.exists(file_name):
+        os.remove(file_name)
+        logging.info(f"Deleted file: {file_name}")
+    else:
+        logging.warning(f"File not found: {file_name}")
+    
+    # Rebuild the FAISS database
+    rebuild_faiss_database()
+    
+    return gr.update(value=[doc["name"] for doc in uploaded_documents], choices=[doc["name"] for doc in uploaded_documents])
+
+def rebuild_faiss_database():
+    logging.info("Rebuilding FAISS database")
+    embed = get_embeddings()
+    all_data = []
+    
+    for doc in uploaded_documents:
+        try:
+            data = load_document(NamedTemporaryFile(delete=False, suffix=".pdf", mode="wb"), "llamaparse")
+            all_data.extend(data)
+        except Exception as e:
+            logging.error(f"Error processing file {doc['name']}: {str(e)}")
+    
+    if all_data:
+        database = FAISS.from_documents(all_data, embed)
+        database.save_local("faiss_database")
+        logging.info("FAISS database rebuilt and saved")
+    else:
+        if os.path.exists("faiss_database"):
+            os.remove("faiss_database")
+        logging.info("No documents left, removed FAISS database")
+
 
 def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
@@ -531,11 +564,36 @@ with demo:
     
     update_output = gr.Textbox(label="Update Status")
     
+    # Create a new row for displaying uploaded files with delete buttons
+    with gr.Row():
+        uploaded_files = gr.CheckboxGroup(label="Uploaded Documents", interactive=True)
+        delete_buttons = gr.Dataset(
+            components=[gr.Button(value="Delete")],
+            label="Actions",
+            samples=[],
+        )
+    
     # Update both the output text and the document selector
-    update_button.click(update_vectors, 
-                        inputs=[file_input, parser_dropdown], 
-                        outputs=[update_output, document_selector])
+    update_button.click(
+        update_vectors, 
+        inputs=[file_input, parser_dropdown], 
+        outputs=[update_output, uploaded_files]
+    )
+    
+    # Handle file deletion
+    delete_buttons.click(
+        delete_file,
+        inputs=[delete_buttons],
+        outputs=[uploaded_files]
+    )
 
+    # Update the document selector in the chat interface
+    uploaded_files.change(
+        lambda x: gr.update(choices=x, value=x),
+        inputs=[uploaded_files],
+        outputs=[document_selector]
+    )
+    
     gr.Markdown(
     """
     ## How to use