Update app.py

app.py

@@ -57,13 +57,39 @@ def extract_section(text: str, heading: str) -> str:
         return ""
 
 def structure_medical_response(text: str) -> Dict:
+    """Improved version that handles both markdown and plain text formats"""
+    def extract_improved(text: str, heading: str) -> str:
+        # Try multiple patterns to match different heading formats
+        patterns = [
+            rf"{re.escape(heading)}:\s*\n(.*?)(?=\n\s*\n|\Z)",  # Heading followed by content until double newline
+            rf"\*\*{re.escape(heading)}\*\*:\s*\n(.*?)(?=\n\s*\n|\Z)",  # Markdown bold heading
+            rf"{re.escape(heading)}[\s\-]+(.*?)(?=\n\s*\n|\Z)",  # Heading with dashes
+            rf"\n{re.escape(heading)}\s*\n(.*?)(?=\n\s*\n|\Z)"  # Heading on its own line
+        ]
+        
+        for pattern in patterns:
+            match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
+            if match:
+                content = match.group(1).strip()
+                # Clean up any remaining markdown or special characters
+                content = re.sub(r'^\s*[\-\*]\s*', '', content, flags=re.MULTILINE)
+                return content
+        
+        return ""
+
+    # Normalize the text first
+    text = text.replace('**', '').replace('__', '')
+    
     return {
-        "summary": extract_section(text, "Summarize the patient's medical history"),
-        "risks": extract_section(text, "Risks or Red Flags"),
-        "missed_issues": extract_section(text, "What the doctor might have missed"),
-        "recommendations": extract_section(text, "Suggested Clinical Actions")
+        "summary": extract_improved(text, "Summary of Patient's Medical History") or 
+                  extract_improved(text, "Summarize the patient's medical history"),
+        "risks": extract_improved(text, "Identify Risks or Red Flags") or 
+                extract_improved(text, "Risks or Red Flags"),
+        "missed_issues": extract_improved(text, "Missed Diagnoses or Treatments") or 
+                       extract_improved(text, "What the doctor might have missed"),
+        "recommendations": extract_improved(text, "Suggest Next Clinical Steps") or 
+                         extract_improved(text, "Suggested Clinical Actions")
     }
-
 def serialize_patient(patient: dict) -> dict:
     patient_copy = patient.copy()
     if "_id" in patient_copy: