Refactoring StreamProcessor -> ReasoningHandler

d8eeab0c0068665bb7626c164a5bf7b6fa71ecdf

Wolfsblvt <wolfsblvt@gmail.com>

Parents: 8b4414b799b602c9c66d84d7f3a409a74936ae3b

public/script.js+14 -79 Viewed

 import { initBulkEdit } from './scripts/bulk-edit.js';
 import { deriveTemplatesFromChatTemplate } from './scripts/chat-templates.js';
 import { getContext } from './scripts/st-context.js';
-import { extractReasoningFromData, initReasoning, isHiddenReasoningModel, PromptReasoning, updateReasoningUI } from './scripts/reasoning.js';
+import { extractReasoningFromData, initReasoning, PromptReasoning, ReasoningHandler, updateReasoningUI } from './scripts/reasoning.js';
 // API OBJECT FOR EXTERNAL WIRING
 globalThis.SillyTavern = {

         this.messageTimerDom = null;
         /** @type {HTMLElement} */
         this.messageTokenCounterDom = null;
-        /** @type {HTMLElement} */
-        this.messageReasoningDom = null;
-        /** @type {HTMLElement} */
-        this.messageReasoningHeaderDom = null;
         /** @type {HTMLTextAreaElement} */
         this.sendTextarea = document.querySelector('#send_textarea');
         this.type = type;

         /** @type {import('./scripts/logprobs.js').TokenLogprobs[]} */
         this.messageLogprobs = [];
         this.toolCalls = [];
-        this.reasoning = '';
+        // Initialize reasoning in its own handler
-        /** @type {Date} */
+        this.reasoningHandler = new ReasoningHandler(type, timeStarted);
-        this.reasoningStartTime = null;
-        /** @type {Date} */
-        this.reasoningEndTime = null;
-        this.isHiddenReasoning = isHiddenReasoningModel();
-    }
-    /** @type {() => number} Reasoning duration in milliseconds */
-    #reasoningDuration() {
-        if (this.reasoningStartTime && this.reasoningEndTime) {
-            return (this.reasoningEndTime.getTime() - this.reasoningStartTime.getTime());
-        }
-        return null;
+    }
     #checkDomElements(messageId) {

             this.messageTextDom = this.messageDom?.querySelector('.mes_text');
             this.messageTimerDom = this.messageDom?.querySelector('.mes_timer');
             this.messageTokenCounterDom = this.messageDom?.querySelector('.tokenCounterDisplay');
-            this.messageReasoningDom = this.messageDom?.querySelector('.mes_reasoning');
-            this.messageReasoningHeaderDom = this.messageDom?.querySelector('.mes_reasoning_header_title');
+        }
+        this.reasoningHandler.checkDomElements(messageId);
-        this.messageDom.classList.toggle('reasoning_hidden', this.isHiddenReasoning);
+    }
     #updateMessageBlockVisibility() {

