Update app.py

app.py

@@ -1,4 +1,7 @@
-
+import sys
+import os
+import pandas as pd
+import pdfplumber
 import json
 import gradio as gr
 from typing import List
@@ -9,9 +12,7 @@ import re
 import psutil
 import subprocess
 
-# ---------------------------------------------------------------------------------------
-# Persistent directory setup
-# ---------------------------------------------------------------------------------------
+# Persistent directory
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
@@ -36,13 +37,8 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# ---------------------------------------------------------------------------------------
-# Utilities
-# ---------------------------------------------------------------------------------------
-MEDICAL_KEYWORDS = {
-    'diagnosis', 'assessment', 'plan', 'results', 'medications',
-    'allergies', 'summary', 'impression', 'findings', 'recommendations'
-}
+MEDICAL_KEYWORDS = {'diagnosis', 'assessment', 'plan', 'results', 'medications',
+                    'allergies', 'summary', 'impression', 'findings', 'recommendations'}
 
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
@@ -60,7 +56,7 @@ def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
                 text_chunks.append(f"=== Page {i+1} ===\n{text.strip()}")
             for i, page in enumerate(pdf.pages[3:max_pages], start=4):
                 page_text = page.extract_text() or ""
-                if any(re.search(rf'\b{kw}\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
+                if any(re.search(rf'\\b{kw}\\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
                     text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
         return "\n\n".join(text_chunks)
     except Exception as e:
@@ -78,7 +74,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             text = extract_priority_pages(file_path)
             result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"})
         elif file_type == "csv":
-            df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str, 
+            df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str,
                              skip_blank_lines=False, on_bad_lines="skip")
             content = df.fillna("").astype(str).values.tolist()
             result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
@@ -110,12 +106,11 @@ def log_system_usage(tag=""):
             used, total, util = result.stdout.strip().split(", ")
             print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
     except Exception as e:
-        print(f"[{tag}] ⚠️ GPU/CPU monitor failed: {e}")
+        print(f"[{tag}] GPU/CPU monitor failed: {e}")
 
 def init_agent():
     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")
     if not os.path.exists(target_tool_path):
@@ -134,12 +129,8 @@ def init_agent():
     agent.init_model()
     log_system_usage("After Load")
     print("✅ Agent Ready")
-
     return agent
 
-# ---------------------------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------------------------
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
@@ -174,7 +165,7 @@ Medical Records:
 
 ### Potential Oversights:
 """
-            response = ""
+            response_chunks = []
             try:
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
@@ -185,21 +176,23 @@ Medical Records:
                     call_agent=False,
                     conversation=[]
                 ):
-                    if chunk is None: continue
+                    if not chunk:
+                        continue
                     if isinstance(chunk, str):
-                        response += chunk
+                        response_chunks.append(chunk)
                     elif isinstance(chunk, list):
-                        response += "".join([c.content for c in chunk if hasattr(c, 'content')])
-                    cleaned = response.strip()
-                    cleaned = re.sub(r"\[TOOL_CALLS\].*?$", "", cleaned, flags=re.DOTALL).strip()
-                    history[-1] = {"role": "assistant", "content": cleaned or "⏳ Processing..."}
+                        response_chunks.extend([c.content for c in chunk if hasattr(c, 'content')])
+                    partial_response = "".join(response_chunks)
+                    cleaned_partial = partial_response.split("[TOOL_CALLS]")[0].strip()
+                    history[-1] = {"role": "assistant", "content": cleaned_partial or "⏳ Working..."}
                     yield history, None
             except Exception as e:
                 history[-1] = {"role": "assistant", "content": f"❌ Error: {str(e)}"}
                 yield history, None
                 return
 
-            final_output = re.sub(r"\[TOOL_CALLS\].*?$", "", response, flags=re.DOTALL).strip()
+            full_response = "".join(response_chunks)
+            final_output = full_response.split("[TOOL_CALLS]")[0].strip()
             if not final_output:
                 final_output = "No clear oversights identified. Recommend comprehensive review."
             history[-1] = {"role": "assistant", "content": final_output}
@@ -211,9 +204,6 @@ Medical Records:
         msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
     return demo
 
-# ---------------------------------------------------------------------------------------
-# Launch
-# ---------------------------------------------------------------------------------------
 if __name__ == "__main__":
     print("🚀 Launching app...")
     agent = init_agent()