app.py

import os
import json
import logging
import torch
from txagent import TxAgent
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download
from tooluniverse import ToolUniverse

# Configuration
CONFIG = {
    "model_name": "mims-harvard/TxAgent-T1-Llama-3.1-8B",
    "rag_model_name": "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding_e27fb393f3144ec28f620f33d4d79911.pt",
    "local_dir": "./models",
    "tool_files": {
        "new_tool": "./data/new_tool.json"
    }
}

# Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def prepare_tool_files():
    os.makedirs("./data", exist_ok=True)
    if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
        tu = ToolUniverse()
        tools = tu.get_all_tools()
        with open(CONFIG["tool_files"]["new_tool"], "w") as f:
            json.dump(tools, f, indent=2)
        logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")

def download_model_files():
    os.makedirs(CONFIG["local_dir"], exist_ok=True)
    logger.info("Downloading model files...")

    snapshot_download(
        repo_id=CONFIG["model_name"],
        local_dir=os.path.join(CONFIG["local_dir"], CONFIG["model_name"]),
        resume_download=True
    )
    snapshot_download(
        repo_id=CONFIG["rag_model_name"],
        local_dir=os.path.join(CONFIG["local_dir"], CONFIG["rag_model_name"]),
        resume_download=True
    )

def generate_embeddings(agent):
    """Generates and assigns embeddings manually"""
    path = os.path.join(CONFIG["local_dir"], CONFIG["embedding_filename"])
    try:
        if os.path.exists(path):
            logger.info("Embedding file already exists, loading...")
            agent.rag_model.tool_desc_embedding = torch.load(path)
            return

        logger.info("Generating embeddings and saving...")
        tools = agent.tooluniverse.get_all_tools()
        descs = [t["description"] for t in tools]
        embeddings = agent.rag_model.generate_embeddings(descs)
        torch.save(embeddings, path)
        agent.rag_model.tool_desc_embedding = embeddings
        logger.info("✅ Embeddings generated and saved")

    except Exception as e:
        logger.error(f"Failed to generate embeddings: {e}")
        raise

class TxAgentApp:
    def __init__(self):
        self.agent = None
        self.is_initialized = False

    def initialize(self):
        if self.is_initialized:
            return "Already initialized"

        try:
            self.agent = TxAgent(
                CONFIG["model_name"],
                CONFIG["rag_model_name"],
                tool_files_dict=CONFIG["tool_files"],
                force_finish=True,
                enable_checker=True,
                step_rag_num=10,
                seed=100,
                additional_default_tools=["DirectResponse", "RequireClarification"]
            )
            self.agent.init_model()
            generate_embeddings(self.agent)
            self.is_initialized = True
            return "✅ TxAgent initialized successfully"
        except Exception as e:
            return f"❌ Initialization failed: {str(e)}"

    def chat(self, message, history):
        if not self.is_initialized:
            return history + [(message, "⚠️ Error: Model not initialized")]

        try:
            response = ""
            for chunk in self.agent.run_gradio_chat(
                message=message,
                history=history,
                temperature=0.3,
                max_new_tokens=1024,
                max_tokens=8192,
                multi_agent=False,
                conversation=[],
                max_round=30
            ):
                response += chunk
            return history + [(message, response)]
        except Exception as e:
            return history + [(message, f"Error: {str(e)}")]

def create_interface():
    app = TxAgentApp()
    with gr.Blocks(title="TxAgent") as demo:
        gr.Markdown("# 🧠 TxAgent: Therapeutic Reasoning AI")

        with gr.Row():
            init_btn = gr.Button("Initialize Model", variant="primary")
            init_status = gr.Textbox(label="Initialization Status")

        chatbot = gr.Chatbot(height=600, label="Conversation")
        msg = gr.Textbox(label="Your Question")
        submit_btn = gr.Button("Submit")

        gr.Examples(
            examples=[
                "How to adjust Journavx dosage for hepatic impairment?",
                "Is Xolremdi safe with Prozac for WHIM syndrome?",
                "Warfarin-Amiodarone contraindications?"
            ],
            inputs=msg
        )

        init_btn.click(fn=app.initialize, outputs=init_status)
        msg.submit(fn=app.chat, inputs=[msg, chatbot], outputs=chatbot)
        submit_btn.click(fn=app.chat, inputs=[msg, chatbot], outputs=chatbot)

    return demo

if __name__ == "__main__":
    prepare_tool_files()
    download_model_files()
    interface = create_interface()
    interface.launch(server_name="0.0.0.0", server_port=7860, share=False)