Update utils/database.py

utils/database.py

@@ -138,7 +138,124 @@ def verify_vector_store(vector_store):
     except Exception as e:
         st.error(f"Vector store verification failed: {e}")
         return False
-    
+
+def handle_document_upload(uploaded_files):
+    """Handle document upload with progress tracking."""
+    try:
+        # Create a progress container
+        progress_container = st.empty()
+        status_container = st.empty()
+        details_container = st.empty()
+        
+        # Initialize progress bar
+        progress_bar = progress_container.progress(0)
+        status_container.info("🔄 Initializing document processing...")
+        
+        # Reset existing states
+        if 'vector_store' in st.session_state:
+            del st.session_state.vector_store
+        if 'qa_system' in st.session_state:
+            del st.session_state.qa_system
+            
+        # Initialize embeddings (10% progress)
+        status_container.info("🔄 Initializing embeddings model...")
+        embeddings = get_embeddings_model()
+        if not embeddings:
+            status_container.error("❌ Failed to initialize embeddings model")
+            return
+        progress_bar.progress(10)
+        
+        # Process documents
+        documents = []
+        document_names = []
+        
+        # Calculate progress steps per file
+        progress_per_file = 70 / len(uploaded_files)  # 70% of progress for file processing
+        current_progress = 10
+        
+        for idx, uploaded_file in enumerate(uploaded_files):
+            file_name = uploaded_file.name
+            status_container.info(f"🔄 Processing document {idx + 1}/{len(uploaded_files)}: {file_name}")
+            details_container.text(f"📄 Current file: {file_name}")
+            
+            # Read file content
+            content = uploaded_file.read()
+            
+            # Store in database
+            details_container.text(f"💾 Storing {file_name} in database...")
+            doc_id = insert_document(st.session_state.db_conn, file_name, content.decode('utf-8'))
+            if not doc_id:
+                status_container.error(f"❌ Failed to store document: {file_name}")
+                continue
+                
+            documents.append(content.decode('utf-8'))
+            document_names.append(file_name)
+            
+            # Update progress
+            current_progress += progress_per_file
+            progress_bar.progress(int(current_progress))
+            
+        if not documents:
+            status_container.error("❌ No documents were successfully processed")
+            return
+            
+        # Initialize vector store (80-90% progress)
+        status_container.info("🔄 Initializing vector store...")
+        details_container.text("🔍 Creating vector embeddings...")
+        vector_store = initialize_faiss(embeddings, documents, document_names)
+        if not vector_store:
+            status_container.error("❌ Failed to initialize vector store")
+            return
+        progress_bar.progress(90)
+        
+        # Verify vector store
+        status_container.info("🔄 Verifying document indexing...")
+        details_container.text("✨ Performing final checks...")
+        if not verify_vector_store(vector_store):
+            status_container.error("❌ Vector store verification failed")
+            return
+            
+        # Store in session state
+        st.session_state.vector_store = vector_store
+        
+        # Initialize QA system (90-100% progress)
+        status_container.info("🔄 Setting up QA system...")
+        qa_system = initialize_qa_system(vector_store)
+        if not qa_system:
+            status_container.error("❌ Failed to initialize QA system")
+            return
+            
+        st.session_state.qa_system = qa_system
+        
+        # Complete!
+        progress_bar.progress(100)
+        status_container.success("✅ Documents processed successfully!")
+        details_container.markdown("""
+        🎉 **Ready to chat!**
+        - Documents loaded: {}
+        - Total content size: {:.2f} KB
+        - Vector store initialized
+        - QA system ready
+        
+        You can now start asking questions about your documents!
+        """.format(
+            len(documents),
+            sum(len(doc) for doc in documents) / 1024
+        ))
+        
+        # Add notification
+        st.balloons()
+        
+    except Exception as e:
+        status_container.error(f"❌ Error processing documents: {e}")
+        details_container.error(traceback.format_exc())
+        
+    finally:
+        # Clean up progress display after 5 seconds if successful
+        if st.session_state.get('qa_system'):
+            time.sleep(5)
+            progress_container.empty()    
+            
 def initialize_qa_system(vector_store):
     """Initialize QA system with proper chat handling."""
     try: