Update app.py

app.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import streamlit as st
 import os
+import json
+import time
 from pathlib import Path
-from typing import Dict, List, Any, Optional
+from typing import Dict, List, Any, Optional, Tuple
 from datetime import datetime
 from dataclasses import dataclass
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -12,6 +14,7 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.chat_models import ChatOpenAI
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
+import tempfile
 from utils.database import (
     create_connection,
     create_tables,
@@ -22,13 +25,15 @@ from utils.database import (
     initialize_qa_system,
     get_embeddings_model,
     verify_database_tables,
-    handle_document_upload,
     create_collection,
     add_document_to_collection,
     get_recent_documents,
     save_chat_message,
     create_new_chat,
-    get_chat_messages)
+    get_chat_messages,
+    get_document_tags,
+    add_document_tags,
+    delete_collection)
 
 
 @dataclass
@@ -44,84 +49,6 @@ class SessionState:
     reinitialize_chat: bool = False
 
 
-def initialize_session_state():
-    """Initialize session state with default values."""
-    defaults = SessionState()
-    if 'initialized' not in st.session_state:
-        # Setup data paths
-        data_path = Path('/data' if os.path.exists('/data') else 'data')
-        vector_store_path = data_path / 'vector_stores'
-
-        # Create necessary directories
-        data_path.mkdir(parents=True, exist_ok=True)
-        vector_store_path.mkdir(parents=True, exist_ok=True)
-
-        # Initialize session state
-        st.session_state.update({
-            'show_collection_dialog': defaults.show_collection_dialog,
-            'selected_collection': defaults.selected_collection,
-            'chat_ready': defaults.chat_ready,
-            'messages': [] if defaults.messages is None else defaults.messages,
-            'current_chat_id': defaults.current_chat_id,
-            'vector_store': defaults.vector_store,
-            'qa_system': defaults.qa_system,
-            'reinitialize_chat': defaults.reinitialize_chat,
-            'initialized': True,
-            'data_path': data_path,
-            'vector_store_path': vector_store_path,
-            'show_explorer': False
-        })
-
-
-def initialize_chat_system(collection_id=None) -> bool:
-    """Initialize chat system with persistent vector stores."""
-    try:
-        # Get documents based on collection or all documents
-        documents = (get_collection_documents(st.session_state.db_conn, collection_id)
-                     if collection_id else get_all_documents(st.session_state.db_conn))
-
-        if not documents:
-            st.error("No documents found.")
-            return False
-
-        # Initialize new vector store
-        with st.spinner("Processing documents..."):
-            embeddings = get_embeddings_model()
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=500,
-                chunk_overlap=50,
-                length_function=len,
-            )
-
-            chunks = []
-            for doc in documents:
-                doc_chunks = text_splitter.split_text(doc['content'])
-                chunks.extend([{
-                    'content': chunk,
-                    'metadata': {
-                        'source': doc['name'],
-                        'document_id': doc['id'],
-                        'collection_id': collection_id
-                    }
-                } for chunk in doc_chunks])
-
-            # Create new vector store with allow_dangerous_deserialization
-            vector_store = FAISS.from_texts(
-                [chunk['content'] for chunk in chunks],
-                embeddings,
-                [chunk['metadata'] for chunk in chunks]
-            )
-
-            st.session_state.vector_store = vector_store
-            st.session_state.qa_system = initialize_qa_system(vector_store)
-            st.session_state.chat_ready = True
-            return True
-
-    except Exception as e:
-        st.error(f"Error initializing chat system: {e}")
-        return False
-
-
 def display_header():
     """Display the application header with navigation."""
     # Add custom CSS for header styling
@@ -188,67 +115,195 @@ def display_header():
     st.divider()
 
 
-def display_welcome_screen():
-    """Display enhanced welcome screen with quick actions."""
-    col1, col2 = st.columns([3, 2])
+async def process_document(file_path: str, collection_id: Optional[int] = None) -> Tuple[List, str]:
+    """Process a document with automatic tagging."""
+    try:
+        # Load PDF
+        loader = PyPDFLoader(file_path)
+        documents = loader.load()
+        
+        # Extract full content
+        full_content = "\n".join(doc.page_content for doc in documents)
+        
+        # Generate tags
+        tags = await generate_document_tags(full_content)
+        
+        # Create text splitter for chunks
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200,
+            length_function=len,
+            separators=["\n\n", "\n", " ", ""]
+        )
+        
+        # Split documents into chunks
+        chunks = text_splitter.split_documents(documents)
+        
+        # Add metadata to chunks
+        for chunk in chunks:
+            chunk.metadata.update({
+                'collection_id': collection_id,
+                'tags': tags
+            })
+        
+        return chunks, full_content, tags
+        
+    except Exception as e:
+        st.error(f"Error processing document: {e}")
+        return [], "", []
 
