749c5df2 · permissionBRICK/SillyTavern

 export async function generateQuietPrompt(quiet_prompt, quietToLoud, skipWIAN, quietImage = null, quietName = null, responseLength = null, force_chid = null) {
     console.log('got into genQuietPrompt');
     const responseLengthCustomized = typeof responseLength === 'number' && responseLength > 0;
-    let originalResponseLength = -1;
+    let eventHook = () => {};
     try {
         /** @type {GenerateOptions} */
         const options = {

             quietName: quietName,
             force_chid: force_chid,
         };
-        originalResponseLength = responseLengthCustomized ? saveResponseLength(main_api, responseLength) : -1;
+        if (responseLengthCustomized) {
+            TempResponseLength.save(main_api, responseLength);
+            eventHook = TempResponseLength.setupEventHook(main_api);
+        }
         return await Generate('quiet', options);
     } finally {
-        if (responseLengthCustomized) {
+        if (responseLengthCustomized && TempResponseLength.isCustomized()) {
-            restoreResponseLength(main_api, originalResponseLength);
+            TempResponseLength.restore(main_api);
+            TempResponseLength.removeEventHook(main_api, eventHook);
+        }
+    }
+}

     const abortController = new AbortController();
     const responseLengthCustomized = typeof responseLength === 'number' && responseLength > 0;
-    let originalResponseLength = -1;
     const isInstruct = power_user.instruct.enabled && api !== 'openai' && api !== 'novel' && !instructOverride;
     const isQuiet = true;
+    let eventHook = () => {};
     if (systemPrompt) {
         systemPrompt = substituteParams(systemPrompt);

     prompt = isInstruct ? (prompt + formatInstructModePrompt(name2, false, '', name1, name2, isQuiet, quietToLoud)) : (prompt + '\n');
     try {
-        originalResponseLength = responseLengthCustomized ? saveResponseLength(api, responseLength) : -1;
+        if (responseLengthCustomized) {
+            TempResponseLength.save(api, responseLength);
+        }
         let generateData = {};
         switch (api) {

                     const koboldSettings = koboldai_settings[koboldai_setting_names[preset_settings]];
                     generateData = getKoboldGenerationData(prompt, koboldSettings, amount_gen, max_context, isHorde, 'quiet');
+                }
+                TempResponseLength.restore(api);
                 break;
             case 'novel': {
                 const novelSettings = novelai_settings[novelai_setting_names[nai_settings.preset_settings_novel]];
                 generateData = getNovelGenerationData(prompt, novelSettings, amount_gen, false, false, null, 'quiet');
+                TempResponseLength.restore(api);
                 break;
+            }
             case 'textgenerationwebui':
                 generateData = getTextGenGenerationData(prompt, amount_gen, false, false, null, 'quiet');
+                TempResponseLength.restore(api);
                 break;
             case 'openai': {
                 generateData = [{ role: 'user', content: prompt.trim() }];
                 if (systemPrompt) {
                     generateData.unshift({ role: 'system', content: systemPrompt.trim() });
+                }
+                eventHook = TempResponseLength.setupEventHook(api);
             } break;
+        }

         return message;
     } finally {
-        if (responseLengthCustomized) {
+        if (responseLengthCustomized && TempResponseLength.isCustomized()) {
-            restoreResponseLength(api, originalResponseLength);
+            TempResponseLength.restore(api);
+            TempResponseLength.removeEventHook(api, eventHook);
+        }
+    }
+}
-/**
+class TempResponseLength {
- * Temporarily change the response length for the specified API.
+    static #originalResponseLength = -1;
- * @param {string} api API to use.
+    static #lastApi = null;
- * @param {number} responseLength Target response length.
- * @returns {number} The original response length.
+    static isCustomized() {
- */
+        return this.#originalResponseLength > -1;
-function saveResponseLength(api, responseLength) {
-    let oldValue = -1;
-    if (api === 'openai') {
-        oldValue = oai_settings.openai_max_tokens;
-        oai_settings.openai_max_tokens = responseLength;
-    } else {
-        oldValue = amount_gen;
-        amount_gen = responseLength;
+    }
-    return oldValue;
-}
-/**
+    /**
- * Restore the original response length for the specified API.
+     * Save the current response length for the specified API.
- * @param {string} api API to use.
+     * @param {string} api API identifier
- * @param {number} responseLength Target response length.
+     * @param {number} responseLength New response length
- * @returns {void}
+     */
- */
+    static save(api, responseLength) {
-function restoreResponseLength(api, responseLength) {
+        if (api === 'openai') {
-    if (api === 'openai') {
+            this.#originalResponseLength = oai_settings.openai_max_tokens;
-        oai_settings.openai_max_tokens = responseLength;
+            oai_settings.openai_max_tokens = responseLength;
-    } else {
+        } else {
-        amount_gen = responseLength;
+            this.#originalResponseLength = amount_gen;
+            amount_gen = responseLength;
+        }
+        this.#lastApi = api;
+        console.log('[TempResponseLength] Saved original response length:', TempResponseLength.#originalResponseLength);
+    }
+    /**
+     * Restore the original response length for the specified API.
+     * @param {string|null} api API identifier
+     * @returns {void}
+     */
+    static restore(api) {
+        if (this.#originalResponseLength === -1) {
+            return;
+        }
+        if (!api && this.#lastApi) {
+            api = this.#lastApi;
+        }
+        if (api === 'openai') {
+            oai_settings.openai_max_tokens = this.#originalResponseLength;
+        } else {
+            amount_gen = this.#originalResponseLength;
+        }
+        console.log('[TempResponseLength] Restored original response length:', this.#originalResponseLength);
+        this.#originalResponseLength = -1;
+        this.#lastApi = null;
+    }
+    /**
+     * Sets up an event hook to restore the original response length when the event is emitted.
+     * @param {string} api API identifier
+     * @returns {function(): void} Event hook function
+     */
+    static setupEventHook(api) {
+        const eventHook = () => {
+            if (this.isCustomized()) {
+                this.restore(api);
+            }
+        };
+        switch (api) {
+            case 'openai':
+                eventSource.once(event_types.CHAT_COMPLETION_SETTINGS_READY, eventHook);
+                break;
+            default:
+                eventSource.once(event_types.GENERATE_AFTER_DATA, eventHook);
+                break;
+        }
+        return eventHook;
+    }
+    /**
+     * Removes the event hook for the specified API.
+     * @param {string} api API identifier
+     * @param {function(): void} eventHook Previously set up event hook
+     */
+    static removeEventHook(api, eventHook) {
+        switch (api) {
+            case 'openai':
+                eventSource.removeListener(event_types.CHAT_COMPLETION_SETTINGS_READY, eventHook);
+                break;
+            default:
+                eventSource.removeListener(event_types.GENERATE_AFTER_DATA, eventHook);
+                break;
+        }
+    }
+}

         return;
+    }
+    if (TempResponseLength.isCustomized()) {
+        console.warn('Response length is currently being overridden. Restoring previous value before saving.');
+        TempResponseLength.restore(null);
+    }
     //console.log('Entering settings with name1 = '+name1);
     return jQuery.ajax({
         type: 'POST',

Restore temp response length as early as possible Fixes #3297