Update app.py

app.py

@@ -3,10 +3,10 @@ import sys
 import random
 import gradio as gr
 
-# Add `src` to import path
+# Add src to path
 sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
 
-from txagent.txagent import TxAgent
+from txagent import TxAgent
 
 # === Environment setup ===
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -62,8 +62,18 @@ question_examples = [
     ["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
 ]
 
-agent = None  # will be initialized in main
-
+# Initialize agent
+agent = TxAgent(
+    model_name,
+    rag_model_name,
+    tool_files_dict=new_tool_files,
+    force_finish=True,
+    enable_checker=True,
+    step_rag_num=10,
+    seed=100,
+    additional_default_tools=["DirectResponse", "RequireClarification"]
+)
+agent.init_model()
 
 # === Create Gradio UI ===
 def create_ui():
@@ -88,7 +98,7 @@ def create_ui():
 
         # === Chat handler (streaming)
         async def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-            return agent.run_gradio_chat(
+            return await agent.run_gradio_chat(
                 message=message,
                 history=history,
                 temperature=temperature,
@@ -104,7 +114,7 @@ def create_ui():
             agent.update_parameters(seed=random.randint(0, 10000))
             new_history = history[:retry_data.index]
             prompt = history[retry_data.index]["content"]
-            return agent.run_gradio_chat(
+            return await agent.run_gradio_chat(
                 message=prompt,
                 history=new_history,
                 temperature=temperature,
@@ -141,20 +151,6 @@ def create_ui():
         gr.Markdown(LICENSE)
     return demo
 
-
-# === HF Spaces + local launch entrypoint ===
 if __name__ == "__main__":
-    agent = TxAgent(
-        model_name,
-        rag_model_name,
-        tool_files_dict=new_tool_files,
-        force_finish=True,
-        enable_checker=True,
-        step_rag_num=10,
-        seed=100,
-        additional_default_tools=["DirectResponse", "RequireClarification"]
-    )
-    agent.init_model()
-
     demo = create_ui()
-    demo.launch(show_error=True)
+    demo.launch(show_error=True)