Update app.py

app.py

@@ -62,7 +62,7 @@ question_examples = [
     ["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
 ]
 
-agent = None  # global agent initialized later
+agent = None  # will be initialized in main
 
 
 # === Create Gradio UI ===
@@ -86,19 +86,34 @@ def create_ui():
             show_copy_button=True
         )
 
-        # === Retry handler (fixed)
-        def handle_retry(history, retry_data, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
+        # === 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(
+                message=message,
+                history=history,
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                max_token=max_tokens,
+                call_agent=multi_agent,
+                conversation=conversation,
+                max_round=max_round
+            )
+
+        # === Retry handler (also streaming-compatible)
+        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"]
-            result = agent.run_gradio_chat(
-                new_history + [{"role": "user", "content": prompt}],
-                temperature, max_new_tokens, max_tokens,
-                multi_agent, conversation, max_round
+            return agent.run_gradio_chat(
+                message=prompt,
+                history=new_history,
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                max_token=max_tokens,
+                call_agent=multi_agent,
+                conversation=conversation,
+                max_round=max_round
             )
-            if hasattr(result, "__iter__") and not isinstance(result, (str, list, dict)):
-                result = list(result)
-            return result
 
         chatbot.retry(
             handle_retry,
@@ -107,13 +122,6 @@ def create_ui():
             multi_agent, conversation_state, max_round
         )
 
-        # === Chat handler
-        def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-            result = agent.run_gradio_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round)
-            if hasattr(result, "__iter__") and not isinstance(result, (str, list, dict)):
-                result = list(result)
-            return result
-
         # === Chat Interface setup
         gr.ChatInterface(
             fn=handle_chat,
@@ -134,7 +142,7 @@ def create_ui():
     return demo
 
 
-# === HF Spaces + vLLM-safe entrypoint ===
+# === HF Spaces + local launch entrypoint ===
 if __name__ == "__main__":
     agent = TxAgent(
         model_name,
@@ -149,4 +157,4 @@ if __name__ == "__main__":
     agent.init_model()
 
     demo = create_ui()
-    demo.launch()
+    demo.launch(show_error=True)