app.py

import sys
import os
import pandas as pd
import json
import gradio as gr
from typing import List
from concurrent.futures import ThreadPoolExecutor
import hashlib
import shutil
import re
import psutil
import subprocess
import logging
import torch
import gc
from diskcache import Cache
import time
import asyncio

# Try importing pypdfium2 and pytesseract, fall back to pdfplumber
try:
    import pypdfium2 as pdfium
    import pytesseract
    from PIL import Image
    HAS_PYPDFIUM2 = True
except ImportError:
    HAS_PYPDFIUM2 = False
    import pdfplumber

# Configure logging
logging.basicConfig(level=logging.INFO)
logging.getLogger("pdfminer").setLevel(logging.ERROR)
logger = logging.getLogger(__name__)

# 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

# Initialize cache with 10GB limit
cache = Cache(file_cache_dir, size_limit=10 * 1024**3)

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()

async def extract_all_pages_async(file_path: str, progress_callback=None, force_ocr=False) -> str:
    try:
        extracted_text = ""
        total_pages = 0
        text_chunks = []

        if HAS_PYPDFIUM2:
            pdf = pdfium.PdfDocument(file_path)
            total_pages = len(pdf)
            if total_pages == 0:
                return ""

            def extract_page(i):
                page = pdf[i]
                text = page.get_textpage().get_text_range() or ""
                if (not text.strip() or len(text) < 100) and force_ocr and 'pytesseract' in sys.modules:
                    logger.info("Falling back to OCR for page %d", i + 1)
                    bitmap = page.render(scale=2).to_pil()
                    text = pytesseract.image_to_string(bitmap, lang="eng")
                return (i, f"=== Page {i + 1} ===\n{text.strip()}")

            with ThreadPoolExecutor(max_workers=4) as executor:
                futures = [executor.submit(extract_page, i) for i in range(total_pages)]
                for future in as_completed(futures):
                    page_num, text = future.result()
                    text_chunks.append((page_num, text))
                    logger.debug("Page %d extracted: %s...", page_num + 1, text[:50])
                    if progress_callback:
                        progress_callback(page_num + 1, total_pages)

            text_chunks.sort(key=lambda x: x[0])
            extracted_text = "\n\n".join(chunk[1] for chunk in text_chunks if chunk[1].strip())
            pdf.close()
        else:
            with pdfplumber.open(file_path) as pdf:
                total_pages = len(pdf.pages)
                if total_pages == 0:
                    return ""

                for i, page in enumerate(pdf.pages):
                    text = page.extract_text() or ""
                    text_chunks.append((i, f"=== Page {i + 1} ===\n{text.strip()}"))
                    logger.debug("Page %d extracted: %s...", i + 1, text[:50])
                    if progress_callback:
                        progress_callback(i + 1, total_pages)

            extracted_text = "\n\n".join(chunk[1] for chunk in text_chunks if chunk[1].strip())

        logger.info("Extracted %d pages, total length: %d chars", total_pages, len(extracted_text))
        if len(extracted_text) < 1000 and not force_ocr and HAS_PYPDFIUM2 and 'pytesseract' in sys.modules:
            logger.info("Text too short, retrying with OCR")
            return await extract_all_pages_async(file_path, progress_callback, force_ocr=True)
        return extracted_text
    except Exception as e:
        logger.error("PDF processing error: %s", e)
        return f"PDF processing error: {str(e)}"

