Update src/utils.py

src/utils.py

@@ -34,18 +34,25 @@ def has_meaningful_content(text):
 
 def remove_reasoning_and_sources(text):
     """
-    Remove reasoning and sources sections from the main response text.
+    Remove reasoning, follow-up questions, and sources sections from the main response text.
     
     Args:
         text (str): The text to clean
         
     Returns:
-        str: Text without reasoning and sources sections
+        str: Text without reasoning, follow-up questions, and sources sections
     """
+    if not text:
+        return text
+        
     # First, remove any reasoning sections
-    pattern_reasoning = r'(?i)(\n+\s*reasoning:|\n+\s*\*{0,2}reasoning\*{0,2}:?|\n+\s*#{1,3}\s*reasoning).*?(?=\n+\s*(?:#{1,3}|sources:|references:|\Z))'
+    pattern_reasoning = r'(?i)(\n+\s*reasoning:|\n+\s*\*{0,2}reasoning\*{0,2}:?|\n+\s*#{1,3}\s*reasoning).*?(?=\n+\s*(?:#{1,3}|follow[ -]?up questions:|sources:|references:|\Z))'
     cleaned_text = re.sub(pattern_reasoning, '', text, flags=re.DOTALL)
     
+    # Remove any follow-up questions sections
+    pattern_followup = r'(?i)(\n+\s*follow[ -]?up questions:|\n+\s*\*{0,2}follow[ -]?up questions\*{0,2}:?|\n+\s*#{1,3}\s*follow[ -]?up questions).*?(?=\n+\s*(?:#{1,3}|reasoning:|sources:|references:|\Z))'
+    cleaned_text = re.sub(pattern_followup, '', cleaned_text, flags=re.DOTALL)
+    
     # Then, remove any sources/references sections
     pattern_sources = r'(?i)(\n+\s*sources:|\n+\s*references:|\n+\s*\*{0,2}sources\*{0,2}:?|\n+\s*\*{0,2}references\*{0,2}:?|\n+\s*#{1,3}\s*sources|\n+\s*#{1,3}\s*references).*?(?=\n+\s*(?:#{1,3}|\Z))'
     cleaned_text = re.sub(pattern_sources, '', cleaned_text, flags=re.DOTALL)
@@ -59,8 +66,8 @@ def remove_reasoning_and_sources(text):
     skip_section = False
     
     for line in lines:
-        # Check if we should skip this line (part of reasoning or sources)
-        if re.search(r'(?i)^(\s*reasoning:|\s*sources:|\s*references:|\s*\*{0,2}reasoning\*{0,2}:?|\s*\*{0,2}sources\*{0,2}:?|\s*\*{0,2}references\*{0,2}:?|\s*#{1,3}\s*reasoning|\s*#{1,3}\s*sources|\s*#{1,3}\s*references)', line):
+        # Check if we should skip this line (part of reasoning, follow-up questions, or sources)
+        if re.search(r'(?i)^(\s*reasoning:|\s*follow[ -]?up questions:|\s*sources:|\s*references:|\s*\*{0,2}reasoning\*{0,2}:?|\s*\*{0,2}follow[ -]?up questions\*{0,2}:?|\s*\*{0,2}sources\*{0,2}:?|\s*\*{0,2}references\*{0,2}:?|\s*#{1,3}\s*reasoning|\s*#{1,3}\s*follow[ -]?up questions|\s*#{1,3}\s*sources|\s*#{1,3}\s*references)', line):
             skip_section = True
             continue
         # Check if we're entering a new section
@@ -75,6 +82,11 @@ def remove_reasoning_and_sources(text):
     result = '\n'.join(filtered_lines).strip()
     result = re.sub(r'\[([^\]]+)\]\(https?://[^)]+\)', r'\1', result)
     
+    # Also remove any sections starting with the headers Immediate Response or Main Response
+    # We want to preserve this content but remove the header itself
+    result = re.sub(r'(?i)^(\s*#{1,3}\s*)?immediate response:?\s*\n', '', result)
+    result = re.sub(r'(?i)^(\s*#{1,3}\s*)?main response:?\s*\n', '', result)
+    
     return result
 
 
@@ -214,4 +226,55 @@ def format_conversation_history(history, patient_info=None):
                 if explanation:
                     formatted_text += f"REASONING: {explanation}\n\n"
     
-    return formatted_text 
+    return formatted_text
+
+
+def format_follow_up_questions(questions_text):
+    """
+    Format follow-up questions text for display.
+    
+    Args:
+        questions_text (str): Raw follow-up questions text
+        
+    Returns:
+        str: Formatted follow-up questions
+    """
+    if not questions_text:
+        return ""
+        
+    # Clean up any header text
+    cleaned_text = re.sub(r'(?i)^(\s*#{1,3}\s*)?follow[ -]?up questions:?\s*\n', '', questions_text)
+    
+    # Ensure questions are numbered consistently
+    lines = cleaned_text.split('\n')
+    formatted_lines = []
+    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)
+        if question_match:
+            # Replace the existing number/bullet with a consistent format
+            formatted_lines.append(f"{question_num}. {question_match.group(1).strip()}")
+            question_num += 1
+        elif line.strip():
+            # If it's not empty and doesn't look like a numbered question, 
+            # treat it as a continuation of the previous question or a new question
+            if formatted_lines and formatted_lines[-1].endswith('?'):
+                # If the previous line ends with a question mark, this is likely a new question
+                formatted_lines.append(f"{question_num}. {line.strip()}")
+                question_num += 1
+            elif formatted_lines:
+                # Otherwise it's a continuation of the previous question
+                formatted_lines[-1] += " " + line.strip()
+            else:
+                # If there's no previous line, start a new question
+                formatted_lines.append(f"{question_num}. {line.strip()}")
+                question_num += 1
+    
+    # Ensure each question ends with a question mark
+    for i in range(len(formatted_lines)):
+        if not formatted_lines[i].endswith('?'):
+            formatted_lines[i] += '?'
+    
+    return '\n'.join(formatted_lines)