app.py

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
import multiprocessing
from functools import partial
import time

# Persistent directory
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

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_page_range(file_path: str, start_page: int, end_page: int) -> str:
    """Extract text from a range of PDF pages."""
    try:
        text_chunks = []
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages[start_page:end_page]:
                page_text = page.extract_text() or ""
                text_chunks.append(f"=== Page {start_page + pdf.pages.index(page) + 1} ===\n{page_text.strip()}")
        return "\n\n".join(text_chunks)
    except Exception:
        return ""

def extract_all_pages(file_path: str, progress_callback=None) -> str:
    """Extract text from all pages of a PDF using parallel processing."""
    try:
        with pdfplumber.open(file_path) as pdf:
            total_pages = len(pdf.pages)
        
        if total_pages == 0:
            return ""
        
        # Use 6 processes (adjust based on CPU cores)
        num_processes = min(6, multiprocessing.cpu_count())
        pages_per_process = max(1, total_pages // num_processes)
        
        # Create page ranges for parallel processing
        ranges = [(i * pages_per_process, min((i + 1) * pages_per_process, total_pages))
                  for i in range(num_processes)]
        if ranges[-1][1] != total_pages:
            ranges[-1] = (ranges[-1][0], total_pages)
        
        # Process page ranges in parallel
        with multiprocessing.Pool(processes=num_processes) as pool:
            extract_func = partial(extract_page_range, file_path)
            results = []
            for idx, result in enumerate(pool.starmap(extract_func, ranges)):
                results.append(result)
                if progress_callback:
                    processed_pages = min((idx + 1) * pages_per_process, total_pages)
                    progress_callback(processed_pages, total_pages)
        
        return "\n\n".join(filter(None, results))
    except Exception as e:
        return f"PDF processing error: {str(e)}"

def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> 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_all_pages(file_path, progress_callback)
            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:
        return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {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}")

def clean_response(text: str) -> str:
    text = sanitize_utf8(text)
    text = re.sub(r"\[TOOL_CALLS\].*", "", text, flags=re.DOTALL)
    text = re.sub(r"\n{3,}", "\n\n", text).strip()
    return text

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):
        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},
        force_finish=True,
        enable_checker=True,
        step_rag_num=4,
        seed=100,
        additional_default_tools=[],
    )
    agent.init_model()
    log_system_usage("After Load")
    print("✅ Agent Ready")
    return 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>")
        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[dict], files: List):
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": "⏳ Extracting text from files..."})
            yield history, None

            extracted = ""
            file_hash_value = ""
            if files:
                # Progress callback for extraction
                total_pages = 0
                processed_pages = 0
                def update_extraction_progress(current, total):
                    nonlocal processed_pages, total_pages
                    processed_pages = current
                    total_pages = total
                    animation = ["🌀", "🔄", "⚙️", "🔃"][(int(time.time() * 2) % 4)]
                    history[-1] = {"role": "assistant", "content": f"Extracting text... {animation} Page {processed_pages}/{total_pages}"}
                    return history, None

                with ThreadPoolExecutor(max_workers=6) as executor:
                    futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files]
                    results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
                    extracted = "\n".join(results)
                    file_hash_value = file_hash(files[0].name) if files else ""

            history.pop()  # Remove extraction message
            history.append({"role": "assistant", "content": "✅ Text extraction complete."})
            yield history, None

            # Split extracted text into chunks of ~6,000 characters
            chunk_size = 6000
            chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
            combined_response = ""

            prompt_template = """
You are a medical analysis assistant. Analyze the following patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown format under these headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, and Urgent Follow-up. For each finding, include:
- Clinical context (why the issue was missed or relevant details from the record).
- Potential risks if unaddressed (e.g., disease progression, adverse events).
- Actionable recommendations (e.g., tests, referrals, medication adjustments).
If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text, formatted as markdown with bullet points under each heading, and avoids generic or static responses.

Patient Record Excerpt (Chunk {0} of {1}):
{chunk}

### Missed Diagnoses
- ...

### Medication Conflicts
- ...

### Incomplete Assessments
- ...

### Urgent Follow-up
- ...
"""

            try:
                # Process each chunk and stream results in real-time
                for chunk_idx, chunk in enumerate(chunks, 1):
                    # Update UI with chunk progress
                    animation = ["🔍", "📊", "🧠", "🔎"][(int(time.time() * 2) % 4)]
                    history.append({"role": "assistant", "content": f"Analyzing records... {animation} Chunk {chunk_idx}/{len(chunks)}"})
                    yield history, None

                    prompt = prompt_template.format(chunk_idx, len(chunks), chunk=chunk[:4000])  # Truncate to avoid token limits
                    chunk_response = ""
                    for chunk_output in agent.run_gradio_chat(
                        message=prompt,
                        history=[],
                        temperature=0.2,
                        max_new_tokens=1024,
                        max_token=4096,
                        call_agent=False,
                        conversation=[],
                    ):
                        if chunk_output is None:
                            continue
                        if isinstance(chunk_output, list):
                            for m in chunk_output:
                                if hasattr(m, 'content') and m.content:
                                    cleaned = clean_response(m.content)
                                    if cleaned:
                                        chunk_response += cleaned + "\n"
                                        # Update UI with partial response
                                        if history[-1]["content"].startswith("Analyzing"):
                                            history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"}
                                        else:
                                            history[-1]["content"] = f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"
                                        yield history, None
                        elif isinstance(chunk_output, str) and chunk_output.strip():
                            cleaned = clean_response(chunk_output)
                            if cleaned:
                                chunk_response += cleaned + "\n"
                                # Update UI with partial response
                                if history[-1]["content"].startswith("Analyzing"):
                                    history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"}
                                else:
                                    history[-1]["content"] = f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"
                                yield history, None

                    # Append completed chunk response to combined response
                    combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"

                # Finalize UI with complete response
                if combined_response:
                    history[-1]["content"] = combined_response.strip()
                else:
                    history.append({"role": "assistant", "content": "No oversights identified."})

                # Generate report file
                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
                if report_path:
                    with open(report_path, "w", encoding="utf-8") as f:
                        f.write(combined_response)
                yield history, report_path if report_path and os.path.exists(report_path) else None

            except Exception as e:
                print("🚨 ERROR:", e)
                history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
                yield history, None

        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

if __name__ == "__main__":
    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
    )