app.py

import os
import json
import logging
import torch
from txagent import TxAgent
import gradio as gr
from tooluniverse import ToolUniverse

# Configuration with hardcoded embedding file
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_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",
    "tool_files": {
        "new_tool": "./data/new_tool.json"
    }
}

# Logging setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def prepare_tool_files():
    os.makedirs("./data", exist_ok=True)
    if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
        logger.info("Generating tool list using ToolUniverse...")
        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 patch_embedding_loading():
    """Monkey-patch the embedding loading functionality"""
    try:
        # Try to get the RAG model class dynamically
        from txagent.txagent import TxAgent as TxAgentClass
        original_init = TxAgentClass.__init__
        
        def patched_init(self, *args, **kwargs):
            # First let the original initialization happen
            original_init(self, *args, **kwargs)
            
            # Then handle the embeddings our way
            try:
                if os.path.exists(CONFIG["embedding_filename"]):
                    logger.info(f"Loading embeddings from {CONFIG['embedding_filename']}")
                    self.rag_model.tool_desc_embedding = torch.load(CONFIG["embedding_filename"])
                    
                    # Handle tool count mismatch
                    tools = self.tooluniverse.get_all_tools()
                    current_count = len(tools)
                    embedding_count = len(self.rag_model.tool_desc_embedding)
                    
                    if current_count != embedding_count:
                        logger.warning(f"Tool count mismatch (tools: {current_count}, embeddings: {embedding_count})")
                        
                        if current_count < embedding_count:
                            self.rag_model.tool_desc_embedding = self.rag_model.tool_desc_embedding[:current_count]
                            logger.info(f"Truncated embeddings to match {current_count} tools")
                        else:
                            last_embedding = self.rag_model.tool_desc_embedding[-1]
                            padding = [last_embedding] * (current_count - embedding_count)
                            self.rag_model.tool_desc_embedding = torch.cat(
                                [self.rag_model.tool_desc_embedding] + padding
                            )
                            logger.info(f"Padded embeddings to match {current_count} tools")
                else:
                    logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
            
            except Exception as e:
                logger.error(f"Failed to load embeddings: {str(e)}")
        
        # Apply the patch
        TxAgentClass.__init__ = patched_init
        logger.info("Successfully patched embedding loading")
        
    except Exception as e:
        logger.error(f"Failed to patch embedding loading: {str(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:
            # Apply our patch before initialization
            patch_embedding_loading()
            
            logger.info("Initializing TxAgent...")
            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"]
            )
            
            logger.info("Loading models...")
            self.agent.init_model()
            
            self.is_initialized = True
            return "✅ TxAgent initialized successfully"
            
        except Exception as e:
            logger.error(f"Initialization failed: {str(e)}")
            return f"❌ Initialization failed: {str(e)}"

    def chat(self, message, history):
        if not self.is_initialized:
            return history + [(message, "⚠️ Please initialize the model first")]
        
        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
                yield history + [(message, response)]
                
        except Exception as e:
            logger.error(f"Chat error: {str(e)}")
            yield history + [(message, f"Error: {str(e)}")]

def create_interface():
    app = TxAgentApp()
    
    with gr.Blocks(
        title="TxAgent",
        css="""
        .gradio-container {max-width: 900px !important}
        """
    ) as demo:
        gr.Markdown("""
        # 🧠 TxAgent: Therapeutic Reasoning AI
        ### (Using pre-loaded embeddings)
        """)
        
        with gr.Row():
            init_btn = gr.Button("Initialize Model", variant="primary")
            init_status = gr.Textbox(label="Status", interactive=False)
        
        chatbot = gr.Chatbot(
            height=500, 
            label="Conversation",
            type="messages"  # Fixing the deprecation warning
        )
        msg = gr.Textbox(label="Your clinical question")
        clear_btn = gr.Button("Clear Chat")
        
        gr.Examples(
            examples=[
                "How to adjust Journavx for renal impairment?",
                "Xolremdi and Prozac interaction in WHIM syndrome?",
                "Alternative to Warfarin for patient with amiodarone?"
            ],
            inputs=msg
        )
        
        init_btn.click(
            fn=app.initialize,
            outputs=init_status
        )
        
        msg.submit(
            fn=app.chat,
            inputs=[msg, chatbot],
            outputs=chatbot
        )
        
        clear_btn.click(
            fn=lambda: ([], ""),
            outputs=[chatbot, msg]
        )
    
    return demo

if __name__ == "__main__":
    try:
        logger.info("Starting application...")
        
        # Verify embedding file exists
        if not os.path.exists(CONFIG["embedding_filename"]):
            logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
            logger.info("Please ensure the file is in the root directory")
        else:
            logger.info(f"Found embedding file: {CONFIG['embedding_filename']}")
        
        # Prepare tool files
        prepare_tool_files()
        
        # Launch interface
        interface = create_interface()
        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False
        )
    except Exception as e:
        logger.error(f"Application failed to start: {str(e)}")
        raise