Restore temp response length as early as possible Fixes #3297

749c5df29aa352f1f68f7a716ebc8bf2ed42cae6

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: 4322197aba621eb22dadbc0571a96add8176ae5b
1 files changed, +109 -35Ignore whitespace
public/script.js+109 -35
@@ -2724,7 +2724,7 @@ export function getStoppingStrings(isImpersonate, isContinue) {
27242724export async function generateQuietPrompt(quiet_prompt, quietToLoud, skipWIAN, quietImage = null, quietName = null, responseLength = null, force_chid = null) {
27252725 console.log('got into genQuietPrompt');
27262726 const responseLengthCustomized = typeof responseLength === 'number' && responseLength > 0;
27272727 let originalResponseLengtheventHook = -1() => {};
27282728 try {
27292729 /** @type {GenerateOptions} */
27302730 const options = {
@@ -2736,11 +2736,15 @@ export async function generateQuietPrompt(quiet_prompt, quietToLoud, skipWIAN, q
27362736 quietName: quietName,
27372737 force_chid: force_chid,
27382738 };
2739- originalResponseLength = responseLengthCustomized ? saveResponseLength(main_api, responseLength) : -1;
2739+ if (responseLengthCustomized) {
2740+ TempResponseLength.save(main_api, responseLength);
2741+ eventHook = TempResponseLength.setupEventHook(main_api);
2742+ }
27402743 return await Generate('quiet', options);
27412744 } finally {
27422745 if (responseLengthCustomized && TempResponseLength.isCustomized()) {
27432746 restoreResponseLengthTempResponseLength.restore(main_api, originalResponseLength);
2747+ TempResponseLength.removeEventHook(main_api, eventHook);
27442748 }
27452749 }
27462750}
@@ -3384,9 +3388,9 @@ export async function generateRaw(prompt, api, instructOverride, quietToLoud, sy
33843388
33853389 const abortController = new AbortController();
33863390 const responseLengthCustomized = typeof responseLength === 'number' && responseLength > 0;
3387- let originalResponseLength = -1;
33883391 const isInstruct = power_user.instruct.enabled && api !== 'openai' && api !== 'novel' && !instructOverride;
33893392 const isQuiet = true;
3393+ let eventHook = () => {};
33903394
33913395 if (systemPrompt) {
33923396 systemPrompt = substituteParams(systemPrompt);
@@ -3400,7 +3404,9 @@ export async function generateRaw(prompt, api, instructOverride, quietToLoud, sy
34003404 prompt = isInstruct ? (prompt + formatInstructModePrompt(name2, false, '', name1, name2, isQuiet, quietToLoud)) : (prompt + '\n');
34013405
34023406 try {
3403- originalResponseLength = responseLengthCustomized ? saveResponseLength(api, responseLength) : -1;
3407+ if (responseLengthCustomized) {
3408+ TempResponseLength.save(api, responseLength);
3409+ }
34043410 let generateData = {};
34053411
34063412 switch (api) {
@@ -3413,20 +3419,24 @@ export async function generateRaw(prompt, api, instructOverride, quietToLoud, sy
34133419 const koboldSettings = koboldai_settings[koboldai_setting_names[preset_settings]];
34143420 generateData = getKoboldGenerationData(prompt, koboldSettings, amount_gen, max_context, isHorde, 'quiet');
34153421 }
3422+ TempResponseLength.restore(api);
34163423 break;
34173424 case 'novel': {
34183425 const novelSettings = novelai_settings[novelai_setting_names[nai_settings.preset_settings_novel]];
34193426 generateData = getNovelGenerationData(prompt, novelSettings, amount_gen, false, false, null, 'quiet');
3427+ TempResponseLength.restore(api);
34203428 break;
34213429 }
34223430 case 'textgenerationwebui':
34233431 generateData = getTextGenGenerationData(prompt, amount_gen, false, false, null, 'quiet');
3432+ TempResponseLength.restore(api);
34243433 break;
34253434 case 'openai': {
34263435 generateData = [{ role: 'user', content: prompt.trim() }];
34273436 if (systemPrompt) {
34283437 generateData.unshift({ role: 'system', content: systemPrompt.trim() });
34293438 }
3439+ eventHook = TempResponseLength.setupEventHook(api);
34303440 } break;
34313441 }
34323442
@@ -3468,41 +3478,100 @@ export async function generateRaw(prompt, api, instructOverride, quietToLoud, sy
34683478
34693479 return message;
34703480 } finally {
34713481 if (responseLengthCustomized && TempResponseLength.isCustomized()) {
34723482 restoreResponseLengthTempResponseLength.restore(api, originalResponseLength);
3483+ TempResponseLength.removeEventHook(api, eventHook);
34733484 }
34743485 }
34753486}
34763487
3477-/**
3488+class TempResponseLength {
3478- * Temporarily change the response length for the specified API.
3489+ static #originalResponseLength = -1;
3479- * @param {string} api API to use.
3490+ static #lastApi = null;
3480- * @param {number} responseLength Target response length.
3491+
3481- * @returns {number} The original response length.
3492+ static isCustomized() {
3482- */
3493+ return this.#originalResponseLength > -1;
3483-function saveResponseLength(api, responseLength) {
3484- let oldValue = -1;
3485- if (api === 'openai') {
3486- oldValue = oai_settings.openai_max_tokens;
3487- oai_settings.openai_max_tokens = responseLength;
3488- } else {
3489- oldValue = amount_gen;
3490- amount_gen = responseLength;
34913494 }
3492- return oldValue;
3493-}
34943495
34953496 /**
34963497 * RestoreSave the originalcurrent response length for the specified API.
34973498 * @param {string} api API to use.identifier
34983499 * @param {number} responseLength TargetNew response length.
3499- * @returns {void}
3500+ */
3500- */
3501+ static save(api, responseLength) {
35013502function if restoreResponseLength(api, responseLength=== 'openai') {
3502- if (api === 'openai') {
3503+ this.#originalResponseLength = oai_settings.openai_max_tokens;
35033504 oai_settings.openai_max_tokens = responseLength;
35043505 } else {
3505- amount_gen = responseLength;
3506+ this.#originalResponseLength = amount_gen;
3507+ amount_gen = responseLength;
3508+ }
3509+
3510+ this.#lastApi = api;
3511+ console.log('[TempResponseLength] Saved original response length:', TempResponseLength.#originalResponseLength);
3512+ }
3513+
3514+ /**
3515+ * Restore the original response length for the specified API.
3516+ * @param {string|null} api API identifier
3517+ * @returns {void}
3518+ */
3519+ static restore(api) {
3520+ if (this.#originalResponseLength === -1) {
3521+ return;
3522+ }
3523+ if (!api && this.#lastApi) {
3524+ api = this.#lastApi;
3525+ }
3526+ if (api === 'openai') {
3527+ oai_settings.openai_max_tokens = this.#originalResponseLength;
3528+ } else {
3529+ amount_gen = this.#originalResponseLength;
3530+ }
3531+
3532+ console.log('[TempResponseLength] Restored original response length:', this.#originalResponseLength);
3533+ this.#originalResponseLength = -1;
3534+ this.#lastApi = null;
3535+ }
3536+
3537+ /**
3538+ * Sets up an event hook to restore the original response length when the event is emitted.
3539+ * @param {string} api API identifier
3540+ * @returns {function(): void} Event hook function
3541+ */
3542+ static setupEventHook(api) {
3543+ const eventHook = () => {
3544+ if (this.isCustomized()) {
3545+ this.restore(api);
3546+ }
3547+ };
3548+
3549+ switch (api) {
3550+ case 'openai':
3551+ eventSource.once(event_types.CHAT_COMPLETION_SETTINGS_READY, eventHook);
3552+ break;
3553+ default:
3554+ eventSource.once(event_types.GENERATE_AFTER_DATA, eventHook);
3555+ break;
3556+ }
3557+
3558+ return eventHook;
3559+ }
3560+
3561+ /**
3562+ * Removes the event hook for the specified API.
3563+ * @param {string} api API identifier
3564+ * @param {function(): void} eventHook Previously set up event hook
3565+ */
3566+ static removeEventHook(api, eventHook) {
3567+ switch (api) {
3568+ case 'openai':
3569+ eventSource.removeListener(event_types.CHAT_COMPLETION_SETTINGS_READY, eventHook);
3570+ break;
3571+ default:
3572+ eventSource.removeListener(event_types.GENERATE_AFTER_DATA, eventHook);
3573+ break;
3574+ }
35063575 }
35073576}
35083577
@@ -6871,6 +6940,11 @@ export async function saveSettings(type) {
68716940 return;
68726941 }
68736942
6943+ if (TempResponseLength.isCustomized()) {
6944+ console.warn('Response length is currently being overridden. Restoring previous value before saving.');
6945+ TempResponseLength.restore(null);
6946+ }
6947+
68746948 //console.log('Entering settings with name1 = '+name1);
68756949 return jQuery.ajax({
68766950 type: 'POST',