Update src/app.py

src/app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import os
 from pathlib import Path
 from rag_pipeline import RAGPipeline
-import time
+import shutil
 
 # Page configuration
 st.set_page_config(
@@ -20,8 +20,8 @@ if "uploaded_files" not in st.session_state:
     st.session_state.uploaded_files = []
 if "rag_pipeline" not in st.session_state:
     st.session_state.rag_pipeline = None
-if "last_upload_time" not in st.session_state:
-    st.session_state.last_upload_time = 0
+if "needs_reindex" not in st.session_state:
+    st.session_state.needs_reindex = False
 
 # Sidebar configuration
 with st.sidebar:
@@ -33,9 +33,6 @@ with st.sidebar:
         help="Path to directory containing PDF files"
     )
     
-    # Ensure directory exists
-    os.makedirs(pdf_dir, exist_ok=True)
-    
     device = st.selectbox(
         "🖥️ Device",
         ["cpu", "cuda"],
@@ -52,44 +49,36 @@ with st.sidebar:
     
     st.divider()
     
-    # PDF Upload Section with Form
+    # PDF Upload Section
     st.subheader("📤 Upload PDF Files")
     
-    # Use a form to separate file upload from submission
-    with st.form("pdf_upload_form", clear_on_submit=True):
-        uploaded_pdfs = st.file_uploader(
-            "Choose PDF files to upload",
-            type="pdf",
-            accept_multiple_files=True,
-            help="Select one or more PDF files to add to the system"
-        )
+    uploaded_pdfs = st.file_uploader(
+        "Choose PDF files to upload",
+        type="pdf",
+        accept_multiple_files=True,
+        help="Select one or more PDF files to add to the system"
+    )
+    
+    if uploaded_pdfs:
+        # Create PDF directory if not exists
+        os.makedirs(pdf_dir, exist_ok=True)
         
-        submit_button = st.form_submit_button("⬆️ Upload PDFs", use_container_width=True)
+        upload_button = st.button("⬆️ Upload PDFs", use_container_width=True)
         
-        if submit_button and uploaded_pdfs:
-            upload_successful = True
+        if upload_button:
             uploaded_count = 0
-            
             for uploaded_file in uploaded_pdfs:
-                try:
-                    file_path = os.path.join(pdf_dir, uploaded_file.name)
-                    
-                    # Save file to disk
-                    with open(file_path, "wb") as f:
-                        f.write(uploaded_file.getbuffer())
-                    
-                    st.session_state.uploaded_files.append(uploaded_file.name)
-                    uploaded_count += 1
-                    
-                except Exception as e:
-                    st.error(f"Failed to upload {uploaded_file.name}: {str(e)}")
-                    upload_successful = False
+                file_path = os.path.join(pdf_dir, uploaded_file.name)
+                
+                # Save file
+                with open(file_path, "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                
+                st.session_state.uploaded_files.append(uploaded_file.name)
+                uploaded_count += 1
             
-            if upload_successful and uploaded_count > 0:
-                st.session_state.last_upload_time = time.time()
-                st.success(f"✅ Uploaded {uploaded_count} PDF(s) successfully!")
-                st.info("📌 Click 'Reload & Index PDFs' below to process them.")
-                # Don't call st.rerun() here - let form handle clear_on_submit
+            st.success(f"✅ Uploaded {uploaded_count} PDF(s) successfully!")
+            st.session_state.needs_reindex = True
     
     st.divider()
     
@@ -97,94 +86,59 @@ with st.sidebar:
     pdf_files = list(Path(pdf_dir).glob("*.pdf"))
     if pdf_files:
         st.subheader(f"📚 Documents ({len(pdf_files)})")
-        
         for pdf_file in pdf_files:
             col1, col2 = st.columns([4, 1])
             with col1:
                 st.write(f"• {pdf_file.name}")
             with col2:
                 if st.button("🗑️", key=f"delete_{pdf_file.name}", help="Delete this file"):
-                    try:
-                        os.remove(pdf_file)
-                        st.session_state.rag_pipeline = None  # Clear pipeline
-                        st.success(f"Deleted {pdf_file.name}")
-                        time.sleep(0.5)
-                        st.rerun()
-                    except Exception as e:
-                        st.error(f"Failed to delete: {str(e)}")
-    else:
-        st.info("📭 No PDF files in directory yet")
+                    os.remove(pdf_file)
+                    st.session_state.needs_reindex = True
+                    st.rerun()
     
     st.divider()
     
-    # Reload/Index button
-    col1, col2 = st.columns(2)
-    with col1:
-        if st.button("🔄 Reload & Index", use_container_width=True):
-            st.session_state.rag_pipeline = None  # Clear cached pipeline
-            st.rerun()
-    
-    with col2:
-        if st.button("🗑️ Clear All", use_container_width=True):
-            # Delete all PDFs
-            for pdf_file in Path(pdf_dir).glob("*.pdf"):
-                try:
-                    os.remove(pdf_file)
-                except:
-                    pass
-            st.session_state.rag_pipeline = None
-            st.session_state.uploaded_files = []
-            st.success("All PDFs cleared")
-            time.sleep(0.5)
-            st.rerun()
+    # Reindex button
+    if st.button("🔄 Reload & Index PDFs", use_container_width=True):
+        st.session_state.rag_pipeline = None
+        st.session_state.needs_reindex = True
+        st.rerun()
 
 
-# Initialize pipeline
+# Initialize pipeline in session state
 @st.cache_resource
 def init_rag_pipeline(_device, _pdf_dir):
     """Initialize RAG pipeline (cached)"""
+    # Create PDF directory if not exists
     os.makedirs(_pdf_dir, exist_ok=True)
     
+    # Check if PDFs exist
     pdf_files = list(Path(_pdf_dir).glob("*.pdf"))
     if not pdf_files:
-        return None, f"No PDF files found in {_pdf_dir}"
+        return None, f"No PDF files found in {_pdf_dir}. Upload PDFs using the sidebar."
     
     try:
-        with st.spinner("⏳ Initializing models..."):
+        with st.spinner("⏳ Initializing RAG pipeline..."):
             pipeline = RAGPipeline(pdf_dir=_pdf_dir, device=_device)
-        
-        with st.spinner("⏳ Indexing PDFs..."):
-            pipeline.index_pdfs()
-        
+            with st.spinner("⏳ Indexing PDFs..."):
+                pipeline.index_pdfs()
         return pipeline, None
     except Exception as e:
         return None, str(e)
 
 
 # Get or initialize pipeline
-if st.session_state.rag_pipeline is None:
-    pdf_files = list(Path(pdf_dir).glob("*.pdf"))
-    
-    if pdf_files:
-        pipeline, error = init_rag_pipeline(device, pdf_dir)
-        if error:
-            st.error(f"❌ Error: {error}")
-            st.stop()
-        st.session_state.rag_pipeline = pipeline
-    else:
-        st.warning("📭 No PDF files found")
-        st.info("""
-        **How to get started:**
-        1. 📤 Upload PDF files using the sidebar file uploader
-        2. ✅ Click 'Upload PDFs' to save them
-        3. 🔄 Click 'Reload & Index PDFs' to process
-        4. ❓ Ask questions in the Q&A tab
-        """)
+if st.session_state.rag_pipeline is None or st.session_state.needs_reindex:
+    pipeline, error = init_rag_pipeline(device, pdf_dir)
+    if error:
+        st.error(f"❌ Error: {error}")
+        st.info("💡 **How to get started:**\n1. Upload PDF files using the sidebar\n2. Click 'Upload PDFs' to save them\n3. Click 'Reload & Index PDFs' to process them")
         st.stop()
+    st.session_state.rag_pipeline = pipeline
+    st.session_state.needs_reindex = False
 else:
     pipeline = st.session_state.rag_pipeline
 
-
 # Main content
 if pipeline:
     # Tabs
@@ -216,19 +170,22 @@ if pipeline:
                     try:
                         result = pipeline.answer_question(question, n_context_docs=n_context_docs)
                     except Exception as e:
-                        st.error(f"Error: {str(e)}")
+                        st.error(f"Error generating answer: {str(e)}")
                         result = None
                 
                 if result and result.get("answer"):
                     st.success("✓ Answer generated!")
                     
+                    # Display answer
                     st.subheader("📝 Answer")
                     st.write(result["answer"])
                     
+                    # Display sources
                     with st.expander("📚 Sources Used"):
                         for i, source in enumerate(result["sources"], 1):
                             st.write(f"{i}. {source}")
                     
+                    # Display stats
                     col1, col2 = st.columns(2)
                     with col1:
                         st.metric("Documents Used", result.get("context_used", 0))
@@ -249,7 +206,7 @@ if pipeline:
                     st.subheader("📄 Document Summary")
                     st.write(summary)
                 except Exception as e:
-                    st.error(f"Error: {str(e)}")
+                    st.error(f"Error generating summary: {str(e)}")
     
     # Tab 3: Document Retrieval
     with tab3:
@@ -296,11 +253,15 @@ if pipeline:
         info = pipeline.vector_store.get_collection_info()
         col1, col2, col3, col4 = st.columns(4)
         with col1:
-            st.metric("📚 Chunks", info.get("document_count", 0))
+            st.metric("📚 Documents", info.get("document_count", 0))
         with col2:
             st.metric("🖥️ Device", device.upper())
         with col3:
-            st.metric("🔍 Context", n_context_docs)
+            st.metric("🔍 Context Docs", n_context_docs)
         with col4:
             pdf_count = len(list(Path(pdf_dir).glob("*.pdf")))
-            st.metric("📁 PDFs", pdf_count)
+            st.metric("📁 PDF Files", pdf_count)
+
+else:
+    st.error("❌ Failed to initialize RAG pipeline")
+    st.info("💡 **How to get started:**\n1. Upload PDF files using the sidebar\n2. Click 'Upload PDFs' to save them\n3. Click 'Reload & Index PDFs' to process them")