app.py

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

# Directories
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

# Paths
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

# Constants
MAX_MODEL_TOKENS = 131072
MAX_NEW_TOKENS = 4096
MAX_CHUNK_TOKENS = 8192
PROMPT_OVERHEAD = 300
BATCH_SIZE = 2

def estimate_tokens(text: str) -> int:
    return len(text) // 4 + 1

def extract_text_from_excel(path: str) -> str:
    all_text = []
    xls = pd.ExcelFile(path)
    for sheet in xls.sheet_names:
        try:
            df = xls.parse(sheet).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}] {line}")
    return "\n".join(all_text)

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

def batch_chunks(chunks: List[str], batch_size: int = BATCH_SIZE) -> 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 clean_response(text: str) -> str:
    text = re.sub(r"\[.*?\]", "", text, flags=re.DOTALL)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()

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_batches(agent, batches: List[List[str]], max_workers: int = 3) -> List[str]:
    results = []

    def process_single_batch(batch):
        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=4096,
                max_token=131072,
                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:
            response = f"❌ Error: {str(e)}"
        return clean_response(response)

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(process_single_batch, 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"""Summarize the following clinical summaries into a final medical report:\n\n{combined}"""
    response = ""
    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):
            response += r
        elif isinstance(r, list):
            for m in r:
                if hasattr(m, "content"):
                    response += m.content
        elif hasattr(r, "content"):
            response += r.content
    return clean_response(response)

def process_file(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], str]:
    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: {file.name}"})
    try:
        extracted_text = extract_text_from_excel(file.name)
        chunks = split_text(extracted_text)
        batches = batch_chunks(chunks)
        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Analyzing..."})

        batch_outputs = analyze_batches(agent, batches)
        valid_outputs = [res for res in batch_outputs if not res.startswith("❌")]

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

        summary = generate_final_summary(agent, "\n\n".join(valid_outputs))

        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"✅ Saved report: {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 { background-color: #0e1621; color: #e0e0e0; }
    button { background: #007bff; color: white; border-radius: 8px; padding: 8px 16px; }
    .gr-chatbot { background: #1b2533; border: 1px solid #2a2f45; border-radius: 16px; padding: 10px; }
    """) as demo:
        gr.Markdown("""## 🧠 CPS: Clinical Patient Support Assistant""")
        chatbot = gr.Chatbot(label="CPS Assistant", height=700, type="messages")
        upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
        analyze_btn = gr.Button("🧠 Analyze File")
        download = gr.File(label="Download Report", visible=False)

        state = gr.State([])

        def handle_analyze(file, chat_state):
            messages, report_path = process_file(agent, file, chat_state)
            return messages, gr.update(visible=bool(report_path), value=report_path), messages

        analyze_btn.click(fn=handle_analyze, inputs=[upload, state], outputs=[chatbot, download, state])

    return demo

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