Update app.py

app.py

@@ -4,9 +4,11 @@ import random
 import gradio as gr
 from datetime import datetime
 
+# Add src to Python path
 sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
-from txagent.txagent import TxAgent
 
+# Adjust to match your file structure
+from txagent.txagent import TxAgent  # e.g., src/txagent/txagent.py
 
 # ==== Environment Setup ====
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -84,11 +86,12 @@ if __name__ == "__main__":
         seed = random.randint(0, 10000)
         return agent.update_parameters(seed=seed)
 
-    def handle_retry(history, retry_data: gr.RetryData, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
+    # ✅ FIXED: retry must return, not yield
+    def handle_retry(history, retry_data, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
         update_seed()
         new_history = history[:retry_data.index]
         previous_prompt = history[retry_data.index]["content"]
-        yield from agent.run_gradio_chat(
+        return agent.run_gradio_chat(
             new_history + [{"role": "user", "content": previous_prompt}],
             temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round
         )
@@ -113,6 +116,7 @@ if __name__ == "__main__":
             show_copy_button=True
         )
 
+        # ✅ Retry now fixed
         chatbot.retry(
             handle_retry,
             chatbot, chatbot,