Update app.py

app.py

@@ -109,6 +109,40 @@ def save_chat_context(new_context, user_chats, current_chat_title, logged_in_use
     save_chats(logged_in_username, user_chats, auth_globals)
     return user_chats, "✅ Context saved!"
 
+def enforce_cooldown(last_time):
+    now = time.time()
+    if now - last_time < 3:
+        gr.Warning("⏳ Please wait 3 seconds before clicking again.")
+        return False, last_time
+    return True, now
+
+def delete_last_message(user_chats, current_title, logged_in_username, auth_globals):
+    if not current_title or not user_chats:
+        return user_chats, "Nothing to delete"
+    for chat in user_chats:
+        if chat["title"] == current_title and len(chat["history"]) >= 2:
+            chat["history"] = chat["history"][:-2]  # remove last user + assistant
+            break
+    save_chats(logged_in_username, user_chats, auth_globals)
+    return user_chats, "🗑️ Last message deleted"
+
+def edit_last_prompt(user_chats, current_title):
+    if not current_title or not user_chats:
+        return "", "No previous prompt"
+    for chat in user_chats:
+        if chat["title"] == current_title and len(chat["history"]) >= 2:
+            last_user_msg = chat["history"][-2]["content"]
+            return last_user_msg, "Ready to edit & re-run"
+    return "", "No previous prompt"
+
+def update_token_display(user_chats, current_title):
+    total = 0
+    for chat in user_chats or []:
+        if chat["title"] == current_title:
+            for msg in chat.get("history", []):
+                total += msg.get("tokens", 0)
+    return f"**Tokens used this chat:** {total:,}"
+
 # ================== HELPER: No conversion needed (Gradio expects messages format) ==================
 def format_history_for_chatbot(history):
     """Return history as-is (already dict format)"""
@@ -177,27 +211,37 @@ def handle_chat(msg, chatbot, current_model, logged_in_username, user_chats_stat
     if not current_model:
         gr.Warning("⚠️ Load a model first!")
         return chatbot
-    # Ensure the message is being passed as expected
-    if not msg or not isinstance(msg, str):
-        gr.Warning("⚠️ Invalid input message format!")
+
+    # Cooldown check
+    ok, new_time = enforce_cooldown(last_interaction_time.value)
+    if not ok:
         return chatbot
-        
-    if file_upload and len(file_upload) > 5:
-        gr.Warning("⚠️ Max 5 files allowed! Ignoring extras.")
-        file_upload = file_upload[:5]
+    last_interaction_time.value = new_time
+
+    is_generating.value = True
+    stop_btn.visible = True
 
-    # Chatbot is already in correct dict format (no conversion needed)
     internal_history = chatbot[:] if chatbot else []
 
-    # Call chat_stream with internal dict history
     try:
         stream_gen = chat_stream(msg, internal_history, current_model, logged_in_username, user_chats_state, current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context)
 
         for updated_history in stream_gen:
-            yield updated_history  # ← direct yield, no formatter
+            # Add approximate token count to last assistant message
+            if updated_history and updated_history[-1]["role"] == "assistant":
+                content = updated_history[-1]["content"]
+                tokens = len(content) // 4 + 50  # simple but good enough approximation
+                updated_history[-1]["tokens"] = tokens
+            yield updated_history
+
     except Exception as e:
-        gr.Warning(f"⚠️ An error occurred in the chat processing: {str(e)}")
-        return chatbot  # Gracefully return the existing chatbot state
+        gr.Warning(f"⚠️ Error: {str(e)[:200]}")
+    finally:
+        is_generating.value = False
+        stop_btn.visible = False
+
+    # Update token display
+    return token_display.update(value=update_token_display(user_chats_state, current_chat_title))
 
 
 # ================== GUI ==================
@@ -216,6 +260,8 @@ with gr.Blocks(title="NVIDIA AI Chat") as demo:
     auth_globals_state = gr.State(auth_globals)
     coding_files_state = gr.State({})      # { "filename.py": {"content": "...", "language": "python"} }
     current_file_name = gr.State(None)
+    last_interaction_time = gr.State(0.0)   # for 3-second cooldown
+    is_generating = gr.State(False)         # for Stop button
     status_label = gr.Markdown("⚠️ No model loaded. Login + create/select a chat to begin.", elem_id="status-bold")  # Made more prominent
 
     
@@ -289,6 +335,14 @@ with gr.Blocks(title="NVIDIA AI Chat") as demo:
         )
         send_btn = gr.Button("➤", variant="primary", scale=1)
 
+    with gr.Row():
+            delete_last_btn = gr.Button("🗑️ Delete Last Message", variant="stop", size="sm")
+            edit_last_btn = gr.Button("✏️ Edit & Re-run Last Prompt", variant="secondary", size="sm")
+            stop_btn = gr.Button("⏹️ Stop Generation", variant="stop", visible=False)
+    
+        # Token display (optional but useful)
+        token_display = gr.Markdown("**Tokens used this chat:** 0", visible=False)
+
     # Authentication (clean version - no phone/OTP)
     with gr.Row():
         username_input = gr.Textbox(label="Username (will be forced lowercase)", scale=2, placeholder="admin")
@@ -333,6 +387,34 @@ with gr.Blocks(title="NVIDIA AI Chat") as demo:
     confirm_delete_btn.click(delete_account, [confirm_username_input, logged_in_username, auth_globals_state], [status_label, logged_in_username, logout_btn, delete_acc_btn, confirm_row, user_chats_state, current_chat_title, chatbot, chat_dropdown, chat_context_tb, save_context_btn])
     
     save_context_btn.click(save_chat_context, [chat_context_tb, user_chats_state, current_chat_title, logged_in_username], [user_chats_state, status_label])
+
+# Cooldown + action buttons
+    delete_last_btn.click(
+        delete_last_message,
+        [user_chats_state, current_chat_title, logged_in_username, auth_globals_state],
+        [user_chats_state, status_label]
+    ).then(
+        lambda: gr.update(visible=True), None, token_display
+    ).then(update_token_display, [user_chats_state, current_chat_title], token_display)
+
+    edit_last_btn.click(
+        edit_last_prompt,
+        [user_chats_state, current_chat_title],
+        [msg, status_label]
+    )
+
+    stop_btn.click(
+        lambda: (False, gr.update(visible=False)),
+        None,
+        [is_generating, stop_btn]
+    )
+
+    # Show token display in Coding Mode and after first message
+    mode_radio.change(
+        lambda m: gr.update(visible=(m == "💻 Coding Mode")),
+        mode_radio,
+        token_display
+    )
     
     msg.submit(
         handle_chat,