Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import os
 import logging
 from txagent import TxAgent
+from tooluniverse import ToolUniverse
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -15,9 +16,19 @@ class TxAgentApp:
         """Initialize the TxAgent with proper parameters"""
         try:
             logger.info("Initializing TxAgent...")
+            
+            # Initialize default tool files
+            tool_files = {
+                "opentarget": "opentarget_tools.json",
+                "fda_drug_label": "fda_drug_labeling_tools.json",
+                "special_tools": "special_tools.json",
+                "monarch": "monarch_tools.json"
+            }
+            
             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=tool_files,  # This is critical!
                 enable_finish=True,
                 enable_rag=True,
                 enable_summary=False,
@@ -33,6 +44,10 @@ class TxAgentApp:
                 enable_chat=False,
                 additional_default_tools=["DirectResponse", "RequireClarification"]
             )
+            
+            # Explicitly initialize the model
+            agent.init_model()
+            
             logger.info("Model loading complete")
             return agent
         except Exception as e:
@@ -43,10 +58,14 @@ class TxAgentApp:
         """Handle streaming responses with Gradio"""
         try:
             if not isinstance(message, str) or len(message.strip()) <= 10:
-                yield "Please provide a valid message longer than 10 characters."
-                return
+                return chat_history + [("", "Please provide a valid message longer than 10 characters.")]
 
-            response_generator = self.agent.run_gradio_chat(
+            # Convert chat history to list of tuples if needed
+            if chat_history and isinstance(chat_history[0], dict):
+                chat_history = [(h["role"], h["content"]) for h in chat_history if "role" in h and "content" in h]
+
+            response = ""
+            for chunk in self.agent.run_gradio_chat(
                 message=message.strip(),
                 history=chat_history,
                 temperature=temperature,
@@ -56,78 +75,63 @@ class TxAgentApp:
                 conversation=conversation_state,
                 max_round=max_round,
                 seed=42
-            )
-
-            full_response = ""
-            for chunk in response_generator:
-                if isinstance(chunk, dict) and "content" in chunk:
-                    content = chunk["content"]
+            ):
+                if isinstance(chunk, dict):
+                    response += chunk.get("content", "")
                 elif isinstance(chunk, str):
-                    content = chunk
+                    response += chunk
                 else:
-                    content = str(chunk)
+                    response += str(chunk)
                 
-                full_response += content
-                yield full_response
+                yield chat_history + [("user", message), ("assistant", response)]
 
         except Exception as e:
             logger.error(f"Error in respond function: {str(e)}")
-            yield f"⚠️ Error: {str(e)}"
+            yield chat_history + [("", f"⚠️ Error: {str(e)}")]
 
 def create_demo():
     """Create and configure the Gradio interface"""
     app = TxAgentApp()
     
-    with gr.Blocks(
-        title="TxAgent Medical AI",
-        theme=gr.themes.Soft(spacing_size="sm", radius_size="none")
-    ) as demo:
-        gr.Markdown("""<h1 style='text-align: center'>TxAgent Biomedical Assistant</h1>""")
+    with gr.Blocks(title="TxAgent Medical AI") as demo:
+        gr.Markdown("# TxAgent Biomedical Assistant")
         
-        with gr.Row(equal_height=False):
-            with gr.Column(scale=2):
-                chatbot = gr.Chatbot(
-                    height=650,
-                    bubble_full_width=False,
-                    render_markdown=True
-                )
-                msg = gr.Textbox(
-                    label="Your medical query",
-                    placeholder="Enter your biomedical question...",
-                    lines=5,
-                    max_lines=10
-                )
-                
-            with gr.Column(scale=1):
-                with gr.Accordion("⚙️ Parameters", open=False):
-                    temperature = gr.Slider(0, 1, value=0.3, label="Creativity")
-                    max_new_tokens = gr.Slider(128, 4096, value=1024, step=128, label="Max Response Length")
-                    max_tokens = gr.Slider(128, 81920, value=81920, step=1024, label="Max Total Tokens")
-                    max_rounds = gr.Slider(1, 30, value=10, step=1, label="Max Rounds")
-                    multi_agent = gr.Checkbox(value=False, label="Multi-Agent Mode")
-                
-                submit_btn = gr.Button("Submit", variant="primary")
-                clear_btn = gr.Button("Clear History")
+        chatbot = gr.Chatbot(
+            label="Conversation",
+            height=600,
+            bubble_full_width=False
+        )
+        
+        msg = gr.Textbox(
+            label="Your medical query",
+            placeholder="Enter your biomedical question...",
+            lines=3
+        )
+        
+        with gr.Row():
+            temp = gr.Slider(0, 1, value=0.3, label="Temperature")
+            max_new_tokens = gr.Slider(128, 4096, value=1024, label="Max New Tokens")
+            max_tokens = gr.Slider(128, 81920, value=81920, label="Max Total Tokens")
+            max_rounds = gr.Slider(1, 30, value=10, label="Max Rounds")
+            multi_agent = gr.Checkbox(label="Multi-Agent Mode")
+        
+        submit = gr.Button("Submit")
+        clear = gr.Button("Clear")
         
         conversation_state = gr.State([])
         
-        # Chat interface
-        msg.submit(
+        submit.click(
             app.respond,
-            [msg, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
+            [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
             chatbot
-        ).then(lambda: "", None, msg)
+        )
         
-        submit_btn.click(
+        clear.click(lambda: [], None, chatbot)
+        
+        msg.submit(
             app.respond,
-            [msg, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
+            [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
             chatbot
-        ).then(lambda: "", None, msg)
-        
-        clear_btn.click(
-            lambda: ([], []), 
-            None, 
-            [chatbot, conversation_state]
         )
     
     return demo