app.py

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

# 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

# Constants
MAX_MODEL_TOKENS = 32768  # Model's maximum sequence length
MAX_CHUNK_TOKENS = 8192   # Chunk size aligned with max_num_batched_tokens
MAX_NEW_TOKENS = 2048     # Maximum tokens for generation
PROMPT_OVERHEAD = 500     # Estimated tokens for prompt template overhead

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:
    """Estimate the number of tokens based on character length."""
    return len(text) // 3.5 + 1  # Add 1 to avoid zero estimates

def extract_text_from_excel(file_path: str) -> str:
    """Extract text from all sheets in an Excel file."""
    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:
        raise ValueError(f"Failed to extract text from Excel file: {str(e)}")
    return "\n".join(all_text)

def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
    """
    Split text into chunks, ensuring each chunk is within token limits,
    accounting for prompt overhead.
    """
    effective_max_tokens = max_tokens - PROMPT_OVERHEAD
    if effective_max_tokens <= 0:
        raise ValueError(f"Effective max tokens ({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_tokens:
            if current_chunk:  # Save the current chunk if it's not empty
                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))

    return chunks

def build_prompt_from_text(chunk: str) -> str:
    """Build a prompt for analyzing a chunk of clinical data."""
    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:
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations
"""

def init_agent():
    """Initialize the TxAgent with model and tool configurations."""
    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=[]
    )
    agent.init_model()
    return agent

def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
    """Process the Excel file and generate a final report."""
    messages = chatbot_state if chatbot_state else []
    report_path = None

    if file is None or not hasattr(file, "name"):
        messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
        return messages, report_path

    try:
        messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
        messages.append({"role": "assistant", "content": "⏳ Extracting and analyzing data..."})

        # Extract text and split into chunks
        extracted_text = extract_text_from_excel(file.name)
        chunks = split_text_into_chunks(extracted_text, max_tokens=MAX_CHUNK_TOKENS)
        chunk_responses = []

        # Process each chunk
        for i, chunk in enumerate(chunks):
            messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
            
            prompt = build_prompt_from_text(chunk)
            prompt_tokens = estimate_tokens(prompt)
            if prompt_tokens > MAX_MODEL_TOKENS:
                messages.append({"role": "assistant", "content": f"❌ Chunk {i+1} prompt too long ({prompt_tokens} tokens). Skipping..."})
                continue

            response = ""
            try:
                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
            except Exception as e:
                messages.append({"role": "assistant", "content": f"❌ Error analyzing chunk {i+1}: {str(e)}"})
                continue
            
            chunk_responses.append(clean_response(response))
            messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})

        if not chunk_responses:
            messages.append({"role": "assistant", "content": "❌ No valid chunk responses to summarize."})
            return messages, report_path

        # Summarize chunk responses incrementally to avoid token limit
        summary = ""
        current_summary_tokens = 0
        for i, response in enumerate(chunk_responses):
            response_tokens = estimate_tokens(response)
            if current_summary_tokens + response_tokens > MAX_MODEL_TOKENS - PROMPT_OVERHEAD - MAX_NEW_TOKENS:
                # Summarize current summary
                summary_prompt = f"Summarize the following analysis:\n\n{summary}\n\nProvide a concise summary."
                summary_response = ""
                try:
                    for result in agent.run_gradio_chat(
                        message=summary_prompt,
                        history=[],
                        temperature=0.2,
                        max_new_tokens=MAX_NEW_TOKENS,
                        max_token=MAX_MODEL_TOKENS,
                        call_agent=False,
                        conversation=[],
                    ):
                        if isinstance(result, str):
                            summary_response += result
                        elif hasattr(result, "content"):
                            summary_response += result.content
                        elif isinstance(result, list):
                            for r in result:
                                if hasattr(r, "content"):
                                    summary_response += r.content
                    summary = clean_response(summary_response)
                    current_summary_tokens = estimate_tokens(summary)
                except Exception as e:
                    messages.append({"role": "assistant", "content": f"❌ Error summarizing intermediate results: {str(e)}"})
                    return messages, report_path

            summary += f"\n\n### Chunk {i+1} Analysis\n{response}"
            current_summary_tokens += response_tokens

        # Final summarization
        final_prompt = f"Summarize the key findings from the following analyses:\n\n{summary}"
        messages.append({"role": "assistant", "content": "📊 Generating final report..."})

        final_report_text = ""
        try:
            for result in agent.run_gradio_chat(
                message=final_prompt,
                history=[],
                temperature=0.2,
                max_new_tokens=MAX_NEW_TOKENS,
                max_token=MAX_MODEL_TOKENS,
                call_agent=False,
                conversation=[],
            ):
                if isinstance(result, str):
                    final_report_text += result
                elif hasattr(result, "content"):
                    final_report_text += result.content
                elif isinstance(result, list):
                    for r in result:
                        if hasattr(r, "content"):
                            final_report_text += r.content
        except Exception as e:
            messages.append({"role": "assistant", "content": f"❌ Error generating final report: {str(e)}"})
            return messages, report_path

        final_report = f"# \U0001f9e0 Final Patient Report\n\n{clean_response(final_report_text)}"
        messages[-1]["content"] = f"📊 Final Report:\n\n{clean_response(final_report_text)}"

        # Save the report
        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 generated and saved: report_{timestamp}.md"})

    except Exception as e:
        messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})

    return messages, report_path

def create_ui(agent):
    """Create the Gradio UI for the patient history analysis tool."""
    with gr.Blocks(title="Patient History Chat", css=".gradio-container {max-width: 900px !important}") as demo:
        gr.Markdown("## 🏥 Patient History Analysis Tool")
        
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Clinical Assistant",
                    show_copy_button=True,
                    height=600,
                    type="messages",
                    avatar_images=(
                        None,
                        "https://i.imgur.com/6wX7Zb4.png"
                    )
                )
            with gr.Column(scale=1):
                file_upload = gr.File(
                    label="Upload Excel File",
                    file_types=[".xlsx"],
                    height=100
                )
                analyze_btn = gr.Button(
                    "🧠 Analyze Patient History",
                    variant="primary"
                )
                report_output = gr.File(
                    label="Download Report",
                    visible=False,
                    interactive=False
                )

        # State to maintain chatbot messages
        chatbot_state = gr.State(value=[])

        def update_ui(file, current_state):
            messages, report_path = process_final_report(agent, file, current_state)
            report_update = gr.update(visible=report_path is not None, value=report_path)
            return messages, report_update, messages

        analyze_btn.click(
            fn=update_ui,
            inputs=[file_upload, chatbot_state],
            outputs=[chatbot, report_output, chatbot_state],
            api_name="analyze"
        )

    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=["/data/hf_cache/reports"],
            share=False
        )
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)