Update app.py

app.py

@@ -6,63 +6,57 @@ from multiprocessing import freeze_support
 import importlib
 import inspect
 
-# === Path fix
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "src"))
+# === Fix path to include src/txagent
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
 
-# === Reload to avoid stale module
+# === Import and reload to ensure correct file
 import txagent.txagent
 importlib.reload(txagent.txagent)
 from txagent.txagent import TxAgent
-from gradio import ChatMessage
 
-# === Debug
+# === Debug print
 print(">>> TxAgent loaded from:", inspect.getfile(TxAgent))
 print(">>> TxAgent has run_gradio_chat:", hasattr(TxAgent, "run_gradio_chat"))
 
-# === Env vars
-current_dir = os.path.dirname(os.path.abspath(__file__))
+# === Environment
+current_dir = os.path.abspath(os.path.dirname(__file__))
 os.environ["MKL_THREADING_LAYER"] = "GNU"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# === UI text
-DESCRIPTION = '''
-<h1 style="text-align: center;">TxAgent: AI for Therapeutic Reasoning</h1>
-'''
-INTRO = "Ask biomedical or therapeutic questions. Results are powered by tools and reasoning."
-LICENSE = "DISCLAIMER: THIS WEBSITE DOES NOT PROVIDE MEDICAL ADVICE."
-
-# === Model & tool config
+# === Model 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 = {
     "new_tool": os.path.join(current_dir, "data", "new_tool.json")
 }
 
+# === Example prompts
 question_examples = [
     ["Given a patient with WHIM syndrome on prophylactic antibiotics, is it advisable to co-administer Xolremdi with fluconazole?"],
     ["What treatment options exist for HER2+ breast cancer resistant to trastuzumab?"]
 ]
 
-# === Gradio UI
+# === UI creation
 def create_ui(agent):
     with gr.Blocks() as demo:
-        gr.Markdown(DESCRIPTION)
-        gr.Markdown(INTRO)
+        gr.Markdown("<h1 style='text-align: center;'>TxAgent: Therapeutic Reasoning</h1>")
+        gr.Markdown("Ask biomedical or therapeutic questions. Powered by step-by-step reasoning and tools.")
 
-        temperature = gr.Slider(0, 1, step=0.1, value=0.3, label="Temperature")
-        max_new_tokens = gr.Slider(128, 4096, step=1, value=1024, label="Max New Tokens")
-        max_tokens = gr.Slider(128, 32000, step=1, value=8192, label="Max Total Tokens")
-        max_round = gr.Slider(1, 50, step=1, value=30, label="Max Rounds")
+        temperature = 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, 32000, value=8192, label="Max Total Tokens")
+        max_round = gr.Slider(1, 50, value=30, label="Max Rounds")
         multi_agent = gr.Checkbox(label="Enable Multi-agent Reasoning", value=False)
         conversation_state = gr.State([])
 
         chatbot = gr.Chatbot(label="TxAgent", height=600, type="messages")
-        message_input = gr.Textbox(placeholder="Ask a biomedical question...", show_label=False)
-        send_btn = gr.Button("Send", variant="primary")
+        message_input = gr.Textbox(placeholder="Ask your biomedical question...", show_label=False)
+        send_button = gr.Button("Send", variant="primary")
 
+        # === Core handler (streaming generator)
         def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-            # Ensure response is a generator that yields list of {role, content} dictionaries
-            return agent.run_gradio_chat(
+            # Must yield a list of {"role": ..., "content": ...} dicts
+            generator = agent.run_gradio_chat(
                 message=message,
                 history=history,
                 temperature=temperature,
@@ -72,8 +66,18 @@ def create_ui(agent):
                 conversation=conversation,
                 max_round=max_round
             )
-
-        send_btn.click(
+            for update in generator:
+                # Convert to list of dicts if not already
+                formatted = [
+                    {"role": m["role"], "content": m["content"]}
+                    if isinstance(m, dict)
+                    else {"role": m.role, "content": m.content}
+                    for m in update
+                ]
+                yield formatted
+
+        # === Trigger handlers
+        send_button.click(
             fn=handle_chat,
             inputs=[message_input, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round],
             outputs=chatbot
@@ -85,18 +89,16 @@ def create_ui(agent):
             outputs=chatbot
         )
 
-        gr.Examples(
-            examples=question_examples,
-            inputs=message_input
-        )
+        gr.Examples(examples=question_examples, inputs=message_input)
 
-        gr.Markdown(LICENSE)
+        gr.Markdown("**DISCLAIMER**: This demo is for research purposes only and does not provide medical advice.")
 
     return demo
 
-# === App startup
+# === Startup
 if __name__ == "__main__":
     freeze_support()
+
     try:
         agent = TxAgent(
             model_name=model_name,
@@ -117,5 +119,5 @@ if __name__ == "__main__":
         demo.launch(show_error=True)
 
     except Exception as e:
-        print(f"🚨 Startup error: {e}")
+        print(f"❌ Application failed to start: {e}")
         raise