app.py

import sys
import os
import pandas as pd
import json
import gradio as gr
from typing import List, Tuple, Union, Generator, BinaryIO
import hashlib
import shutil
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")
file_cache_dir = os.path.join(persistent_dir, "cache")
report_dir = os.path.join(persistent_dir, "reports")

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

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

# Fixed this line - added missing closing parenthesis
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src"))

try:
    from txagent.txagent import TxAgent
except ImportError as e:
    print(f"Error importing TxAgent: {e}")
    sys.exit(1)

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 estimate_tokens(text: str) -> int:
    return len(text) // 3.5 + 1

def extract_text_from_excel(file_obj: Union[str, os.PathLike, BinaryIO]) -> str:
    """Extract text from Excel file which can be a path, file-like object, or Gradio file object."""
    all_text = []
    try:
        # Handle Gradio file object which has .name attribute
        if hasattr(file_obj, 'name'):
            file_path = file_obj.name
        else:
            file_path = file_obj
            
        xls = pd.ExcelFile(file_path)
    except Exception as e:
        raise ValueError(f"❌ Error reading Excel file: {e}")
    
    for sheet_name in xls.sheet_names:
        try:
            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)
        except Exception as e:
            print(f"Warning: Could not parse sheet {sheet_name}: {e}")
            continue
            
    return "\n".join(all_text)

def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS, max_chunks: int = 30) -> List[str]:
    effective_max = max_tokens - PROMPT_OVERHEAD
    lines, chunks, curr_chunk, curr_tokens = text.split("\n"), [], [], 0
    for line in lines:
        t = estimate_tokens(line)
        if curr_tokens + t > effective_max:
            if curr_chunk:
                chunks.append("\n".join(curr_chunk))
            if len(chunks) >= max_chunks:
                break
            curr_chunk, curr_tokens = [line], t
        else:
            curr_chunk.append(line)
            curr_tokens += t
    if curr_chunk and len(chunks) < max_chunks:
        chunks.append("\n".join(curr_chunk))
    return chunks

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

Analyze the following clinical notes and provide a detailed, concise summary focusing on:
- Diagnostic Patterns
- Medication Issues
- Missed Opportunities
- Inconsistencies
- Follow-up Recommendations

---

{chunk}

---
Respond in well-structured bullet points with medical reasoning.
"""

def init_agent():
    tool_path = os.path.join(tool_cache_dir, "new_tool.json")
    if not os.path.exists(tool_path):
        # Create default tool file if it doesn't exist
        default_tool = {
            "name": "new_tool",
            "description": "Default tool configuration",
            "version": "1.0",
            "tools": []
        }
        with open(tool_path, 'w') as f:
            json.dump(default_tool, f)
            
    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 stream_report(agent, input_file: Union[str, BinaryIO], full_output: str) -> Generator[Tuple[str, Union[str, None], str], None, None]:
    accumulated_text = ""
    try:
        if input_file is None:
            yield "❌ Please upload a valid Excel file.", None, ""
            return

        text = extract_text_from_excel(input_file)
        chunks = split_text_into_chunks(text)

        for i, chunk in enumerate(chunks):
            prompt = build_prompt_from_text(chunk)
            partial = ""
            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):
                    partial += res
                elif hasattr(res, "content"):
                    partial += res.content
            cleaned = clean_response(partial)
            accumulated_text += f"\n\n📄 **Chunk {i+1}**:\n{cleaned}"
            yield accumulated_text, None, ""

        summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + accumulated_text
        final_report = ""
        for res 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(res, str):
                final_report += res
            elif hasattr(res, "content"):
                final_report += res.content

        cleaned = clean_response(final_report)
        accumulated_text += f"\n\n📊 **Final Summary**:\n{cleaned}"
        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(f"# 🧠 Final Patient Report\n\n{cleaned}")

        yield accumulated_text, report_path, cleaned

    except Exception as e:
        yield f"❌ Error: {str(e)}", 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;
        }
        .output-markdown {
            background-color: #131720;
            border-radius: 12px;
            padding: 20px;
            min-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);
        }
    """) as demo:
        gr.Markdown("""# 🧠 Clinical Reasoning Assistant
Upload clinical Excel records below and click **Analyze** to generate a medical summary.
""")
        file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
        analyze_btn = gr.Button("Analyze")
        report_output_markdown = gr.Markdown(elem_classes="output-markdown")
        report_file = gr.File(label="Download Report", visible=False)
        full_output = gr.State(value="")

        analyze_btn.click(
            fn=stream_report,
            inputs=[file_upload, full_output],
            outputs=[report_output_markdown, report_file, full_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, 
            allowed_paths=["/data/hf_cache/reports"], 
            share=True,
            show_error=True
        )
    except Exception as e:
        print(f"Error: {str(e)}", file=sys.stderr)
        sys.exit(1)