app.py

import sys
import os
import pandas as pd
import json
import gradio as gr
from typing import List, Tuple, Dict, Any, Generator, 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_TOKENS = 32768
MAX_NEW_TOKENS = 2048


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


def extract_text_from_excel(file_path: str) -> str:
    all_text = []
    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)
    return "\n".join(all_text)


def split_text_into_chunks(text: str, max_tokens: int = MAX_TOKENS) -> List[str]:
    lines = text.split("\n")
    chunks = []
    current_chunk = []
    current_tokens = 0

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

    if current_chunk:
        chunks.append("\n".join(current_chunk))
    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:
- 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=[]
    )
    agent.init_model()
    return agent


def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
    # Initialize with empty values
    messages = []
    report_path = None
    
    if file is None or not hasattr(file, "name"):
        messages = [{"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."}]
        yield messages, None
        return

    try:
        # Initial processing message
        messages = [{"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"},
                   {"role": "assistant", "content": "⏳ Extracting and analyzing data..."}]
        yield messages, None

        extracted_text = extract_text_from_excel(file.name)
        chunks = split_text_into_chunks(extracted_text)
        chunk_responses = []

        # Process each chunk
        for i, chunk in enumerate(chunks):
            messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
            yield messages, None
            
            prompt = build_prompt_from_text(chunk)
            response = ""
            for result in agent.run_gradio_chat(
                message=prompt,
                history=[],
                temperature=0.2,
                max_new_tokens=MAX_NEW_TOKENS,
                max_token=MAX_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
            
            chunk_responses.append(clean_response(response))
            messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})
            yield messages, None

        # Final summarization
        final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above."
        messages.append({"role": "assistant", "content": "📊 Generating final report..."})
        yield messages, None

        stream_text = ""
        for result in agent.run_gradio_chat(
            message=final_prompt,
            history=[],
            temperature=0.2,
            max_new_tokens=MAX_NEW_TOKENS,
            max_token=MAX_TOKENS,
            call_agent=False,
            conversation=[],
        ):
            if isinstance(result, str):
                stream_text += result
            elif hasattr(result, "content"):
                stream_text += result.content
            elif isinstance(result, list):
                for r in result:
                    if hasattr(r, "content"):
                        stream_text += r.content
            
            messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(stream_text)}"
            yield messages, None

        # Save final report
        final_report = f"# \U0001f9e0 Final Patient Report\n\n{clean_response(stream_text)}"
        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"})
        yield messages, report_path

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


def create_ui(agent):
    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,  # User avatar
                        "https://i.imgur.com/6wX7Zb4.png"  # Bot avatar
                    )
                )
            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
                )

        analyze_btn.click(
            fn=lambda file: stream_final_report(agent, file),
            inputs=[file_upload],
            outputs=[chatbot, report_output],
            api_name="analyze"
        )

        def show_report(report_path):
            if report_path:
                return gr.File(visible=True, value=report_path)
            return gr.File(visible=False)

    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)