app.py

import sys
import os
import pandas as pd
import pdfplumber
import json
import gradio as gr
from typing import List, Dict, Optional, Generator
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
import shutil
import re
import psutil
import subprocess
import logging
import torch
import gc
from diskcache import Cache
import time
from transformers import AutoTokenizer
from functools import lru_cache
import numpy as np
from difflib import SequenceMatcher

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants
MAX_TOKENS = 1800
BATCH_SIZE = 1
MAX_WORKERS = 2
CHUNK_SIZE = 5
MODEL_MAX_TOKENS = 131072
MAX_TEXT_LENGTH = 500000

# Persistent directory setup
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.update({
    "HF_HOME": model_cache_dir,
    "TRANSFORMERS_CACHE": model_cache_dir,
    "VLLM_CACHE_DIR": vllm_cache_dir,
    "TOKENIZERS_PARALLELISM": "false",
    "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)

@lru_cache(maxsize=1)
def get_tokenizer():
    return AutoTokenizer.from_pretrained("mims-harvard/TxAgent-T1-Llama-3.1-8B")

def sanitize_utf8(text: str) -> str:
    return text.encode("utf-8", "ignore").decode("utf-8")

def file_hash(path: str) -> str:
    hash_md5 = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

def extract_pdf_page(page, tokenizer, max_tokens=MAX_TOKENS) -> List[str]:
    try:
        text = page.extract_text() or ""
        text = sanitize_utf8(text)
        if len(text) > MAX_TEXT_LENGTH // 10:
            text = text[:MAX_TEXT_LENGTH // 10]
        
        tokens = tokenizer.encode(text, add_special_tokens=False)
        if len(tokens) > max_tokens:
            chunks = []
            current_chunk = []
            current_length = 0
            for token in tokens:
                if current_length + 1 > max_tokens:
                    chunks.append(tokenizer.decode(current_chunk))
                    current_chunk = [token]
                    current_length = 1
                else:
                    current_chunk.append(token)
                    current_length += 1
            if current_chunk:
                chunks.append(tokenizer.decode(current_chunk))
            return [f"=== Page {page.page_number} ===\n{c}" for c in chunks]
        return [f"=== Page {page.page_number} ===\n{text}"]
    except Exception as e:
        logger.warning(f"Error extracting page {page.page_number}: {str(e)}")
        return []

def extract_all_pages(file_path: str, progress_callback=None) -> List[str]:
    try:
        tokenizer = get_tokenizer()
        with pdfplumber.open(file_path) as pdf:
            total_pages = len(pdf.pages)
            if total_pages == 0:
                logger.error("PDF has 0 pages - may be corrupted or empty")
                return []

        results = []
        total_tokens = 0
        for chunk_start in range(0, total_pages, CHUNK_SIZE):
            chunk_end = min(chunk_start + CHUNK_SIZE, total_pages)
            
            with pdfplumber.open(file_path) as pdf:
                with ThreadPoolExecutor(max_workers=min(CHUNK_SIZE, 2)) as executor:
                    futures = [executor.submit(extract_pdf_page, pdf.pages[i], tokenizer) 
                             for i in range(chunk_start, chunk_end)]
                    
                    for future in as_completed(futures):
                        page_chunks = future.result()
                        for chunk in page_chunks:
                            chunk_tokens = len(tokenizer.encode(chunk, add_special_tokens=False))
                            if total_tokens + chunk_tokens > MODEL_MAX_TOKENS:
                                logger.warning("Total tokens exceed model limit. Stopping.")
                                return results
                            results.append(chunk)
                            total_tokens += chunk_tokens
                    
                    if progress_callback:
                        progress_callback(min(chunk_end, total_pages), total_pages)
            
            del pdf
            gc.collect()
        
        if not results:
            logger.error("No content extracted from PDF - may be scanned or encrypted")
            return ["PDF appears to be empty or unreadable"]
        
        return results
    except Exception as e:
        logger.error(f"PDF processing error: {e}")
        return [f"PDF processing error: {str(e)}"]

def excel_to_json(file_path: str) -> List[Dict]:
    """Enhanced Excel processing with multiple engine support"""
    engines = ['openpyxl', 'xlrd', 'odf']
    last_error = None
    
    for engine in engines:
        try:
            with pd.ExcelFile(file_path, engine=engine) as excel_file:
                sheets = excel_file.sheet_names
                if not sheets:
                    return [{"error": "No sheets found in Excel file"}]
                
                results = []
                for sheet_name in sheets:
                    try:
                        df = pd.read_excel(
                            excel_file,
                            sheet_name=sheet_name,
                            header=None,
                            dtype=str,
                            na_filter=False,
                            engine=engine
                        )
                        if not df.empty:
                            # Convert all cells to string and clean
                            df = df.applymap(lambda x: str(x).strip() if pd.notna(x) else "")
                            results.append({
                                "filename": f"{os.path.basename(file_path)} - {sheet_name}",
                                "rows": df.values.tolist(),
                                "type": "excel",
                                "sheet": sheet_name,
                                "dimensions": f"{len(df)} rows x {len(df.columns)} cols"
                            })
                    except Exception as sheet_error:
                        logger.warning(f"Error processing sheet {sheet_name}: {sheet_error}")
                        continue
                
                if results:
                    logger.info(f"Successfully processed Excel file with {engine} engine")
                    return results
        except Exception as engine_error:
            last_error = engine_error
            continue
    
    return [{"error": f"Failed to process Excel file with all engines. Last error: {str(last_error)}"}]

def csv_to_json(file_path: str) -> List[Dict]:
    try:
        chunks = []
        for chunk in pd.read_csv(
            file_path,
            header=None,
            dtype=str,
            encoding_errors='replace',
            on_bad_lines='skip',
            chunksize=10000,
            na_filter=False
        ):
            chunks.append(chunk)
        
        df = pd.concat(chunks) if chunks else pd.DataFrame()
        if df.empty:
            return [{"error": "CSV file is empty or could not be read"}]
            
        return [{
            "filename": os.path.basename(file_path),
            "rows": df.values.tolist(),
            "type": "csv",
            "dimensions": f"{len(df)} rows x {len(df.columns)} cols"
        }]
    except Exception as e:
        logger.error(f"CSV processing error: {e}")
        return [{"error": f"CSV processing error: {str(e)}"}]

@lru_cache(maxsize=100)
def process_file_cached(file_path: str, file_type: str) -> List[Dict]:
    """Enhanced file processing with detailed logging"""
    try:
        logger.info(f"Processing file: {file_path} (type: {file_type})")
        
        if file_type == "pdf":
            chunks = extract_all_pages(file_path)
            if not chunks or (len(chunks) == 1 and "error" in chunks[0]):
                return [{"error": chunks[0] if chunks else "PDF appears to be empty"}]
            return [{
                "filename": os.path.basename(file_path),
                "content": chunk,
                "status": "initial",
                "type": "pdf",
                "page": i+1
            } for i, chunk in enumerate(chunks)]
            
        elif file_type in ["xls", "xlsx"]:
            result = excel_to_json(file_path)
            if "error" in result[0]:
                logger.error(f"Excel processing failed: {result[0]['error']}")
            else:
                logger.info(f"Excel processing successful - found {len(result)} sheets")
            return result
            
        elif file_type == "csv":
            result = csv_to_json(file_path)
            if "error" in result[0]:
                logger.error(f"CSV processing failed: {result[0]['error']}")
            else:
                logger.info(f"CSV processing successful - found {len(result[0]['rows'])} rows")
            return result
            
        else:
            logger.warning(f"Unsupported file type: {file_type}")
            return [{"error": f"Unsupported file type: {file_type}"}]
            
    except Exception as e:
        logger.error(f"Error processing {file_path}: {str(e)}", exc_info=True)
        return [{"error": f"Error processing file: {str(e)}"}]

def clean_response(text: str) -> str:
    if not text:
        return ""
    
    patterns = [
        (re.compile(r"\[.*?\]|\bNone\b", re.IGNORECASE), ""),
        (re.compile(r"\s+"), " "),
        (re.compile(r"[^\w\s\.\,\(\)\-]"), ""),
    ]
    
    for pattern, repl in patterns:
        text = pattern.sub(repl, text)
    
    sentences = text.split(". ")
    unique_sentences = []
    seen = set()
    
    for s in sentences:
        if not s:
            continue
        is_unique = True
        for seen_s in seen:
            if SequenceMatcher(None, s.lower(), seen_s.lower()).ratio() > 0.9:
                is_unique = False
                break
        if is_unique:
            unique_sentences.append(s)
            seen.add(s)
    
    text = ". ".join(unique_sentences).strip()
    return text if text else "No missed diagnoses identified."

@lru_cache(maxsize=1)
def init_agent():
    logger.info("Initializing model...")
    
    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,
        step_rag_num=4,
        seed=100,
        additional_default_tools=[],
    )
    agent.init_model()
    logger.info("Agent Ready")
    return agent

def create_ui(agent):
    PROMPT_TEMPLATE = """
Analyze the patient record excerpt for missed diagnoses. Provide detailed, evidence-based analysis.
Patient Record Excerpt (Chunk {0} of {1}):
{chunk}
"""

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
        
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(label="Analysis Summary", height=600, value=[])
                msg_input = gr.Textbox(placeholder="Ask about potential oversights...")
                send_btn = gr.Button("Analyze", variant="primary")
                file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
            
            with gr.Column(scale=1):
                final_summary = gr.Markdown(label="Missed Diagnoses Summary")
                download_output = gr.File(label="Download Detailed Report")
                progress_bar = gr.Progress()

        def analyze(message: str, history: List[List[str]], files: List, progress=gr.Progress()):
            """Enhanced analysis with detailed file processing feedback"""
            try:
                if history is None:
                    history = []
                
                history.append([message, None])
                yield history, None, ""

                if not files:
                    history[-1][1] = "❌ Please upload a file to analyze"
                    yield history, None, "No files uploaded"
                    return

                extracted = []
                file_hash_value = ""
                
                for f in files:
                    file_type = f.name.split(".")[-1].lower()
                    logger.info(f"Processing file: {f.name} (type: {file_type})")
                    
                    cache_key = f"{file_hash(f.name)}_{file_type}"
                    if cache_key in cache:
                        cached_data = cache[cache_key]
                        if isinstance(cached_data, list) and len(cached_data) > 0:
                            extracted.extend(cached_data)
                            history[-1][1] = f"✅ Using cached data for {os.path.basename(f.name)}"
                            yield history, None, ""
                            continue
                    
                    try:
                        result = process_file_cached(f.name, file_type)
                        if "error" in result[0]:
                            history[-1][1] = f"❌ Error processing {os.path.basename(f.name)}: {result[0]['error']}"
                            yield history, None, result[0]['error']
                            return
                        
                        cache[cache_key] = result
                        extracted.extend(result)
                        history[-1][1] = f"✅ Processed {os.path.basename(f.name)}"
                        yield history, None, ""
                    except Exception as e:
                        logger.error(f"File processing error: {e}", exc_info=True)
                        history[-1][1] = f"❌ Critical error processing {os.path.basename(f.name)}"
                        yield history, None, str(e)
                        return
                
                file_hash_value = file_hash(files[0].name) if files else ""
                
                # Debug extracted content
                logger.info(f"Extracted content summary:")
                for item in extracted:
                    if "content" in item:
                        logger.info(f"- {item['filename']}: {len(item['content'])} chars")
                    elif "rows" in item:
                        logger.info(f"- {item['filename']}: {len(item['rows'])} rows")
                
                if not extracted:
                    history[-1][1] = "❌ No valid content extracted from files"
                    yield history, None, "No valid content extracted"
                    return

                chunks = []
                for item in extracted:
                    if "content" in item:
                        chunks.append(item["content"])
                    elif "rows" in item:
                        # Convert Excel/CSV rows to text
                        rows_text = "\n".join([", ".join(map(str, row)) for row in item["rows"]])
                        chunks.append(f"=== {item['filename']} ===\n{rows_text}")
                
                if not chunks:
                    history[-1][1] = "❌ No processable content found in files"
                    yield history, None, "No processable content found"
                    return

                combined_response = ""
                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
                
                try:
                    for batch_idx in range(0, len(chunks), BATCH_SIZE):
                        batch_chunks = chunks[batch_idx:batch_idx + BATCH_SIZE]
                        
                        progress(batch_idx / len(chunks), 
                               desc=f"Processing batch {(batch_idx // BATCH_SIZE) + 1}/{(len(chunks) + BATCH_SIZE - 1) // BATCH_SIZE}")
                        
                        with ThreadPoolExecutor(max_workers=min(BATCH_SIZE, MAX_WORKERS)) as executor:
                            futures = {
                                executor.submit(
                                    agent.run_quick_summary,
                                    chunk, 0.2, 256, 1024
                                ): idx
                                for idx, chunk in enumerate(batch_chunks)
                            }
                            
                            for future in as_completed(futures):
                                chunk_idx = futures[future]
                                try:
                                    response = clean_response(future.result())
                                    if response:
                                        combined_response += f"\n--- Analysis for Chunk {batch_idx + chunk_idx + 1} ---\n{response}\n"
                                        history[-1][1] = combined_response.strip()
                                        yield history, None, ""
                                except Exception as e:
                                    logger.error(f"Chunk processing error: {e}")
                                    history[-1][1] = f"Error processing chunk: {str(e)}"
                                    yield history, None, ""
                                finally:
                                    del future
                                    torch.cuda.empty_cache()
                                    gc.collect()

                    summary = "Analysis complete. " + ("Download full report below." if report_path and os.path.exists(report_path) else "")
                    history.append(["Analysis completed", None])
                    history[-1][1] = summary
                    yield history, report_path, summary

                except Exception as e:
                    logger.error(f"Analysis error: {e}")
                    history.append(["Analysis failed", None])
                    history[-1][1] = f"❌ Error occurred: {str(e)}"
                    yield history, None, f"Error occurred: {str(e)}"
                finally:
                    torch.cuda.empty_cache()
                    gc.collect()

            except Exception as e:
                logger.error(f"Unexpected error in analysis: {e}")
                history.append(["System error", None])
                history[-1][1] = f"❌ System error occurred: {str(e)}"
                yield history, None, f"System error: {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().launch(
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True
        )
    except Exception as e:
        logger.error(f"Fatal error: {e}")
        raise