Fix streams getting stuck on regen

This commit is contained in:
Cohee 2024-07-25 08:51:49 +03:00
parent c12a283efc
commit 807487ce85
1 changed files with 20 additions and 5 deletions

View File

@ -2818,7 +2818,7 @@ class StreamingProcessor {
} }
#checkDomElements(messageId) { #checkDomElements(messageId) {
if (this.messageDom === null) { if (this.messageDom === null || this.messageTextDom === null) {
this.messageDom = document.querySelector(`#chat .mes[mesid="${messageId}"]`); this.messageDom = document.querySelector(`#chat .mes[mesid="${messageId}"]`);
this.messageTextDom = this.messageDom?.querySelector('.mes_text'); this.messageTextDom = this.messageDom?.querySelector('.mes_text');
this.messageTimerDom = this.messageDom?.querySelector('.mes_timer'); this.messageTimerDom = this.messageDom?.querySelector('.mes_timer');
@ -2890,6 +2890,7 @@ class StreamingProcessor {
this.sendTextarea.dispatchEvent(new Event('input', { bubbles: true })); this.sendTextarea.dispatchEvent(new Event('input', { bubbles: true }));
} }
else { else {
this.#checkDomElements(messageId);
const currentTime = new Date(); const currentTime = new Date();
// Don't waste time calculating token count for streaming // Don't waste time calculating token count for streaming
const currentTokenCount = isFinal && power_user.message_token_count_enabled ? getTokenCount(processedText, 0) : 0; const currentTokenCount = isFinal && power_user.message_token_count_enabled ? getTokenCount(processedText, 0) : 0;
@ -2921,7 +2922,6 @@ class StreamingProcessor {
chat[messageId].is_user, chat[messageId].is_user,
messageId, messageId,
); );
this.#checkDomElements(messageId);
if (this.messageTextDom instanceof HTMLElement) { if (this.messageTextDom instanceof HTMLElement) {
this.messageTextDom.innerHTML = formattedText; this.messageTextDom.innerHTML = formattedText;
} }
@ -3214,6 +3214,23 @@ function restoreResponseLength(api, responseLength) {
} }
} }
/**
* Removes last message from the chat DOM.
* @returns {Promise<void>} Resolves when the message is removed.
*/
function removeLastMessage() {
return new Promise((resolve) => {
const lastMes = $('#chat').children('.mes').last();
if (lastMes.length === 0) {
return resolve();
}
lastMes.hide(animation_duration, function () {
$(this).remove();
resolve();
});
});
}
/** /**
* Runs a generation using the current chat context. * Runs a generation using the current chat context.
* @param {string} type Generation type * @param {string} type Generation type
@ -3346,9 +3363,7 @@ export async function Generate(type, { automatic_trigger, force_name2, quiet_pro
} }
else if (type !== 'quiet' && type !== 'swipe' && !isImpersonate && !dryRun && chat.length) { else if (type !== 'quiet' && type !== 'swipe' && !isImpersonate && !dryRun && chat.length) {
chat.length = chat.length - 1; chat.length = chat.length - 1;
$('#chat').children().last().hide(250, function () { await removeLastMessage();
$(this).remove();
});
await eventSource.emit(event_types.MESSAGE_DELETED, chat.length); await eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
} }
} }