app.py

import sys
import os
import gradio as gr
from typing import List
import hashlib
import time
import json
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Thread
import pandas as pd
import pdfplumber

# Optimized environment setup
os.environ.update({
    "HF_HOME": "/data/hf_cache",
    "VLLM_CACHE_DIR": "/data/vllm_cache",
    "TOKENIZERS_PARALLELISM": "false",
    "CUDA_LAUNCH_BLOCKING": "1"
})

# Create cache directories if they don't exist
os.makedirs("/data/hf_cache", exist_ok=True)
os.makedirs("/data/tool_cache", exist_ok=True)
os.makedirs("/data/file_cache", exist_ok=True)
os.makedirs("/data/reports", exist_ok=True)
os.makedirs("/data/vllm_cache", exist_ok=True)

# Lazy loading of heavy dependencies
def lazy_load_agent():
    from txagent.txagent import TxAgent
    
    # Initialize agent with optimized settings
    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": "/data/tool_cache/new_tool.json"},
        force_finish=True,
        enable_checker=True,
        step_rag_num=8,
        seed=100,
        additional_default_tools=[],
    )
    agent.init_model()
    return agent

# Pre-load the agent in a separate thread
agent = None
def preload_agent():
    global agent
    agent = lazy_load_agent()

Thread(target=preload_agent).start()

# File processing functions
def file_hash(path: str) -> str:
    with open(path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()

def extract_priority_pages(file_path: str, max_pages: int = 10) -> str:
    try:
        with pdfplumber.open(file_path) as pdf:
            return "\n\n".join(
                f"=== Page {i+1} ===\n{(page.extract_text() or '').strip()}"
                for i, page in enumerate(pdf.pages[:max_pages])
            )
    except Exception as e:
        return f"PDF processing error: {str(e)}"

def process_file(file_path: str, file_type: str) -> str:
    try:
        h = file_hash(file_path)
        cache_path = f"/data/file_cache/{h}.json"
        
        if os.path.exists(cache_path):
            with open(cache_path, "r", encoding="utf-8") as f:
                return f.read()
                
        if file_type == "pdf":
            content = extract_priority_pages(file_path)
            result = json.dumps({"filename": os.path.basename(file_path), "content": content})
        elif file_type == "csv":
            df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str)
            result = json.dumps({"filename": os.path.basename(file_path), "rows": df.fillna("").values.tolist()})
        elif file_type in ["xls", "xlsx"]:
            df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str)
            result = json.dumps({"filename": os.path.basename(file_path), "rows": df.fillna("").values.tolist()})
        else:
            return json.dumps({"error": f"Unsupported file type: {file_type}"})

        with open(cache_path, "w", encoding="utf-8") as f:
            f.write(result)
        return result
        
    except Exception as e:
        return json.dumps({"error": str(e)})

def format_response(response: str) -> str:
    response = response.replace("[TOOL_CALLS]", "").strip()
    if "Based on the medical records provided" in response:
        parts = response.split("Based on the medical records provided")
        response = "Based on the medical records provided" + parts[-1]
    
    replacements = {
        "1. **Missed Diagnoses**:": "### 🔍 Missed Diagnoses",
        "2. **Medication Conflicts**:": "\n### 💊 Medication Conflicts",
        "3. **Incomplete Assessments**:": "\n### 📋 Incomplete Assessments",
        "4. **Abnormal Results Needing Follow-up**:": "\n### ⚠️ Abnormal Results Needing Follow-up",
        "Overall, the patient's medical records": "\n### 📝 Overall Assessment"
    }
    
    for old, new in replacements.items():
        response = response.replace(old, new)
    
    return response

def analyze_files(message: str, history: List, files: List):
    try:
        # Wait for agent to load if not ready
        while agent is None:
            time.sleep(0.1)
            
        # Append user message to history in correct format
        history.append([message, None])
        yield history, None
        
        # Process files in parallel
        extracted_data = ""
        if files:
            with ThreadPoolExecutor(max_workers=4) as executor:
                futures = [executor.submit(process_file, f.name, f.name.split(".")[-1].lower()) 
                          for f in files if hasattr(f, 'name')]
                extracted_data = "\n".join(f.result() for f in as_completed(futures))
        
        prompt = f"""Review these medical records:
{extracted_data[:10000]}

Identify:
1. Potential missed diagnoses
2. Medication conflicts
3. Incomplete assessments
4. Abnormal results needing follow-up

Analysis:"""
        
        response = ""
        for chunk in agent.run_gradio_chat(
            message=prompt,
            history=[],
            temperature=0.2,
            max_new_tokens=800,
            max_token=3000
        ):
            if isinstance(chunk, str):
                response += chunk
            elif isinstance(chunk, list):
                response += "".join(getattr(c, 'content', '') for c in chunk)
            
            formatted = format_response(response)
            if formatted.strip():
                history[-1][1] = formatted
                yield history, None
        
        final_output = format_response(response) or "No clear oversights identified."
        history[-1][1] = final_output
        yield history, None
        
    except Exception as e:
        history[-1][1] = f"❌ Error: {str(e)}"
        yield history, None

# Create optimized UI with better layout
with gr.Blocks(title="Clinical Oversight Assistant", css="""
    .gradio-container {
        max-width: 1200px !important;
        margin: auto;
    }
    .container {
        max-width: 1200px !important;
    }
    .chatbot {
        min-height: 500px;
    }
""") as demo:
    gr.Markdown("""
    <div style='text-align: center; margin-bottom: 20px;'>
        <h1 style='margin-bottom: 10px;'>🩺 Clinical Oversight Assistant</h1>
        <p>Upload medical records to analyze for potential oversights in patient care</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1, min_width=400):
            file_upload = gr.File(
                label="Upload Medical Records",
                file_types=[".pdf", ".csv", ".xls", ".xlsx"],
                file_count="multiple",
                height=100
            )
            query = gr.Textbox(
                label="Your Query",
                placeholder="Ask about potential oversights...",
                lines=3
            )
            submit = gr.Button("Analyze", variant="primary")
            
            gr.Examples(
                examples=[
                    ["What potential diagnoses might have been missed?"],
                    ["Are there any medication conflicts I should be aware of?"],
                    ["What assessments appear incomplete in these records?"]
                ],
                inputs=query,
                label="Example Queries"
            )
            
        with gr.Column(scale=2, min_width=600):
            chatbot = gr.Chatbot(
                label="Analysis Results",
                height=600,
                bubble_full_width=False,
                show_copy_button=True
            )
    
    submit.click(
        analyze_files,
        inputs=[query, chatbot, file_upload],
        outputs=[chatbot, gr.File(visible=False)]
    )
    
    query.submit(
        analyze_files,
        inputs=[query, chatbot, file_upload],
        outputs=[chatbot, gr.File(visible=False)]
    )

if __name__ == "__main__":
    demo.queue(concurrency_count=1).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )