From d0abba23dca061bb1875a2bb62319810a1f70976 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Wed, 29 Jan 2025 21:38:08 +0200
Subject: [PATCH] Add reasoning auto-parsing, always show reasoning
 prefix/suffix in Markdown, respect space trim preference

---
 public/index.html            |  6 +++
 public/script.js             | 18 ++++++-
 public/scripts/power-user.js |  1 +
 public/scripts/reasoning.js  | 95 +++++++++++++++++++++++++++++++++++-
 4 files changed, 117 insertions(+), 3 deletions(-)

diff --git a/public/index.html b/public/index.html
index 67ca38777..ff999b2f2 100644
--- a/public/index.html
+++ b/public/index.html
@@ -3791,6 +3791,12 @@
                                 <span data-i18n="Reasoning">Reasoning</span>
                             </h4>
                             <div>
+                                <label class="checkbox_label" for="reasoning_auto_parse" title="Automatically parse reasoning blocks from main content between the reasoning prefix/suffix. Both fields must be defined and non-empty." data-i18n="[title]reasoning_auto_parse">
+                                    <input id="reasoning_auto_parse" type="checkbox" />
+                                    <small data-i18n="Auto-Parse Reasoning">
+                                        Auto-Parse Reasoning
+                                    </small>
+                                </label>
                                 <label class="checkbox_label" for="reasoning_add_to_prompts" title="Add existing reasoning blocks to prompts. To add a new reasoning block, use the message edit menu." data-i18n="[title]reasoning_add_to_prompts">
                                     <input id="reasoning_add_to_prompts" type="checkbox" />
                                     <small data-i18n="Add Reasoning to Prompts">
diff --git a/public/script.js b/public/script.js
index b5fadcdc1..181437247 100644
--- a/public/script.js
+++ b/public/script.js
@@ -169,6 +169,7 @@ import {
     toggleDrawer,
     isElementInViewport,
     copyText,
+    escapeHtml,
 } from './scripts/utils.js';
 import { debounce_timeout } from './scripts/constants.js';
 
@@ -2066,6 +2067,17 @@ export function messageFormatting(mes, ch_name, isSystem, isUser, messageId, san
         mes = mes.replaceAll('<', '&lt;').replaceAll('>', '&gt;');
     }
 
+    // Make sure reasoning strings are always shown, even if they include "<" or ">"
+    [power_user.reasoning.prefix, power_user.reasoning.suffix].forEach((reasoningString) => {
+        if (!reasoningString || !reasoningString.trim().length) {
+            return;
+        }
+        // Only replace the first occurrence of the reasoning string
+        if (mes.includes(reasoningString)) {
+            mes = mes.replace(reasoningString, escapeHtml(reasoningString));
+        }
+    });
+
     if (!isSystem) {
         // Save double quotes in tags as a special character to prevent them from being encoded
         if (!power_user.encode_tags) {
@@ -3208,7 +3220,7 @@ class StreamingProcessor {
             }
 
             if (this.reasoning) {
-                chat[messageId]['extra']['reasoning'] = this.reasoning;
+                chat[messageId]['extra']['reasoning'] = power_user.trim_spaces ? this.reasoning.trim() : this.reasoning;
                 if (this.messageReasoningDom instanceof HTMLElement) {
                     const formattedReasoning = messageFormatting(this.reasoning, '', false, false, messageId, {}, true);
                     this.messageReasoningDom.innerHTML = formattedReasoning;
@@ -4805,6 +4817,10 @@ export async function Generate(type, { automatic_trigger, force_name2, quiet_pro
         messageChunk = cleanUpMessage(getMessage, isImpersonate, isContinue, false);
         reasoning = getRegexedString(reasoning, regex_placement.REASONING);
 
+        if (power_user.trim_spaces) {
+            reasoning = reasoning.trim();
+        }
+
         if (isContinue) {
             getMessage = continue_mag + getMessage;
         }
diff --git a/public/scripts/power-user.js b/public/scripts/power-user.js
index 637fdbffb..5305fec34 100644
--- a/public/scripts/power-user.js
+++ b/public/scripts/power-user.js
@@ -254,6 +254,7 @@ let power_user = {
     },
 
     reasoning: {
+        auto_parse: false,
         add_to_prompts: false,
         prefix: '<think>\n',
         suffix: '\n</think>',
diff --git a/public/scripts/reasoning.js b/public/scripts/reasoning.js
index 888276f8e..c79e03629 100644
--- a/public/scripts/reasoning.js
+++ b/public/scripts/reasoning.js
@@ -1,4 +1,4 @@
-import { chat, closeMessageEditor, saveChatConditional, saveSettingsDebounced, substituteParams, updateMessageBlock } from '../script.js';
+import { chat, closeMessageEditor, event_types, eventSource, saveChatConditional, saveSettingsDebounced, substituteParams, updateMessageBlock } from '../script.js';
 import { getRegexedString, regex_placement } from './extensions/regex/engine.js';
 import { t } from './i18n.js';
 import { MacrosParser } from './macros.js';
@@ -8,7 +8,7 @@ import { SlashCommand } from './slash-commands/SlashCommand.js';
 import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from './slash-commands/SlashCommandArgument.js';
 import { commonEnumProviders } from './slash-commands/SlashCommandCommonEnumsProvider.js';
 import { SlashCommandParser } from './slash-commands/SlashCommandParser.js';
-import { copyText } from './utils.js';
+import { copyText, escapeRegex } from './utils.js';
 
 /**
  * Gets a message from a jQuery element.
@@ -106,6 +106,12 @@ function loadReasoningSettings() {
         power_user.reasoning.max_additions = Number($(this).val());
         saveSettingsDebounced();
     });
+
+    $('#reasoning_auto_parse').prop('checked', power_user.reasoning.auto_parse);
+    $('#reasoning_auto_parse').on('change', function () {
+        power_user.reasoning.auto_parse = !!$(this).prop('checked');
+        saveSettingsDebounced();
+    });
 }
 
 function registerReasoningSlashCommands() {
@@ -290,9 +296,94 @@ function setReasoningEventHandlers(){
     });
 }
 
+/**
+ * Parses reasoning from a string using the power user reasoning settings.
+ * @typedef {Object} ParsedReasoning
+ * @property {string} reasoning Reasoning block
+ * @property {string} content Message content
+ * @param {string} str Content of the message
+ * @returns {ParsedReasoning|null} Parsed reasoning block and message content
+ */
+function parseReasoningFromString(str) {
+    // Both prefix and suffix must be defined
+    if (!power_user.reasoning.prefix || !power_user.reasoning.suffix) {
+        return null;
+    }
+
+    try {
+        let reasoning = '';
+
+        const regex = new RegExp(`${escapeRegex(power_user.reasoning.prefix)}(.*?)${escapeRegex(power_user.reasoning.suffix)}`, 's');
+        const content = String(str).replace(regex, (_match, captureGroup) => {
+            reasoning = power_user.trim_spaces ? captureGroup.trim() : captureGroup;
+            return '';
+        });
+
+        return { reasoning, content };
+    } catch (error) {
+        console.error('[Reasoning] Error parsing reasoning block', error);
+        return null;
+    }
+}
+
+function registerReasoningAppEvents() {
+    eventSource.makeFirst(event_types.MESSAGE_RECEIVED, (/** @type {number} */ idx) => {
+        if (!power_user.reasoning.auto_parse) {
+            return;
+        }
+
+        console.debug('[Reasoning] Auto-parsing reasoning block for message', idx);
+        const message = chat[idx];
+
+        if (!message) {
+            console.warn('[Reasoning] Message not found', idx);
+            return null;
+        }
+
+        if (!message.mes || message.mes === '...') {
+            console.debug('[Reasoning] Message content is empty or a placeholder', idx);
+            return null;
+        }
+
+        const parsedReasoning = parseReasoningFromString(message.mes);
+
+        // No reasoning block found
+        if (!parsedReasoning) {
+            return;
+        }
+
+        // Make sure the message has an extra object
+        if (!message.extra || typeof message.extra !== 'object') {
+            message.extra = {};
+        }
+
+        const contentUpdated = !!parsedReasoning.reasoning || parsedReasoning.content !== message.mes;
+
+        // If reasoning was found, add it to the message
+        if (parsedReasoning.reasoning) {
+            message.extra.reasoning = parsedReasoning.reasoning;
+        }
+
+        // Update the message text if it was changed
+        if (parsedReasoning.content !== message.mes) {
+            message.mes = parsedReasoning.content;
+        }
+
+        // Find if a message already exists in DOM and must be updated
+        if (contentUpdated) {
+            const messageRendered = document.querySelector(`.mes[mesid="${idx}"]`) !== null;
+            if (messageRendered) {
+                console.debug('[Reasoning] Updating message block', idx);
+                updateMessageBlock(idx, message);
+            }
+        }
+    });
+}
+
 export function initReasoning() {
     loadReasoningSettings();
     setReasoningEventHandlers();
     registerReasoningSlashCommands();
     registerReasoningMacros();
+    registerReasoningAppEvents();
 }