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_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 process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Dict[str, Any]]: messages = chatbot_state if chatbot_state else [] report_output = {"visible": False, "value": 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_output try: messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"}) messages.append({"role": "assistant", "content": "⏳ Extracting and analyzing data..."}) extracted_text = extract_text_from_excel(file.name) chunks = split_text_into_chunks(extracted_text) chunk_responses = [] for i, chunk in enumerate(chunks): messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."}) 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"}) final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above." messages.append({"role": "assistant", "content": "📊 Generating final report..."}) 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 final_report = f"# \U0001f9e0 Final Patient Report\n\n{clean_response(stream_text)}" messages[-1]["content"] = f"📊 Final 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"}) report_output = {"visible": True, "value": report_path} except Exception as e: messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"}) return messages, report_output 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, "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_output = process_final_report(agent, file, current_state) return messages, report_output, 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)