-    with col1:
-        st.header("Quick Start")
 
-        # Upload new documents with collection linking
-        st.markdown("### Upload Documents")
-        collection_id = None
-        collections = get_collections(st.session_state.db_conn)
+async def handle_document_upload(uploaded_files: List, collection_id: Optional[int] = None) -> bool:
+    """Handle document upload with progress tracking and auto-tagging."""
+    try:
+        progress_container = st.empty()
+        status_container = st.empty()
+        progress_bar = progress_container.progress(0)
         
-        if collections:
-            selected_collection = st.selectbox(
-                "Select Collection (Optional)",
-                options=[("None", None)] + [(c["name"], c["id"]) for c in collections],
-                format_func=lambda x: x[0]
-            )
-            collection_id = selected_collection[1] if selected_collection[0] != "None" else None
+        # Initialize embeddings
+        embeddings = get_embeddings_model()
+        if not embeddings:
+            status_container.error("Failed to initialize embeddings model")
+            return False
+        
+        progress_bar.progress(10)
+        all_chunks = []
+        documents = []
+        
+        # Process each document
+        progress_per_file = 70 / len(uploaded_files)
+        current_progress = 10
+        
+        for idx, uploaded_file in enumerate(uploaded_files):
+            status_container.info(f"Processing {uploaded_file.name}...")
             
-            # Add new collection button
-            if st.button("Create New Collection", use_container_width=True):
-                st.session_state.show_collection_dialog = True
-                st.rerun()
+            # Create temporary file
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
+                tmp_file.write(uploaded_file.getvalue())
+                tmp_file.flush()
+                
+                # Process document with tagging
+                chunks, content, tags = await process_document(tmp_file.name, collection_id)
+                
+                # Store in database
+                doc_id = insert_document(st.session_state.db_conn, uploaded_file.name, content)
+                if not doc_id:
+                    status_container.error(f"Failed to store document: {uploaded_file.name}")
+                    continue
+                
+                # Add tags
+                if tags:
+                    add_document_tags(st.session_state.db_conn, doc_id, tags)
+                
+                # Add to collection if specified
+                if collection_id:
+                    add_document_to_collection(st.session_state.db_conn, doc_id, collection_id)
+                
+                all_chunks.extend(chunks)
+                documents.append(content)
+            
+            current_progress += progress_per_file
+            progress_bar.progress(int(current_progress))
+        
+        # Initialize vector store
+        status_container.info("Creating document index...")
+        vector_store = FAISS.from_documents(all_chunks, embeddings)
+        
+        st.session_state.vector_store = vector_store
+        st.session_state.qa_system = initialize_qa_system(vector_store)
+        st.session_state.chat_ready = True
+        
+        progress_bar.progress(100)
+        status_container.success("Documents processed successfully!")
+        
+        # Clean up progress display
+        time.sleep(2)
+        progress_container.empty()
+        status_container.empty()
+        
+        return True
+        
+    except Exception as e:
+        st.error(f"Error uploading documents: {e}")
+        return False
 
-        uploaded_files = st.file_uploader(
-            "Upload Documents",
-            type=['pdf'],
-            accept_multiple_files=True,
-            help="Upload PDF documents to start analyzing"
-        )
 
-        if uploaded_files:
-            with st.spinner("Processing documents..."):
-                if handle_document_upload(uploaded_files, collection_id=collection_id):
-                    initialize_chat_system(collection_id)
+def display_collection_management():
+    """Display collection management interface."""
+    st.header("📁 Collection Management")
+    
+    # Get existing collections
+    collections = get_collections(st.session_state.db_conn)
+    
+    col1, col2 = st.columns([2, 1])
+    
+    with col1:
+        # Create new collection form
+        with st.form("create_collection_form"):
+            st.subheader("Create New Collection")
+            name = st.text_input("Collection Name")
+            description = st.text_area("Description")
+            submit = st.form_submit_button("Create Collection", use_container_width=True)
+            
+            if submit and name:
+                collection_id = create_collection(st.session_state.db_conn, name, description)
+                if collection_id:
+                    st.success(f"Collection '{name}' created successfully!")
                     st.rerun()
