Update components/chat.py

components/chat.py

@@ -1,118 +1,87 @@
-# components_chat.py
+# components/chat.py
 
 import streamlit as st
 from langchain_core.messages import HumanMessage, AIMessage
 from utils.database import verify_vector_store
 from threading import Lock
+from typing import Optional
+import traceback
 
 # Create a lock for QA system access
 qa_lock = Lock()
 
 def display_chat_interface():
-    """Display modern chat interface with clean formatting."""
-    
-    # Add custom CSS for modern chat styling
-    st.markdown("""
-    <style>
-        /* Clean, modern chat container */
-        .chat-container {
-            max-width: 800px;
-            margin: auto;
-        }
-        
-        /* Message styling */
-        .user-message {
-            background-color: #f0f2f6;
-            padding: 1rem;
-            border-radius: 10px;
-            margin: 1rem 0;
-        }
-        
-        .assistant-message {
-            background-color: #ffffff;
-            border: 1px solid #e0e0e0;
-            padding: 1.5rem;
-            border-radius: 10px;
-            margin: 1rem 0;
-        }
-        
-        /* Section styling */
-        .section-header {
-            font-weight: bold;
-            color: #0f52ba;
-            margin-top: 1rem;
-        }
-        
-        .source-info {
-            background-color: #f8f9fa;
-            padding: 0.5rem;
-            border-left: 3px solid #0f52ba;
-            margin-top: 1rem;
-            font-size: 0.9rem;
-        }
-        
-        /* Content sections */
-        .content-section {
-            margin: 0.5rem 0;
-        }
-        
-        /* Clean bullets */
-        .clean-bullet {
-            margin: 0.3rem 0;
-        }
-    </style>
-    """, unsafe_allow_html=True)
-    
-    # Check if QA system is initialized
-    if 'qa_system' not in st.session_state or st.session_state.qa_system is None:
-        st.warning("Please upload documents first to initialize the chat system.")
+    """
+    Display a modern chat interface with error handling and clean formatting.
+    Manages conversation state and handles document context.
+    """
+    # Verify QA system initialization
+    if not _verify_chat_ready():
         return
-        
-    # Initialize chat history
+
+    # Initialize chat history if needed
     if 'messages' not in st.session_state:
         st.session_state.messages = []
-    
-    # Display chat history
+
+    # Display existing chat history
+    _display_chat_history()
+
+    # Handle new user input
+    if prompt := st.chat_input("Ask about your documents..."):
+        _process_user_message(prompt)
+
+def _verify_chat_ready() -> bool:
+    """Check if the chat system is properly initialized."""
+    if 'qa_system' not in st.session_state or st.session_state.qa_system is None:
+        st.warning("Please upload documents first to initialize the chat system.")
+        return False
+    return True
+
+def _display_chat_history():
+    """Display all messages in the chat history."""
     for message in st.session_state.messages:
         if isinstance(message, HumanMessage):
-            st.markdown(f"""
-            <div class="user-message">
-                🧑‍💼 <strong>You:</strong><br>{message.content}
-            </div>
-            """, unsafe_allow_html=True)
+            with st.chat_message("user"):
+                st.write(message.content)
         elif isinstance(message, AIMessage):
-            content = str(message.content)
-            st.markdown(f"""
-            <div class="assistant-message">
-                🤖 <strong>Assistant:</strong><br>
-                {content}
-            </div>
-            """, unsafe_allow_html=True)
+            with st.chat_message("assistant"):
+                st.write(message.content)
+
+def _process_user_message(prompt: str):
+    """
+    Process a new user message and generate AI response.
     
-    # Chat input at the bottom
-    if prompt := st.chat_input("Ask about your documents..."):
-        try:
-            with st.spinner("Analyzing..."):
-                # Create message objects
-                human_message = HumanMessage(content=prompt)
-                st.session_state.messages.append(human_message)
+    Args:
+        prompt (str): User's input message
+    """
+    try:
+        with st.spinner("Analyzing..."):
+            # Create and display user message
+            human_message = HumanMessage(content=prompt)
+            st.session_state.messages.append(human_message)
+            with st.chat_message("user"):
+                st.write(prompt)
+
+            # Generate AI response with thread safety
+            with qa_lock:
+                response = st.session_state.qa_system.invoke({
+                    "input": prompt,
+                    "chat_history": st.session_state.messages
+                })
 
-                # Use lock to access QA system
-                with qa_lock:
-                    response = st.session_state.qa_system.invoke({
-                        "input": prompt,
-                        "chat_history": st.session_state.messages
-                    })
+            if response:
+                # Create and display AI message
+                ai_message = AIMessage(content=str(response))
+                st.session_state.messages.append(ai_message)
+                with st.chat_message("assistant"):
+                    st.write(response)
+                
+                # Rerun to update UI cleanly
+                st.rerun()
+            else:
+                st.error("No response received. Please try again.")
 
-                if response:
-                    content = str(response)
-                    ai_message = AIMessage(content=content)
-                    st.session_state.messages.append(ai_message)
-                    st.rerun()
-                else:
-                    st.error("No valid response received")
-                    
-        except Exception as e:
-            st.error(f"Error: {e}")
-            import traceback
-            st.error(traceback.format_exc())
+    except Exception as e:
+        st.error(f"An error occurred while processing your message: {str(e)}")
+        if st.session_state.get('debug_mode'):
+            st.error(traceback.format_exc())