Update app.py

app.py

@@ -139,9 +139,25 @@ 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."""
     text = sanitize_utf8(text)
-    text = re.sub(r"\[TOOL_CALLS\].*", "", text, flags=re.DOTALL)
-    text = re.sub(r"\n{3,}", "\n\n", text).strip()
+    # 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)
+    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."
     return text
 
 def init_agent():
@@ -211,11 +227,21 @@ def create_ui(agent):
             combined_response = ""
 
             prompt_template = """
-You are a medical analysis assistant. Analyze the following patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown format under these headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, and Urgent Follow-up. For each finding, include:
+You are a medical analysis assistant. Analyze the following patient record excerpt for clinical oversights and provide a concise, evidence-based summary in markdown format under these headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, and Urgent Follow-up. For each finding, include:
 - Clinical context (why the issue was missed or relevant details from the record).
 - Potential risks if unaddressed (e.g., disease progression, adverse events).
 - Actionable recommendations (e.g., tests, referrals, medication adjustments).
-If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text, formatted as markdown with bullet points under each heading, and avoids generic or static responses.
+Output ONLY the markdown-formatted findings, with bullet points under each heading. Do NOT include reasoning, tool calls, or intermediate steps. If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text and avoids generic responses.
+
+Example Output:
+### Missed Diagnoses
+- Elevated BP noted without diagnosis. Missed due to inconsistent visits. Risks: stroke. Recommend: BP monitoring, antihypertensives.
+### Medication Conflicts
+- No issues identified.
+### Incomplete Assessments
+- Chest pain not evaluated. Time constraints likely cause. Risks: cardiac issues. Recommend: ECG, stress test.
+### Urgent Follow-up
+- Abnormal creatinine not addressed. Delayed lab review. Risks: renal failure. Recommend: nephrology referral.
 
 Patient Record Excerpt (Chunk {0} of {1}):
 {chunk}
@@ -258,7 +284,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                             for m in chunk_output:
                                 if hasattr(m, 'content') and m.content:
                                     cleaned = clean_response(m.content)
-                                    if cleaned:
+                                    if cleaned and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
                                         chunk_response += cleaned + "\n"
                                         # Update UI with partial response
                                         if history[-1]["content"].startswith("Analyzing"):
@@ -268,7 +294,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                                         yield history, None
                         elif isinstance(chunk_output, str) and chunk_output.strip():
                             cleaned = clean_response(chunk_output)
-                            if cleaned:
+                            if cleaned and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
                                 chunk_response += cleaned + "\n"
                                 # Update UI with partial response
                                 if history[-1]["content"].startswith("Analyzing"):
@@ -278,7 +304,10 @@ Patient Record Excerpt (Chunk {0} of {1}):
                                 yield history, None
 
                     # Append completed chunk response to combined response
-                    combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
+                    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"
 
                 # Finalize UI with complete response
                 if combined_response: