Update utils/database.py

utils/database.py

@@ -150,6 +150,12 @@ def verify_vector_store(vector_store):
 def handle_document_upload(uploaded_files):
     """Handle document upload with progress tracking."""
     try:
+        # Initialize session state variables if they don't exist
+        if 'qa_system' not in st.session_state:
+            st.session_state.qa_system = None
+        if 'vector_store' not in st.session_state:
+            st.session_state.vector_store = None
+        
         # Create a progress container
         progress_container = st.empty()
         status_container = st.empty()
@@ -160,10 +166,10 @@ def handle_document_upload(uploaded_files):
         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
+        if st.session_state.vector_store is not None:
+            st.session_state.vector_store = None
+        if st.session_state.qa_system is not None:
+            st.session_state.qa_system = None
             
         # Initialize embeddings (10% progress)
         status_container.info("🔄 Initializing embeddings model...")
@@ -224,6 +230,9 @@ def handle_document_upload(uploaded_files):
         if not vector_store:
             status_container.error("❌ Failed to initialize vector store")
             return
+        
+        # Store vector store in session state
+        st.session_state.vector_store = vector_store
         progress_bar.progress(90)
         
         # Verify vector store
@@ -233,9 +242,6 @@ def handle_document_upload(uploaded_files):
             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)
@@ -243,6 +249,7 @@ def handle_document_upload(uploaded_files):
             status_container.error("❌ Failed to initialize QA system")
             return
             
+        # Store QA system in session state
         st.session_state.qa_system = qa_system
         
         # Complete!
@@ -264,13 +271,20 @@ def handle_document_upload(uploaded_files):
         # Add notification
         st.balloons()
         
+        # Set chat ready flag
+        st.session_state.chat_ready = True
+        
     except Exception as e:
         status_container.error(f"❌ Error processing documents: {e}")
         details_container.error(traceback.format_exc())
+        # Reset states on error
+        st.session_state.vector_store = None
+        st.session_state.qa_system = None
+        st.session_state.chat_ready = False
         
     finally:
         # Clean up progress display after 5 seconds if successful
-        if st.session_state.get('qa_system'):
+        if st.session_state.get('qa_system') is not None:
             time.sleep(5)
             progress_container.empty()