Update app.py

app.py

@@ -139,25 +139,37 @@ def log_system_usage(tag=""):
         print(f"[{tag}] GPU/CPU monitor failed: {e}")
 
 def clean_response(text: str) -> str:
-    """Clean TxAgent response to remove tool calls, reasoning, and non-markdown content."""
+    """Clean TxAgent response to keep only markdown sections with valid findings."""
     text = sanitize_utf8(text)
     # Remove tool call artifacts, None, and reasoning
     text = re.sub(r"\[.*?\]|\bNone\b|To analyze the patient record excerpt.*?medications\.|Since the previous attempts.*?\.|I need to.*?medications\.|Retrieving tools.*?\.", "", text, flags=re.DOTALL)
     # Remove extra whitespace and non-markdown content
     text = re.sub(r"\n{3,}", "\n\n", text)
     text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)  # Keep markdown-relevant characters
-    # Extract only markdown sections
-    markdown_pattern = r"(###\s*(?:Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up).*?(?=\n###|\Z))|(-.*?)(?=\n-|\n###|\Z)"
-    matches = re.findall(markdown_pattern, text, re.DOTALL | re.MULTILINE)
+    
+    # Extract markdown sections with valid findings
+    sections = []
+    current_section = None
+    lines = text.splitlines()
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        if re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
+            current_section = line
+            sections.append([current_section])
+        elif current_section and re.match(r"-\s*.+", line) and not re.match(r"-\s*No issues identified", line):
+            sections[-1].append(line)
+    
+    # Combine only non-empty sections
     cleaned = []
-    for match in matches:
-        section = match[0].strip() or match[1].strip()
-        if section:
-            cleaned.append(section)
-    text = "\n".join(cleaned).strip()
-    # Ensure consistent markdown formatting
-    if not text.startswith("###"):
-        text = "### Missed Diagnoses\n- No issues identified.\n\n### Medication Conflicts\n- No issues identified.\n\n### Incomplete Assessments\n- No issues identified.\n\n### Urgent Follow-up\n- No issues identified."
+    for section in sections:
+        if len(section) > 1:  # Section has at least one finding
+            cleaned.append("\n".join(section))
+    
+    text = "\n\n".join(cleaned).strip()
+    if not text:
+        text = ""  # Return empty string if no valid findings
     return text
 
 def init_agent():
@@ -285,7 +297,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                                 if hasattr(m, 'content') and m.content:
                                     cleaned = clean_response(m.content)
                                     if cleaned and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
-                                        chunk_response += cleaned + "\n"
+                                        chunk_response += cleaned + "\n\n"
                                         # Update UI with partial response
                                         if history[-1]["content"].startswith("Analyzing"):
                                             history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"}
@@ -295,7 +307,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                         elif isinstance(chunk_output, str) and chunk_output.strip():
                             cleaned = clean_response(chunk_output)
                             if cleaned and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
-                                chunk_response += cleaned + "\n"
+                                chunk_response += cleaned + "\n\n"
                                 # Update UI with partial response
                                 if history[-1]["content"].startswith("Analyzing"):
                                     history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"}
@@ -307,13 +319,13 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     if chunk_response:
                         combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
                     else:
-                        combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n### Missed Diagnoses\n- No issues identified.\n\n### Medication Conflicts\n- No issues identified.\n\n### Incomplete Assessments\n- No issues identified.\n\n### Urgent Follow-up\n- No issues identified.\n\n"
+                        combined_response += f"--- Analysis for Chunk {chunk_idx} ---\nNo oversights identified for this chunk.\n\n"
 
                 # Finalize UI with complete response
-                if combined_response:
+                if combined_response.strip() and not all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
                     history[-1]["content"] = combined_response.strip()
                 else:
-                    history.append({"role": "assistant", "content": "No oversights identified."})
+                    history.append({"role": "assistant", "content": "No oversights identified in the provided records."})
 
                 # Generate report file
                 report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None