Update components/chat.py

components/chat.py

@@ -13,24 +13,42 @@ from utils.response_formatter import EnhancedResponseFormatter, display_enhanced
 qa_lock = Lock()
 
 def display_chat_interface():
-    """
-    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 if needed
-    if 'messages' not in st.session_state:
+    """Display chat interface with proper formatting."""
+    if not st.session_state.get('messages'):
         st.session_state.messages = []
 
-    # Display existing chat history
-    _display_chat_history()
+    # Display chat messages
+    for message in st.session_state.messages:
+        if isinstance(message, HumanMessage):
+            with st.chat_message("user"):
+                st.write(message.content)
+        elif isinstance(message, AIMessage):
+            with st.chat_message("assistant"):
+                display_enhanced_response(message.content, message.additional_kwargs.get('sources'))
 
-    # Handle new user input
+    # Chat input
     if prompt := st.chat_input("Ask about your documents..."):
-        _process_user_message(prompt)
+        with st.chat_message("user"):
+            st.write(prompt)
+            st.session_state.messages.append(HumanMessage(content=prompt))
+
+        with st.chat_message("assistant"):
+            with st.spinner("Analyzing documents..."):
+                response = st.session_state.qa_system.invoke({
+                    "input": prompt,
+                    "chat_history": st.session_state.messages
+                })
+                
+                if response:
+                    ai_message = AIMessage(
+                        content=str(response),
+                        additional_kwargs={'sources': response.metadata.get('sources', [])}
+                    )
+                    st.session_state.messages.append(ai_message)
+                    display_enhanced_response(
+                        str(response),
+                        response.metadata.get('sources', [])
+                    )
 
 def _verify_chat_ready() -> bool:
     """Check if the chat system is properly initialized."""