Update app.py

app.py

@@ -460,27 +460,42 @@ def create_chatbot_demo():
 
         # --- Connect UI elements ---
 
-        # User Input Submission (Textbox Enter or Send Button Click)
-        submit_triggers = [user_input.submit, send_btn.click]
-
-        # 1. Add user message to UI immediately
-        for trigger in submit_triggers:
-            trigger.then(
-                add_user_message_to_history,
-                inputs=[user_input, chat_history],
-                outputs=[chat_history, chatbot_ui, user_input, output_vis] # Update chat, clear input, clear vis
-            ).then( # 2. Trigger bot response generation (as a generator)
-                generate_dream_response,
-                inputs=[
-                    chat_history, gen_length, steps, constraints_input,
-                    temperature, top_p, top_k, remasking_strategy, alg_temp,
-                    visualization_delay
-                ],
-                outputs=[chatbot_ui, output_vis] # Stream updates to chatbot and visualization
-                # Note: The final text response is implicitly handled by updating chatbot_ui
-            )
-
-        # Clear Button Action
+        # Define the inputs for the generation function once
+        generation_inputs = [
+            chat_history, gen_length, steps, constraints_input,
+            temperature, top_p, top_k, remasking_strategy, alg_temp,
+            visualization_delay
+        ]
+        # Define the outputs for the generation function
+        generation_outputs = [chatbot_ui, output_vis]
+
+        # Handle Textbox Submission (Enter key)
+        submit_listener = user_input.submit(
+            fn=add_user_message_to_history,
+            inputs=[user_input, chat_history],
+            outputs=[chat_history, chatbot_ui, user_input, output_vis] # Step 1: Add user msg
+        )
+        # Chain the bot response generation after the user message is added
+        submit_listener.then(
+            fn=generate_dream_response,
+            inputs=generation_inputs,
+            outputs=generation_outputs # Step 2: Generate response and stream vis
+        )
+
+        # Handle Send Button Click
+        click_listener = send_btn.click(
+            fn=add_user_message_to_history,
+            inputs=[user_input, chat_history],
+            outputs=[chat_history, chatbot_ui, user_input, output_vis] # Step 1: Add user msg
+        )
+        # Chain the bot response generation after the user message is added
+        click_listener.then(
+            fn=generate_dream_response,
+            inputs=generation_inputs,
+            outputs=generation_outputs # Step 2: Generate response and stream vis
+        )
+
+        # Clear Button Action remains the same
         clear_btn.click(
             clear_conversation,
             inputs=[],