Update app.py

app.py

@@ -1,16 +1,25 @@
+import random
+import datetime
+import sys
+from txagent import TxAgent
+import spaces
+import gradio as gr
 import os
-import json
-import logging
 import torch
-import gradio as gr
-from tooluniverse import ToolUniverse
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import warnings
-from typing import List, Dict, Any
+import logging
 from importlib.resources import files
+import traceback
 
-# Suppress specific warnings
-warnings.filterwarnings("ignore", category=UserWarning)
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Determine the directory where the current file is located
+current_dir = os.path.dirname(os.path.abspath(__file__))
+os.environ["MKL_THREADING_LAYER"] = "GNU"
 
 # Configuration
 CONFIG = {
@@ -22,27 +31,76 @@ CONFIG = {
         "fda_drug_label": str(files('tooluniverse.data').joinpath('fda_drug_labeling_tools.json')),
         "special_tools": str(files('tooluniverse.data').joinpath('special_tools.json')),
         "monarch": str(files('tooluniverse.data').joinpath('monarch_tools.json')),
-        "new_tool": "./data/new_tool.json"
+        "new_tool": os.path.join(current_dir, 'data', 'new_tool.json')
     }
 }
 
-# Logging setup
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+# Set an environment variable
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-def prepare_tool_files():
-    """Ensure tool files exist and are populated"""
-    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() if hasattr(tu, 'get_all_tools') else []
-        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']}")
+DESCRIPTION = '''
+<div>
+<h1 style="text-align: center;">TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools </h1>
+</div>
+'''
+INTRO = """
+Precision therapeutics require multimodal adaptive models that provide personalized treatment recommendations. We introduce TxAgent, an AI agent that leverages multi-step reasoning and real-time biomedical knowledge retrieval across a toolbox of 211 expert-curated tools to navigate complex drug interactions, contraindications, and patient-specific treatment strategies, delivering evidence-grounded therapeutic decisions. TxAgent executes goal-oriented tool selection and iterative function calls to solve therapeutic tasks that require deep clinical understanding and cross-source validation. The ToolUniverse consolidates 211 tools linked to trusted sources, including all US FDA-approved drugs since 1939 and validated clinical insights from Open Targets.
+"""
+
+LICENSE = """
+We welcome your feedback and suggestions to enhance your experience with TxAgent, and if you're interested in collaboration, please email Marinka Zitnik and Shanghua Gao.
+
+### Medical Advice Disclaimer
+DISCLAIMER: THIS WEBSITE DOES NOT PROVIDE MEDICAL ADVICE
+The information, including but not limited to, text, graphics, images and other material contained on this website are for informational purposes only. No material on this site is intended to be a substitute for professional medical advice, diagnosis or treatment. Always seek the advice of your physician or other qualified health care provider with any questions you may have regarding a medical condition or treatment and before undertaking a new health care regimen, and never disregard professional medical advice or delay in seeking it because of something you have read on this website.
+"""
+
+PLACEHOLDER = """
+<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
+   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">TxAgent</h1>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Tips before using TxAgent:</p>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Please click clear🗑️
+ (top-right) to remove previous context before sumbmitting a new question.</p>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Click retry🔄 (below message)  to get multiple versions of the answer.</p>
+</div>
+"""
+
+css = """
+h1 {
+  text-align: center;
+  display: block;
+}
+
+#duplicate-button {
+  margin: auto;
+  color: white;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+.small-button button {
+    font-size: 12px !important;
+    padding: 4px 8px !important;
+    height: 6px !important;
+    width: 4px !important;
+}
+.gradio-accordion {
+    margin-top: 0px !important;
+    margin-bottom: 0px !important;
+}
+"""
+
+chat_css = """
+.gr-button { font-size: 20px !important; }  /* Enlarges button icons */
+.gr-button svg { width: 32px !important; height: 32px !important; } /* Enlarges SVG icons */
+"""
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+question_examples = [
+    ['Given a 50-year-old patient experiencing severe acute pain and considering the use of the newly approved medication, Journavx, how should the dosage be adjusted considering the presence of moderate hepatic impairment?'],
+    ['Given a 50-year-old patient experiencing severe acute pain and considering the use of the newly approved medication, Journavx, how should the dosage be adjusted considering the presence of severe hepatic impairment?'],
+    ['A 30-year-old patient is taking Prozac to treat their depression. They were recently diagnosed with WHIM syndrome and require a treatment for that condition as well. Is Xolremdi suitable for this patient, considering contraindications?'],
+]
 
 def safe_load_embeddings(filepath: str) -> Any:
     """Safely load embeddings with proper weights_only handling"""
