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, isHiddenReasoningModelPromptReasoning, PromptReasoningReasoningHandler, 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.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.#reasoningDurationreasoningHandler.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	3377	if (logprobs) {
3444	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	3382	await eventSource.emit(event_types.STREAM_TOKEN_RECEIVED, text);
3448	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 } = {}) {
     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.