Update app.py

app.py

@@ -17,7 +17,7 @@ class TxAgentApp:
         """Initialize the TxAgent with proper tool file paths"""
         try:
             logger.info("Initializing TxAgent...")
-            
+
             # Get absolute paths to tool files from package installation
             tool_files = {
                 "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
@@ -25,9 +25,9 @@ class TxAgentApp:
                 "special_tools": str(files('tooluniverse.data').joinpath('special_tools.json')),
                 "monarch": str(files('tooluniverse.data').joinpath('monarch_tools.json'))
             }
-            
+
             logger.info(f"Using tool files at: {tool_files}")
-            
+
             agent = TxAgent(
                 model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
                 rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -47,12 +47,11 @@ 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:
             logger.error(f"Initialization failed: {str(e)}")
             raise
@@ -63,7 +62,6 @@ class TxAgentApp:
             if not isinstance(message, str) or len(message.strip()) <= 10:
                 return chat_history + [("", "Please provide a valid message longer than 10 characters.")]
 
-            # 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]
 
@@ -85,74 +83,55 @@ class TxAgentApp:
                     response += chunk
                 else:
                     response += str(chunk)
-                
+
                 yield chat_history + [("user", message), ("assistant", response)]
 
         except Exception as e:
             logger.error(f"Error in respond function: {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") as demo:
-        gr.Markdown("# TxAgent Biomedical Assistant")
-        
-        chatbot = gr.Chatbot(
-            label="Conversation",
-            height=600
-        )
-        
-        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([])
-        
-        submit.click(
-            app.respond,
-            [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
-            chatbot
-        )
-        
-        clear.click(lambda: [], None, chatbot)
-        
-        msg.submit(
-            app.respond,
-            [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
-            chatbot
-        )
-    
-    return demo
-
-def main():
-    """Main entry point for the application"""
-    try:
-        logger.info("Starting TxAgent application...")
-        demo = create_demo()
-        
-        logger.info("Launching Gradio interface...")
-        demo.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
-
-if __name__ == "__main__":
-    main()
+# Initialize app instance
+tx_app = TxAgentApp()
+
+# Define Gradio UI blocks
+with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
+    gr.Markdown("# 🧠 TxAgent Biomedical Assistant")
+
+    chatbot = gr.Chatbot(
+        label="Conversation",
+        height=600
+    )
+
+    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([])
+
+    submit.click(
+        tx_app.respond,
+        [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
+        chatbot
+    )
+
+    clear.click(lambda: [], None, chatbot)
+
+    msg.submit(
+        tx_app.respond,
+        [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
+        chatbot
+    )
+
+# `app` will be auto-discovered and served by Hugging Face Spaces