app.py

import sys
import os
import gc
import json
import shutil
import re
import time
import pandas as pd
import gradio as gr
import torch
from typing import List, Tuple, Dict, Union
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

# Constants
MAX_MODEL_TOKENS = 131072
MAX_NEW_TOKENS = 4096
MAX_CHUNK_TOKENS = 8192
PROMPT_OVERHEAD = 300
BATCH_SIZE = 4  # 4 chunks per batch
MAX_WORKERS = 6  # 6 parallel batches

# Paths
persistent_dir = "/data/hf_cache"
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")

for d in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
    os.makedirs(d, exist_ok=True)

os.environ["HF_HOME"] = model_cache_dir
os.environ["TRANSFORMERS_CACHE"] = model_cache_dir

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 estimate_tokens(text: str) -> int:
    return len(text) // 4 + 1

def clean_response(text: str) -> str:
    text = re.sub(r"\[.*?\]|\bNone\b", "", text, flags=re.DOTALL)
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
    return text.strip()

def extract_text_from_excel(path: str) -> str:
    all_text = []
    xls = pd.ExcelFile(path)
    for sheet_name in xls.sheet_names:
        try:
            df = xls.parse(sheet_name).astype(str).fillna("")
        except Exception:
            continue
        for _, row in df.iterrows():
            non_empty = [cell.strip() for cell in row if cell.strip()]
            if len(non_empty) >= 2:
                line = " | ".join(non_empty)
                if len(line) > 15:
                    all_text.append(f"[{sheet_name}] {line}")
    return "\n".join(all_text)

def split_text(text: str, max_tokens=MAX_CHUNK_TOKENS) -> List[str]:
    effective_limit = max_tokens - PROMPT_OVERHEAD
    chunks, current, current_tokens = [], [], 0
    for line in text.split("\n"):
        tokens = estimate_tokens(line)
        if current_tokens + tokens > effective_limit:
            if current:
                chunks.append("\n".join(current))
            current, current_tokens = [line], tokens
        else:
            current.append(line)
            current_tokens += tokens
    if current:
        chunks.append("\n".join(current))
    return chunks

def batch_chunks(chunks: List[str], batch_size: int = 4) -> List[List[str]]:
    return [chunks[i:i+batch_size] for i in range(0, len(chunks), batch_size)]

def build_prompt(chunk: str) -> str:
    return f"""### Unstructured Clinical Records\n\nAnalyze the clinical notes below and summarize with:\n- Diagnostic Patterns\n- Medication Issues\n- Missed Opportunities\n- Inconsistencies\n- Follow-up Recommendations\n\n---\n\n{chunk}\n\n---\nRespond concisely in bullet points with clinical reasoning."""

def init_agent() -> TxAgent:
    tool_path = os.path.join(tool_cache_dir, "new_tool.json")
    if not os.path.exists(tool_path):
        shutil.copy(os.path.abspath("data/new_tool.json"), 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": tool_path},
        force_finish=True,
        enable_checker=True,
        step_rag_num=4,
        seed=100
    )
    agent.init_model()
    return agent

def analyze_batch(agent, batch: List[str]) -> str:
    prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
    response = ""
    try:
        for r in agent.run_gradio_chat(
            message=prompt,
            history=[],
            temperature=0.0,
            max_new_tokens=MAX_NEW_TOKENS,
            max_token=MAX_MODEL_TOKENS,
            call_agent=False,
            conversation=[]
        ):
            if isinstance(r, str):
                response += r
            elif isinstance(r, list):
                for m in r:
                    if hasattr(m, "content"):
                        response += m.content
            elif hasattr(r, "content"):
                response += r.content
    except Exception as e:
        return f"❌ Error in batch: {str(e)}"
    finally:
        torch.cuda.empty_cache()
        gc.collect()
    return clean_response(response)

def analyze_batches_parallel(agent, batches: List[List[str]]) -> List[str]:
    results = []
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(analyze_batch, agent, batch) for batch in batches]
        for future in as_completed(futures):
            results.append(future.result())
    return results

def generate_final_summary(agent, combined: str) -> str:
    final_prompt = f"""Provide a structured medical report based on the following summaries:\n\n{combined}\n\nRespond in detailed medical bullet points."""
    full_report = ""
    for r in agent.run_gradio_chat(
        message=final_prompt,
        history=[],
        temperature=0.0,
        max_new_tokens=MAX_NEW_TOKENS,
        max_token=MAX_MODEL_TOKENS,
        call_agent=False,
        conversation=[]
    ):
        if isinstance(r, str):
            full_report += r
        elif isinstance(r, list):
            for m in r:
                if hasattr(m, "content"):
                    full_report += m.content
        elif hasattr(r, "content"):
            full_report += r.content
    return clean_response(full_report)

def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
    if not file or not hasattr(file, "name"):
        messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file."})
        return messages, None

    messages.append({"role": "user", "content": f"📂 Processing file: {os.path.basename(file.name)}"})
    try:
        extracted = extract_text_from_excel(file.name)
        chunks = split_text(extracted)
        batches = batch_chunks(chunks, batch_size=BATCH_SIZE)
        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Analyzing in parallel..."})

        batch_results = analyze_batches_parallel(agent, batches)
        valid = [res for res in batch_results if not res.startswith("❌")]

        if not valid:
            messages.append({"role": "assistant", "content": "❌ No valid batch outputs."})
            return messages, None

        summary = generate_final_summary(agent, "\n\n".join(valid))
        report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(f"# 🧠 Final Medical Report\n\n{summary}")

        messages.append({"role": "assistant", "content": f"📊 Final Report:\n\n{summary}"})
        messages.append({"role": "assistant", "content": f"✅ Report saved: {os.path.basename(report_path)}"})
        return messages, report_path

    except Exception as e:
        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
        return messages, None

def create_ui(agent):
    with gr.Blocks(css="""
    html, body, .gradio-container {background-color: #0e1621; color: #e0e0e0; font-family: 'Inter', sans-serif;}
    h2, h3, h4 {color: #89b4fa; font-weight: 600;}
    button.gr-button-primary {background-color: #007bff !important; color: white !important;}
    .gr-chatbot, .gr-markdown, .gr-file-upload {border-radius: 16px; background-color: #1b2533;}
    .gr-chatbot .message {font-size: 16px; padding: 12px; border-radius: 18px;}
    .gr-chatbot .message.user {background-color: #334155;}
    .gr-chatbot .message.assistant {background-color: #1e293b;}
    """) as demo:
        gr.Markdown("""<h2>📄 CPS: Clinical Patient Support System</h2><p>Upload a file and analyze medical notes.</p>""")
        with gr.Column():
            chatbot = gr.Chatbot(label="CPS Assistant", height=700, type="messages")
            upload = gr.File(label="Upload Medical File", file_types=[".xlsx"])
            analyze = gr.Button("🧠 Analyze", variant="primary")
            download = gr.File(label="Download Report", visible=False, interactive=False)
        state = gr.State(value=[])

        def handle_analysis(file, chat):
            messages, report_path = process_report(agent, file, chat)
            return messages, gr.update(visible=bool(report_path), value=report_path), messages

        analyze.click(fn=handle_analysis, inputs=[upload, state], outputs=[chatbot, download, state])

    return demo

if __name__ == "__main__":
    try:
        agent = init_agent()
        ui = create_ui(agent)
        ui.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)
    except Exception as err:
        print(f"Startup failed: {err}")
        sys.exit(1)