+    }
     showMessageButtons(messageId) {
-        if (messageId == -1) {
-            return;
-        }
         showStopButton();
         $(`#chat .mes[mesid="${messageId}"] .mes_buttons`).css({ 'display': 'none' });
+    }
     hideMessageButtons(messageId) {
-        if (messageId == -1) {
-            return;
-        }
         hideStopButton();
         $(`#chat .mes[mesid="${messageId}"] .mes_buttons`).css({ 'display': 'flex' });
+    }

         if (this.type == 'impersonate') {
             this.sendTextarea.value = '';
             this.sendTextarea.dispatchEvent(new Event('input', { bubbles: true }));
-        }
+        } else {
-        else {
             await saveReply(this.type, text, true, '', [], '');
             messageId = chat.length - 1;
             this.#checkDomElements(messageId);
             this.showMessageButtons(messageId);
+        }
         hideSwipeButtons();
         scrollChatToBottom();
         return messageId;

         let processedText = cleanUpMessage(text, isImpersonate, isContinue, !isFinal, this.stoppingStrings);
-        // Predict unbalanced asterisks / quotes during streaming
         const charsToBalance = ['*', '"', '```'];
         for (const char of charsToBalance) {
             if (!isFinal && isOdd(countOccurrences(processedText, char))) {
-                // Add character at the end to balance it
                 const separator = char.length > 1 ? '\n' : '';
                 processedText = processedText.trimEnd() + separator + char;
+            }

         if (isImpersonate) {
             this.sendTextarea.value = processedText;
             this.sendTextarea.dispatchEvent(new Event('input', { bubbles: true }));
-        }
+        } else {
-        else {
             const mesChanged = chat[messageId]['mes'] !== processedText;
             this.#checkDomElements(messageId);
             this.#updateMessageBlockVisibility();
             const currentTime = new Date();
             chat[messageId]['mes'] = processedText;
             chat[messageId]['gen_started'] = this.timeStarted;
             chat[messageId]['gen_finished'] = currentTime;
             if (!chat[messageId]['extra']) {
                 chat[messageId]['extra'] = {};
+            }
-            if (this.reasoning || this.isHiddenReasoning) {
+            // Update reasoning
-                const reasoning = power_user.trim_spaces ? this.reasoning.trim() : this.reasoning;
+            await this.reasoningHandler.process(messageId, mesChanged, currentTime);
-                const reasoningChanged = chat[messageId]['extra']['reasoning'] !== reasoning;
-                chat[messageId]['extra']['reasoning'] = reasoning;
-                if ((this.isHiddenReasoning || reasoningChanged) && this.reasoningStartTime === null) {
-                    this.reasoningStartTime = this.timeStarted;
-                }
-                if ((this.isHiddenReasoning || !reasoningChanged) && mesChanged && this.reasoningStartTime !== null && this.reasoningEndTime === null) {
-                    this.reasoningEndTime = currentTime;
-                    await eventSource.emit(event_types.STREAM_REASONING_DONE, this.reasoning, this.#reasoningDuration);
-                }
-                await this.#updateReasoningTime(messageId);
-                if (this.messageReasoningDom instanceof HTMLElement) {
-                    const formattedReasoning = messageFormatting(this.reasoning, '', false, false, messageId, {}, true);
-                    this.messageReasoningDom.innerHTML = formattedReasoning;
-                }
-                if (this.messageDom instanceof HTMLElement) {
-                    this.messageDom.classList.add('reasoning');
-                }
-            }
-            // Don't waste time calculating token count for streaming
+            // Token count update.
-            const tokenCountText = (this.reasoning || '') + processedText;
+            const tokenCountText = this.reasoningHandler.reasoning + processedText;
             const currentTokenCount = isFinal && power_user.message_token_count_enabled ? getTokenCount(tokenCountText, 0) : 0;
             if (currentTokenCount) {
                 chat[messageId]['extra']['token_count'] = currentTokenCount;
                 if (this.messageTokenCounterDom instanceof HTMLElement) {

                 this.messageTextDom.innerHTML = formattedText;
+            }
-            const timePassed = formatGenerationTimer(this.timeStarted, currentTime, currentTokenCount, this.#reasoningDuration());
+            const timePassed = formatGenerationTimer(this.timeStarted, currentTime, currentTokenCount, this.reasoningHandler.getDuration());
             if (this.messageTimerDom instanceof HTMLElement) {
                 this.messageTimerDom.textContent = timePassed.timerValue;
                 this.messageTimerDom.title = timePassed.timerTitle;

+        }
+    }
-    async #updateReasoningTime(messageId, { forceEnd = false } = {}) {
-        const duration = this.#reasoningDuration();
-        chat[messageId]['extra']['reasoning_duration'] = duration;
-        updateReasoningUI(this.messageDom, this.reasoning, duration, { forceEnd: forceEnd });
-    }
     async onFinishStreaming(messageId, text) {
         this.hideMessageButtons(this.messageId);
         await this.onProgressStreaming(messageId, text, true);
         addCopyToCodeBlocks($(`#chat .mes[mesid="${messageId}"]`));
-        // Ensure reasoning finish time is recorded if not already
+        await this.reasoningHandler.finish(messageId);
-        if (this.reasoningStartTime !== null && this.reasoningEndTime === null) {
-            this.reasoningEndTime = new Date();
-            await eventSource.emit(event_types.STREAM_REASONING_DONE, this.reasoning, this.#reasoningDuration);
-            await this.#updateReasoningTime(messageId, { forceEnd: true });
-        }
         if (Array.isArray(this.swipes) && this.swipes.length > 0) {
             const message = chat[messageId];

@@ -3443,7 +3377,8 @@ class StreamingProcessor {
3443	if (logprobs) {	3377	if (logprobs) {
3444	this.messageLogprobs.push(...(Array.isArray(logprobs) ? logprobs : [logprobs]));	3378	this.messageLogprobs.push(...(Array.isArray(logprobs) ? logprobs : [logprobs]));
3445	}	3379	}
3446	this.reasoning = getRegexedString(state?.reasoning ?? '', regex_placement.REASONING);	3380	// Get the updated reasoning string into the handler
		3381	this.reasoningHandler.updateReasoning(state?.reasoning ?? '');
3447	await eventSource.emit(event_types.STREAM_TOKEN_RECEIVED, text);	3382	await eventSource.emit(event_types.STREAM_TOKEN_RECEIVED, text);
3448	await sw.tick(async () => await this.onProgressStreaming(this.messageId, this.continueMessage + text));	3383	await sw.tick(async () => await this.onProgressStreaming(this.messageId, this.continueMessage + text));
3449	}	3384	}

public/scripts/reasoning.js+155 -1 Viewed

  * @param {object} [options={}] Options for the function
  * @param {boolean} [options.forceEnd=false] If true, there will be no "Thinking..." when no duration exists
  */
-export function updateReasoningTimeUI(element, duration, { forceEnd = false } = {}) {
+function updateReasoningTimeUI(element, duration, { forceEnd = false } = {}) {
     if (duration) {
         const durationStr = moment.duration(duration).locale(getCurrentLocale()).humanize({ s: 50, ss: 3 });
         const secondsStr = moment.duration(duration).asSeconds();

+    }
+}
+/** @enum {string} */
+export const ReasoningState = {
+    None: 'none',
+    Thinking: 'thinking',
+    Done: 'done',
+    Hidden: 'hidden',
+};
+/**
+ * Handles reasoning-specific logic and DOM updates for messages.
+ * Used inside the @see {StreamingProcessor}
+ */
+export class ReasoningHandler {
+    #isHidden;
+    /**
+     * @param {string} type - The streaming type
+     * @param {Date} timeStarted - When the generation started
+     */
+    constructor(type, timeStarted) {
+        /** @type {ReasoningState} The current state of the reasoning process */
+        this.state = ReasoningState.None;
+        /** @type {string} The reasoning output */
+        this.reasoning = '';
+        /** @type {Date} When the reasoning started */
+        this.startTime = null;
+        /** @type {Date} When the reasoning ended */
+        this.endTime = null;
+        /** @type {string} Generation type (normal, continue, impersonation, etc) */
+        this.type = type;
+        /** @type {Date} Initial starting time of the generation */
+        this.initialTime = timeStarted;
+        /** @type {boolean} True if the model supports reasoning, but hides the reasoning output */
+        this.#isHidden = isHiddenReasoningModel();
+        // Cached DOM elements for reasoning
+        /** @type {HTMLElement} Main message DOM element `.mes` */
+        this.messageDom = null;
+        /** @type {HTMLElement} Reasoning details DOM element `.mes_reasoning_details` */
+        this.messageReasoningDetailsDom = null;
+        /** @type {HTMLElement} Reasoning content DOM element `.mes_reasoning_content` */
+        this.messageReasoningContentDom = null;
+        /** @type {HTMLElement} Reasoning header DOM element `.mes_reasoning_header` */
+        this.messageReasoningHeaderDom = null;
+    }
+    /**
+     * Gets the duration of the reasoning in milliseconds.
+     * @returns {number|null} The duration in milliseconds, or null if the start or end time is not set.
+     */
+    getDuration() {
+        if (this.startTime && this.endTime) {
+            return this.endTime.getTime() - this.startTime.getTime();
+        }
+        return null;
+    }
+    /**
+     * Finds and caches reasoning-related DOM elements for the given message.
+     * @param {number} messageId The message ID
+     */
+    checkDomElements(messageId) {
+        // Make sure we reset dom elements if we are checking for a different message (shouldn't happen, but be sure)
+        if (this.messageDom !== null && this.messageDom.getAttribute('mesid') !== messageId.toString()) {
+            this.messageDom = null;
+        }
+        // Cache the DOM elements once
+        if (this.messageDom === null) {
+            this.messageDom = document.querySelector(`#chat .mes[mesid="${messageId}"]`);
+            this.messageReasoningDetailsDom = this.messageDom.querySelector('.mes_reasoning_details');
+            this.messageReasoningContentDom = this.messageDom.querySelector('.mes_reasoning');
+            this.messageReasoningHeaderDom = this.messageDom.querySelector('.mes_reasoning_header_title');
+            // Update the DOM with the current reasoning state.
+            this.messageDom.dataset.state = this.state;
+            this.messageDom.classList.toggle('reasoning_hidden', this.#isHidden);
+        }
+        // Update main DOM state
+        this.#updateDomState();
+    }
+    #updateDomState() {
+        this.messageDom.dataset.state = this.state;
+        this.messageDom.classList.toggle('reasoning_hidden', this.#isHidden);
+    }
+    updateReasoning(reasoning = null) {
+        reasoning = reasoning ?? this.reasoning;
+        this.reasoning = getRegexedString(reasoning ?? '', regex_placement.REASONING);
+    }
+    /**
+     * Processes and updates reasoning info for the message.
+     * @param {number} messageId - The ID of the message.
+     * @param {boolean} mesChanged - True if the message text changed.
+     * @param {Date} currentTime - The current time.
+     */
+    async process(messageId, mesChanged, currentTime) {
+        if (!this.reasoning && !this.#isHidden) return;
+        this.updateReasoning();
+        // Ensure the chat extra exists.
+        if (!chat[messageId]['extra']) {
+            chat[messageId]['extra'] = {};
+        }
+        const extra = chat[messageId]['extra'];
+        const finalReasoning = power_user.trim_spaces ? this.reasoning.trim() : this.reasoning;
+        const reasoningChanged = extra['reasoning'] !== finalReasoning;
+        extra['reasoning'] = finalReasoning;
+        if ((this.#isHidden || reasoningChanged) && this.startTime === null) {
+            this.startTime = this.initialTime;
+        }
+        if ((this.#isHidden || !reasoningChanged) && mesChanged && this.startTime !== null && this.endTime === null) {
+            this.endTime = currentTime;
+            await eventSource.emit(event_types.STREAM_REASONING_DONE, finalReasoning, () => this.getDuration());
+        }
+        await this.updateTime(messageId);
+        if (this.messageReasoningContentDom instanceof HTMLElement) {
+            const formattedReasoning = messageFormatting(finalReasoning, '', false, false, messageId, {}, true);
+            this.messageReasoningContentDom.innerHTML = formattedReasoning;
+        }
+        if (this.messageDom instanceof HTMLElement) {
+            this.messageDom.classList.add('reasoning');
+        }
+    }
+    async finish(messageId) {
+        // Make sure the finish time is recorded if a reasoning was in process and it wasn't ended correctly during streaming
+        if (this.startTime !== null && this.endTime === null) {
+            this.endTime = new Date();
+            const finalReasoning = power_user.trim_spaces ? this.reasoning.trim() : this.reasoning;
+            await eventSource.emit(event_types.STREAM_REASONING_DONE, finalReasoning, () => this.getDuration());
+            await this.updateTime(messageId);
+        }
+    }
+    /**
+     * Updates the reasoning duration in the UI.
+     * @param {number} messageId - The ID of the message
+     * @param {object} [options={}] - Optional argument
+     * @param {boolean} [options.forceEnd=false] - If true, there will be no "Thinking..." when no duration exists
+     */
+    async updateTime(messageId, { forceEnd = false } = {}) {
+        const duration = this.getDuration();
+        chat[messageId]['extra']['reasoning_duration'] = duration;
+        updateReasoningUI(this.messageDom, this.reasoning, duration, { forceEnd });
+    }
+}
 /**
  * Helper class for adding reasoning to messages.
  * Keeps track of the number of reasoning additions.