def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str:
    try:
        file_h = file_hash(file_path)
        cache_key = f"{file_h}_{file_type}"
        if cache_key in cache:
            logger.info("Using cached extraction for %s", file_path)
            return cache[cache_key]

        if file_type == "pdf":
            text = asyncio.run(extract_all_pages_async(file_path, progress_callback, force_ocr=False))
            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}"})

        cache[cache_key] = result
        logger.info("Cached extraction for %s, size: %d bytes", file_path, len(result))
        return result
    except Exception as e:
        logger.error("Error processing %s: %s", os.path.basename(file_path), 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()
        logger.info("[%s] CPU: %.1f%% | RAM: %dMB / %dMB", tag, cpu, mem.used // (1024**2), mem.total // (1024**2))
        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(", ")
            logger.info("[%s] GPU: %sMB / %sMB | Utilization: %s%%", tag, used, total, util)
    except Exception as e:
        logger.error("[%s] GPU/CPU monitor failed: %s", tag, e)

def clean_response(text: str) -> str:
    text = sanitize_utf8(text)
    text = text.replace("[", "").replace("]", "").replace("None", "")
    text = text.replace("\n\n\n", "\n\n")
    sections = {}
    current_section = None
    seen_lines = set()
    for line in text.splitlines():
        line = line.strip()
        if not line or line in seen_lines:
            continue
        seen_lines.add(line)
        section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
        if section_match:
            current_section = section_match.group(1)
            sections.setdefault(current_section, [])
            continue
        if current_section and line.startswith("- "):
            sections[current_section].append(line)
    cleaned = [f"### {heading}\n" + "\n".join(findings) for heading, findings in sections.items() if findings]
    result = "\n\n".join(cleaned).strip()
    logger.debug("Cleaned response length: %d chars", len(result))
    return result or "No oversights identified"

def summarize_findings(all_responses: List[str]) -> str:
    combined_response = "\n\n".join(all_responses)
    if not combined_response or all("No oversights identified" in resp.lower() for resp in all_responses):
        return "### Comprehensive Clinical Oversight Summary\nNo critical oversights were identified across the provided patient records after thorough analysis."

    sections = {
        "Missed Diagnoses": [],
        "Medication Conflicts": [],
        "Incomplete Assessments": [],
        "Urgent Follow-up": []
    }
    current_section = None
    seen_findings = set()
    for line in combined_response.splitlines():
        line = line.strip()
        if not line:
            continue
        section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
        if section_match:
            current_section = section_match.group(1)
            continue
        if current_section and line.startswith("- ") and line not in seen_findings:
            sections[current_section].append(line)
            seen_findings.add(line)

    summary_lines = []
    for heading, findings in sections.items():
        if findings:
            summary_lines.append(f"### {heading}")
            for finding in findings:
                summary_lines.append(f"{finding}\n  - **Risks**: Potential adverse outcomes if not addressed.\n  - **Recommendation**: Immediate clinical review and follow-up.")

    result = "### Comprehensive Clinical Oversight Summary\n" + "\n".join(summary_lines) if summary_lines else "### Comprehensive Clinical Oversight Summary\nNo critical oversights identified."
    logger.debug("Summary length: %d chars", len(result))
    return result

def init_agent():
    logger.info("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=False,
        enable_rag=False,
        init_rag_num=0,
        step_rag_num=0,
        seed=100,
        additional_default_tools=[],
    )
    agent.init_model()
    log_system_usage("After Load")
    logger.info("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="Detailed Analysis", height=600, type="messages", visible=False)
        final_summary = gr.Markdown(label="Comprehensive Clinical Oversight Summary")
        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")
        progress_bar = gr.Progress()

        prompt_template = """
Analyze the patient record excerpt for clinical oversights. Provide a detailed, evidence-based summary in markdown with findings grouped under headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, Urgent Follow-up. For each finding, include clinical context, risks, and recommendations. Output only markdown bullet points under headings. If no issues, state "No oversights identified" once.

Patient Record Excerpt:
{chunk}
"""

        async def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
            history.append({"role": "user", "content": message})
            yield history, None, ""

            extracted = ""
            file_hash_value = ""
            if files:
                def update_extraction_progress(current, total):
                    progress(current / total, desc=f"Extracting text... Page {current}/{total}")
                    return history, None, ""

                futures = [convert_file_to_json(f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files]
                results = [sanitize_utf8(future) for future in futures]
                extracted = "\n".join(results)
                file_hash_value = file_hash(files[0].name) if files else ""

                history.append({"role": "assistant", "content": "✅ Text extraction complete."})
                yield history, None, ""
                logger.info("Extracted text length: %d chars", len(extracted))

            chunk_size = 3000
            chunks = [extracted[i:i + chunk_size] for i in range(0, max(len(extracted), 1), chunk_size)] or [""]
            logger.info("Created %d chunks", len(chunks))
            for i, chunk in enumerate(chunks):
                logger.debug("Chunk %d content: %s...", i + 1, chunk[:100])
            all_responses = []
            batch_size = 2

            try:
                for batch_idx in range(0, len(chunks), batch_size):
                    batch_chunks = chunks[batch_idx:batch_idx + batch_size]
                    batch_prompts = [prompt_template.format(chunk=chunk[:2000]) for chunk in batch_chunks]
                    batch_responses = []

                    progress((batch_idx + 1) / len(chunks), desc=f"Analyzing chunks {batch_idx + 1}-{min(batch_idx + batch_size, len(chunks))}/{len(chunks)}")

                    async def process_chunk(prompt):
                        chunk_response = ""
                        raw_outputs = []
                        for chunk_output in agent.run_gradio_chat(
                            message=prompt, history=[], temperature=0.2, max_new_tokens=512, max_token=1024, 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:
                                        raw_outputs.append(m.content)
                                        cleaned = clean_response(m.content)
                                        chunk_response += cleaned + "\n\n"
                            elif isinstance(chunk_output, str) and chunk_output.strip():
                                raw_outputs.append(chunk_output)
                                cleaned = clean_response(chunk_output)
                                chunk_response += cleaned + "\n\n"
                        logger.debug("Raw outputs: %s", raw_outputs[:100])
                        logger.debug("Chunk response length: %d chars", len(chunk_response))
                        return chunk_response

                    futures = [process_chunk(prompt) for prompt in batch_prompts]
                    batch_responses = await asyncio.gather(*futures)
                    all_responses.extend([resp.strip() for resp in batch_responses if resp.strip()])
                    torch.cuda.empty_cache()
                    gc.collect()

                summary = summarize_findings(all_responses)
                history.append({"role": "assistant", "content": "Analysis complete. See summary below."})
                yield history, None, summary

                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(summary)
                yield history, report_path if report_path and os.path.exists(report_path) else None, summary

            except Exception as e:
                logger.error("Analysis error: %s", e)
                history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
                yield history, None, f"### Comprehensive Clinical Oversight Summary\nError occurred during analysis: {str(e)}"

        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
    return demo

if __name__ == "__main__":
    try:
        logger.info("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
        )
    finally:
        if torch.distributed.is_initialized():
            torch.distributed.destroy_process_group()