Update app.py

app.py

@@ -46,19 +46,19 @@ llama_parser = LlamaParse(
     language="en",
 )
 
-def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
+def load_document(file_path: str, parser: str = "llamaparse") -> List[Document]:
     """Loads and splits the document into pages."""
     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'.")
@@ -85,7 +85,12 @@ def update_vectors(files, parser):
     for file in files:
         logging.info(f"Processing file: {file.name}")
         try:
-            data = load_document(file, parser)
+            # Save the uploaded file
+            os.makedirs("uploaded_files", exist_ok=True)
+            file_path = os.path.join("uploaded_files", file.name)
+            file.save(file_path)
+            
+            data = load_document(file_path, parser)
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
@@ -100,16 +105,10 @@ 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")
+    if all_data:
         database = FAISS.from_documents(all_data, embed)
-    
-    database.save_local("faiss_database")
-    logging.info("FAISS database saved")
+        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],
@@ -117,6 +116,52 @@ def update_vectors(files, parser):
         label="Select documents to query"
     )
 
+def delete_selected_files(selected_files):
+    global uploaded_documents
+    if not selected_files:
+        return "No files selected for deletion.", document_selector
+
+    deleted_files = []
+    for file_name in selected_files:
+        # Remove the file from uploaded_documents
+        uploaded_documents = [doc for doc in uploaded_documents if doc["name"] != file_name]
+        
+        # Delete the file from the file system if it exists
+        file_path = os.path.join("uploaded_files", file_name)
+        if os.path.exists(file_path):
+            os.remove(file_path)
+        
+        deleted_files.append(file_name)
+
+    # Update the FAISS database
+    update_faiss_database()
+
+    remaining_files = [doc["name"] for doc in uploaded_documents]
+    return f"Deleted files: {', '.join(deleted_files)}", gr.CheckboxGroup(choices=remaining_files, value=remaining_files, label="Select documents to query")
+
+def update_faiss_database():
+    global uploaded_documents
+    
+    embed = get_embeddings()
+    all_data = []
+
+    for doc in uploaded_documents:
+        file_path = os.path.join("uploaded_files", doc["name"])
+        if os.path.exists(file_path):
+            data = load_document(file_path, parser="llamaparse")  # or use your preferred parser
+            all_data.extend(data)
+
+    if all_data:
+        database = FAISS.from_documents(all_data, embed)
+        database.save_local("faiss_database")
+        logging.info("FAISS database updated after deletion")
+    else:
+        # If no documents left, remove the FAISS database
+        if os.path.exists("faiss_database"):
+            import shutil
+            shutil.rmtree("faiss_database")
+        logging.info("All documents deleted, FAISS database removed")
+
 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 = ""
@@ -247,7 +292,6 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
-
     logging.info(f"Selected Documents: {selected_docs}")
 
     try:
@@ -265,7 +309,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
                 
                 # Filter relevant documents based on user selection
                 all_relevant_docs = retriever.get_relevant_documents(message)
-                relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
+                relevant_docs = [doc for doc in all_relevant_docs if os.path.basename(doc.metadata["source"]) in selected_docs]
                 
                 if not relevant_docs:
                     yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."