Update app.py

app.py

@@ -1,5 +1,19 @@
-
-
+import sys
+import os
+import pandas as pd
+import pdfplumber
+import json
+import gradio as gr
+from typing import List
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import hashlib
+import shutil
+import re
+import psutil
+import subprocess
+
+# Persistent directory
+persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
 model_cache_dir = os.path.join(persistent_dir, "txagent_models")
@@ -119,21 +133,24 @@ def init_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)
+    # First remove all tool call blocks
+    response = re.sub(r'\[TOOL_CALLS\].*?(\[TOOL_CALLS\]|$)', '', response, flags=re.DOTALL)
     
-    # Remove duplicate sentences (simple approach)
-    sentences = [s.strip() for s in response.split('.') if s.strip()]
-    unique_sentences = []
-    seen_sentences = set()
+    # Then remove any remaining standalone tool call markers
+    response = response.replace('[TOOL_CALLS]', '')
     
-    for sentence in sentences:
-        if sentence not in seen_sentences:
-            seen_sentences.add(sentence)
-            unique_sentences.append(sentence)
+    # Remove duplicate sections (looking for repeated identical paragraphs)
+    paragraphs = [p.strip() for p in response.split('\n\n') if p.strip()]
+    unique_paragraphs = []
+    seen_paragraphs = set()
+    
+    for para in paragraphs:
+        if para not in seen_paragraphs:
+            seen_paragraphs.add(para)
+            unique_paragraphs.append(para)
     
     # Reconstruct the response
-    cleaned = '. '.join(unique_sentences) + '.' if unique_sentences else response
+    cleaned = '\n\n'.join(unique_paragraphs)
     
     # Remove any remaining JSON-like artifacts
     cleaned = re.sub(r'\{.*?\}', '', cleaned)
@@ -177,7 +194,7 @@ Medical Records:
 """
 
             try:
-                response = ""
+                full_response = ""
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -190,24 +207,29 @@ Medical Records:
                     if chunk is None:
                         continue
                     if isinstance(chunk, str):
-                        response += chunk
+                        full_response += chunk
                     elif isinstance(chunk, list):
-                        response += "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-
-                # Clean the response before displaying
-                clean_response_text = clean_response(response)
-                
-                if not clean_response_text:
-                    clean_response_text = "⚠️ No clear oversights identified or model output was invalid."
-                
+                        full_response += "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
+
+                    # Clean the current response for display
+                    current_cleaned = clean_response(full_response)
+                    if current_cleaned:
+                        history[-1] = {"role": "assistant", "content": current_cleaned}
+                        yield history, None
+
+                # Final cleaning and processing
+                final_cleaned = clean_response(full_response)
+                if not final_cleaned:
+                    final_cleaned = "⚠️ No clear oversights identified or model output was invalid."
+
                 # Save the full report
                 report_path = None
                 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(clean_response_text)
-                
-                history[-1] = {"role": "assistant", "content": clean_response_text}
+                        f.write(final_cleaned)
+
+                history[-1] = {"role": "assistant", "content": final_cleaned}
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e: