From d35818f4e1f51b86d07c7b3c92af9c96facb3e8a Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 25 Apr 2025 18:02:02 -0300 Subject: [PATCH] UI: Add a collapsible thinking block to messages with `` steps (#6902) --- css/main.css | 82 ++++++++++++++++++++++++++++++++-- js/global_scope_js.js | 76 ++++++++++++++++++++++++++++++-- modules/chat.py | 10 +---- modules/html_generator.py | 93 ++++++++++++++++++++++++++++++++++++--- modules/shared.py | 1 - modules/ui.py | 1 - modules/ui_parameters.py | 1 - settings-template.yaml | 1 - 8 files changed, 238 insertions(+), 27 deletions(-) diff --git a/css/main.css b/css/main.css index 8ca5b33f..9dce4d0e 100644 --- a/css/main.css +++ b/css/main.css @@ -625,19 +625,19 @@ div.svelte-362y77>*, div.svelte-362y77>.form>* { width: 100%; overflow-y: visible; } - + .message { break-inside: avoid; } - + .gradio-container { overflow: visible; } - + .tab-nav { display: none !important; } - + #chat-tab > :first-child { max-width: unset; } @@ -1308,3 +1308,77 @@ div.svelte-362y77>*, div.svelte-362y77>.form>* { padding-left: 1rem; padding-right: 1rem; } + +/* Thinking blocks styling */ +.thinking-block { + margin-bottom: 12px; + border-radius: 8px; + border: 1px solid rgba(0, 0, 0, 0.1); + background-color: var(--light-theme-gray); + overflow: hidden; +} + +.dark .thinking-block { + background-color: var(--darker-gray); +} + +.thinking-header { + display: flex; + align-items: center; + padding: 10px 16px; + cursor: pointer; + user-select: none; + font-size: 14px; + color: rgba(0, 0, 0, 0.7); + transition: background-color 0.2s; +} + +.thinking-header:hover { + background-color: rgba(0, 0, 0, 0.03); +} + +.thinking-header::-webkit-details-marker { + display: none; +} + +.thinking-icon { + margin-right: 8px; + color: rgba(0, 0, 0, 0.5); +} + +.thinking-title { + font-weight: 500; +} + +.thinking-content { + padding: 12px 16px; + border-top: 1px solid rgba(0, 0, 0, 0.07); + color: rgba(0, 0, 0, 0.7); + font-size: 14px; + line-height: 1.5; + overflow-wrap: break-word; + max-height: 300px; + overflow-y: scroll; + contain: layout; +} + +/* Animation for opening thinking blocks */ +@keyframes fadeIn { + from { opacity: 0; } + to { opacity: 1; } +} + +.thinking-block[open] .thinking-content { + animation: fadeIn 0.3s ease-out; +} + +/* Additional style for in-progress thinking */ +.thinking-block[data-streaming="true"] .thinking-title { + animation: pulse 1.5s infinite; +} + +@keyframes pulse { + 0% { opacity: 0.6; } + 50% { opacity: 1; } + 100% { opacity: 0.6; } +} diff --git a/js/global_scope_js.js b/js/global_scope_js.js index f308edb9..e808c473 100644 --- a/js/global_scope_js.js +++ b/js/global_scope_js.js @@ -31,24 +31,94 @@ function removeLastClick() { } function handleMorphdomUpdate(text) { + // Track closed blocks + const closedBlocks = new Set(); + document.querySelectorAll(".thinking-block").forEach(block => { + const blockId = block.getAttribute("data-block-id"); + // If block exists and is not open, add to closed set + if (blockId && !block.hasAttribute("open")) { + closedBlocks.add(blockId); + } + }); + + // Store scroll positions for any open blocks + const scrollPositions = {}; + document.querySelectorAll(".thinking-block[open]").forEach(block => { + const content = block.querySelector(".thinking-content"); + const blockId = block.getAttribute("data-block-id"); + if (content && blockId) { + const isAtBottom = Math.abs((content.scrollHeight - content.scrollTop) - content.clientHeight) < 5; + scrollPositions[blockId] = { + position: content.scrollTop, + isAtBottom: isAtBottom + }; + } + }); + morphdom( document.getElementById("chat").parentNode, "
" + text + "
", { onBeforeElUpdated: function(fromEl, toEl) { + // Preserve code highlighting if (fromEl.tagName === "PRE" && fromEl.querySelector("code[data-highlighted]")) { const fromCode = fromEl.querySelector("code"); const toCode = toEl.querySelector("code"); if (fromCode && toCode && fromCode.textContent === toCode.textContent) { - // If the content is the same, preserve the entire 
 element
             toEl.className = fromEl.className;
             toEl.innerHTML = fromEl.innerHTML;
-            return false; // Skip updating the 
 element
+            return false;
+          }
+        }
+
+        // For thinking blocks, respect closed state
+        if (fromEl.classList && fromEl.classList.contains("thinking-block") &&
+            toEl.classList && toEl.classList.contains("thinking-block")) {
+          const blockId = toEl.getAttribute("data-block-id");
+          // If this block was closed by user, keep it closed
+          if (blockId && closedBlocks.has(blockId)) {
+            toEl.removeAttribute("open");
+          }
+        }
+
+        return !fromEl.isEqualNode(toEl);
+      },
+
+      onElUpdated: function(el) {
+        // Restore scroll positions for open thinking blocks
+        if (el.classList && el.classList.contains("thinking-block") && el.hasAttribute("open")) {
+          const blockId = el.getAttribute("data-block-id");
+          const content = el.querySelector(".thinking-content");
+
+          if (content && blockId && scrollPositions[blockId]) {
+            setTimeout(() => {
+              if (scrollPositions[blockId].isAtBottom) {
+                content.scrollTop = content.scrollHeight;
+              } else {
+                content.scrollTop = scrollPositions[blockId].position;
+              }
+            }, 0);
           }
         }
-        return !fromEl.isEqualNode(toEl); // Update only if nodes differ
       }
     }
   );
