Update app.py

app.py

@@ -139,7 +139,7 @@ def log_system_usage(tag=""):
         print(f"[{tag}] GPU/CPU monitor failed: {e}")
 
 def clean_response(text: str) -> str:
-    """Clean TxAgent response to keep only markdown sections with valid findings."""
+    """Clean TxAgent response to group findings under tool-derived headings."""
     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)
@@ -147,25 +147,53 @@ def clean_response(text: str) -> str:
     text = re.sub(r"\n{3,}", "\n\n", text)
     text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)  # Keep markdown-relevant characters
     
-    # Extract markdown sections with valid findings
-    sections = []
+    # Define tool-to-heading mapping
+    tool_to_heading = {
+        "get_abuse_info_by_drug_name": "Drugs",
+        "get_dependence_info_by_drug_name": "Drugs",
+        "get_abuse_types_and_related_adverse_reactions_and_controlled_substance_status_by_drug_name": "Drugs",
+        "get_info_for_patients_by_drug_name": "Drugs",
+        # Add other tools from new_tool.json if applicable
+    }
+    
+    # Parse sections and findings
+    sections = {}
     current_section = None
+    current_tool = 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)
+        # Detect tool tag
+        tool_match = re.match(r"\[TOOL:\s*(\w+)\]", line)
+        if tool_match:
+            current_tool = tool_match.group(1)
+            continue
+        # Detect section heading
+        section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
+        if section_match:
+            current_section = section_match.group(1)
+            if current_section not in sections:
+                sections[current_section] = []
+            continue
+        # Detect finding
+        finding_match = re.match(r"-\s*.+", line)
+        if finding_match and current_section and not re.match(r"-\s*No issues identified", line):
+            # Assign to tool-derived heading if tool is specified
+            if current_tool and current_tool in tool_to_heading:
+                heading = tool_to_heading[current_tool]
+                if heading not in sections:
+                    sections[heading] = []
+                sections[heading].append(line)
+            else:
+                sections[current_section].append(line)
     
-    # Combine only non-empty sections
+    # Combine non-empty sections
     cleaned = []
-    for section in sections:
-        if len(section) > 1:  # Section has at least one finding
-            cleaned.append("\n".join(section))
+    for heading, findings in sections.items():
+        if findings:  # Only include sections with findings
+            cleaned.append(f"### {heading}\n" + "\n".join(findings))
     
     text = "\n\n".join(cleaned).strip()
     if not text:
@@ -239,17 +267,18 @@ def create_ui(agent):
             combined_response = ""
 
             prompt_template = """
-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:
+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. Group findings under appropriate headings based on the tool used (e.g., drug-related findings under 'Drugs'). 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).
-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.
+Output ONLY the markdown-formatted findings, with bullet points under each heading. Precede each finding with a tool tag (e.g., [TOOL: get_abuse_info_by_drug_name]) to indicate the tool used. Do NOT include reasoning, tool calls, or intermediate steps. If no issues are found for a tool or category, state "No issues identified" for that section. Ensure the output is specific to the provided text and avoids generic responses.
 
 Example Output:
+### Drugs
+[TOOL: get_abuse_info_by_drug_name]
+- Opioid use disorder not addressed. Missed due to lack of screening. Risks: overdose. Recommend: addiction specialist referral.
 ### 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
@@ -257,18 +286,6 @@ Example Output:
 
 Patient Record Excerpt (Chunk {0} of {1}):
 {chunk}
-
-### Missed Diagnoses
-- ...
-
-### Medication Conflicts
-- ...
-
-### Incomplete Assessments
-- ...
-
-### Urgent Follow-up
-- ...
 """
 
             try:
@@ -296,7 +313,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 and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
+                                    if cleaned and re.search(r"###\s*\w+", cleaned):
                                         chunk_response += cleaned + "\n\n"
                                         # Update UI with partial response
                                         if history[-1]["content"].startswith("Analyzing"):
@@ -306,7 +323,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 and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
+                            if cleaned and re.search(r"###\s*\w+", cleaned):
                                 chunk_response += cleaned + "\n\n"
                                 # Update UI with partial response
                                 if history[-1]["content"].startswith("Analyzing"):