Update utils/database.py

utils/database.py

@@ -148,29 +148,20 @@ def verify_vector_store(vector_store):
 
 
 def handle_document_upload(uploaded_files):
-    """Handle document upload with progress tracking."""
+    """Handle document upload with improved chunking and progress tracking."""
     try:
-        # Initialize session state variables if they don't exist
+        # Initialize session state variables
         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
+        # Create progress containers
         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 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...")
         embeddings = get_embeddings_model()
@@ -179,12 +170,16 @@ def handle_document_upload(uploaded_files):
             return
         progress_bar.progress(10)
         
-        # Process documents
-        documents = []
-        document_names = []
+        # Initialize document chunker
+        chunker = DocumentChunker(
+            chunk_size=1000,  # Adjust these parameters based on your needs
+            chunk_overlap=200,
+            max_tokens_per_chunk=2000
+        )
         
-        # Calculate progress steps per file
-        progress_per_file = 70 / len(uploaded_files)  # 70% of progress for file processing
+        # Process documents
+        document_pairs = []  # List to store (content, filename) pairs
+        progress_per_file = 70 / len(uploaded_files)
         current_progress = 10
         
         for idx, uploaded_file in enumerate(uploaded_files):
@@ -192,98 +187,89 @@ def handle_document_upload(uploaded_files):
             status_container.info(f"🔄 Processing document {idx + 1}/{len(uploaded_files)}: {file_name}")
             details_container.text(f"📄 Current file: {file_name}")
             
-            # Create a temporary file to save the PDF
+            # Create temporary file for PDF processing
             with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
-                # Write the uploaded file content to the temporary file
                 tmp_file.write(uploaded_file.getvalue())
                 tmp_file.flush()
                 
-                # Use PyPDFLoader to load the PDF
+                # Load PDF content
                 loader = PyPDFLoader(tmp_file.name)
                 pdf_documents = loader.load()
-                
-                # Extract text content from the PDF
                 content = "\n".join(doc.page_content for doc in pdf_documents)
                 
-                # Store in database
-                details_container.text(f"💾 Storing {file_name} in database...")
+                # 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
                 
-                documents.append(content)
-                document_names.append(file_name)
+                document_pairs.append((content, file_name))
             
-            # Update progress
             current_progress += progress_per_file
             progress_bar.progress(int(current_progress))
-            
-        if not documents:
+        
+        if not document_pairs:
             status_container.error("❌ No documents were successfully processed")
             return
-            
-        # Initialize vector store (80-90% progress)
+        
+        # Chunk documents (80% progress)
+        status_container.info("🔄 Chunking documents...")
+        details_container.text("📑 Splitting documents into manageable chunks...")
+        chunks, chunk_metadatas = chunker.process_documents(document_pairs)
+        
+        if not chunks:
+            status_container.error("❌ Failed to chunk documents")
+            return
+        
+        progress_bar.progress(80)
+        
+        # Initialize vector store (90% progress)
         status_container.info("🔄 Initializing vector store...")
         details_container.text("🔍 Creating vector embeddings...")
-        vector_store = initialize_faiss(embeddings, documents, document_names)
+        vector_store = initialize_faiss(embeddings, chunks, chunk_metadatas)
+        
         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
-        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
-            
-        # Initialize QA system (90-100% progress)
+        # Initialize QA system (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
-            
-        # Store QA system in session state
-        st.session_state.qa_system = qa_system
         
-        # Complete!
+        st.session_state.qa_system = qa_system
         progress_bar.progress(100)
+        
+        # Success message
         status_container.success("✅ Documents processed successfully!")
-        details_container.markdown("""
+        details_container.markdown(f"""
         🎉 **Ready to chat!**
-        - Documents loaded: {}
-        - Total content size: {:.2f} KB
+        - 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
         - 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()
-        
-        # 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') is not None:
             time.sleep(5)
             progress_container.empty()