Update app.py

app.py

@@ -3,36 +3,35 @@ import sys
 import random
 import gradio as gr
 from multiprocessing import freeze_support
+import importlib
+import inspect
 
-# ✅ Fix path first (before importing anything custom)
+# === Fix: Add src directory to import path BEFORE importing anything else
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "src"))
 
-# ✅ Now import the correct module
-import importlib
+# === Reload to avoid stale imports in Hugging Face
 import txagent.txagent
 importlib.reload(txagent.txagent)
 from txagent.txagent import TxAgent
 
-# ✅ Confirm
-import inspect
+# === Debug: Confirm we're loading the right class
 print(">>> TxAgent loaded from:", inspect.getfile(TxAgent))
 print(">>> TxAgent has run_gradio_chat:", hasattr(TxAgent, "run_gradio_chat"))
+print(">>> TxAgent methods:", [m for m in dir(TxAgent) if not m.startswith("_")])
 
-# === Environment setup ===
+# === Environment setup
 current_dir = os.path.dirname(os.path.abspath(__file__))
 os.environ["MKL_THREADING_LAYER"] = "GNU"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# === UI constants ===
+# === UI constants
 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..."
 LICENSE = "DISCLAIMER: THIS WEBSITE DOES NOT PROVIDE MEDICAL ADVICE..."
-
 PLACEHOLDER = '''
 <div style="padding: 30px; text-align: center;">
    <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">TxAgent</h1>
@@ -40,7 +39,6 @@ PLACEHOLDER = '''
    <p style="font-size: 18px;">Click retry 🔄 to see another answer.</p>
 </div>
 '''
-
 css = """
 h1 { text-align: center; }
 #duplicate-button {
@@ -54,13 +52,12 @@ h1 { text-align: center; }
     margin-bottom: 0px !important;
 }
 """
-
 chat_css = """
 .gr-button { font-size: 20px !important; }
 .gr-button svg { width: 32px !important; height: 32px !important; }
 """
 
-# === Model and tool config ===
+# === Model and tool config
 model_name = "mims-harvard/TxAgent-T1-Llama-3.1-8B"
 rag_model_name = "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B"
 new_tool_files = {
@@ -72,6 +69,7 @@ question_examples = [
     ["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
 ]
 
+# === UI constructor
 def create_ui(agent):
     with gr.Blocks(css=css) as demo:
         gr.Markdown(DESCRIPTION)
@@ -92,9 +90,9 @@ def create_ui(agent):
             show_copy_button=True
         )
 
-        # === Chat handler (streaming) ===
+        # === Chat handler
         async def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-            response = await agent.run_gradio_chat(
+            return agent.run_gradio_chat(
                 message=message,
                 history=history,
                 temperature=temperature,
@@ -104,14 +102,13 @@ def create_ui(agent):
                 conversation=conversation,
                 max_round=max_round
             )
-            return response
 
-        # === Retry handler ===
+        # === Retry handler
         async def handle_retry(history, retry_data, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
             agent.update_parameters(seed=random.randint(0, 10000))
             new_history = history[:retry_data.index]
             prompt = history[retry_data.index]["content"]
-            return await agent.run_gradio_chat(
+            return agent.run_gradio_chat(
                 message=prompt,
                 history=new_history,
                 temperature=temperature,
@@ -122,14 +119,14 @@ def create_ui(agent):
                 max_round=max_round
             )
 
-        # Configure retry button
+        # === Retry button config
         chatbot.retry(
             handle_retry,
             inputs=[chatbot, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round],
             outputs=chatbot
         )
 
-        # === Chat Interface setup ===
+        # === Main interface
         gr.ChatInterface(
             fn=handle_chat,
             chatbot=chatbot,
@@ -148,14 +145,14 @@ def create_ui(agent):
         gr.Markdown(LICENSE)
     return demo
 
+# === App start
 if __name__ == "__main__":
     freeze_support()
-    
+
     try:
-        # Initialize agent
         agent = TxAgent(
-            model_name,
-            rag_model_name,
+            model_name=model_name,
+            rag_model_name=rag_model_name,
             tool_files_dict=new_tool_files,
             force_finish=True,
             enable_checker=True,
@@ -164,15 +161,13 @@ if __name__ == "__main__":
             additional_default_tools=["DirectResponse", "RequireClarification"]
         )
         agent.init_model()
-        
-        # Verify the agent has the required method
+
         if not hasattr(agent, 'run_gradio_chat'):
-            raise AttributeError("The TxAgent instance is missing the run_gradio_chat method!")
-        
-        # Create and launch UI
+            raise AttributeError("TxAgent is missing `run_gradio_chat`. Make sure txagent.py is up to date.")
+
         demo = create_ui(agent)
         demo.launch(show_error=True)
-        
+
     except Exception as e:
-        print(f"Application failed to start: {str(e)}")
-        raise
+        print(f"🚨 Application failed to start: {e}")
+        raise