Update utils/database.py

utils/database.py

@@ -146,24 +146,24 @@ def verify_vector_store(vector_store):
 
 def handle_document_upload(uploaded_files):
     """Handle document upload with improved chunking and progress tracking."""
+    # Initialize containers first - before any processing
+    progress_container = st.empty()
+    status_container = st.empty()
+    details_container = st.empty()
+    progress_bar = progress_container.progress(0)
+
     try:
         # Initialize session state variables
-        # Initialize persistence manager
-        persistence = PersistenceManager()
-        
-        # Generate a session ID based on timestamp and files
-        session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-        
         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 progress containers
-        progress_container = st.empty()
-        status_container = st.empty()
-        details_container = st.empty()
-        progress_bar = progress_container.progress(0)
+        # Initialize persistence manager
+        persistence = PersistenceManager()
+        
+        # Generate a session ID based on timestamp and files
+        session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
         
         # Initialize embeddings (10% progress)
         status_container.info("🔄 Initializing embeddings model...")
@@ -175,7 +175,7 @@ def handle_document_upload(uploaded_files):
         
         # Initialize document chunker
         chunker = DocumentChunker(
-            chunk_size=1000,  # Adjust these parameters based on your needs
+            chunk_size=1000,
             chunk_overlap=200,
             max_tokens_per_chunk=2000
         )
@@ -195,18 +195,26 @@ def handle_document_upload(uploaded_files):
                 tmp_file.write(uploaded_file.getvalue())
                 tmp_file.flush()
                 
-                # Load PDF content
-                loader = PyPDFLoader(tmp_file.name)
-                pdf_documents = loader.load()
-                content = "\n".join(doc.page_content for doc in pdf_documents)
-                
-                # Store original content in database
-                doc_id = insert_document(st.session_state.db_conn, file_name, content)
-                if not doc_id:
-                    status_container.error(f"❌ Failed to store document: {file_name}")
-                    continue
-                
-                document_pairs.append((content, file_name))
+                try:
+                    # Load PDF content
+                    loader = PyPDFLoader(tmp_file.name)
+                    pdf_documents = loader.load()
+                    content = "\n".join(doc.page_content for doc in pdf_documents)
+                    
+                    # Store original content in database
+                    doc_id = insert_document(st.session_state.db_conn, file_name, content)
+                    if not doc_id:
+                        status_container.error(f"❌ Failed to store document: {file_name}")
+                        continue
+                    
+                    document_pairs.append((content, file_name))
+                    
+                finally:
+                    # Ensure temporary file is cleaned up
+                    try:
+                        os.unlink(tmp_file.name)
+                    except Exception as e:
+                        st.warning(f"Could not delete temporary file: {e}")
             
             current_progress += progress_per_file
             progress_bar.progress(int(current_progress))
@@ -225,7 +233,8 @@ def handle_document_upload(uploaded_files):
             return
         
         progress_bar.progress(80)
-        # After chunking documents
+        
+        # Save chunks for persistence
         persistence.save_chunks(chunks, chunk_metadatas, session_id)
         
         # Initialize vector store (90% progress)
@@ -237,11 +246,12 @@ def handle_document_upload(uploaded_files):
             status_container.error("❌ Failed to initialize vector store")
             return
         
-        st.session_state.vector_store = vector_store
-        progress_bar.progress(90)
+        # Save vector store and update session state
         persistence.save_vector_store(vector_store, session_id)
-        # Store session ID in state
+        st.session_state.vector_store = vector_store
         st.session_state.current_session_id = session_id
+        progress_bar.progress(90)
+        
         # Initialize QA system (100% progress)
         status_container.info("🔄 Setting up QA system...")
         qa_system = initialize_qa_system(vector_store)
@@ -260,8 +270,9 @@ def handle_document_upload(uploaded_files):
         - Documents processed: {len(document_pairs)}
         - Total chunks created: {len(chunks)}
         - Average chunk size: {sum(len(chunk) for chunk in chunks) / len(chunks):.0f} characters
-        - Vector store initialized
+        - Vector store initialized and saved
         - QA system ready
+        - Session ID: {session_id}
         
         You can now start asking questions about your documents!
         """)
@@ -270,14 +281,14 @@ def handle_document_upload(uploaded_files):
         st.session_state.chat_ready = True
         
     except Exception as e:
-        status_container.error(f"❌ Error processing documents: {e}")
+        status_container.error(f"❌ Error processing documents: {str(e)}")
         details_container.error(traceback.format_exc())
         st.session_state.vector_store = None
         st.session_state.qa_system = None
         st.session_state.chat_ready = False
-        st.error(f"Error in document upload: {str(e)}")
         
     finally:
+        # Clean up progress display after successful processing
         if st.session_state.get('qa_system') is not None:
             time.sleep(5)
             progress_container.empty()