Update app.py

app.py

@@ -144,19 +144,106 @@ def initialize_database():
         st.error(f"Database initialization error: {e}")
         return False
 
-def initialize_session_state():
-    """Initialize all session state variables."""
+@dataclass
+class SessionStateDefaults:
+    """Default values for session state variables."""
+    # Collection Management
+    show_collection_dialog: bool = False
+    selected_collection: Optional[Dict] = None
+    current_collection: Optional[Dict] = None
     
-    default_states = {
-        'show_collection_dialog': False,
-        'selected_collection': None,
-        'chat_ready': False,
-        'current_collection': None,
-        'vector_store': None,
-        'qa_system': None,
-        'messages': [],
-        'current_chat_id': None,
-        'reinitialize_chat': False
+    # Chat System
+    chat_ready: bool = False
+    messages: list[BaseMessage] = None  # Will be initialized as empty list
+    current_chat_id: Optional[int] = None
+    
+    # Document Processing
+    vector_store: Optional[FAISS] = None
+    qa_system: Optional[ConversationalRetrievalChain] = None
+    
+    # State Management
+    reinitialize_chat: bool = False
+    debug_mode: bool = False
+
+def initialize_session_state() -> None:
+    """
+    Initialize all session state variables with proper typing and documentation.
+    
+    This function ensures all necessary session state variables are properly initialized
+    with appropriate default values and types.
+    """
+    defaults = SessionStateDefaults()
+    
+    # Dictionary mapping state variables to their descriptions
+    state_vars: Dict[str, Dict[str, Any]] = {
+        'show_collection_dialog': {
+            'default': defaults.show_collection_dialog,
+            'desc': 'Controls visibility of collection creation dialog'
+        },
+        'selected_collection': {
+            'default': defaults.selected_collection,
+            'desc': 'Currently selected collection for viewing/editing'
+        },
+        'current_collection': {
+            'default': defaults.current_collection,
+            'desc': 'Active collection for chat/analysis'
+        },
+        'chat_ready': {
+            'default': defaults.chat_ready,
+            'desc': 'Indicates if chat system is initialized and ready'
+        },
+        'messages': {
+            'default': [] if defaults.messages is None else defaults.messages,
+            'desc': 'List of chat messages in the current session'
+        },
+        'current_chat_id': {
+            'default': defaults.current_chat_id,
+            'desc': 'ID of the current chat session'
+        },
+        'vector_store': {
+            'default': defaults.vector_store,
+            'desc': 'FAISS vector store for document embeddings'
+        },
+        'qa_system': {
+            'default': defaults.qa_system,
+            'desc': 'Initialized QA system for chat'
+        },
+        'reinitialize_chat': {
+            'default': defaults.reinitialize_chat,
+            'desc': 'Flag to trigger chat system reinitialization'
+        },
+        'debug_mode': {
+            'default': defaults.debug_mode,
+            'desc': 'Enable/disable debug information'
+        }
+    }
+    
+    # Initialize each state variable if not already present
+    for var_name, config in state_vars.items():
+        if var_name not in st.session_state:
+            st.session_state[var_name] = config['default']
+
+def reset_chat_state() -> None:
+    """Reset chat-related session state variables to defaults."""
+    st.session_state.messages = []
+    st.session_state.current_chat_id = None
+    st.session_state.chat_ready = False
+    st.session_state.qa_system = None
+
+def reset_collection_state() -> None:
+    """Reset collection-related session state variables to defaults."""
+    st.session_state.selected_collection = None
+    st.session_state.current_collection = None
+    st.session_state.show_collection_dialog = False
+
+def get_current_state() -> Dict[str, Any]:
+    """
+    Get the current state of all session variables.
+    Useful for debugging and state management.
+    """
+    return {
+        key: value for key, value in st.session_state.items()
+        if not key.startswith('_')  # Exclude internal Streamlit states
     }
     
     for key, default_value in default_states.items():
@@ -258,44 +345,44 @@ def initialize_chat_from_collection():
         return False
 
 def display_collection_sidebar():
-    """Display the document management sidebar."""
+    """Display enhanced sidebar with collection management."""
     with st.sidebar:
-        st.title("📚 Document Manager")
-
-        # Collection Creation Button
-        if st.button("➕ Create New Collection", use_container_width=True):
-            st.session_state.show_collection_dialog = True
-
-        # Upload Section
-        st.header("Upload Documents")
+        st.title("💬 Chat Controls")
         
-        # Get collections for dropdown
+        if st.button("🔄 Start New Chat", use_container_width=True):
+            st.session_state.messages = []
+            st.session_state.current_chat_id = None
+            if st.session_state.chat_ready:
+                st.rerun()
+        
+        st.divider()
+        
+        # Collection Management
+        st.title("📚 Collections")
         collections = get_collections(st.session_state.db_conn)
-        if collections:
-            upload_to = st.selectbox(
-                "Upload to",
-                options=["Select Collection..."] + [c['name'] for c in collections],
-                key="upload_destination"
-            )
         
-        uploaded_files = st.file_uploader(
-            "Upload PDF documents",
-            type=['pdf'],
-            accept_multiple_files=True,
-            help="Limit 200MB per file • PDF"
-        )
-
-        if uploaded_files and upload_to != "Select Collection...":
-            collection_id = next(
-                (c['id'] for c in collections if c['name'] == upload_to),
-                None
+        if st.button("➕ Create Collection", use_container_width=True):
+            st.session_state.show_collection_dialog = True
+        
+        # Collection selection with document preview
+        if collections:
+            selected = st.selectbox(
+                "Select Collection",
+                options=["All Documents"] + [c['name'] for c in collections],
+                key="collection_select"
             )
-            if collection_id:
-                handle_document_upload(uploaded_files, collection_id=collection_id)
-
-    # Show collection creation dialog if triggered
-    if st.session_state.get('show_collection_dialog', False):
-        show_collection_creation_dialog()
+            
+            if selected != "All Documents":
+                collection = next((c for c in collections if c['name'] == selected), None)
+                if collection:
+                    documents = get_collection_documents(st.session_state.db_conn, collection['id'])
+                    if documents:
+                        st.markdown("### Documents in Collection")
+                        for doc in documents:
+                            with st.expander(f"📄 {doc['name']}", expanded=False):
+                                st.caption(f"Uploaded: {doc['upload_date']}")
+                                if st.button("Use for Chat", key=f"use_{doc['id']}"):
+                                    initialize_chat_for_document(doc['id'])
 
 def display_collection_dialog():
     """Display the create collection dialog."""