Update app.py

app.py

@@ -1,281 +1,18 @@
-import sys
 import os
-import pandas as pd
-import pdfplumber
-import json
+import sys
 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
-import traceback
-import torch
-import copy
-from gradio import ChatMessage
-
-os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"
-if not torch.cuda.is_available():
-    print("No GPU detected. Forcing CPU mode by setting CUDA_VISIBLE_DEVICES to an empty string.")
-    os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
-persistent_dir = "/data/hf_cache"
-os.makedirs(persistent_dir, exist_ok=True)
-
-model_cache_dir = os.path.join(persistent_dir, "txagent_models")
-tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
-file_cache_dir = os.path.join(persistent_dir, "cache")
-report_dir = os.path.join(persistent_dir, "reports")
-vllm_cache_dir = os.path.join(persistent_dir, "vllm_cache")
-
-for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir, vllm_cache_dir]:
-    os.makedirs(directory, exist_ok=True)
-
-os.environ["HF_HOME"] = model_cache_dir
-os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
-os.environ["VLLM_CACHE_DIR"] = vllm_cache_dir
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
-
-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)
-
-from txagent.txagent import TxAgent
-
-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")
-
-def file_hash(path: str) -> str:
-    with open(path, "rb") as f:
-        return hashlib.md5(f.read()).hexdigest()
-
-def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
-    try:
-        text_chunks = []
-        with pdfplumber.open(file_path) as pdf:
-            for i, page in enumerate(pdf.pages[:3]):
-                text = page.extract_text() or ""
-                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):
-                    text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
-        return "\n\n".join(text_chunks)
-    except Exception as e:
-        print("PDF processing error:", str(e))
-        traceback.print_exc()
-        return str(e)
-
-def convert_file_to_json(file_path: str, file_type: str) -> str:
-    try:
-        h = file_hash(file_path)
-        cache_path = os.path.join(file_cache_dir, f"{h}.json")
-        if os.path.exists(cache_path):
-            with open(cache_path, "r", encoding="utf-8") as f:
-                return f.read()
-
-        if file_type == "pdf":
-            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, 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})
-        elif file_type in ["xls", "xlsx"]:
-            try:
-                df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str)
-            except Exception:
-                df = pd.read_excel(file_path, engine="xlrd", header=None, dtype=str)
-            content = df.fillna("").astype(str).values.tolist()
-            result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
-        else:
-            result = json.dumps({"error": f"Unsupported file type: {file_type}"})
-
-        with open(cache_path, "w", encoding="utf-8") as f:
-            f.write(result)
-        return result
-    except Exception as e:
-        print("Error processing", file_path, str(e))
-        traceback.print_exc()
-        return json.dumps({"error": str(e)})
-
-def log_system_usage(tag=""):
-    try:
-        cpu = psutil.cpu_percent(interval=1)
-        mem = psutil.virtual_memory()
-        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
-        )
-        if result.returncode == 0:
-            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}")
-        traceback.print_exc()
-
-def init_agent():
-    try:
-        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):
-            shutil.copy(default_tool_path, target_tool_path)
-
-        agent = TxAgent(
-            model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
-            rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-            tool_files_dict={"new_tool": target_tool_path},
-            enable_finish=True,
-            enable_rag=True,
-            enable_summary=False,
-            init_rag_num=0,
-            step_rag_num=8,
-            summary_mode='step',
-            summary_skip_last_k=0,
-            summary_context_length=None,
-            force_finish=True,
-            avoid_repeat=True,
-            seed=100,
-            enable_checker=True,
-            enable_chat=False,
-            additional_default_tools=[],
-        )
-        agent.init_model()
-        log_system_usage("After Load")
-        print("✅ Agent Ready")
-        return agent
-    except Exception as e:
-        print("❌ Error initializing agent:", str(e))
-        traceback.print_exc()
-        raise e
-
-def create_ui(agent):
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
-        chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
-        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(message: str, history: List[List[str]], files: list):
-            try:
-                # Initialize with loading message
-                history.append([message, None])
-                yield history, None
-
-                extracted = ""
-                file_hash_value = ""
-                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]
-                        results = []
-                        for future in as_completed(futures):
-                            try:
-                                res = future.result()
-                                results.append(sanitize_utf8(res))
-                            except Exception as e:
-                                print("❌ Error in file processing:", str(e))
-                                traceback.print_exc()
-                    extracted = "\n".join(results)
-                    file_hash_value = file_hash(files[0].name)
-
-                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:
-{extracted[:8000]}
-
-### Potential Oversights:
-"""
-                print("🔎 Generated prompt:")
-                print(prompt)
-
-                # Remove loading message before streaming actual response
-                history[-1][1] = "⏳ Analyzing records..."
-                yield history, None
-
-                # Initialize conversation state
-                conversation = []
-                full_response = ""
-
-                # Stream responses from the agent
-                for chunk in agent.run_gradio_chat(
-                    message=prompt,
-                    history=[],
-                    temperature=0.2,
-                    max_new_tokens=2048,
-                    max_token=4096,
-                    call_agent=False,
-                    conversation=conversation
-                ):
-                    if isinstance(chunk, str):
-                        # Update the last message in history with the new chunk
-                        full_response = chunk
-                        history[-1][1] = full_response
-                        yield history, None
-                    elif isinstance(chunk, list):
-                        # Handle tool calls or other structured responses
-                        for item in chunk:
-                            if isinstance(item, ChatMessage):
-                                # Add tool call messages to history
-                                if item.role == "assistant":
-                                    history.append([None, item.content])
-                                else:
-                                    history.append([None, f"⚒️ {item.content}"])
-                        yield history, None
-
-                # Final cleanup and report generation
-                full_response = full_response.replace('[TxAgent]', '').strip()
-                full_response = re.sub(r"\[TOOL_CALLS\].*?\n*", "", full_response, flags=re.DOTALL).strip()
-
-                # Update the final response
-                history[-1][1] = full_response
-
-                # Generate report file if we have files
-                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(full_response)
-
-                yield history, report_path if report_path and os.path.exists(report_path) else None
-
-            except Exception as e:
-                error_msg = f"❌ An error occurred: {str(e)}"
-                print(error_msg)
-                traceback.print_exc()
-                history[-1][1] = error_msg
-                yield history, None
+from multiprocessing import freeze_support
 
-        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
+from ui.ui_core import create_ui
+from backend.agent_instance import init_agent
 
 if __name__ == "__main__":
-    try:
-        print("🚀 Launching app...")
-        agent = init_agent()
-        demo = create_ui(agent)
-        demo.queue(api_open=False).launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_error=True,
-            allowed_paths=[report_dir],
-            share=False
-        )
-    except Exception as e:
-        print("❌ Fatal error during launch:", str(e))
-        traceback.print_exc()
+    freeze_support()
+    agent = init_agent()
+    demo = create_ui(agent)
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        share=True
+    )