Merge pull request #3634 from SillyTavern/continue-from-reasoning Fix auto-parsing of continue from reasoning

a392593e534551998116c2faf4da349088d02f88

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed

Parents: cabd6de85b5db765c3efba0371e45086f4d6e6f3 c03da65821d729014858a4478b1f1e290fd643b3

public/script.js+15 -6 Viewed

         this.promptReasoning = promptReasoning;
+    }
-    #checkDomElements(messageId) {
+    /**
+     * Initializes DOM elements for the current message.
+     * @param {number} messageId Current message ID
+     * @param {boolean?} continueOnReasoning If continuing on reasoning
+     */
+    async #checkDomElements(messageId, continueOnReasoning = null) {
         if (this.messageDom === null || this.messageTextDom === null) {
             this.messageDom = document.querySelector(`#chat .mes[mesid="${messageId}"]`);
             this.messageTextDom = this.messageDom?.querySelector('.mes_text');
             this.messageTimerDom = this.messageDom?.querySelector('.mes_timer');
             this.messageTokenCounterDom = this.messageDom?.querySelector('.tokenCounterDisplay');
+        }
+        if (continueOnReasoning) {
+            await this.reasoningHandler.process(messageId, false, this.promptReasoning);
+        }
         this.reasoningHandler.updateDom(messageId);
+    }

@@ -3230,7 +3238,8 @@ class StreamingProcessor {
3230	}	3238	}
3231		3239
3232	async onStartStreaming(text) {	3240	async onStartStreaming(text) {
3233	if (this.type === 'continue' && this.promptReasoning.prefixReasoning) {	3241	const continueOnReasoning = !!(this.type === 'continue' && this.promptReasoning.prefixReasoning);
		3242	if (continueOnReasoning) {
3234	this.reasoningHandler.initContinue(this.promptReasoning);	3243	this.reasoningHandler.initContinue(this.promptReasoning);
3235	}	3244	}
3236		3245

         } else {
             await saveReply(this.type, text, true, '', [], '');
             messageId = chat.length - 1;
-            this.#checkDomElements(messageId);
+            await this.#checkDomElements(messageId, continueOnReasoning);
             this.markUIGenStarted();
+        }
         hideSwipeButtons();

             this.sendTextarea.dispatchEvent(new Event('input', { bubbles: true }));
         } else {
             const mesChanged = chat[messageId]['mes'] !== processedText;
-            this.#checkDomElements(messageId);
+            await this.#checkDomElements(messageId);
             this.#updateMessageBlockVisibility();
             const currentTime = new Date();
             chat[messageId]['mes'] = processedText;

             chat[messageId]['extra']['time_to_first_token'] = this.timeToFirstToken;
             // Update reasoning
-            await this.reasoningHandler.process(messageId, mesChanged);
+            await this.reasoningHandler.process(messageId, mesChanged, this.promptReasoning);
             processedText = chat[messageId]['mes'];
             // Token count update.

@@ -5957,7 +5966,7 @@ export function cleanUpMessage(getMessage, isImpersonate, isContinue, displayInc
5957	getMessage = trimToEndSentence(getMessage);	5966	getMessage = trimToEndSentence(getMessage);
5958	}	5967	}
5959		5968
5960	if (power_user.trim_spaces) {	5969	if (power_user.trim_spaces && !PromptReasoning.getLatestPrefix()) {
5961	getMessage = getMessage.trim();	5970	getMessage = getMessage.trim();
5962	}	5971	}
5963		5972

public/scripts/reasoning.js+75 -14 Viewed

      */
     initContinue(promptReasoning) {
         this.reasoning = promptReasoning.prefixReasoning;
-        this.state = ReasoningState.Done;
+        this.state = promptReasoning.prefixIncomplete ? ReasoningState.None : ReasoningState.Done;
         this.startTime = this.initialTime;
         this.endTime = promptReasoning.prefixDuration ? new Date(this.initialTime.getTime() + promptReasoning.prefixDuration) : null;
+    }

