app.py

import sys
import os
import pandas as pd
import gradio as gr
from typing import List, Tuple, Dict, Any, Union, Generator
import shutil
import re
from datetime import datetime
import time
import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# Configuration and 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")

for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
    os.makedirs(directory, 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

# Updated token limits as specified
MAX_MODEL_TOKENS = 131072  # TxAgent's max token limit
MAX_CHUNK_TOKENS = 32768   # Larger chunks to reduce number of chunks
MAX_NEW_TOKENS = 512       # Optimized for fast generation
PROMPT_OVERHEAD = 500      # Estimated tokens for prompt template
MAX_CONCURRENT = 8         # High concurrency for A100 80GB

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def clean_response(text: str) -> str:
    try:
        text = text.encode('utf-8', 'surrogatepass').decode('utf-8')
    except UnicodeError:
        text = text.encode('utf-8', 'replace').decode('utf-8')
    text = re.sub(r"\[.*?\]|\bNone\b", "", text, flags=re.DOTALL)
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
    return text.strip()

def estimate_tokens(text: str) -> int:
    return len(text) // 3.5 + 1  # More conservative estimate

def extract_text_from_excel(file_path: str) -> str:
    all_text = []
    try:
        xls = pd.ExcelFile(file_path)
        for sheet_name in xls.sheet_names:
            df = xls.parse(sheet_name)
            df = df.astype(str).fillna("")
            rows = df.apply(lambda row: " | ".join(row), axis=1)
            sheet_text = [f"[{sheet_name}] {line}" for line in rows]
            all_text.extend(sheet_text)
    except Exception as e:
        logger.error(f"Error extracting Excel: {str(e)}")
        raise ValueError(f"Failed to process Excel file: {str(e)}")
    return "\n".join(all_text)

def split_text_into_chunks(text: str) -> List[str]:
    """Split text into chunks respecting MAX_CHUNK_TOKENS and PROMPT_OVERHEAD"""
    effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD
    if effective_max <= 0:
        raise ValueError("Effective max tokens must be positive")
    
    lines = text.split("\n")
    chunks = []
    current_chunk = []
    current_tokens = 0

    for line in lines:
        line_tokens = estimate_tokens(line)
        if current_tokens + line_tokens > effective_max:
            if current_chunk:
                chunks.append("\n".join(current_chunk))
            current_chunk = [line]
            current_tokens = line_tokens
        else:
            current_chunk.append(line)
            current_tokens += line_tokens

    if current_chunk:
        chunks.append("\n".join(current_chunk))
    
    logger.info(f"Split text into {len(chunks)} chunks")
    return chunks

def build_prompt_from_text(chunk: str) -> str:
    return f"""
### Unstructured Clinical Records

You are reviewing unstructured, mixed-format clinical documentation from various forms, tables, and sheets.

**Objective:** Identify patterns, missed diagnoses, inconsistencies, and follow-up gaps.

Here is the extracted content chunk:

{chunk}

Please analyze the above and provide concise responses (max {MAX_NEW_TOKENS} tokens):
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations
"""

def init_agent():
    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=True,
        step_rag_num=4,
        seed=100,
        additional_default_tools=[],
        max_model_tokens=MAX_MODEL_TOKENS  # Pass the updated token limit
    )
    agent.init_model()
    return agent

async def process_chunk(agent: TxAgent, chunk: str, chunk_idx: int) -> Tuple[int, str]:
    """Process a single chunk with error handling"""
    try:
        prompt = build_prompt_from_text(chunk)
        prompt_tokens = estimate_tokens(prompt)
        
        if prompt_tokens > MAX_MODEL_TOKENS:
            logger.warning(f"Chunk {chunk_idx} prompt too long ({prompt_tokens} tokens)")
            return chunk_idx, ""
        
        response = ""
        for result in agent.run_gradio_chat(
            message=prompt,
            history=[],
            temperature=0.2,
            max_new_tokens=MAX_NEW_TOKENS,
            max_token=MAX_MODEL_TOKENS,
            call_agent=False,
            conversation=[],
        ):
            if isinstance(result, str):
                response += result
            elif hasattr(result, "content"):
                response += result.content
            elif isinstance(result, list):
                for r in result:
                    if hasattr(r, "content"):
                        response += r.content
        
        return chunk_idx, clean_response(response)
    
    except Exception as e:
        logger.error(f"Error processing chunk {chunk_idx}: {str(e)}")
        return chunk_idx, ""

