Update app.py

app.py

@@ -431,17 +431,32 @@ with gr.Blocks(title="NVIDIA AI Chat") as demo:
         token_display
     )
     
+# Send button + Enter
     msg.submit(
         handle_chat, 
-        [msg, chatbot, current_model, logged_in_username, user_chats_state, current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context], 
+        [msg, chatbot, current_model, logged_in_username, user_chats_state, 
+         current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context], 
         chatbot
-    ).then(lambda: "", None, msg)
-    
+    ).then(
+        lambda: "", None, msg   # clear input box
+    ).then(
+        update_token_display, 
+        [user_chats_state, current_chat_title], 
+        token_display
+    )
+
     send_btn.click(
         handle_chat, 
-        [msg, chatbot, current_model, logged_in_username, user_chats_state, current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context], 
+        [msg, chatbot, current_model, logged_in_username, user_chats_state, 
+         current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context], 
         chatbot
-    ).then(lambda: "", None, msg)
+    ).then(
+        lambda: "", None, msg
+    ).then(
+        update_token_display, 
+        [user_chats_state, current_chat_title], 
+        token_display
+    )
 
 # File manager events
     add_files_btn.click(