+     *
      * @param {number} messageId - The ID of the message to process
      * @param {boolean} mesChanged - Whether the message has changed
+     * @param {PromptReasoning} promptReasoning - Prompt reasoning object
      * @returns {Promise<void>}
      */
-    async process(messageId, mesChanged) {
+    async process(messageId, mesChanged, promptReasoning) {
-        mesChanged = this.#autoParseReasoningFromMessage(messageId, mesChanged);
+        mesChanged = this.#autoParseReasoningFromMessage(messageId, mesChanged, promptReasoning);
         if (!this.reasoning && !this.#isHiddenReasoningModel)
             return;

+        }
+    }
-    #autoParseReasoningFromMessage(messageId, mesChanged) {
+    /**
+     * Parse reasoning from a message during streaming.
+     * @param {number} messageId Message ID
+     * @param {boolean} mesChanged Whether the message has changed before reasoning parsing
+     * @param {PromptReasoning} promptReasoning Prompt reasoning object
+     * @returns {boolean} Whether the message has changed after reasoning parsing
+     */
+    #autoParseReasoningFromMessage(messageId, mesChanged, promptReasoning) {
         if (!power_user.reasoning.auto_parse)
             return;
         if (!power_user.reasoning.prefix || !power_user.reasoning.suffix)

         const message = chat[messageId];
         if (!message) return mesChanged;
+        const parseTarget = promptReasoning?.prefixIncomplete ? (promptReasoning.prefixReasoningFormatted + message.mes) : message.mes;
         // If we are done with reasoning parse, we just split the message correctly so the reasoning doesn't show up inside of it.
         if (this.#parsingReasoningMesStartIndex) {
-            message.mes = trimSpaces(message.mes.slice(this.#parsingReasoningMesStartIndex));
+            message.mes = trimSpaces(parseTarget.slice(this.#parsingReasoningMesStartIndex));
             return mesChanged;
+        }
         if (this.state === ReasoningState.None || this.#isHiddenReasoningModel) {
             // If streamed message starts with the opening, cut it out and put all inside reasoning
-            if (message.mes.startsWith(power_user.reasoning.prefix) && message.mes.length > power_user.reasoning.prefix.length) {
+            if (parseTarget.startsWith(power_user.reasoning.prefix) && parseTarget.length > power_user.reasoning.prefix.length) {
                 this.#isParsingReasoning = true;
                 // Manually set starting state here, as we might already have received the ending suffix

             return mesChanged;
         // If we are in manual parsing mode, all currently streaming mes tokens will go the the reasoning block
-        const originalMes = message.mes;
+        this.reasoning = parseTarget.slice(power_user.reasoning.prefix.length);
-        this.reasoning = originalMes.slice(power_user.reasoning.prefix.length);
         message.mes = '';
         // If the reasoning contains the ending suffix, we cut that off and continue as message streaming
         if (this.reasoning.includes(power_user.reasoning.suffix)) {
             this.reasoning = this.reasoning.slice(0, this.reasoning.indexOf(power_user.reasoning.suffix));
-            this.#parsingReasoningMesStartIndex = originalMes.indexOf(power_user.reasoning.suffix) + power_user.reasoning.suffix.length;
+            this.#parsingReasoningMesStartIndex = parseTarget.indexOf(power_user.reasoning.suffix) + power_user.reasoning.suffix.length;
-            message.mes = trimSpaces(originalMes.slice(this.#parsingReasoningMesStartIndex));
+            message.mes = trimSpaces(parseTarget.slice(this.#parsingReasoningMesStartIndex));
             this.#isParsingReasoning = false;
+        }

  * Keeps track of the number of reasoning additions.
  */
 export class PromptReasoning {
+    /**
+     * An instance initiated during the latest prompt processing.
+     * @type {PromptReasoning}
+     * */
+    static #LATEST = null;
+    /**
+     * @readonly Zero-width space character used as a placeholder for reasoning.
+     * @type {string}
+    */
     static REASONING_PLACEHOLDER = '\u200B';
+    /**
+     * Returns the latest formatted reasoning prefix if the prefix is incomplete.
+     * @returns {string} Formatted reasoning prefix
+     */
+    static getLatestPrefix() {
+        if (!PromptReasoning.#LATEST) {
+            return '';
+        }
+        if (!PromptReasoning.#LATEST.prefixIncomplete) {
+            return '';
+        }
+        return PromptReasoning.#LATEST.prefixReasoningFormatted;
+    }
+    /**
+     * Free the latest reasoning instance.
+     * To be called when the generation has ended or stopped.
+     */
+    static clearLatest() {
+        PromptReasoning.#LATEST = null;
+    }
     constructor() {
+        PromptReasoning.#LATEST = this;
+        /** @type {number} */
         this.counter = 0;
+        /** @type {number} */
         this.prefixLength = -1;
+        /** @type {string} */
         this.prefixReasoning = '';
+        /** @type {string} */
+        this.prefixReasoningFormatted = '';
+        /** @type {number?} */
         this.prefixDuration = null;
+        /** @type {boolean} */
+        this.prefixIncomplete = false;
+    }
     /**

             const formattedReasoning = `${prefix}${reasoning}`;
             if (isPrefix) {
                 this.prefixReasoning = reasoning;
+                this.prefixReasoningFormatted = formattedReasoning;
                 this.prefixLength = formattedReasoning.length;
                 this.prefixDuration = duration;
+                this.prefixIncomplete = true;
+            }
             return formattedReasoning;
+        }

         const formattedReasoning = `${prefix}${reasoning}${suffix}${separator}`;
         if (isPrefix) {
             this.prefixReasoning = reasoning;
+            this.prefixReasoningFormatted = formattedReasoning;
             this.prefixLength = formattedReasoning.length;
             this.prefixDuration = duration;
+            this.prefixIncomplete = false;
+        }
         return `${formattedReasoning}${content}`;
+    }

+}
 function registerReasoningAppEvents() {
-    const eventHandler = (/** @type {number} */ idx) => {
+    const eventHandler = (/** @type {string} */ type, /** @type {number} */ idx) => {
         if (!power_user.reasoning.auto_parse) {
             return;
+        }
         console.debug('[Reasoning] Auto-parsing reasoning block for message', idx);
+        const prefix = type === event_types.MESSAGE_RECEIVED ? PromptReasoning.getLatestPrefix() : '';
         const message = chat[idx];
         if (!message) {

@@ -1093,12 +1150,12 @@ function registerReasoningAppEvents() {
1093	return null;	1150	return null;
1094	}	1151	}
1095		1152
1096	if (message.extra?.reasoning) {	1153	if (message.extra?.reasoning && !prefix) {
1097	console.debug('[Reasoning] Message already has reasoning', idx);	1154	console.debug('[Reasoning] Message already has reasoning', idx);
1098	return null;	1155	return null;
1099	}	1156	}
1100		1157
1101	const parsedReasoning = parseReasoningFromString(message.mes);	1158	const parsedReasoning = parseReasoningFromString(prefix + message.mes);
1102		1159
1103	// No reasoning block found	1160	// No reasoning block found
1104	if (!parsedReasoning) {	1161	if (!parsedReasoning) {

@@ -1137,7 +1194,11 @@ function registerReasoningAppEvents() {
1137	};	1194	};
1138		1195
1139	for (const event of [event_types.MESSAGE_RECEIVED, event_types.MESSAGE_UPDATED]) {	1196	for (const event of [event_types.MESSAGE_RECEIVED, event_types.MESSAGE_UPDATED]) {
1140	eventSource.on(event, eventHandler);	1197	eventSource.on(event, (/** @type {number} */ idx) => eventHandler(event, idx));
		1198	}
		1199
		1200	for (const event of [event_types.GENERATION_STOPPED, event_types.GENERATION_ENDED, event_types.CHAT_CHANGED]) {
		1201	eventSource.on(event, () => PromptReasoning.clearLatest());
1141	}	1202	}
1142	}	1203	}
1143		1204