import sys import os import pandas as pd import json import gradio as gr from typing import List, Tuple, Union, Generator 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 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 estimate_tokens(text: str) -> int: return len(text) // 3.5 + 1 def extract_text_from_excel(file_obj: Union[str, os.PathLike, 'file']) -> str: all_text = [] try: xls = pd.ExcelFile(file_obj) except Exception as e: raise ValueError(f"❌ Error reading Excel file: {e}") 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, 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 validate_tool_file(file_path): try: with open(file_path, 'r') as f: data = json.load(f) if isinstance(data, list): assert all(isinstance(t, dict) and "name" in t for t in data), "Invalid list format" elif isinstance(data, dict): assert "tools" in data and isinstance(data["tools"], list), "'tools' field missing or invalid" assert all(isinstance(t, dict) and "name" in t for t in data["tools"]), "Invalid item in 'tools'" else: raise ValueError("Unexpected structure") return True except Exception as e: print(f"❌ Tool validation failed for {file_path}: {e}") return False def init_agent(): all_tool_paths = { "opentarget": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/opentarget_tools.json", "fda_drug_label": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/fda_drug_labeling_tools.json", "special_tools": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/special_tools.json", "monarch": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/monarch_tools.json", "new_tool": os.path.join(tool_cache_dir, "new_tool.json"), } if not os.path.exists(all_tool_paths["new_tool"]): shutil.copy(os.path.abspath("data/new_tool.json"), all_tool_paths["new_tool"]) valid_tool_paths = {} for key, path in all_tool_paths.items(): if validate_tool_file(path): valid_tool_paths[key] = path else: print(f"⚠️ Skipping invalid tool file: {path}") if not valid_tool_paths: raise RuntimeError("❌ No valid tool files found to load into TxAgent.") 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=valid_tool_paths, 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, 'file'], 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 if hasattr(input_file, "read"): text = extract_text_from_excel(input_file) elif isinstance(input_file, str) and os.path.exists(input_file): text = extract_text_from_excel(input_file) else: raise ValueError("❌ Invalid or missing 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) except Exception as e: print(f"Error: {str(e)}") sys.exit(1)