Update src/chat_display.py

src/chat_display.py

@@ -4,9 +4,9 @@ This module handles the chat display components and rendering.
 
 import streamlit as st
 import os
+import re
 from utils import has_meaningful_content, remove_reasoning_and_sources, clean_explanation, get_image_base64
 from session_state import get_full_history
-import re
 
 
 def get_avatars():
@@ -31,6 +31,50 @@ def get_avatars():
     return user_avatar, assistant_avatar
 
 
+def extract_follow_up_questions(content):
+    """
+    Extract follow-up questions from the main content.
+    
+    Args:
+        content (str): The message content
+        
+    Returns:
+        str: Extracted follow-up questions or empty string if none found
+    """
+    # Look for various forms of follow-up questions sections
+    patterns = [
+        r'(?i)follow[ -]?up questions:?\s*(.*?)(?=\n+\s*(?:#{1,3}|reasoning:|sources:|\Z))',
+        r'(?i)#{1,3}\s*follow[ -]?up questions:?\s*(.*?)(?=\n+\s*(?:#{1,3}|reasoning:|sources:|\Z))',
+        r'(?i)important questions to ask:?\s*(.*?)(?=\n+\s*(?:#{1,3}|reasoning:|sources:|\Z))',
+        r'(?i)clarifying questions:?\s*(.*?)(?=\n+\s*(?:#{1,3}|reasoning:|sources:|\Z))'
+    ]
+    
+    # Try each pattern until we find a match
+    for pattern in patterns:
+        follow_up_match = re.search(pattern, content, re.DOTALL)
+        if follow_up_match:
+            follow_up_text = follow_up_match.group(1).strip()
+            # Format the questions into a numbered list
+            questions = []
+            lines = follow_up_text.split('\n')
+            question_num = 1
+            
+            for line in lines:
+                # Check if this is a question line (starts with a number or bullet)
+                question_match = re.match(r'^\s*(?:\d+\.|\-|\•|\*)\s*(.*)', line.strip())
+                if question_match and line.strip():
+                    questions.append(f"{question_num}. {question_match.group(1).strip()}")
+                    question_num += 1
+                elif line.strip():
+                    # If not formatted as a list but has content
+                    questions.append(f"{question_num}. {line.strip()}")
+                    question_num += 1
+            
+            return "\n".join(questions)
+    
+    return ""
+
+
 def display_chat_history():
     """
     Display the chat history from the database.
@@ -40,9 +84,6 @@ def display_chat_history():
     # Get full history from database
     history = get_full_history()
     
-    # Show message count for debugging
-    # st.caption(f"Displaying {len(history)} messages")
-    
     # Display all messages instead of limiting to a fixed number
     display_history = history
     
@@ -64,6 +105,12 @@ 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:
+                            with st.expander("Show Follow-up Questions"):
+                                st.markdown(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")):
                             # Clean up the explanation text