Merge pull request #3634 from SillyTavern/continue-from-reasoning

Fix auto-parsing of continue from reasoning
This commit is contained in:
Cohee 2025-03-09 16:22:02 +02:00 committed by GitHub
commit a392593e53
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 90 additions and 20 deletions

View File

@ -3204,13 +3204,21 @@ class StreamingProcessor {
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 {
}
async onStartStreaming(text) {
if (this.type === 'continue' && this.promptReasoning.prefixReasoning) {
const continueOnReasoning = !!(this.type === 'continue' && this.promptReasoning.prefixReasoning);
if (continueOnReasoning) {
this.reasoningHandler.initContinue(this.promptReasoning);
}
@ -3242,7 +3251,7 @@ class StreamingProcessor {
} else {
await saveReply(this.type, text, true, '', [], '');
messageId = chat.length - 1;
this.#checkDomElements(messageId);
await this.#checkDomElements(messageId, continueOnReasoning);
this.markUIGenStarted();
}
hideSwipeButtons();
@ -3275,7 +3284,7 @@ class StreamingProcessor {
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;
@ -3287,7 +3296,7 @@ class StreamingProcessor {
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
getMessage = trimToEndSentence(getMessage);
}
if (power_user.trim_spaces) {
if (power_user.trim_spaces && !PromptReasoning.getLatestPrefix()) {
getMessage = getMessage.trim();
}

View File

@ -196,7 +196,7 @@ export class ReasoningHandler {
*/
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;
}
@ -324,10 +324,11 @@ export class ReasoningHandler {
*
* @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) {
mesChanged = this.#autoParseReasoningFromMessage(messageId, mesChanged);
async process(messageId, mesChanged, promptReasoning) {
mesChanged = this.#autoParseReasoningFromMessage(messageId, mesChanged, promptReasoning);
if (!this.reasoning && !this.#isHiddenReasoningModel)
return;
@ -345,7 +346,14 @@ export class ReasoningHandler {
}
}
#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)
@ -355,15 +363,17 @@ export class ReasoningHandler {
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
@ -377,15 +387,14 @@ export class ReasoningHandler {
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 = originalMes.slice(power_user.reasoning.prefix.length);
this.reasoning = parseTarget.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;
message.mes = trimSpaces(originalMes.slice(this.#parsingReasoningMesStartIndex));
this.#parsingReasoningMesStartIndex = parseTarget.indexOf(power_user.reasoning.suffix) + power_user.reasoning.suffix.length;
message.mes = trimSpaces(parseTarget.slice(this.#parsingReasoningMesStartIndex));
this.#isParsingReasoning = false;
}
@ -525,13 +534,56 @@ export class ReasoningHandler {
* 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;
}
/**
@ -578,8 +630,10 @@ export class PromptReasoning {
const formattedReasoning = `${prefix}${reasoning}`;
if (isPrefix) {
this.prefixReasoning = reasoning;
this.prefixReasoningFormatted = formattedReasoning;
this.prefixLength = formattedReasoning.length;
this.prefixDuration = duration;
this.prefixIncomplete = true;
}
return formattedReasoning;
}
@ -588,8 +642,10 @@ export class PromptReasoning {
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}`;
}
@ -1075,12 +1131,13 @@ export function parseReasoningInSwipes(swipes, swipeInfoArray, duration) {
}
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() {
return null;
}
if (message.extra?.reasoning) {
if (message.extra?.reasoning && !prefix) {
console.debug('[Reasoning] Message already has reasoning', idx);
return null;
}
const parsedReasoning = parseReasoningFromString(message.mes);
const parsedReasoning = parseReasoningFromString(prefix + message.mes);
// No reasoning block found
if (!parsedReasoning) {
@ -1137,7 +1194,11 @@ function registerReasoningAppEvents() {
};
for (const event of [event_types.MESSAGE_RECEIVED, event_types.MESSAGE_UPDATED]) {
eventSource.on(event, eventHandler);
eventSource.on(event, (/** @type {number} */ idx) => eventHandler(event, idx));
}
for (const event of [event_types.GENERATION_STOPPED, event_types.GENERATION_ENDED, event_types.CHAT_CHANGED]) {
eventSource.on(event, () => PromptReasoning.clearLatest());
}
}