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
2 files changed, +90 -20Ignore whitespace
public/script.js+15 -6
@@ -3204,13 +3204,21 @@ class StreamingProcessor {
32043204 this.promptReasoning = promptReasoning;
32053205 }
32063206
3207- #checkDomElements(messageId) {
3207+ /**
3208+ * Initializes DOM elements for the current message.
3209+ * @param {number} messageId Current message ID
3210+ * @param {boolean?} continueOnReasoning If continuing on reasoning
3211+ */
3212+ async #checkDomElements(messageId, continueOnReasoning = null) {
32083213 if (this.messageDom === null || this.messageTextDom === null) {
32093214 this.messageDom = document.querySelector(`#chat .mes[mesid="${messageId}"]`);
32103215 this.messageTextDom = this.messageDom?.querySelector('.mes_text');
32113216 this.messageTimerDom = this.messageDom?.querySelector('.mes_timer');
32123217 this.messageTokenCounterDom = this.messageDom?.querySelector('.tokenCounterDisplay');
32133218 }
3219+ if (continueOnReasoning) {
3220+ await this.reasoningHandler.process(messageId, false, this.promptReasoning);
3221+ }
32143222 this.reasoningHandler.updateDom(messageId);
32153223 }
32163224
@@ -3230,7 +3238,8 @@ class StreamingProcessor {
32303238 }
32313239
32323240 async onStartStreaming(text) {
32333241 ifconst continueOnReasoning = !!(this.type === 'continue' && this.promptReasoning.prefixReasoning) {;
3242+ if (continueOnReasoning) {
32343243 this.reasoningHandler.initContinue(this.promptReasoning);
32353244 }
32363245
@@ -3242,7 +3251,7 @@ class StreamingProcessor {
32423251 } else {
32433252 await saveReply(this.type, text, true, '', [], '');
32443253 messageId = chat.length - 1;
32453254 await this.#checkDomElements(messageId, continueOnReasoning);
32463255 this.markUIGenStarted();
32473256 }
32483257 hideSwipeButtons();
@@ -3275,7 +3284,7 @@ class StreamingProcessor {
32753284 this.sendTextarea.dispatchEvent(new Event('input', { bubbles: true }));
32763285 } else {
32773286 const mesChanged = chat[messageId]['mes'] !== processedText;
32783287 await this.#checkDomElements(messageId);
32793288 this.#updateMessageBlockVisibility();
32803289 const currentTime = new Date();
32813290 chat[messageId]['mes'] = processedText;
@@ -3287,7 +3296,7 @@ class StreamingProcessor {
32873296 chat[messageId]['extra']['time_to_first_token'] = this.timeToFirstToken;
32883297
32893298 // Update reasoning
32903299 await this.reasoningHandler.process(messageId, mesChanged, this.promptReasoning);
32913300 processedText = chat[messageId]['mes'];
32923301
32933302 // Token count update.
@@ -5957,7 +5966,7 @@ export function cleanUpMessage(getMessage, isImpersonate, isContinue, displayInc
59575966 getMessage = trimToEndSentence(getMessage);
59585967 }
59595968
59605969 if (power_user.trim_spaces && !PromptReasoning.getLatestPrefix()) {
59615970 getMessage = getMessage.trim();
59625971 }
59635972
public/scripts/reasoning.js+75 -14
@@ -196,7 +196,7 @@ export class ReasoningHandler {
196196 */
197197 initContinue(promptReasoning) {
198198 this.reasoning = promptReasoning.prefixReasoning;
199199 this.state = promptReasoning.prefixIncomplete ? ReasoningState.None : ReasoningState.Done;
200200 this.startTime = this.initialTime;
201201 this.endTime = promptReasoning.prefixDuration ? new Date(this.initialTime.getTime() + promptReasoning.prefixDuration) : null;
202202 }
@@ -324,10 +324,11 @@ export class ReasoningHandler {
324324 *
325325 * @param {number} messageId - The ID of the message to process
326326 * @param {boolean} mesChanged - Whether the message has changed
327+ * @param {PromptReasoning} promptReasoning - Prompt reasoning object
327328 * @returns {Promise<void>}
328329 */
329330 async process(messageId, mesChanged, promptReasoning) {
330331 mesChanged = this.#autoParseReasoningFromMessage(messageId, mesChanged, promptReasoning);
331332
332333 if (!this.reasoning && !this.#isHiddenReasoningModel)
333334 return;
@@ -345,7 +346,14 @@ export class ReasoningHandler {
345346 }
346347 }
347348
348- #autoParseReasoningFromMessage(messageId, mesChanged) {
349+ /**
350+ * Parse reasoning from a message during streaming.
351+ * @param {number} messageId Message ID
352+ * @param {boolean} mesChanged Whether the message has changed before reasoning parsing
353+ * @param {PromptReasoning} promptReasoning Prompt reasoning object
354+ * @returns {boolean} Whether the message has changed after reasoning parsing
355+ */
356+ #autoParseReasoningFromMessage(messageId, mesChanged, promptReasoning) {
349357 if (!power_user.reasoning.auto_parse)
350358 return;
351359 if (!power_user.reasoning.prefix || !power_user.reasoning.suffix)
@@ -355,15 +363,17 @@ export class ReasoningHandler {
355363 const message = chat[messageId];
356364 if (!message) return mesChanged;
357365
366+ const parseTarget = promptReasoning?.prefixIncomplete ? (promptReasoning.prefixReasoningFormatted + message.mes) : message.mes;
367+
358368 // If we are done with reasoning parse, we just split the message correctly so the reasoning doesn't show up inside of it.
359369 if (this.#parsingReasoningMesStartIndex) {
360370 message.mes = trimSpaces(message.mesparseTarget.slice(this.#parsingReasoningMesStartIndex));
361371 return mesChanged;
362372 }
363373
364374 if (this.state === ReasoningState.None || this.#isHiddenReasoningModel) {
365375 // If streamed message starts with the opening, cut it out and put all inside reasoning
366376 if (message.mesparseTarget.startsWith(power_user.reasoning.prefix) && message.mesparseTarget.length > power_user.reasoning.prefix.length) {
367377 this.#isParsingReasoning = true;
368378
369379 // Manually set starting state here, as we might already have received the ending suffix
@@ -377,15 +387,14 @@ export class ReasoningHandler {
377387 return mesChanged;
378388
379389 // If we are in manual parsing mode, all currently streaming mes tokens will go the the reasoning block
380- const originalMes = message.mes;
390+ this.reasoning = parseTarget.slice(power_user.reasoning.prefix.length);
381- this.reasoning = originalMes.slice(power_user.reasoning.prefix.length);
382391 message.mes = '';
383392
384393 // If the reasoning contains the ending suffix, we cut that off and continue as message streaming
385394 if (this.reasoning.includes(power_user.reasoning.suffix)) {
386395 this.reasoning = this.reasoning.slice(0, this.reasoning.indexOf(power_user.reasoning.suffix));
387396 this.#parsingReasoningMesStartIndex = originalMesparseTarget.indexOf(power_user.reasoning.suffix) + power_user.reasoning.suffix.length;
388397 message.mes = trimSpaces(originalMesparseTarget.slice(this.#parsingReasoningMesStartIndex));
389398 this.#isParsingReasoning = false;
390399 }
391400
@@ -525,13 +534,56 @@ export class ReasoningHandler {
525534 * Keeps track of the number of reasoning additions.
526535 */
527536export class PromptReasoning {
537+ /**
538+ * An instance initiated during the latest prompt processing.
539+ * @type {PromptReasoning}
540+ * */
541+ static #LATEST = null;
542+ /**
543+ * @readonly Zero-width space character used as a placeholder for reasoning.
544+ * @type {string}
545+ */
528546 static REASONING_PLACEHOLDER = '\u200B';
529547
548+ /**
549+ * Returns the latest formatted reasoning prefix if the prefix is incomplete.
550+ * @returns {string} Formatted reasoning prefix
551+ */
552+ static getLatestPrefix() {
553+ if (!PromptReasoning.#LATEST) {
554+ return '';
555+ }
556+
557+ if (!PromptReasoning.#LATEST.prefixIncomplete) {
558+ return '';
559+ }
560+
561+ return PromptReasoning.#LATEST.prefixReasoningFormatted;
562+ }
563+
564+ /**
565+ * Free the latest reasoning instance.
566+ * To be called when the generation has ended or stopped.
567+ */
568+ static clearLatest() {
569+ PromptReasoning.#LATEST = null;
570+ }
571+
530572 constructor() {
573+ PromptReasoning.#LATEST = this;
574+
575+ /** @type {number} */
531576 this.counter = 0;
577+ /** @type {number} */
532578 this.prefixLength = -1;
579+ /** @type {string} */
533580 this.prefixReasoning = '';
581+ /** @type {string} */
582+ this.prefixReasoningFormatted = '';
583+ /** @type {number?} */
534584 this.prefixDuration = null;
585+ /** @type {boolean} */
586+ this.prefixIncomplete = false;
535587 }
536588
537589 /**
@@ -578,8 +630,10 @@ export class PromptReasoning {
578630 const formattedReasoning = `${prefix}${reasoning}`;
579631 if (isPrefix) {
580632 this.prefixReasoning = reasoning;
633+ this.prefixReasoningFormatted = formattedReasoning;
581634 this.prefixLength = formattedReasoning.length;
582635 this.prefixDuration = duration;
636+ this.prefixIncomplete = true;
583637 }
584638 return formattedReasoning;
585639 }
@@ -588,8 +642,10 @@ export class PromptReasoning {
588642 const formattedReasoning = `${prefix}${reasoning}${suffix}${separator}`;
589643 if (isPrefix) {
590644 this.prefixReasoning = reasoning;
645+ this.prefixReasoningFormatted = formattedReasoning;
591646 this.prefixLength = formattedReasoning.length;
592647 this.prefixDuration = duration;
648+ this.prefixIncomplete = false;
593649 }
594650 return `${formattedReasoning}${content}`;
595651 }
@@ -1075,12 +1131,13 @@ export function parseReasoningInSwipes(swipes, swipeInfoArray, duration) {
10751131}
10761132
10771133function registerReasoningAppEvents() {
10781134 const eventHandler = (/** @type {string} */ type, /** @type {number} */ idx) => {
10791135 if (!power_user.reasoning.auto_parse) {
10801136 return;
10811137 }
10821138
10831139 console.debug('[Reasoning] Auto-parsing reasoning block for message', idx);
1140+ const prefix = type === event_types.MESSAGE_RECEIVED ? PromptReasoning.getLatestPrefix() : '';
10841141 const message = chat[idx];
10851142
10861143 if (!message) {
@@ -1093,12 +1150,12 @@ function registerReasoningAppEvents() {
10931150 return null;
10941151 }
10951152
10961153 if (message.extra?.reasoning && !prefix) {
10971154 console.debug('[Reasoning] Message already has reasoning', idx);
10981155 return null;
10991156 }
11001157
11011158 const parsedReasoning = parseReasoningFromString(prefix + message.mes);
11021159
11031160 // No reasoning block found
11041161 if (!parsedReasoning) {
@@ -1137,7 +1194,11 @@ function registerReasoningAppEvents() {
11371194 };
11381195
11391196 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());
11411202 }
11421203}
11431204