async def process_file(agent: TxAgent, file_path: str) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
    """Process the entire file and yield progress updates"""
    messages = []
    report_path = None
    
    try:
        # Initial messages
        messages.append({"role": "user", "content": f"Processing file: {os.path.basename(file_path)}"})
        messages.append({"role": "assistant", "content": "⏳ Extracting data from Excel..."})
        yield messages, None
        
        # Extract and chunk text
        start_time = time.time()
        text = extract_text_from_excel(file_path)
        chunks = split_text_into_chunks(text)
        messages.append({"role": "assistant", "content": f"✅ Extracted {len(chunks)} chunks in {time.time()-start_time:.1f}s"})
        yield messages, None
        
        # Process chunks in parallel
        chunk_responses = [None] * len(chunks)
        with ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor:
            futures = []
            for idx, chunk in enumerate(chunks):
                future = executor.submit(
                    lambda c, i: asyncio.run(process_chunk(agent, c, i)),
                    chunk, idx
                )
                futures.append(future)
                messages.append({"role": "assistant", "content": f"🔍 Processing chunk {idx+1}/{len(chunks)}..."})
                yield messages, None
            
            for future in as_completed(futures):
                idx, response = future.result()
                chunk_responses[idx] = response
                messages.append({"role": "assistant", "content": f"✅ Chunk {idx+1} processed"})
                yield messages, None
        
        # Combine and summarize
        combined = "\n\n".join([r for r in chunk_responses if r])
        messages.append({"role": "assistant", "content": "📊 Generating final report..."})
        yield messages, None
        
        final_response = ""
        for result in agent.run_gradio_chat(
            message=f"Summarize these clinical findings:\n\n{combined}",
            history=[],
            temperature=0.2,
            max_new_tokens=MAX_NEW_TOKENS*2,  # Allow more tokens for summary
            max_token=MAX_MODEL_TOKENS,
            call_agent=False,
            conversation=[],
        ):
            if isinstance(result, str):
                final_response += result
            elif hasattr(result, "content"):
                final_response += result.content
            elif isinstance(result, list):
                for r in result:
                    if hasattr(r, "content"):
                        final_response += r.content
            
            messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(final_response)}"
            yield messages, None
        
        # Save report
        final_report = f"# Final Clinical Report\n\n{clean_response(final_response)}"
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        report_path = os.path.join(report_dir, f"report_{timestamp}.md")
        
        with open(report_path, 'w') as f:
            f.write(final_report)
        
        messages.append({"role": "assistant", "content": f"✅ Report saved: report_{timestamp}.md"})
        yield messages, report_path
    
    except Exception as e:
        logger.error(f"Processing failed: {str(e)}")
        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
        yield messages, None

def create_ui(agent: TxAgent):
    """Create the Gradio interface"""
    with gr.Blocks(title="Clinical Analysis", css=".gradio-container {max-width: 900px}") as demo:
        gr.Markdown("## 🏥 Clinical Data Analysis (TxAgent)")
        
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Analysis Progress",
                    show_copy_button=True,
                    height=600,
                    type="messages"
                )
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Upload Excel File",
                    file_types=[".xlsx"],
                    height=100
                )
                analyze_btn = gr.Button(
                    "🧠 Analyze Data",
                    variant="primary"
                )
                report_output = gr.File(
                    label="Download Report",
                    visible=False
                )
        
        analyze_btn.click(
            fn=lambda file: process_file(agent, file.name) if file else ([{"role": "assistant", "content": "❌ Please upload a file"}], None),
            inputs=[file_input],
            outputs=[chatbot, report_output]
        )
    
    return demo

if __name__ == "__main__":
    try:
        agent = init_agent()
        demo = create_ui(agent)
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True,
            allowed_paths=[report_dir],
            share=False
        )
    except Exception as e:
        logger.error(f"Application failed: {str(e)}")
        sys.exit(1)