-
+    
     with col2:
-        # Display existing collections
-        st.header("Collections")
+        # Upload to existing collection
         if collections:
-            for collection in collections:
-                with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
-                    st.write(collection.get('description', ''))
-                    if st.button("Start Chat", key=f"chat_{collection['id']}", use_container_width=True):
-                        st.session_state.selected_collection = collection
-                        if initialize_chat_system(collection['id']):
-                            st.rerun()
-
-        # Show recent documents
-        st.header("Recent Documents")
-        recent_docs = get_recent_documents(st.session_state.db_conn, limit=5)
-        for doc in recent_docs:
-            with st.expander(f"📄 {doc['name']}"):
-                st.caption(f"Upload date: {doc['upload_date']}")
-                if doc['collections']:
-                    st.caption(f"Collections: {', '.join(doc['collections'])}")
-                if st.button("Start Chat", key=f"doc_{doc['id']}", use_container_width=True):
-                    if initialize_chat_system():
+            st.subheader("Add to Collection")
+            selected_collection = st.selectbox(
+                "Select Collection",
+                options=[(c["name"], c["id"]) for c in collections],
+                format_func=lambda x: x[0]
+            )
+            
+            uploaded_files = st.file_uploader(
+                "Upload Documents",
+                type=['pdf'],
+                accept_multiple_files=True,
+                key="collection_uploader"
+            )
+            
+            if uploaded_files:
+                collection_id = selected_collection[1]
+                with st.spinner("Processing documents..."):
+                    if await handle_document_upload(uploaded_files, collection_id=collection_id):
+                        st.success("Documents added to collection successfully!")
                         st.rerun()
+    
+    # Display existing collections
+    if collections:
+        st.markdown("### Existing Collections")
+        for collection in collections:
+            with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
+                col1, col2 = st.columns([3, 1])
+                
+                with col1:
+                    st.write(f"**Description:** {collection.get('description', 'No description')}")
+                    st.write(f"**Created:** {collection['created_at']}")
+                    
+                    # Display documents in collection
+                    docs = get_collection_documents(st.session_state.db_conn, collection['id'])
+                    if docs:
+                        st.write("**Documents:**")
+                        for doc in docs:
+                            st.write(f"- {doc['name']}")
+                            tags = get_document_tags(st.session_state.db_conn, doc['id'])
+                            if tags:
+                                st.write(f"  Tags: {', '.join(tags)}")
+                
+                with col2:
+                    st.button("Start Chat", key=f"chat_{collection['id']}", use_container_width=True)
+                    st.button("Upload Files", key=f"upload_{collection['id']}", use_container_width=True)
+                    if st.button("Delete", key=f"delete_{collection['id']}", use_container_width=True):
+                        if st.warning("Are you sure you want to delete this collection?"):
+                            if delete_collection(st.session_state.db_conn, collection['id']):
+                                st.success("Collection deleted successfully!")
+                                st.rerun()
 
 
 def display_chat_interface():
@@ -294,7 +349,6 @@ def display_chat_interface():
 
         st.rerun()
 
-# Rest of the code remains the same...
 
 def main():
     """Main application function with improved state management."""
@@ -304,7 +358,7 @@ def main():
         initial_sidebar_state="collapsed"
     )
 
-    # Initialize session state with paths
+    # Initialize session state
     initialize_session_state()
 
     # Initialize database connection
@@ -318,29 +372,13 @@ def main():
     # Display header
     display_header()
 
-    # Show collection creation dialog if triggered
+    # Show collection management if triggered
     if st.session_state.show_collection_dialog:
-        with st.form("create_collection"):
-            st.subheader("Create New Collection")
-            name = st.text_input("Collection Name")
-            description = st.text_area("Description")
-
-            if st.form_submit_button("Create", use_container_width=True):
-                if name:
-                    if create_collection(st.session_state.db_conn, name, description):
-                        st.success(f"Collection '{name}' created successfully!")
-                        st.session_state.show_collection_dialog = False
-                        st.rerun()
-                    else:
-                        st.error("Failed to create collection.")
-                else:
-                    st.warning("Please enter a collection name.")
-
-    # Display different views based on application state
-    if st.session_state.chat_ready:
+        display_collection_management()
+    # Display chat interface if ready
+    elif st.session_state.chat_ready:
         display_chat_interface()
-    elif st.session_state.show_explorer:
-        display_document_chunks()
+    # Show welcome screen
     else:
         display_welcome_screen()