Update app.py

app.py

@@ -15,6 +15,7 @@ from langchain_core.documents import Document
 from huggingface_hub import InferenceClient
 import inspect
 import logging
+import shutil
 
 
 # Set up basic configuration for logging
@@ -99,6 +100,9 @@ def update_vectors(files, parser):
         logging.info(f"Processing file: {file.name}")
         try:
             data = load_document(file, parser)
+            if not data:
+                logging.warning(f"No chunks loaded from {file.name}")
+                continue
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
@@ -112,22 +116,69 @@ def update_vectors(files, parser):
     
     logging.info(f"Total chunks processed: {total_chunks}")
     
-    if os.path.exists("faiss_database"):
-        logging.info("Updating existing FAISS database")
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        database.add_documents(all_data)
-    else:
-        logging.info("Creating new FAISS database")
-        database = FAISS.from_documents(all_data, embed)
+    if not all_data:
+        logging.warning("No valid data extracted from uploaded files")
+        return "No valid data could be extracted from the uploaded files. Please check the file contents and try again.", display_documents()
     
-    database.save_local("faiss_database")
-    logging.info("FAISS database saved")
+    try:
+        if os.path.exists("faiss_database"):
+            logging.info("Updating existing FAISS database")
+            database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+            database.add_documents(all_data)
+        else:
+            logging.info("Creating new FAISS database")
+            database = FAISS.from_documents(all_data, embed)
+        
+        database.save_local("faiss_database")
+        logging.info("FAISS database saved")
+    except Exception as e:
+        logging.error(f"Error updating FAISS database: {str(e)}")
+        return f"Error updating vector store: {str(e)}", display_documents()
 
     # Save the updated list of documents
     save_documents(uploaded_documents)
 
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", display_documents()
 
+def delete_documents(selected_docs):
+    global uploaded_documents
+    
+    if not selected_docs:
+        return "No documents selected for deletion.", display_documents()
+    
+    embed = get_embeddings()
+    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+    
+    deleted_docs = []
+    docs_to_keep = []
+    for doc in database.docstore._dict.values():
+        if doc.metadata.get("source") not in selected_docs:
+            docs_to_keep.append(doc)
+        else:
+            deleted_docs.append(doc.metadata.get("source", "Unknown"))
+    
+    # Print debugging information
+    logging.info(f"Total documents before deletion: {len(database.docstore._dict)}")
+    logging.info(f"Documents to keep: {len(docs_to_keep)}")
+    logging.info(f"Documents to delete: {len(deleted_docs)}")
+    
+    if not docs_to_keep:
+        # If all documents are deleted, remove the FAISS database directory
+        if os.path.exists("faiss_database"):
+            shutil.rmtree("faiss_database")
+        logging.info("All documents deleted. Removed FAISS database directory.")
+    else:
+        # Create new FAISS index with remaining documents
+        new_database = FAISS.from_documents(docs_to_keep, embed)
+        new_database.save_local("faiss_database")
+        logging.info(f"Created new FAISS index with {len(docs_to_keep)} documents.")
+    
+    # Update uploaded_documents list
+    uploaded_documents = [doc for doc in uploaded_documents if doc["name"] not in deleted_docs]
+    save_documents(uploaded_documents)
+    
+    return f"Deleted documents: {', '.join(deleted_docs)}", display_documents()
+
 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")
     full_response = ""
@@ -536,7 +587,7 @@ def display_documents():
     return 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"
+        label="Select documents to query or delete"
     )
 
 # Add this new function
@@ -623,6 +674,7 @@ with demo:
         refresh_button = gr.Button("Refresh Document List")
     
     update_output = gr.Textbox(label="Update Status")
+    delete_button = gr.Button("Delete Selected Documents")
     
     # Update both the output text and the document selector
     update_button.click(update_vectors, 
@@ -633,6 +685,11 @@ with demo:
     refresh_button.click(refresh_documents, 
                          inputs=[], 
                          outputs=[document_selector])
+    
+    # Add the delete button functionality
+    delete_button.click(delete_documents,
+                        inputs=[document_selector],
+                        outputs=[update_output, document_selector])
 
     gr.Markdown(
     """