Update app.py

app.py

@@ -124,10 +124,21 @@ def update_vectors(files, parser):
     for file in files:
         logging.info(f"Processing file: {file.name}")
         try:
-            if file.name.lower().endswith(('.xlsx', '.xls', '.docx')):
+            file_extension = os.path.splitext(file.name)[1].lower()
+            
+            if file_extension in ['.xlsx', '.xls', '.docx']:
+                if parser != "office":
+                    logging.warning(f"Using office parser for {file.name} regardless of selected parser")
                 data = load_office_document(file)
+            elif file_extension == '.pdf':
+                if parser == "office":
+                    logging.warning(f"Cannot use office parser for PDF file {file.name}. Using llamaparse.")
+                    data = load_document(file, "llamaparse")
+                else:
+                    data = load_document(file, parser)
             else:
-                data = load_document(file, parser)
+                logging.warning(f"Unsupported file type: {file_extension}")
+                continue
             
             if not data:
                 logging.warning(f"No chunks loaded from {file.name}")
@@ -150,19 +161,30 @@ def update_vectors(files, parser):
         return "No valid data could be extracted from the uploaded files. Please check the file contents and try again.", display_documents()
 
     try:
-        # Create or update the office documents vector store
-        if os.path.exists("office_faiss_database"):
-            logging.info("Updating existing office FAISS database")
-            office_database = FAISS.load_local("office_faiss_database", embed, allow_dangerous_deserialization=True)
-            office_database.add_documents(all_data)
-        else:
-            logging.info("Creating new office FAISS database")
-            office_database = FAISS.from_documents(all_data, embed)
+        # Update the appropriate vector store based on file type
+        pdf_data = [doc for doc in all_data if doc.metadata["source"].lower().endswith('.pdf')]
+        office_data = [doc for doc in all_data if not doc.metadata["source"].lower().endswith('.pdf')]
+
+        if pdf_data:
+            if os.path.exists("faiss_database"):
+                pdf_database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+                pdf_database.add_documents(pdf_data)
+            else:
+                pdf_database = FAISS.from_documents(pdf_data, embed)
+            pdf_database.save_local("faiss_database")
+            logging.info("PDF FAISS database updated and saved")
+
+        if office_data:
+            if os.path.exists("office_faiss_database"):
+                office_database = FAISS.load_local("office_faiss_database", embed, allow_dangerous_deserialization=True)
+                office_database.add_documents(office_data)
+            else:
+                office_database = FAISS.from_documents(office_data, embed)
+            office_database.save_local("office_faiss_database")
+            logging.info("Office FAISS database updated and saved")
 
-        office_database.save_local("office_faiss_database")
-        logging.info("Office FAISS database saved")
     except Exception as e:
-        logging.error(f"Error updating office FAISS database: {str(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
@@ -667,7 +689,7 @@ with demo:
     gr.Markdown("## Upload and Manage PDF Documents")
     with gr.Row():
         file_input = gr.Files(label="Upload your documents", file_types=[".pdf", ".docx", ".xlsx", ".xls"])
-        parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="llamaparse")
+        parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse", "office"], label="Select PDF Parser", value="llamaparse")
         update_button = gr.Button("Upload Document")
         refresh_button = gr.Button("Refresh Document List")