Update src/chat_display.py

src/chat_display.py

@@ -136,11 +136,10 @@ def display_chat_history():
                         cleaned_response = remove_reasoning_and_sources(message["content"])
                         st.markdown(cleaned_response)
                         
-                        # Extract follow-up questions
-                        follow_up_questions = extract_follow_up_questions(message["content"])
-                        if follow_up_questions:
+                        # Use follow-up questions from the message object directly
+                        if message.get("follow_up_questions") and message["follow_up_questions"].strip():
                             with st.expander("Additional Questions"):
-                                st.markdown(follow_up_questions)
+                                st.markdown(message["follow_up_questions"])
                         
                         # Only display the explanation in an expander if it exists AND has actual content
                         if message.get("explanation") and has_meaningful_content(message.get("explanation")):