@@ -59,209 +117,263 @@ def safe_load_embeddings(filepath: str) -> Any:
             logger.error(f"Failed to load embeddings even with safe_globals: {str(e)}")
             return None
 
-class TxAgentWrapper:
-    def __init__(self):
-        self.model = None
-        self.tokenizer = None
-        self.rag_model = None
-        self.tooluniverse = None
-        self.is_initialized = False
-        self.special_tools = ['Finish', 'Tool_RAG', 'DirectResponse', 'RequireClarification']
-
-    def initialize(self) -> str:
-        """Initialize the model from Hugging Face"""
-        if self.is_initialized:
-            return "✅ Already initialized"
-        
-        try:
-            logger.info("Loading models from Hugging Face Hub...")
-            
-            # Verify tool files exist
-            for tool_name, tool_path in CONFIG["tool_files"].items():
-                if tool_name != "new_tool" and not os.path.exists(tool_path):
-                    raise FileNotFoundError(f"Tool file not found: {tool_path}")
-            
-            # Initialize ToolUniverse with verified paths
-            self.tooluniverse = ToolUniverse(tool_files=CONFIG["tool_files"])
-            if hasattr(self.tooluniverse, 'load_tools'):
-                self.tooluniverse.load_tools()
-                logger.info(f"Loaded {len(self.tooluniverse.tools)} tools")
-            else:
-                logger.error("ToolUniverse doesn't have load_tools method")
-                return "❌ Failed to load tools"
-            
-            # Load main model
-            self.tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_name"])
-            self.model = AutoModelForCausalLM.from_pretrained(
-                CONFIG["model_name"],
-                device_map="auto",
-                torch_dtype=torch.float16
-            )
-            
-            # Load embeddings if file exists
-            if os.path.exists(CONFIG["embedding_filename"]):
-                self.rag_model = safe_load_embeddings(CONFIG["embedding_filename"])
-                if self.rag_model is None:
-                    return "❌ Failed to load embeddings"
-            
-            self.is_initialized = True
-            return "✅ Model initialized successfully"
-            
-        except Exception as e:
-            logger.error(f"Initialization failed: {str(e)}")
-            return f"❌ Initialization failed: {str(e)}"
-
-    def chat(self, message: str, history: List[List[str]]) -> List[List[str]]:
-        """Handle chat interactions with the model"""
-        if not self.is_initialized:
-            return history + [["", "⚠️ Please initialize the model first"]]
-        
-        try:
-            if len(message) <= 10:
-                return history + [["", "Please provide a more detailed question (at least 10 characters)"]]
-            
-            # Prepare tools prompt
-            tools_prompt = self._prepare_tools_prompt(message)
-            
-            # Format conversation
-            conversation = [
-                {"role": "system", "content": "You are a helpful assistant that will solve problems through detailed, step-by-step reasoning." + tools_prompt},
-                *self._format_history(history),
-                {"role": "user", "content": message}
-            ]
-            
-            # Generate response
-            inputs = self.tokenizer.apply_chat_template(
-                conversation,
-                add_generation_prompt=True,
-                return_tensors="pt"
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                inputs,
-                max_new_tokens=1024,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-            
-            # Decode and clean response
-            response = self.tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
-            response = response.split("[TOOL_CALLS]")[0].strip()
-            
-            return history + [[message, response]]
-            
-        except Exception as e:
-            logger.error(f"Chat error: {str(e)}")
-            return history + [["", f"Error: {str(e)}"]]
-
-    def _prepare_tools_prompt(self, message: str) -> str:
-        """Prepare the tools prompt section"""
-        if not hasattr(self.tooluniverse, 'tools'):
-            return ""
-            
-        tools_prompt = "\n\nYou have access to the following tools:\n"
-        for tool in self.tooluniverse.tools:
-            if tool['name'] not in self.special_tools:
-                tools_prompt += f"- {tool['name']}: {tool['description']}\n"
-        
-        # Add special tools
-        tools_prompt += "\nSpecial tools:\n"
-        tools_prompt += "- Finish: Use when you have the final answer\n"
-        tools_prompt += "- Tool_RAG: Search for additional tools when needed\n"
+def patch_embedding_loading():
+    """Monkey-patch the embedding loading functionality"""
+    try:
+        from txagent.toolrag import ToolRAGModel
         
-        return tools_prompt
-
-    def _format_history(self, history: List[List[str]]) -> List[Dict[str, str]]:
-        """Format chat history for the model"""
-        formatted = []
-        for user_msg, bot_msg in history:
-            formatted.append({"role": "user", "content": user_msg})
-            if bot_msg:
-                formatted.append({"role": "assistant", "content": bot_msg})
-        return formatted
-
-def create_interface() -> gr.Blocks:
-    """Create the Gradio interface"""
-    agent = TxAgentWrapper()
-    
-    with gr.Blocks(
-        title="TxAgent",
-        css="""
-        .gradio-container {max-width: 900px !important}
-        """
-    ) as demo:
-        gr.Markdown("""
-        # 🧠 TxAgent: Therapeutic Reasoning AI
-        ### (Loading from Hugging Face Hub)
-        """)
+        original_load = ToolRAGModel.load_tool_desc_embedding
         
-        with gr.Row():
-            init_btn = gr.Button("Initialize Model", variant="primary")
-            init_status = gr.Textbox(label="Status", interactive=False)
+        def patched_load(self, tooluniverse):
+            try:
+                if not os.path.exists(CONFIG["embedding_filename"]):
+                    logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
+                    return False
+                
+                # Load embeddings safely
+                self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
+                
+                # Handle tool count mismatch
+                tools = tooluniverse.get_all_tools()
+                current_count = len(tools)
+                embedding_count = len(self.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.tool_desc_embedding = self.tool_desc_embedding[:current_count]
+                        logger.info(f"Truncated embeddings to match {current_count} tools")
+                    else:
+                        last_embedding = self.tool_desc_embedding[-1]
+                        padding = [last_embedding] * (current_count - embedding_count)
+                        self.tool_desc_embedding = torch.cat(
+                            [self.tool_desc_embedding] + padding
+                        )
+                        logger.info(f"Padded embeddings to match {current_count} tools")
+                
+                return True
+                
+            except Exception as e:
+                logger.error(f"Failed to load embeddings: {str(e)}")
+                return False
         
-        chatbot = gr.Chatbot(
-            height=500, 
-            label="Conversation"
-        )
-        msg = gr.Textbox(label="Your clinical question")
-        clear_btn = gr.Button("Clear Chat")
+        # Apply the patch
+        ToolRAGModel.load_tool_desc_embedding = patched_load
+        logger.info("Successfully patched embedding loading")
         
-        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
+    except Exception as e:
+        logger.error(f"Failed to patch embedding loading: {str(e)}")
+        raise
+
+def prepare_tool_files():
+    """Ensure tool files exist and are populated"""
+    os.makedirs(os.path.join(current_dir, '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() if hasattr(tu, 'get_all_tools') else []
+        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']}")
+
+# Apply the embedding patch before creating the agent
+patch_embedding_loading()
+prepare_tool_files()
+
+# Initialize the agent
+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']
+)
+agent.init_model()
+
+def update_model_parameters(enable_finish, enable_rag, enable_summary,
+                            init_rag_num, step_rag_num, skip_last_k,
+                            summary_mode, summary_skip_last_k, summary_context_length, force_finish, seed):
+    # Update model instance parameters dynamically
+    updated_params = agent.update_parameters(
+        enable_finish=enable_finish,
+        enable_rag=enable_rag,
+        enable_summary=enable_summary,
+        init_rag_num=init_rag_num,
+        step_rag_num=step_rag_num,
+        skip_last_k=skip_last_k,
+        summary_mode=summary_mode,
+        summary_skip_last_k=summary_skip_last_k,
+        summary_context_length=summary_context_length,
+        force_finish=force_finish,
+        seed=seed,
+    )
+
+    return updated_params
+
+def update_seed():
+    # Update model instance parameters dynamically
+    seed = random.randint(0, 10000)
+    updated_params = agent.update_parameters(
+        seed=seed,
+    )
+    return updated_params
+
+def handle_retry(history, retry_data: gr.RetryData, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
+    print("Updated seed:", update_seed())
+    new_history = history[:retry_data.index]
+    previous_prompt = history[retry_data.index]['content']
+
+    print("previous_prompt", previous_prompt)
+
+    yield from agent.run_gradio_chat(new_history + [{"role": "user", "content": previous_prompt}], temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round)
+
+PASSWORD = "mypassword"
+
+def check_password(input_password):
+    if input_password == PASSWORD:
+        return gr.update(visible=True), ""
+    else:
+        return gr.update(visible=False), "Incorrect password, try again!"
+
+conversation_state = gr.State([])
+
+# Gradio block
+chatbot = gr.Chatbot(height=800, placeholder=PLACEHOLDER,
+                     label='TxAgent', type="messages",  show_copy_button=True)
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.Markdown(INTRO)
+    default_temperature = 0.3
+    default_max_new_tokens = 1024
+    default_max_tokens = 81920
+    default_max_round = 30
+    temperature_state = gr.State(value=default_temperature)
+    max_new_tokens_state = gr.State(value=default_max_new_tokens)
+    max_tokens_state = gr.State(value=default_max_tokens)
+    max_round_state = gr.State(value=default_max_round)
+    chatbot.retry(handle_retry, chatbot, chatbot, temperature_state, max_new_tokens_state,
+                  max_tokens_state, gr.Checkbox(value=False, render=False), conversation_state, max_round_state)
+
+    gr.ChatInterface(
+        fn=agent.run_gradio_chat,
+        chatbot=chatbot,
+        fill_height=True, fill_width=True, stop_btn=True,
+        additional_inputs_accordion=gr.Accordion(
+            label="⚙️ Inference Parameters", open=False, render=False),
+        additional_inputs=[
+            temperature_state, max_new_tokens_state, max_tokens_state,
+            gr.Checkbox(
+                label="Activate multi-agent reasoning mode (it requires additional time but offers a more comprehensive analysis).", value=False, render=False),
+            conversation_state,
+            max_round_state,
+            gr.Number(label="Seed", value=100, render=False)
+        ],
+        examples=question_examples,
+        cache_examples=False,
+        css=chat_css,
+    )
+
+    with gr.Accordion("Settings", open=False):
+        # Define the sliders
+        temperature_slider = gr.Slider(
+            minimum=0,
+            maximum=1,
+            step=0.1,
+            value=default_temperature,
+            label="Temperature"
         )
-        
-        def wrapper_initialize():
-            status = agent.initialize()
-            return status, gr.update(interactive=False)
-        
-        init_btn.click(
-            fn=wrapper_initialize,
-            outputs=[init_status, init_btn]
+        max_new_tokens_slider = gr.Slider(
+            minimum=128,
+            maximum=4096,
+            step=1,
+            value=default_max_new_tokens,
+            label="Max new tokens"
         )
-        
-        msg.submit(
-            fn=agent.chat,
-            inputs=[msg, chatbot],
-            outputs=chatbot
-        ).then(
-            lambda: "",  # Clear message box
-            outputs=msg
+        max_tokens_slider = gr.Slider(
+            minimum=128,
+            maximum=32000,
+            step=1,
+            value=default_max_tokens,
+            label="Max tokens"
         )
-        
-        clear_btn.click(
-            fn=lambda: ([], ""),
-            outputs=[chatbot, msg]
+        max_round_slider = gr.Slider(
+            minimum=0,
+            maximum=50,
+            step=1,
+            value=default_max_round,
+            label="Max round")
+
+        # Automatically update states when slider values change
+        temperature_slider.change(
+            lambda x: x, inputs=temperature_slider, outputs=temperature_state)
+        max_new_tokens_slider.change(
+            lambda x: x, inputs=max_new_tokens_slider, outputs=max_new_tokens_state)
+        max_tokens_slider.change(
+            lambda x: x, inputs=max_tokens_slider, outputs=max_tokens_state)
+        max_round_slider.change(
+            lambda x: x, inputs=max_round_slider, outputs=max_round_state)
+
+        password_input = gr.Textbox(
+            label="Enter Password for More Settings", type="password")
+        incorrect_message = gr.Textbox(visible=False, interactive=False)
+        with gr.Accordion("⚙️ Settings", open=False, visible=False) as protected_accordion:
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Accordion("⚙️ Model Loading", open=False):
+                        model_name_input = gr.Textbox(
+                            label="Enter model path", value=CONFIG["model_name"])
+                        load_model_btn = gr.Button(value="Load Model")
+                        load_model_btn.click(
+                            agent.load_models, inputs=model_name_input, outputs=gr.Textbox(label="Status"))
+                with gr.Column(scale=1):
+                    with gr.Accordion("⚙️ Functional Parameters", open=False):
+                        # Create Gradio components for parameter inputs
+                        enable_finish = gr.Checkbox(
+                            label="Enable Finish", value=True)
+                        enable_rag = gr.Checkbox(
+                            label="Enable RAG", value=True)
+                        enable_summary = gr.Checkbox(
+                            label="Enable Summary", value=False)
+                        init_rag_num = gr.Number(
+                            label="Initial RAG Num", value=0)
+                        step_rag_num = gr.Number(
+                            label="Step RAG Num", value=10)
+                        skip_last_k = gr.Number(label="Skip Last K", value=0)
+                        summary_mode = gr.Textbox(
+                            label="Summary Mode", value='step')
+                        summary_skip_last_k = gr.Number(
+                            label="Summary Skip Last K", value=0)
+                        summary_context_length = gr.Number(
+                            label="Summary Context Length", value=None)
+                        force_finish = gr.Checkbox(
+                            label="Force FinalAnswer", value=True)
+                        seed = gr.Number(label="Seed", value=100)
+                        # Button to submit and update parameters
+                        submit_btn = gr.Button("Update Parameters")
+
+                        # Display the updated parameters
+                        updated_parameters_output = gr.JSON()
+
+                        # When button is clicked, update parameters
+                        submit_btn.click(fn=update_model_parameters,
+                                         inputs=[enable_finish, enable_rag, enable_summary, init_rag_num, step_rag_num, skip_last_k,
+                                                 summary_mode, summary_skip_last_k, summary_context_length, force_finish, seed],
+                                         outputs=updated_parameters_output)
+        # Button to submit the password
+        submit_button = gr.Button("Submit")
+
+        # When the button is clicked, check if the password is correct
+        submit_button.click(
+            check_password,
+            inputs=password_input,
+            outputs=[protected_accordion, incorrect_message]
         )
-    
-    return demo
+    gr.Markdown(LICENSE)
 
 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
+    demo.launch(share=True)