Update app.py

app.py

@@ -47,18 +47,20 @@ llama_parser = LlamaParse(
 )
 
 def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
-    """Loads and splits the document into pages."""
+    file_path = os.path.join(UPLOAD_FOLDER, file.name)
+    shutil.copy(file.name, file_path)
+    
     if parser == "pypdf":
-        loader = PyPDFLoader(file.name)
+        loader = PyPDFLoader(file_path)
         return loader.load_and_split()
     elif parser == "llamaparse":
         try:
-            documents = llama_parser.load_data(file.name)
-            return [Document(page_content=doc.text, metadata={"source": file.name}) for doc in documents]
+            documents = llama_parser.load_data(file_path)
+            return [Document(page_content=doc.text, metadata={"source": file_path}) for doc in documents]
         except Exception as e:
             print(f"Error using Llama Parse: {str(e)}")
             print("Falling back to PyPDF parser")
-            loader = PyPDFLoader(file.name)
+            loader = PyPDFLoader(file_path)
             return loader.load_and_split()
     else:
         raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
@@ -108,6 +110,10 @@ def update_vectors(files, parser):
     
     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"]])
 
+UPLOAD_FOLDER = "uploaded_files"
+if not os.path.exists(UPLOAD_FOLDER):
+    os.makedirs(UPLOAD_FOLDER)
+
 # Add this new function to handle file deletion
 def delete_file(file_name):
     global uploaded_documents
@@ -117,11 +123,12 @@ def delete_file(file_name):
     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}")
+    file_path = os.path.join(UPLOAD_FOLDER, file_name)
+    if os.path.exists(file_path):
+        os.remove(file_path)
+        logging.info(f"Deleted file: {file_path}")
     else:
-        logging.warning(f"File not found: {file_name}")
+        logging.warning(f"File not found: {file_path}")
     
     # Rebuild the FAISS database
     rebuild_faiss_database()
@@ -135,7 +142,8 @@ def rebuild_faiss_database():
     
     for doc in uploaded_documents:
         try:
-            data = load_document(NamedTemporaryFile(delete=False, suffix=".pdf", mode="wb"), "llamaparse")
+            file_path = os.path.join(UPLOAD_FOLDER, doc["name"])
+            data = load_document(NamedTemporaryFile(delete=False, suffix=".pdf", mode="wb", dir=UPLOAD_FOLDER, name=doc["name"]), "llamaparse")
             all_data.extend(data)
         except Exception as e:
             logging.error(f"Error processing file {doc['name']}: {str(e)}")
@@ -146,10 +154,9 @@ def rebuild_faiss_database():
         logging.info("FAISS database rebuilt and saved")
     else:
         if os.path.exists("faiss_database"):
-            os.remove("faiss_database")
+            shutil.rmtree("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")
     full_response = ""
@@ -567,11 +574,7 @@ with demo:
     # 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=[],
-        )
+        delete_button = gr.Button("Delete Selected")
     
     # Update both the output text and the document selector
     update_button.click(
@@ -581,9 +584,9 @@ with demo:
     )
     
     # Handle file deletion
-    delete_buttons.click(
-        delete_file,
-        inputs=[delete_buttons],
+    delete_button.click(
+        lambda selected: [delete_file(file) for file in selected],
+        inputs=[uploaded_files],
         outputs=[uploaded_files]
     )