Patch to allow multiple <think> UI blocks in oobabooga/text-generation-webui

ooba_multi_thinking.patch

diff --git a/modules/html_generator.py b/modules/html_generator.py
index 79237f7f..a85003e5 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -1,5 +1,6 @@
 import datetime
 import functools
+from itertools import zip_longest
 import html
 import os
 import re
@@ -108,8 +109,8 @@ 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."""
+def extract_thinking_blocks(string):
+    """Extract thinking blocks from a string."""
     if not string:
         return None, string
 
@@ -120,22 +121,44 @@ def extract_thinking_block(string):
     start_pos = string.find(THINK_START_TAG)
     end_pos = string.find(THINK_END_TAG)
 
-    # If think tags found, use existing logic
-    if start_pos != -1 or end_pos != -1:
-        # handle missing start or end tags
-        if start_pos == -1:
-            thought_start = 0
-        else:
-            thought_start = start_pos + len(THINK_START_TAG)
+    thinking_contents = []
+    remaining_contents = []
+
+    # no start tag, only end tag
+    if start_pos == -1 and end_pos != -1:
+        think_end_pos = end_pos + len(THINK_END_TAG)
+        thinking_contents.append(string[:think_end_pos])
+        remaining_contents.append(string[think_end_pos:])
+        return thinking_contents, remaining_contents
+
+    think_start_re = re.compile(THINK_START_TAG)
+    # Adjust start position to account for any leading whitespace
+    start_pos_matches = think_start_re.finditer(string)
+    end_pos = 0
+    for start_pos_match in start_pos_matches:
+        tag_start, tag_end = start_pos_match.span()
+        # If there are multiple start tags before an end tag, only use the first
+        if end_pos > tag_start:
+            continue
+        if end_pos > 0:
+            remaining_contents.append(string[end_pos+len(THINK_END_TAG):tag_start])
+        # Find the content after the opening tag
+        content_start = tag_end
+
+        # Look for closing tag
+        end_pos = string.find(THINK_END_TAG, content_start)
+
         if end_pos == -1:
-            thought_end = len(string)
-            content_start = len(string)
+            # Only opening tag found - everything else is thinking content
+            thinking_contents.append(string[content_start:])
+            return thinking_contents, remaining_contents
         else:
-            thought_end = end_pos
-            content_start = end_pos + len(THINK_END_TAG)
-        thinking_content = string[thought_start:thought_end]
-        remaining_content = string[content_start:]
-        return thinking_content, remaining_content
+            # Both tags found - extract content between them
+            thinking_contents.append(string[content_start:end_pos])
+
+    if thinking_contents:
+        remaining_contents.append(string[end_pos + len(THINK_END_TAG):])
+        return thinking_contents, remaining_contents
 
     # If think tags not found, try alternative format
     ALT_START = "<|channel|>analysis<|message|>"
@@ -162,14 +185,17 @@ def extract_thinking_block(string):
                 content_start = len(string)
         else:
             thought_end = alt_end_pos
-            content_start = alt_content_pos + len(ALT_CONTENT_START) if alt_content_pos != -1 else alt_end_pos + len(ALT_END)
+            content_start = alt_content_pos + len(ALT_CONTENT_START) if alt_content_pos != -1 else alt_end_pos + len(
+                ALT_END)
 
         thinking_content = string[thought_start:thought_end]
         remaining_content = string[content_start:]
-        return thinking_content, remaining_content
+        return [thinking_content], [remaining_content]
 
     # Return if neither format is found
-    return None, string
+    return [], [string]
+
+
 
 
 @functools.lru_cache(maxsize=None)
@@ -182,34 +208,37 @@ def convert_to_markdown(string, message_id=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'''
-        <details class="thinking-block" data-block-id="{block_id}" data-streaming="{str(is_streaming).lower()}">
-            <summary class="thinking-header">
-                {info_svg_small}
-                <span class="thinking-title">{title_text}</span>
-            </summary>
-            <div class="thinking-content pretty_scrollbar">{thinking_html}</div>
-        </details>
-        '''
-
-        # Prepend the thinking block to the message HTML
-        html_output = thinking_block + html_output
+    thinking_contents, remaining_contents = extract_thinking_blocks(string)
+    html_output = ""
+    for i, (thinking_content, remaining_content) in enumerate(zip_longest(thinking_contents, remaining_contents)):
+
+        # If thinking content was found, process it using the same function
+        if thinking_content:
+            thinking_html = process_markdown_content(thinking_content)
+
+            # Generate unique ID for the thinking block
+            block_id = f"thinking-{message_id}-{i}"
+
+            # 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'''
+            <details class="thinking-block" data-block-id="{block_id}" data-streaming="{str(is_streaming).lower()}">
+                <summary class="thinking-header">
+                    {info_svg_small}
+                    <span class="thinking-title">{title_text}</span>
+                </summary>
+                <div class="thinking-content pretty_scrollbar">{thinking_html}</div>
+            </details>
+            '''
+
+            # Prepend the thinking block to the message HTML
+            html_output += thinking_block
+
+        # Process the main content
+        if remaining_content:
+            html_output += process_markdown_content(remaining_content)
 
     return html_output