Update app.py

app.py

@@ -12,7 +12,7 @@ import re
 import psutil
 import subprocess
 
-# Persistent directory
+# Persistent directory setup
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
@@ -132,20 +132,25 @@ def init_agent():
     return agent
 
 def clean_response(response: str) -> str:
-    """Clean the response by removing tool calls and duplicate content."""
-    # Remove all tool call blocks
-    response = re.sub(r'\[TOOL_CALLS\].*?$', '', response, flags=re.DOTALL)
+    """Enhanced response cleaner that handles duplicates and tool calls."""
+    # First extract the main analysis content
+    analysis_match = re.search(
+        r'(Based on the medical records provided.*?)(?=\[TOOL_CALLS\]|Based on|$)',
+        response,
+        flags=re.DOTALL
+    )
     
-    # Remove duplicate content by keeping only the last occurrence
-    if "Based on the medical records provided" in response:
-        last_occurrence = response.rfind("Based on the medical records provided")
-        response = response[last_occurrence:]
+    if analysis_match:
+        cleaned = analysis_match.group(1).strip()
+    else:
+        # Fallback if pattern not found
+        cleaned = re.sub(r'\[TOOL_CALLS\].*?$', '', response, flags=re.DOTALL).strip()
     
     # Remove any remaining JSON artifacts
-    response = re.sub(r'\{.*?\}', '', response)
-    response = re.sub(r'\[.*?\]', '', response)
+    cleaned = re.sub(r'\{.*?\}', '', cleaned)
+    cleaned = re.sub(r'\[.*?\]', '', cleaned)
     
-    return response.strip()
+    return cleaned
 
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -207,20 +212,20 @@ Medical Records:
                             finish_detected = True
 
                     # Display intermediate response
-                    clean_intermediate = clean_response(full_response)
-                    if clean_intermediate:
-                        history[-1] = {"role": "assistant", "content": clean_intermediate}
+                    current_cleaned = clean_response(full_response)
+                    if current_cleaned:
+                        history[-1] = {"role": "assistant", "content": current_cleaned}
                         yield history, None
 
-                # Final processing after all chunks are received
+                # Final processing
                 final_cleaned = clean_response(full_response)
                 
                 if not final_cleaned:
                     final_cleaned = "⚠️ No clear oversights identified or model output was invalid."
                 
-                # Only save report if we got a proper finish signal
+                # Save report
                 report_path = None
-                if finish_detected and file_hash_value:
+                if file_hash_value:
                     report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
                     with open(report_path, "w", encoding="utf-8") as f:
                         f.write(final_cleaned)