Update app.py

app.py

@@ -2,10 +2,10 @@ import os
 import sys
 import random
 import gradio as gr
-from multiprocessing import freeze_support  # Added for multiprocessing support
+from multiprocessing import freeze_support
 
-# Add src to path
-sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
+# Fix path to include src directory
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "src"))
 
 from txagent import TxAgent
 
@@ -63,8 +63,7 @@ question_examples = [
     ["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
 ]
 
-# === Create Gradio UI ===
-def create_ui():
+def create_ui(agent):
     with gr.Blocks(css=css) as demo:
         gr.Markdown(DESCRIPTION)
         gr.Markdown(INTRO)
@@ -84,9 +83,9 @@ def create_ui():
             show_copy_button=True
         )
 
-        # === Chat handler (streaming)
+        # === Chat handler (streaming) ===
         async def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-            return await agent.run_gradio_chat(
+            response = await agent.run_gradio_chat(
                 message=message,
                 history=history,
                 temperature=temperature,
@@ -96,8 +95,9 @@ def create_ui():
                 conversation=conversation,
                 max_round=max_round
             )
+            return response
 
-        # === Retry handler (also streaming-compatible)
+        # === 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]
@@ -113,13 +113,14 @@ def create_ui():
                 max_round=max_round
             )
 
+        # Configure retry button
         chatbot.retry(
             handle_retry,
             inputs=[chatbot, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round],
             outputs=chatbot
         )
 
-        # === Chat Interface setup
+        # === Chat Interface setup ===
         gr.ChatInterface(
             fn=handle_chat,
             chatbot=chatbot,
@@ -139,21 +140,30 @@ def create_ui():
     return demo
 
 if __name__ == "__main__":
-    freeze_support()  # Initialize multiprocessing support
-    
-    # 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()
+    freeze_support()
     
-    # Create and launch UI
-    demo = create_ui()
-    demo.launch(show_error=True)
+    try:
+        # 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()
+        
+        # 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
+        demo = create_ui(agent)
+        demo.launch(show_error=True)
+        
+    except Exception as e:
+        print(f"Application failed to start: {str(e)}")
+        raise