Update app.py

app.py

@@ -8,13 +8,12 @@ from typing import List
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
-import time
 import re
 import psutil
 import subprocess
 
 # ---------------------------------------------------------------------------------------
-# Persistent directory for Hugging Face Spaces
+# Persistent directory setup
 # ---------------------------------------------------------------------------------------
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
@@ -34,9 +33,6 @@ os.environ["VLLM_CACHE_DIR"] = vllm_cache_dir
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
 
-# ---------------------------------------------------------------------------------------
-# Add src to path
-# ---------------------------------------------------------------------------------------
 current_dir = os.path.dirname(os.path.abspath(__file__))
 src_path = os.path.abspath(os.path.join(current_dir, "src"))
 sys.path.insert(0, src_path)
@@ -44,7 +40,7 @@ sys.path.insert(0, src_path)
 from txagent.txagent import TxAgent
 
 # ---------------------------------------------------------------------------------------
-# Helper functions
+# Utilities
 # ---------------------------------------------------------------------------------------
 MEDICAL_KEYWORDS = {
     'diagnosis', 'assessment', 'plan', 'results', 'medications',
@@ -106,25 +102,22 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
 
 def log_system_usage(tag=""):
     try:
-        cpu_percent = psutil.cpu_percent(interval=1)
+        cpu = psutil.cpu_percent(interval=1)
         mem = psutil.virtual_memory()
-        print(f"[{tag}] 🧠 CPU: {cpu_percent}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB")
+        print(f"[{tag}] CPU: {cpu}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB")
         result = subprocess.run(
             ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,nounits,noheader"],
-            capture_output=True,
-            text=True,
+            capture_output=True, text=True
         )
         if result.returncode == 0:
-            mem_used, mem_total, util = result.stdout.strip().split(", ")
-            print(f"[{tag}] ⚡ GPU: {mem_used}MB / {mem_total}MB | Utilization: {util}%")
-        else:
-            print(f"[{tag}] ⚡ GPU info not available.")
+            used, total, util = result.stdout.strip().split(", ")
+            print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
     except Exception as e:
-        print(f"[{tag}] ⚠️ Failed to log system usage: {e}")
+        print(f"[{tag}] ⚠️ GPU/CPU monitor failed: {e}")
 
 def init_agent():
-    print("🔁 Initializing TxAgent...")
-    log_system_usage("Before Model Load")
+    print("🔁 Initializing model...")
+    log_system_usage("Before Load")
 
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
@@ -142,13 +135,8 @@ def init_agent():
         additional_default_tools=[],
     )
     agent.init_model()
-    log_system_usage("After Model Load")
-
-    print("✅ TxAgent is ready.")
-    print("📦 Cached model files:")
-    for root, _, files in os.walk(model_cache_dir):
-        for file in files:
-            print(os.path.join(root, file))
+    log_system_usage("After Load")
+    print("✅ Agent Ready")
 
     return agent
 
@@ -157,53 +145,42 @@ def init_agent():
 # ---------------------------------------------------------------------------------------
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        <h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>
-        <h3 style='text-align: center;'>Identify potential oversights in patient care</h3>
-        """)
+        gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
         chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
-        file_upload = gr.File(label="Upload Medical Records", 
-                              file_types=[".pdf", ".csv", ".xls", ".xlsx"], 
-                              file_count="multiple")
+        file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
         msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
         send_btn = gr.Button("Analyze", variant="primary")
         download_output = gr.File(label="Download Full Report")
 
-        def analyze_potential_oversights(message: str, history: list, files: list):
+        def analyze(message: str, history: list, files: list):
             history = history + [{"role": "user", "content": message},
                                  {"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."}]
             yield history, None
 
-            extracted_data = ""
+            extracted = ""
             file_hash_value = ""
-            if files and isinstance(files, list):
+            if files:
                 with ThreadPoolExecutor(max_workers=4) as executor:
-                    futures = [
-                        executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower())
-                        for f in files if hasattr(f, 'name')
-                    ]
+                    futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files]
                     results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
-                    extracted_data = "\n".join(results)
-                    file_hash_value = file_hash(files[0].name) if hasattr(files[0], 'name') else ""
+                    extracted = "\n".join(results)
+                    file_hash_value = file_hash(files[0].name)
 
-            max_extracted_chars = 12000
-            truncated_data = extracted_data[:max_extracted_chars]
-
-            analysis_prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
+            prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
 3. Note incomplete assessments
 4. Highlight abnormal results needing follow-up
 
 Medical Records:
-{truncated_data}
+{extracted[:12000]}
 
 ### Potential Oversights:
 """
             response = ""
             try:
                 for chunk in agent.run_gradio_chat(
-                    message=analysis_prompt,
+                    message=prompt,
                     history=[],
                     temperature=0.2,
                     max_new_tokens=1024,
@@ -216,46 +193,32 @@ Medical Records:
                         response += chunk
                     elif isinstance(chunk, list):
                         response += "".join([c.content for c in chunk if hasattr(c, 'content')])
-                    cleaned = response.replace("[TOOL_CALLS]", "").strip()
-                    history[-1] = {"role": "assistant", "content": cleaned}
+                    cleaned = response.strip()
+                    cleaned = re.sub(r"\[TOOL_CALLS\].*?$", "", cleaned, flags=re.DOTALL).strip()
+                    history[-1] = {"role": "assistant", "content": cleaned or "⏳ Processing..."}
                     yield history, None
-            except Exception as agent_error:
-                history[-1] = {"role": "assistant", "content": f"❌ Analysis failed: {str(agent_error)}"}
+            except Exception as e:
+                history[-1] = {"role": "assistant", "content": f"❌ Error: {str(e)}"}
                 yield history, None
                 return
 
-            final_output = response.replace("[TOOL_CALLS]", "").strip()
+            final_output = re.sub(r"\[TOOL_CALLS\].*?$", "", response, flags=re.DOTALL).strip()
             if not final_output:
                 final_output = "No clear oversights identified. Recommend comprehensive review."
-
             history[-1] = {"role": "assistant", "content": final_output}
 
-            report_path = None
-            if file_hash_value:
-                possible_report = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                if os.path.exists(possible_report):
-                    report_path = possible_report
-
-            yield history, report_path
+            report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
+            yield history, report_path if report_path and os.path.exists(report_path) else None
 
-        send_btn.click(analyze_potential_oversights,
-                       inputs=[msg_input, gr.State([]), file_upload],
-                       outputs=[chatbot, download_output])
-        msg_input.submit(analyze_potential_oversights,
-                         inputs=[msg_input, gr.State([]), file_upload],
-                         outputs=[chatbot, download_output])
-        gr.Examples([
-            ["What might have been missed in this patient's treatment?"],
-            ["Are there any medication conflicts in these records?"],
-            ["What abnormal results require follow-up?"]],
-            inputs=msg_input)
+        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
+        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
     return demo
 
 # ---------------------------------------------------------------------------------------
 # Launch
 # ---------------------------------------------------------------------------------------
 if __name__ == "__main__":
-    print("🚀 Starting TxAgent App...")
+    print("🚀 Launching app...")
     agent = init_agent()
     demo = create_ui(agent)
     demo.queue(api_open=False).launch(