Update components/chat.py

components/chat.py

@@ -6,7 +6,8 @@ from utils.database import verify_vector_store
 from threading import Lock
 from typing import Optional
 import traceback
-from utils.response_formatter import ResponseFormatter, display_formatted_response
+from utils.response_formatter import ResponseFormatter, display_enhanced_response
+
 
 # Create a lock for QA system access
 qa_lock = Lock()
@@ -38,15 +39,18 @@ def _verify_chat_ready() -> bool:
         return False
     return True
 
+
 def _process_user_message(prompt: str):
     """Process a new user message and generate AI response."""
     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 source tracking
             with qa_lock:
                 response = st.session_state.qa_system.invoke({
                     "input": prompt,
@@ -54,14 +58,25 @@ def _process_user_message(prompt: str):
                 })
 
             if response:
-                # Format the response
-                ai_message = AIMessage(content=str(response))
+                # Extract sources from response metadata
+                sources = response.metadata.get('sources', []) if hasattr(response, 'metadata') else []
+                
+                # Create and display AI message with enhanced formatting
+                ai_message = AIMessage(
+                    content=str(response),
+                    additional_kwargs={'sources': sources}
+                )
                 st.session_state.messages.append(ai_message)
+                
                 with st.chat_message("assistant"):
-                    display_formatted_response(
+                    display_enhanced_response(
                         str(response),
-                        metadata=getattr(response, 'metadata', None)
+                        sources=sources
                     )
+                
+                st.rerun()
+            else:
+                st.error("No response received. Please try again.")
 
     except Exception as e:
         st.error(f"An error occurred while processing your message: {str(e)}")