UI: Add message version navigation (#6947)

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>

css/main.css

@@ -1456,3 +1456,42 @@ strong {
 .dark .attachment-icon {
     color: #ccc;
 }
+
+/* --- Simple Version Navigation --- */
+.version-navigation {
+    position: absolute;
+    bottom: -23px;
+    right: 0;
+    display: flex;
+    align-items: center;
+    gap: 5px;
+    opacity: 0;
+    transition: opacity 0.2s;
+}
+
+.message:hover .version-navigation,
+.user-message:hover .version-navigation,
+.assistant-message:hover .version-navigation {
+    opacity: 1;
+}
+
+.version-nav-button {
+    padding: 2px 6px;
+    font-size: 12px;
+    min-width: auto;
+}
+
+.version-nav-button[disabled] {
+    opacity: 0.3;
+    cursor: not-allowed;
+}
+
+.version-position {
+    font-size: 11px;
+    color: currentColor;
+    font-family: monospace;
+    min-width: 35px;
+    text-align: center;
+    opacity: 0.8;
+    user-select: none;
+}

js/global_scope_js.js

@@ -49,6 +49,44 @@ function branchHere(element) {
 
 }
 
+function navigateVersion(element, direction) {
+  if (!element) return;
+
+  const messageElement = element.closest(".message, .user-message, .assistant-message");
+  if (!messageElement) return;
+
+  const index = messageElement.getAttribute("data-index");
+  if (!index) return;
+
+  const indexInput = document.getElementById("Navigate-message-index").querySelector("input");
+  if (!indexInput) {
+    console.error("Element with ID 'Navigate-message-index' not found.");
+    return;
+  }
+
+  const directionInput = document.getElementById("Navigate-direction").querySelector("textarea");
+  if (!directionInput) {
+    console.error("Element with ID 'Navigate-direction' not found.");
+    return;
+  }
+
+  const navigateButton = document.getElementById("Navigate-version");
+  if (!navigateButton) {
+    console.error("Required element 'Navigate-version' not found.");
+    return;
+  }
+
+  indexInput.value = index;
+  directionInput.value = direction;
+
+  // Trigger any 'change' or 'input' events Gradio might be listening for
+  const event = new Event("input", { bubbles: true });
+  indexInput.dispatchEvent(event);
+  directionInput.dispatchEvent(event);
+
+  navigateButton.click();
+}
+
 function regenerateClick() {
   document.getElementById("Regenerate").click();
 }

js/main.js

@@ -39,9 +39,24 @@ document.querySelector(".header_bar").addEventListener("click", function(event)
 //------------------------------------------------
 // Keyboard shortcuts
 //------------------------------------------------
+
+// --- Helper functions --- //
+function isModifiedKeyboardEvent() {
+  return (event instanceof KeyboardEvent &&
+    event.shiftKey ||
+    event.ctrlKey ||
+    event.altKey ||
+    event.metaKey);
+}
+
+function isFocusedOnEditableTextbox() {
+  if (event.target.tagName === "INPUT" || event.target.tagName === "TEXTAREA") {
+    return !!event.target.value;
+  }
+}
+
 let previousTabId = "chat-tab-button";
 document.addEventListener("keydown", function(event) {
-
   // Stop generation on Esc pressed
   if (event.key === "Escape") {
     // Find the element with id 'stop' and click it
@@ -49,10 +64,15 @@ document.addEventListener("keydown", function(event) {
     if (stopButton) {
       stopButton.click();
     }
+    return;
+  }
+
+  if (!document.querySelector("#chat-tab").checkVisibility() ) {
+    return;
   }
 
   // Show chat controls on Ctrl + S
-  else if (event.ctrlKey && event.key == "s") {
+  if (event.ctrlKey && event.key == "s") {
     event.preventDefault();
 
     var showControlsElement = document.getElementById("show-controls");
@@ -100,6 +120,23 @@ document.addEventListener("keydown", function(event) {
     document.getElementById("Impersonate").click();
   }
 
+  // --- Simple version navigation --- //
+  if (!isFocusedOnEditableTextbox()) {
+    // Version navigation on Arrow keys (horizontal)
+    if (!isModifiedKeyboardEvent() && event.key === "ArrowLeft") {
+      event.preventDefault();
+      navigateLastAssistantMessage("left");
+    }
+
+    else if (!isModifiedKeyboardEvent() && event.key === "ArrowRight") {
+      event.preventDefault();
+      if (!navigateLastAssistantMessage("right")) {
+        // If can't navigate right (last version), regenerate
+        document.getElementById("Regenerate").click();
+      }
+    }
+  }
+
 });
 
 //------------------------------------------------
@@ -789,3 +826,55 @@ function createMobileTopBar() {
 }
 
 createMobileTopBar();
+
+//------------------------------------------------
+// Simple Navigation Functions
+//------------------------------------------------
+
+function navigateLastAssistantMessage(direction) {
+  const chat = document.querySelector("#chat");
+  if (!chat) return false;
+
+  const messages = chat.querySelectorAll("[data-index]");
+  if (messages.length === 0) return false;
+
+  // Find the last assistant message (starting from the end)
+  let lastAssistantMessage = null;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const msg = messages[i];
+    if (
+      msg.classList.contains("assistant-message") ||
+      msg.querySelector(".circle-bot") ||
+      msg.querySelector(".text-bot")
+    ) {
+      lastAssistantMessage = msg;
+      break;
+    }
+  }
+
+  if (!lastAssistantMessage) return false;
+
+  const buttons = lastAssistantMessage.querySelectorAll(".version-nav-button");
+
+  for (let i = 0; i < buttons.length; i++) {
+    const button = buttons[i];
+    const onclick = button.getAttribute("onclick");
+    const disabled = button.hasAttribute("disabled");
+
+    const isLeft = onclick && onclick.includes("'left'");
+    const isRight = onclick && onclick.includes("'right'");
+
+    if (!disabled) {
+      if (direction === "left" && isLeft) {
+        navigateVersion(button, direction);
+        return true;
+      }
+      if (direction === "right" && isRight) {
+        navigateVersion(button, direction);
+        return true;
+      }
+    }
+  }
+
+  return false;
+}

modules/chat.py

@@ -414,10 +414,20 @@ def add_message_version(history, row_idx, is_current=True):
     if "versions" not in history['metadata'][key]:
         history['metadata'][key]["versions"] = []
 
+    # Check if this version already exists
+    current_content = history['internal'][row_idx][1]
+    current_visible = history['visible'][row_idx][1]
+
+    for i, version in enumerate(history['metadata'][key]["versions"]):
+        if version['content'] == current_content and version['visible_content'] == current_visible:
+            if is_current:
+                history['metadata'][key]["current_version_index"] = i
+            return
+
     # Add current message as a version
     history['metadata'][key]["versions"].append({
-        "content": history['internal'][row_idx][1],
+        "content": current_content,
-        "visible_content": history['visible'][row_idx][1],
+        "visible_content": current_visible,
         "timestamp": get_current_timestamp()
     })
 
@@ -540,7 +550,8 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess
         if regenerate:
             row_idx = len(output['internal']) - 1
 
-            # Store the existing response as a version before regenerating
+            # Store the first response as a version before regenerating
+            if not output['metadata'].get(f"assistant_{row_idx}", {}).get('versions'):
                 add_message_version(output, row_idx, is_current=False)
 
             if loading_message:
@@ -1414,6 +1425,46 @@ def handle_branch_chat_click(state):
     return [history, html, past_chats_update, -1]
 
 
+def handle_navigate_version_click(state):
+    history = state['history']
+    message_index = int(state['navigate_message_index'])
+    direction = state['navigate_direction']
+
+    # Get assistant message metadata
+    key = f"assistant_{message_index}"
+    if key not in history['metadata'] or 'versions' not in history['metadata'][key]:
+        # No versions to navigate
+        html = redraw_html(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
+        return [history, html]
+
+    metadata = history['metadata'][key]
+    current_idx = metadata.get('current_version_index', 0)
+    versions = metadata['versions']
+
+    # Calculate new index
+    if direction == 'left':
+        new_idx = max(0, current_idx - 1)
+    else:  # right
+        new_idx = min(len(versions) - 1, current_idx + 1)
+
+    if new_idx == current_idx:
+        # No change needed
+        html = redraw_html(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
+        return [history, html]
+
+    # Update history with new version
+    version = versions[new_idx]
+    history['internal'][message_index][1] = version['content']
+    history['visible'][message_index][1] = version['visible_content']
+    metadata['current_version_index'] = new_idx
+
+    # Redraw and save
+    html = redraw_html(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
+    save_history(history, state['unique_id'], state['character_menu'], state['mode'])
+
+    return [history, html]
+
+
 def handle_rename_chat_click():
     return [
         gr.update(value="My New Chat"),

modules/html_generator.py

@@ -380,6 +380,30 @@ def format_message_attachments(history, role, index):
     return ""
 
 
+def get_version_navigation_html(history, i):
+    """Generate simple navigation arrows for message versions"""
+    key = f"assistant_{i}"
+    metadata = history.get('metadata', {})
+
+    if key not in metadata or 'versions' not in metadata[key]:
+        return ""
+
+    versions = metadata[key]['versions']
+    current_idx = metadata[key].get('current_version_index', 0)
+
+    if len(versions) <= 1:
+        return ""
+
+    left_disabled = ' disabled' if current_idx == 0 else ''
+    right_disabled = ' disabled' if current_idx >= len(versions) - 1 else ''
+
+    left_arrow = f'<button class="footer-button version-nav-button"{left_disabled} onclick="navigateVersion(this, \'left\')" title="Previous version">&lt;</button>'
+    right_arrow = f'<button class="footer-button version-nav-button"{right_disabled} onclick="navigateVersion(this, \'right\')" title="Next version">&gt;</button>'
+    position = f'<span class="version-position">{current_idx + 1}/{len(versions)}</span>'
+
+    return f'<div class="version-navigation">{left_arrow}{position}{right_arrow}</div>'
+
+
 def actions_html(history, i, info_message=""):
     return (f'<div class="message-actions">'
             f'{copy_button}'
@@ -388,7 +412,8 @@ def actions_html(history, i, info_message=""):
             f'{remove_button if i == len(history["visible"]) - 1 else ""}'
             f'{branch_button}'
             f'{info_message}'
-            f'</div>')
+            f'</div>'
+            f'{get_version_navigation_html(history, i)}')
 
 
 def generate_instruct_html(history):

modules/ui.py

@@ -157,6 +157,8 @@ def list_model_elements():
 
 def list_interface_input_elements():
     elements = [
+        'navigate_message_index',
+        'navigate_direction',
         'temperature',
         'dynatemp_low',
         'dynatemp_high',

modules/ui_chat.py

@@ -97,6 +97,12 @@ def create_ui():
                 with gr.Row():
                     shared.gradio['chat-instruct_command'] = gr.Textbox(value=shared.settings['chat-instruct_command'], lines=12, label='Command for chat-instruct mode', info='<|character|> and <|prompt|> get replaced with the bot name and the regular chat prompt respectively.', visible=shared.settings['mode'] == 'chat-instruct', elem_classes=['add_scrollbar'])
 
+        # Hidden elements for version navigation (similar to branch)
+        with gr.Row(visible=False):
+            shared.gradio['navigate_message_index'] = gr.Number(value=-1, precision=0, elem_id="Navigate-message-index")
+            shared.gradio['navigate_direction'] = gr.Textbox(value="", elem_id="Navigate-direction")
+            shared.gradio['navigate_version'] = gr.Button(elem_id="Navigate-version")
+
 
 def create_chat_settings_ui():
     mu = shared.args.multi_user
@@ -293,6 +299,10 @@ def create_event_handlers():
     shared.gradio['chat_style'].change(chat.redraw_html, gradio(reload_arr), gradio('display'), show_progress=False)
     shared.gradio['Copy last reply'].click(chat.send_last_reply_to_input, gradio('history'), gradio('textbox'), show_progress=False)
 
+    shared.gradio['navigate_version'].click(
+        ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
+        chat.handle_navigate_version_click, gradio('interface_state'), gradio('history', 'display'), show_progress=False)
+
     # Save/delete a character
     shared.gradio['save_character'].click(chat.handle_save_character_click, gradio('name2'), gradio('save_character_filename', 'character_saver'), show_progress=False)
     shared.gradio['delete_character'].click(lambda: gr.update(visible=True), None, gradio('character_deleter'), show_progress=False)