Update src/streamlit_app.py

src/streamlit_app.py

@@ -426,23 +426,31 @@ def show_login_page():
 
 def upload_and_poll_files(api_key: str, uploaded_files: list):
     """
-    Upload PDFs, display status, and clear messages + reset uploader after processing.
+    Upload PDFs to backend, poll for processing, and clean up uploader after completion.
+    Works for both local and HF Spaces deployments.
     """
     if not api_key:
         st.error("Missing API key. Please login first.")
         return
 
-    # Create placeholders for overall messages
-    progress_bar = st.empty()
+    progress_bar = st.progress(0)
     status_text = st.empty()
     
-    # Track which files were successfully processed
-    processed_files = []
+    # Make a copy of uploaded_files so we can remove processed ones
+    remaining_files = uploaded_files.copy()
 
     for i, file in enumerate(uploaded_files):
         status_text.text(f"Uploading {file.name}...")
-        result = upload_pdf(api_key, file)
+        
+        # Upload
+        try:
+            result = upload_pdf(api_key, file)
+        except Exception as e:
+            st.error(f"❌ {file.name} upload failed: {e}")
+            progress_bar.progress((i + 1) / len(uploaded_files))
+            continue
 
+        # If upload request itself failed
         if not result["success"]:
             st.error(f"❌ {file.name}: {result['error']}")
             progress_bar.progress((i + 1) / len(uploaded_files))
@@ -450,15 +458,15 @@ def upload_and_poll_files(api_key: str, uploaded_files: list):
 
         data = result.get("data") or {}
         job_id = data.get("document_id") or data.get("job_id")
-        status = data.get("status")
+        status = data.get("status", "processing")
 
-        if status == "processing" and job_id:
-            # Polling messages
+        # Poll backend for processing status
+        if job_id and status == "processing":
             poll_text = st.empty()
-            poll_progress = st.empty()
-            elapsed = 0
+            poll_progress = st.progress(0)
             max_poll_seconds = 600
             poll_interval = 2
+            elapsed = 0
 
             while elapsed < max_poll_seconds:
                 try:
@@ -467,47 +475,49 @@ def upload_and_poll_files(api_key: str, uploaded_files: list):
                         headers=get_headers(api_key),
                         timeout=10
                     )
-                    if resp.status_code == 200:
+                    # Always attempt JSON parse, fallback if HF wraps response
+                    try:
                         job = resp.json()
-                        final_status = job.get("status")
-                        display_name = job.get("username") or job.get("user_id")
-                        if final_status not in ["processing"]:
-                            if final_status in ["success", "completed"]:
-                                res = job.get("result", {})
-                                st.success(f"✅ {file.name} processed for {display_name}: {res.get('filename', file.name)}")
-                                processed_files.append(file.name)
-                            else:
-                                error_msg = job.get('error', 'unknown error')
-                                st.error(f"❌ {file.name} failed for {display_name}: {error_msg}")
-                            break
+                    except Exception:
+                        job = {}
+
+                    final_status = job.get("status", "processing")
+                    display_name = job.get("username") or job.get("user_id") or "You"
+
+                    if final_status not in ["processing"]:
+                        if final_status in ["success", "completed"]:
+                            res = job.get("result", {})
+                            st.success(f"✅ {file.name} processed for {display_name}: {res.get('filename', file.name)}")
                         else:
-                            poll_text.text(f"Processing {file.name} for {display_name}... elapsed {elapsed}s")
+                            error_msg = job.get('error', 'unknown error')
+                            st.error(f"❌ {file.name} failed for {display_name}: {error_msg}")
+                        break
                     else:
-                        poll_text.text(f"Waiting for processing... (status {resp.status_code})")
+                        poll_text.text(f"Processing {file.name} for {display_name}... elapsed {elapsed}s")
                 except Exception as e:
                     poll_text.text(f"Waiting... ({str(e)})")
+                
                 time.sleep(poll_interval)
                 elapsed += poll_interval
                 poll_progress.progress(min(1.0, elapsed / max_poll_seconds))
 
-            # Clear polling messages after processing
-            poll_text.empty()
-            poll_progress.empty()
+            if elapsed >= max_poll_seconds:
+                st.error(f"❌ {file.name}: Processing timed out after {max_poll_seconds}s")
+
         elif status in ["success", "completed"]:
             st.success(f"✅ {file.name} uploaded and processed successfully")
-            processed_files.append(file.name)
+        else:
+            # Unknown status — assume processing and continue
+            st.info(f"ℹ️ {file.name}: Upload in progress")
 
-        progress_bar.progress((i + 1) / len(uploaded_files))
+        # Remove the file from remaining_files to clear uploader
+        remaining_files.remove(file)
+        # Force re-render of file uploader with remaining files
+        st.session_state["pdf_uploader"] = remaining_files
 
-    # Clear overall status
-    status_text.empty()
-    progress_bar.empty()
+        progress_bar.progress((i + 1) / len(uploaded_files))
 
-    # ✅ Remove successfully processed files from uploader
-    if processed_files and "pdf_uploader" in st.session_state:
-        # Streamlit doesn't allow direct modification, so delete and rerun
-        del st.session_state["pdf_uploader"]
-        st.experimental_rerun()
+    status_text.text("✅ Upload complete!")
 
 # ============================================
 # MAIN APPLICATION UI