+
+  // Add toggle listeners for new blocks
+  document.querySelectorAll(".thinking-block").forEach(block => {
+    if (!block._hasToggleListener) {
+      block.addEventListener("toggle", function(e) {
+        if (this.open) {
+          const content = this.querySelector(".thinking-content");
+          if (content) {
+            setTimeout(() => {
+              content.scrollTop = content.scrollHeight;
+            }, 0);
+          }
+        }
+      });
+      block._hasToggleListener = true;
+    }
+  });
 }
diff --git a/modules/chat.py b/modules/chat.py
index fd949907..94d90bdc 100644
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -417,16 +417,8 @@ def generate_chat_reply(text, state, regenerate=False, _continue=False, loading_
             yield history
             return
 
-    show_after = html.escape(state.get("show_after")) if state.get("show_after") else None
     for history in chatbot_wrapper(text, state, regenerate=regenerate, _continue=_continue, loading_message=loading_message, for_ui=for_ui):
-        if show_after:
-            after = history["visible"][-1][1].partition(show_after)[2] or "*Is thinking...*"
-            yield {
-                'internal': history['internal'],
-                'visible': history['visible'][:-1] + [[history['visible'][-1][0], after]]
-            }
-        else:
-            yield history
+        yield history
 
 
 def character_is_loaded(state, raise_exception=False):
diff --git a/modules/html_generator.py b/modules/html_generator.py
index 5227e87e..a72e4859 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -107,8 +107,87 @@ def replace_blockquote(m):
     return m.group().replace('\n', '\n> ').replace('\\begin{blockquote}', '').replace('\\end{blockquote}', '')
 
 
+def extract_thinking_block(string):
+    """Extract thinking blocks from the beginning of a string."""
+    if not string:
+        return None, string
+
+    THINK_START_TAG = "<think>"
+    THINK_END_TAG = "</think>"
+
+    # Look for opening tag
+    start_pos = string.lstrip().find(THINK_START_TAG)
+    if start_pos == -1:
+        return None, string
+
+    # Adjust start position to account for any leading whitespace
+    start_pos = string.find(THINK_START_TAG)
+
+    # Find the content after the opening tag
+    content_start = start_pos + len(THINK_START_TAG)
+
+    # Look for closing tag
+    end_pos = string.find(THINK_END_TAG, content_start)
+
+    if end_pos != -1:
+        # Both tags found - extract content between them
+        thinking_content = string[content_start:end_pos]
+        remaining_content = string[end_pos + len(THINK_END_TAG):]
+        return thinking_content, remaining_content
+    else:
+        # Only opening tag found - everything else is thinking content
+        thinking_content = string[content_start:]
+        return thinking_content, ""
+
+
 @functools.lru_cache(maxsize=None)
-def convert_to_markdown(string):
+def convert_to_markdown(string, message_id=None):
+    if not string:
+        return ""
+
+    # Use a default message ID if none provided
+    if message_id is None:
+        message_id = "unknown"
+
+    # Extract thinking block if present
+    thinking_content, remaining_content = extract_thinking_block(string)
+
+    # Process the main content
+    html_output = process_markdown_content(remaining_content)
+
+    # If thinking content was found, process it using the same function
+    if thinking_content is not None:
+        thinking_html = process_markdown_content(thinking_content)
+
+        # Generate unique ID for the thinking block
+        block_id = f"thinking-{message_id}-0"
+
+        # Check if thinking is complete or still in progress
+        is_streaming = not remaining_content
+        title_text = "Thinking..." if is_streaming else "Thought"
+
+        thinking_block = f'''
+        

+            
+                
+                    
+                    
+                    
+                
+                {title_text}
+            
+            
{thinking_html}

+        

+        '''
+
+        # Prepend the thinking block to the message HTML
+        html_output = thinking_block + html_output
+
+    return html_output
+
+
+def process_markdown_content(string):
+    """Process a string through the markdown conversion pipeline."""
     if not string:
         return ""
 
@@ -209,15 +288,15 @@ def convert_to_markdown(string):
     return html_output
 
 
-def convert_to_markdown_wrapped(string, use_cache=True):
+def convert_to_markdown_wrapped(string, message_id=None, use_cache=True):
     '''
     Used to avoid caching convert_to_markdown calls during streaming.
     '''
 
     if use_cache:
-        return convert_to_markdown(string)
+        return convert_to_markdown(string, message_id=message_id)
 
-    return convert_to_markdown.__wrapped__(string)
+    return convert_to_markdown.__wrapped__(string, message_id=message_id)
 
 
 def generate_basic_html(string):
@@ -273,7 +352,7 @@ def generate_instruct_html(history):
     for i in range(len(history['visible'])):
         row_visible = history['visible'][i]
         row_internal = history['internal'][i]
-        converted_visible = [convert_to_markdown_wrapped(entry, use_cache=i != len(history['visible']) - 1) for entry in row_visible]
+        converted_visible = [convert_to_markdown_wrapped(entry, message_id=i, use_cache=i != len(history['visible']) - 1) for entry in row_visible]
 
         if converted_visible[0]:  # Don't display empty user messages
             output += (
@@ -320,7 +399,7 @@ def generate_cai_chat_html(history, name1, name2, style, character, reset_cache=
     for i in range(len(history['visible'])):
         row_visible = history['visible'][i]
         row_internal = history['internal'][i]
-        converted_visible = [convert_to_markdown_wrapped(entry, use_cache=i != len(history['visible']) - 1) for entry in row_visible]
+        converted_visible = [convert_to_markdown_wrapped(entry, message_id=i, use_cache=i != len(history['visible']) - 1) for entry in row_visible]
 
         if converted_visible[0]:  # Don't display empty user messages
             output += (
@@ -360,7 +439,7 @@ def generate_chat_html(history, name1, name2, reset_cache=False):
     for i in range(len(history['visible'])):
         row_visible = history['visible'][i]
         row_internal = history['internal'][i]
-        converted_visible = [convert_to_markdown_wrapped(entry, use_cache=i != len(history['visible']) - 1) for entry in row_visible]
+        converted_visible = [convert_to_markdown_wrapped(entry, message_id=i, use_cache=i != len(history['visible']) - 1) for entry in row_visible]
 
         if converted_visible[0]:  # Don't display empty user messages
             output += (
diff --git a/modules/shared.py b/modules/shared.py
index b4f87082..5177ac67 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -59,7 +59,6 @@ settings = {
     'seed': -1,
     'custom_stopping_strings': '',
     'custom_token_bans': '',
-    'show_after': '',
     'negative_prompt': '',
     'autoload_model': False,
     'dark_theme': True,
diff --git a/modules/ui.py b/modules/ui.py
index 19b76cee..3e1bf6d8 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -207,7 +207,6 @@ def list_interface_input_elements():
         'sampler_priority',
         'custom_stopping_strings',
         'custom_token_bans',
-        'show_after',
         'negative_prompt',
         'dry_sequence_breakers',
         'grammar_string',
diff --git a/modules/ui_parameters.py b/modules/ui_parameters.py
index c3245a9d..b494a758 100644
--- a/modules/ui_parameters.py
+++ b/modules/ui_parameters.py
@@ -93,7 +93,6 @@ def create_ui(default_preset):
                             shared.gradio['sampler_priority'] = gr.Textbox(value=generate_params['sampler_priority'], lines=12, label='Sampler priority', info='Parameter names separated by new lines or commas.', elem_classes=['add_scrollbar'])
                             shared.gradio['custom_stopping_strings'] = gr.Textbox(lines=2, value=shared.settings["custom_stopping_strings"] or None, label='Custom stopping strings', info='Written between "" and separated by commas.', placeholder='"\\n", "\\nYou:"')
                             shared.gradio['custom_token_bans'] = gr.Textbox(value=shared.settings['custom_token_bans'] or None, label='Token bans', info='Token IDs to ban, separated by commas. The IDs can be found in the Default or Notebook tab.')
-                            shared.gradio['show_after'] = gr.Textbox(value=shared.settings['show_after'] or None, label='Show after', info='Hide the reply before this text.', placeholder="")
                             shared.gradio['negative_prompt'] = gr.Textbox(value=shared.settings['negative_prompt'], label='Negative prompt', info='For CFG. Only used when guidance_scale is different than 1.', lines=3, elem_classes=['add_scrollbar'])
                             shared.gradio['dry_sequence_breakers'] = gr.Textbox(value=generate_params['dry_sequence_breakers'], label='dry_sequence_breakers', info='Tokens across which sequence matching is not continued. Specified as a comma-separated list of quoted strings.')
                             with gr.Row() as shared.gradio['grammar_file_row']:
diff --git a/settings-template.yaml b/settings-template.yaml
index 94a5c034..83764f97 100644
--- a/settings-template.yaml
+++ b/settings-template.yaml
@@ -29,7 +29,6 @@ truncation_length: 8192
 seed: -1
 custom_stopping_strings: ''
 custom_token_bans: ''
-show_after: ''
 negative_prompt: ''
 autoload_model: false
 dark_theme: true