app.py

import sys
import os
import pandas as pd
import gradio as gr
from typing import List, Tuple
import re
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed

# Setup directories
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")
report_dir = os.path.join(persistent_dir, "reports")

for d in [model_cache_dir, tool_cache_dir, report_dir]:
    os.makedirs(d, exist_ok=True)

os.environ["HF_HOME"] = model_cache_dir
os.environ["TRANSFORMERS_CACHE"] = model_cache_dir

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src"))
from txagent.txagent import TxAgent

MAX_MODEL_TOKENS = 32768
MAX_CHUNK_TOKENS = 8192
MAX_NEW_TOKENS = 2048
PROMPT_OVERHEAD = 500

def clean_response(text: str) -> str:
    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 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).astype(str).fillna("")
        rows = df.apply(lambda row: " | ".join([cell for cell in row if cell.strip()]), axis=1)
        sheet_text = [f"[{sheet_name}] {line}" for line in rows if line.strip()]
        all_text.extend(sheet_text)
    return "\n".join(all_text)

def split_text_into_chunks(text: str) -> List[str]:
    effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD
    lines, chunks, curr_chunk = text.split("\n"), [], []
    curr_tokens = sum(len(line.split()) for line in curr_chunk)
    
    for line in lines:
        line_tokens = len(line.split())
        if curr_tokens + line_tokens > effective_max:
            if curr_chunk:
                chunks.append("\n".join(curr_chunk))
            curr_chunk, curr_tokens = [line], line_tokens
        else:
            curr_chunk.append(line)
            curr_tokens += line_tokens
    if curr_chunk:
        chunks.append("\n".join(curr_chunk))
    return chunks

def build_prompt_from_text(chunk: str) -> str:
    return f"""Analyze these clinical notes and provide:
- Diagnostic patterns
- Medication issues 
- Missed opportunities
- Inconsistencies
- Follow-up recommendations

Respond with clear bullet points:

{chunk}"""

def init_agent():
    tool_path = os.path.join(tool_cache_dir, "new_tool.json")
    if not os.path.exists(tool_path):
        import shutil
        shutil.copy("data/new_tool.json", 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": tool_path},
        force_finish=True,
        enable_checker=True,
        step_rag_num=4,
        seed=100
    )
    agent.init_model()
    return agent

def process_final_report(agent, file, chatbot_state: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], str]:
    messages = chatbot_state.copy() if chatbot_state else []
    
    if file is None:
        messages.append(("assistant", "❌ Please upload a valid Excel file."))
        return messages, None

    messages.append(("user", f"Processing Excel file: {os.path.basename(file.name)}"))
    yield messages, None
    
    try:
        text = extract_text_from_excel(file.name)
        chunks = split_text_into_chunks(text)
        
        messages.append(("assistant", "🔍 Analyzing clinical data..."))
        yield messages, None
        
        full_report = []
        for i, chunk in enumerate(chunks, 1):
            prompt = build_prompt_from_text(chunk)
            response = ""
            
            for res 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(res, str):
                    response += res
                elif hasattr(res, "content"):
                    response += res.content
            
            cleaned = clean_response(response)
            full_report.append(cleaned)
            
            # Update progress in chat
            progress_msg = f"✅ Analyzed section {i}/{len(chunks)}"
            if len(messages) > 2 and "Analyzed section" in messages[-1][1]:
                messages[-1] = ("assistant", progress_msg)
            else:
                messages.append(("assistant", progress_msg))
            yield messages, None
        
        # Generate final report
        final_report = "## 🧠 Final Clinical Report\n\n" + "\n\n".join(full_report)
        report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
        with open(report_path, 'w') as f:
            f.write(final_report)
        
        messages.append(("assistant", f"✅ Report generated and saved: {os.path.basename(report_path)}"))
        messages.append(("assistant", final_report))
        yield messages, report_path
        
    except Exception as e:
        messages.append(("assistant", f"❌ Error: {str(e)}"))
        yield messages, None

def create_ui(agent):
    with gr.Blocks(css="""
        body {
            background: #10141f;
            color: #ffffff;
            font-family: 'Inter', sans-serif;
            margin: 0;
            padding: 0;
        }
        .gradio-container {
            padding: 30px;
            width: 100vw;
            max-width: 100%;
            border-radius: 0;
            background-color: #1a1f2e;
        }
        .chatbot {
            background-color: #131720;
            border-radius: 12px;
            padding: 20px;
            height: 600px;
            overflow-y: auto;
            border: 1px solid #2c3344;
        }
        .gr-button {
            background: linear-gradient(135deg, #4b4ced, #37b6e9);
            color: white;
            font-weight: 500;
            border: none;
            padding: 10px 20px;
            border-radius: 8px;
            transition: background 0.3s ease;
        }
        .gr-button:hover {
            background: linear-gradient(135deg, #37b6e9, #4b4ced);
        }
        .report-content {
            background-color: #1a1f2e;
            padding: 15px;
            border-radius: 8px;
            margin-top: 10px;
            border: 1px solid #2c3344;
        }
        .bullet-points {
            margin-left: 20px;
        }
    """) as demo:
        gr.Markdown("""# Clinical Reasoning Assistant
Upload clinical Excel records below and click **Analyze** to generate a medical summary.
""")
        
        chatbot = gr.Chatbot(label="Chatbot", elem_classes="chatbot")
        file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
        analyze_btn = gr.Button("Analyze")
        report_output = gr.File(label="Download Report", visible=False)
        chatbot_state = gr.State([])
        
        analyze_btn.click(
            fn=process_final_report,
            inputs=[file_upload, chatbot_state, gr.State(agent)],
            outputs=[chatbot, report_output],
            show_progress="hidden"
        )

    return demo

if __name__ == "__main__":
    try:
        agent = init_agent()
        demo = create_ui(agent)
        demo.launch(
            server_name="0.0.0.0", 
            server_port=7860, 
            allowed_paths=["/data/hf_cache/reports"], 
            share=False
        )
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)