app.py

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

# === Configuration ===
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

# === Constants ===
MAX_MODEL_TOKENS = 131072
MAX_NEW_TOKENS = 4096
MAX_CHUNK_TOKENS = 8192
BATCH_SIZE = 2
PROMPT_OVERHEAD = 300
SAFE_SLEEP = 0.5  # seconds between batches

# === Utility Functions ===
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)
    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 idx, row in df.iterrows():
            non_empty = [cell.strip() for cell in row if cell.strip()]
            if len(non_empty) >= 2:
                text_line = " | ".join(non_empty)
                if len(text_line) > 15:
                    all_text.append(f"[{sheet_name}] {text_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 = 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 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

# === Main Processing ===
def analyze_batches(agent, batches: List[List[str]]) -> List[str]:
    results = []
    for batch in batches:
        prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
        try:
            batch_response = ""
            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):
                    batch_response += r
                elif isinstance(r, list):
                    for m in r:
                        if hasattr(m, "content"):
                            batch_response += m.content
                elif hasattr(r, "content"):
                    batch_response += r.content
            results.append(clean_response(batch_response))
            time.sleep(SAFE_SLEEP)
        except Exception as e:
            results.append(f"❌ Batch failed: {str(e)}")
            time.sleep(SAFE_SLEEP * 2)  # longer sleep on error
    torch.cuda.empty_cache()
    gc.collect()
    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."
    final_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):
            final_response += r
        elif isinstance(r, list):
            for m in r:
                if hasattr(m, "content"):
                    final_response += m.content
        elif hasattr(r, "content"):
            final_response += r.content
    return clean_response(final_response)

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..."})

        batch_results = analyze_batches(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: #0e1621; color: #e0e0e0;
        }
        button.svelte-1ipelgc {
            background: linear-gradient(to right, #1e88e5, #0d47a1) !important;
            border: 1px solid #0d47a1 !important;
            color: white !important;
            font-weight: bold !important;
        }
        button.svelte-1ipelgc:hover {
            background: linear-gradient(to right, #2196f3, #1565c0) !important;
            border: 1px solid #1565c0 !important;
            color: white !important;
        }
    """) as demo:
        gr.Markdown("""
        <h2>📄 CPS: Clinical Patient Support System</h2>
        <p>Analyze and summarize unstructured medical files using AI (optimized for A100 GPU).</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")
            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

# === Main ===
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)