Update app.py

app.py

@@ -1,14 +1,14 @@
 import random
 import datetime
 import sys
-from txagent import TxAgent
-import spaces
-import gradio as gr
 import os
 import torch
 import logging
+import json
 from importlib.resources import files
-import traceback
+from txagent import TxAgent
+from tooluniverse import ToolUniverse
+import gradio as gr
 
 # Set up logging
 logging.basicConfig(
@@ -20,6 +20,7 @@ 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"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # Configuration
 CONFIG = {
@@ -35,33 +36,36 @@ CONFIG = {
     }
 }
 
-# Set an environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
 DESCRIPTION = '''
 <div>
-<h1 style="text-align: center;">TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools </h1>
+<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.
+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.
 """
 
 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.
+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.
+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.
 """
 
 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>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Please click clear🗑️ (top-right) to remove previous context before submitting 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>
 """
 
@@ -90,27 +94,17 @@ h1 {
 """
 
 chat_css = """
-.gr-button { font-size: 20px !important; }  /* Enlarges button icons */
-.gr-button svg { width: 32px !important; height: 32px !important; } /* Enlarges SVG icons */
+.gr-button { font-size: 20px !important; }
+.gr-button svg { width: 32px !important; height: 32px !important; }
 """
 
-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"""
     try:
-        # First try with weights_only=True (secure mode)
         return torch.load(filepath, weights_only=True)
     except Exception as e:
         logger.warning(f"Secure load failed, trying with weights_only=False: {str(e)}")
         try:
-            # Try with the safe_globals context manager
             with torch.serialization.safe_globals([torch.serialization._reconstruct]):
                 return torch.load(filepath, weights_only=False)
         except Exception as e:
@@ -130,10 +124,8 @@ def patch_embedding_loading():
                     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)
@@ -147,9 +139,7 @@ def patch_embedding_loading():
                     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
-                        )
+                        self.tool_desc_embedding = torch.cat([self.tool_desc_embedding] + padding)
                         logger.info(f"Padded embeddings to match {current_count} tools")
                 
                 return True
@@ -158,7 +148,6 @@ def patch_embedding_loading():
                 logger.error(f"Failed to load embeddings: {str(e)}")
                 return False
         
-        # Apply the patch
         ToolRAGModel.load_tool_desc_embedding = patched_load
         logger.info("Successfully patched embedding loading")
         
@@ -177,27 +166,31 @@ def prepare_tool_files():
             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
+def create_agent():
+    """Create and initialize the TxAgent"""
+    # 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()
+    return agent
+
+def update_model_parameters(agent, 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 parameters"""
     updated_params = agent.update_parameters(
         enable_finish=enable_finish,
         enable_rag=enable_rag,
@@ -211,169 +204,180 @@ def update_model_parameters(enable_finish, enable_rag, enable_summary,
         force_finish=force_finish,
         seed=seed,
     )
-
     return updated_params
 
-def update_seed():
-    # Update model instance parameters dynamically
+def update_seed(agent):
+    """Update random seed"""
     seed = random.randint(0, 10000)
-    updated_params = agent.update_parameters(
-        seed=seed,
-    )
+    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())
+def handle_retry(agent, history, retry_data: gr.RetryData, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
+    """Handle retry functionality"""
+    print("Updated seed:", update_seed(agent))
     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)
+    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):
+    """Check password for protected settings"""
     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)
+def create_demo(agent):
+    """Create the Gradio interface"""
     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"
-        )
-        max_new_tokens_slider = gr.Slider(
-            minimum=128,
-            maximum=4096,
-            step=1,
-            value=default_max_new_tokens,
-            label="Max new tokens"
-        )
-        max_tokens_slider = gr.Slider(
-            minimum=128,
-            maximum=32000,
-            step=1,
-            value=default_max_tokens,
-            label="Max tokens"
+    
+    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?'],
+    ]
+
+    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)
+        
+        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(
+            lambda *args: handle_retry(agent, *args), 
+            inputs=[chatbot, chatbot, temperature_state, max_new_tokens_state,
+                   max_tokens_state, gr.Checkbox(value=False, render=False), 
+                   gr.State([]), max_round_state]
         )
-        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()
+        gr.ChatInterface(
+            fn=lambda *args: agent.run_gradio_chat(*args),
+            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", 
+                    value=False, 
+                    render=False),
+                gr.State([]),
+                max_round_state,
+                gr.Number(label="Seed", value=100, render=False)
+            ],
+            examples=question_examples,
+            cache_examples=False,
+            css=chat_css,
+        )
 
-                        # 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")
+        with gr.Accordion("Settings", open=False):
+            temperature_slider = gr.Slider(
+                minimum=0,
+                maximum=1,
+                step=0.1,
+                value=default_temperature,
+                label="Temperature"
+            )
+            max_new_tokens_slider = gr.Slider(
+                minimum=128,
+                maximum=4096,
+                step=1,
+                value=default_max_new_tokens,
+                label="Max new tokens"
+            )
+            max_tokens_slider = gr.Slider(
+                minimum=128,
+                maximum=32000,
+                step=1,
+                value=default_max_tokens,
+                label="Max tokens"
+            )
+            max_round_slider = gr.Slider(
+                minimum=0,
+                maximum=50,
+                step=1,
+                value=default_max_round,
+                label="Max round")
+
+            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):
+                            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)
+                            submit_btn = gr.Button("Update Parameters")
+                            updated_parameters_output = gr.JSON()
+                            submit_btn.click(
+                                lambda *args: update_model_parameters(agent, *args),
+                                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
+                            )
+            
+            submit_button = gr.Button("Submit")
+            submit_button.click(
+                check_password,
+                inputs=password_input,
+                outputs=[protected_accordion, incorrect_message]
+            )
+        
+        gr.Markdown(LICENSE)
+    
+    return demo
 
-        # When the button is clicked, check if the password is correct
-        submit_button.click(
-            check_password,
-            inputs=password_input,
-            outputs=[protected_accordion, incorrect_message]
-        )
-    gr.Markdown(LICENSE)
+def main():
+    """Main function to run the application"""
+    agent = create_agent()
+    demo = create_demo(agent)
+    demo.launch(share=True)
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    main()