Clean up Generate(), part 2 (#1578)

* Move StreamingProcessor constructor to the top

Typical code style is to declare the constructor at the top of the class
definition.

* Remove removePrefix

cleanupMessage does this already.

* Make message_already_generated local

We can pass it into StreamingProcessor so it doesn't have to be a global
variable.

* Consolidate setting isStopped and abort signal

Various places were doing some combination of setting isStopped, calling
abort on the streaming processor's abort controller, and calling
onStopStreaming. Let's consolidate all that functionality into
onStopStreaming/onErrorStreaming.

* More cleanly separate streaming/nonstreaming paths

* Replace promise with async function w/ handlers

By using onSuccess and onError as promise handlers, we can use normal
control flow and don't need to remember to use try/catch blocks or call
onSuccess every time.

* Remove runGenerate

Placing the rest of the code in a separate function doesn't really do
anything for its structure.

* Move StreamingProcessor() into streaming code path

* Fix return from circuit breaker

* Fix non-streaming chat completion request

* Fix Horde generation and quiet unblocking

---------

Co-authored-by: Cohee <18619528+Cohee1207@users.noreply.github.com>

public/script.js

@@ -701,8 +701,6 @@ export let user_avatar = 'you.png';
 export var amount_gen = 80; //default max length of AI generated responses
 var max_context = 2048;
 
-var message_already_generated = '';
-
 var swipes = true;
 let extension_prompts = {};
 
@@ -2600,6 +2598,21 @@ function hideStopButton() {
 }
 
 class StreamingProcessor {
+    constructor(type, force_name2, timeStarted, messageAlreadyGenerated) {
+        this.result = '';
+        this.messageId = -1;
+        this.type = type;
+        this.force_name2 = force_name2;
+        this.isStopped = false;
+        this.isFinished = false;
+        this.generator = this.nullStreamingGeneration;
+        this.abortController = new AbortController();
+        this.firstMessageText = '...';
+        this.timeStarted = timeStarted;
+        this.messageAlreadyGenerated = messageAlreadyGenerated;
+        this.swipes = [];
+    }
+
     showMessageButtons(messageId) {
         if (messageId == -1) {
             return;
@@ -2635,32 +2648,16 @@ class StreamingProcessor {
         return messageId;
     }
 
-    removePrefix(text) {
-        const name1Marker = `${name1}: `;
-        const name2Marker = `${name2}: `;
-
-        if (text) {
-            if (text.startsWith(name1Marker)) {
-                text = text.replace(name1Marker, '');
-            }
-            if (text.startsWith(name2Marker)) {
-                text = text.replace(name2Marker, '');
-            }
-        }
-        return text;
-    }
-
     onProgressStreaming(messageId, text, isFinal) {
         const isImpersonate = this.type == 'impersonate';
         const isContinue = this.type == 'continue';
 
         if (!isImpersonate && !isContinue && Array.isArray(this.swipes) && this.swipes.length > 0) {
             for (let i = 0; i < this.swipes.length; i++) {
-                this.swipes[i] = cleanUpMessage(this.removePrefix(this.swipes[i]), false, false, true, this.stoppingStrings);
+                this.swipes[i] = cleanUpMessage(this.swipes[i], false, false, true, this.stoppingStrings);
             }
         }
 
-        text = this.removePrefix(text);
         let processedText = cleanUpMessage(text, isImpersonate, isContinue, !isFinal, this.stoppingStrings);
 
         // Predict unbalanced asterisks / quotes during streaming
@@ -2786,6 +2783,9 @@ class StreamingProcessor {
     }
 
     onErrorStreaming() {
+        this.abortController.abort();
+        this.isStopped = true;
+
         this.hideMessageButtons(this.messageId);
         $('#send_textarea').removeAttr('disabled');
         is_send_press = false;
@@ -2811,20 +2811,6 @@ class StreamingProcessor {
         throw new Error('Generation function for streaming is not hooked up');
     }
 
-    constructor(type, force_name2, timeStarted) {
-        this.result = '';
-        this.messageId = -1;
-        this.type = type;
-        this.force_name2 = force_name2;
-        this.isStopped = false;
-        this.isFinished = false;
-        this.generator = this.nullStreamingGeneration;
-        this.abortController = new AbortController();
-        this.firstMessageText = '...';
-        this.timeStarted = timeStarted;
-        this.swipes = [];
-    }
-
     async generate() {
         if (this.messageId == -1) {
             this.messageId = await this.onStartStreaming(this.firstMessageText);
@@ -2844,13 +2830,12 @@ class StreamingProcessor {
             for await (const { text, swipes } of this.generator()) {
                 timestamps.push(Date.now());
                 if (this.isStopped) {
-                    this.onStopStreaming();
                     return;
                 }
 
                 this.result = text;
                 this.swipes = swipes;
-                await sw.tick(() => this.onProgressStreaming(this.messageId, message_already_generated + text));
+                await sw.tick(() => this.onProgressStreaming(this.messageId, this.messageAlreadyGenerated + text));
             }
             const seconds = (timestamps[timestamps.length - 1] - timestamps[0]) / 1000;
             console.warn(`Stream stats: ${timestamps.length} tokens, ${seconds.toFixed(2)} seconds, rate: ${Number(timestamps.length / seconds).toFixed(2)} TPS`);
@@ -2858,7 +2843,6 @@ class StreamingProcessor {
         catch (err) {
             console.error(err);
             this.onErrorStreaming();
-            this.isStopped = true;
             return;
         }
 
@@ -2965,7 +2949,7 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
     const isInstruct = power_user.instruct.enabled && main_api !== 'openai';
     const isImpersonate = type == 'impersonate';
 
-    message_already_generated = isImpersonate ? `${name1}: ` : `${name2}: `;
+    let message_already_generated = isImpersonate ? `${name1}: ` : `${name2}: `;
 
     const interruptedByCommand = await processCommands($('#send_textarea').val(), type, dryRun);
 
@@ -3378,10 +3362,6 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
         let mesSend = [];
         console.debug('calling runGenerate');
 
-        if (!dryRun) {
-            streamingProcessor = isStreamingEnabled() && type !== 'quiet' ? new StreamingProcessor(type, force_name2, generation_started) : false;
-        }
-
         if (isContinue) {
             // Coping mechanism for OAI spacing
             const isForceInstruct = isOpenRouterWithInstruct();
@@ -3389,21 +3369,16 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                 cyclePrompt += ' ';
                 continue_mag += ' ';
             }
-
-            // Save reply does add cycle text to the prompt, so it's not needed here
-            streamingProcessor && (streamingProcessor.firstMessageText = '');
             message_already_generated = continue_mag;
         }
 
         const originalType = type;
-        return runGenerate(cyclePrompt);
 
-        async function runGenerate(cycleGenerationPrompt = '') {
         if (!dryRun) {
             is_send_press = true;
         }
 
-            generatedPromptCache += cycleGenerationPrompt;
+        generatedPromptCache += cyclePrompt;
         if (generatedPromptCache.length == 0 || type === 'continue') {
             console.debug('generating prompt');
             chatString = '';
@@ -3771,14 +3746,13 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
             }
         }
 
-            return new Promise(async (resolve, reject) => {
+        async function finishGenerating() {
-                if (true === dryRun) return onSuccess({ error: 'dryRun' });
+            if (dryRun) return { error: 'dryRun' };
 
             if (power_user.console_log_prompts) {
                 console.log(generate_data.prompt);
             }
 
-                let generate_url = getGenerateUrl(main_api);
             console.debug('rungenerate calling API');
 
             showStopButton();
@@ -3825,55 +3799,16 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
             }
 
             console.debug(`pushed prompt bits to itemizedPrompts array. Length is now: ${itemizedPrompts.length}`);
-                /** @type {Promise<any>} */
-                let streamingGeneratorPromise = Promise.resolve();
-
-                if (main_api == 'openai') {
-                    if (isStreamingEnabled() && type !== 'quiet') {
-                        streamingGeneratorPromise = sendOpenAIRequest(type, generate_data.prompt, streamingProcessor.abortController.signal);
-                    }
-                    else {
-                        sendOpenAIRequest(type, generate_data.prompt, abortController.signal).then(onSuccess).catch(onError);
-                    }
-                }
-                else if (main_api == 'koboldhorde') {
-                    generateHorde(finalPrompt, generate_data, abortController.signal, true).then(onSuccess).catch(onError);
-                }
-                else if (main_api == 'textgenerationwebui' && isStreamingEnabled() && type !== 'quiet') {
-                    streamingGeneratorPromise = generateTextGenWithStreaming(generate_data, streamingProcessor.abortController.signal);
-                }
-                else if (main_api == 'novel' && isStreamingEnabled() && type !== 'quiet') {
-                    streamingGeneratorPromise = generateNovelWithStreaming(generate_data, streamingProcessor.abortController.signal);
-                }
-                else if (main_api == 'kobold' && isStreamingEnabled() && type !== 'quiet') {
-                    streamingGeneratorPromise = generateKoboldWithStreaming(generate_data, streamingProcessor.abortController.signal);
-                }
-                else {
-                    try {
-                        const response = await fetch(generate_url, {
-                            method: 'POST',
-                            headers: getRequestHeaders(),
-                            cache: 'no-cache',
-                            body: JSON.stringify(generate_data),
-                            signal: abortController.signal,
-                        });
-
-                        if (!response.ok) {
-                            const error = await response.json();
-                            throw error;
-                        }
-
-                        const data = await response.json();
-                        onSuccess(data);
-                    } catch (error) {
-                        onError(error);
-                    }
-                }
 
             if (isStreamingEnabled() && type !== 'quiet') {
-                    try {
+                streamingProcessor = new StreamingProcessor(type, force_name2, generation_started, message_already_generated);
-                        const streamingGenerator = await streamingGeneratorPromise;
+                if (isContinue) {
-                        streamingProcessor.generator = streamingGenerator;
+                    // Save reply does add cycle text to the prompt, so it's not needed here
+                    streamingProcessor.firstMessageText = '';
+                }
+
+                streamingProcessor.generator = await sendStreamingRequest(type, generate_data);
+
                 hideSwipeButtons();
                 let getMessage = await streamingProcessor.generate();
                 let messageChunk = cleanUpMessage(getMessage, isImpersonate, isContinue, false);
@@ -3887,19 +3822,19 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                     streamingProcessor = null;
                     triggerAutoContinue(messageChunk, isImpersonate);
                 }
-                        resolve();
+            } else {
-                    } catch (err) {
+                return await sendGenerationRequest(type, generate_data);
-                        onError(err);
+            }
         }
 
-                }
+        return finishGenerating().then(onSuccess, onError);
 
         async function onSuccess(data) {
+            if (!data) return;
             let messageChunk = '';
 
             if (data.error == 'dryRun') {
                 generatedPromptCache = '';
-                        resolve();
                 return;
             }
 
@@ -3928,7 +3863,8 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                         await eventSource.emit(event_types.IMPERSONATE_READY, getMessage);
                     }
                     else if (type == 'quiet') {
-                                resolve(getMessage);
+                        unblockGeneration();
+                        return getMessage;
                     }
                     else {
                         // Without streaming we'll be having a full message on continuation. Treat it as a last chunk.
@@ -3948,21 +3884,18 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                     maxLoops ??= MAX_GENERATION_LOOPS;
 
                     if (maxLoops === 0) {
-                                reject(new Error('Generate circuit breaker interruption'));
                         if (type !== 'quiet') {
                             throwCircuitBreakerError();
                         }
-                                return;
+                        throw new Error('Generate circuit breaker interruption');
                     }
 
                     // regenerate with character speech reenforced
                     // to make sure we leave on swipe type while also adding the name2 appendage
-                            delay(1000).then(async () => {
+                    await delay(1000);
                     // The first await is for waiting for the generate to start. The second one is waiting for it to finish
                     const result = await await Generate(type, { automatic_trigger, force_name2: true, quiet_prompt, skipWIAN, force_chid, maxLoops: maxLoops - 1 });
-                                resolve(result);
+                    return result;
-                            });
-                            return;
                 }
 
                 if (power_user.auto_swipe) {
@@ -3989,7 +3922,6 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                         is_send_press = false;
                         swipe_right();
                         // TODO: do we want to resolve after an auto-swipe?
-                                resolve();
                         return;
                     }
                 }
@@ -3999,7 +3931,7 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                 if (data?.response) {
                     toastr.error(data.response, 'API Error');
                 }
-                        reject(data.response);
+                throw data?.response;
             }
 
             console.debug('/api/chats/save called by /Generate');
@@ -4010,7 +3942,6 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
             if (type !== 'quiet') {
                 triggerAutoContinue(messageChunk, isImpersonate);
             }
-                    resolve();
         }
 
         function onError(exception) {
@@ -4018,23 +3949,18 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
                 toastr.error(exception.error.message, 'Error', { timeOut: 10000, extendedTimeOut: 20000 });
             }
 
-                    reject(exception);
             unblockGeneration();
             console.log(exception);
             streamingProcessor = null;
+            throw exception;
         }
-            });
-
-        } //rungenerate ends
     } else {    //generate's primary loop ends, after this is error handling for no-connection or safety-id
         if (this_chid === undefined || this_chid === 'invalid-safety-id') {
             toastr.warning('Сharacter is not selected');
         }
         is_send_press = false;
     }
-
+}
-    //console.log('generate ending');
-} //generate ends
 
 function flushWIDepthInjections() {
     //prevent custom depth WI entries (which have unique random key names) from duplicating
@@ -4481,22 +4407,82 @@ function setInContextMessages(lastmsg, type) {
     }
 }
 
-function getGenerateUrl(api) {
+/**
-    let generate_url = '';
+ * Sends a non-streaming request to the API.
-    if (api == 'kobold') {
+ * @param {string} type Generation type
-        generate_url = '/api/backends/kobold/generate';
+ * @param {object} data Generation data
-    } else if (api == 'textgenerationwebui') {
+ * @returns {Promise<object>} Response data from the API
-        generate_url = '/api/backends/text-completions/generate';
+ */
-    } else if (api == 'novel') {
+async function sendGenerationRequest(type, data) {
-        generate_url = '/api/novelai/generate';
+    if (main_api === 'openai') {
+        return await sendOpenAIRequest(type, data.prompt, abortController.signal);
+    }
+
+    if (main_api === 'koboldhorde') {
+        return await generateHorde(data.prompt, data, abortController.signal, true);
+    }
+
+    const response = await fetch(getGenerateUrl(main_api), {
+        method: 'POST',
+        headers: getRequestHeaders(),
+        cache: 'no-cache',
+        body: JSON.stringify(data),
+        signal: abortController.signal,
+    });
+
+    if (!response.ok) {
+        const error = await response.json();
+        throw error;
+    }
+
+    const responseData = await response.json();
+    return responseData;
+}
+
+/**
+ * Sends a streaming request to the API.
+ * @param {string} type Generation type
+ * @param {object} data Generation data
+ * @returns {Promise<any>} Streaming generator
+ */
+async function sendStreamingRequest(type, data) {
+    switch (main_api) {
+        case 'openai':
+            return await sendOpenAIRequest(type, data.prompt, streamingProcessor.abortController.signal);
+        case 'textgenerationwebui':
+            return await generateTextGenWithStreaming(data, streamingProcessor.abortController.signal);
+        case 'novel':
+            return await generateNovelWithStreaming(data, streamingProcessor.abortController.signal);
+        case 'kobold':
+            return await generateKoboldWithStreaming(data, streamingProcessor.abortController.signal);
+        default:
+            throw new Error('Streaming is enabled, but the current API does not support streaming.');
+    }
+}
+
+/**
+ * Gets the generation endpoint URL for the specified API.
+ * @param {string} api API name
+ * @returns {string} Generation URL
+ */
+function getGenerateUrl(api) {
+    switch (api) {
+        case 'kobold':
+            return '/api/backends/kobold/generate';
+        case 'koboldhorde':
+            return '/api/backends/koboldhorde/generate';
+        case 'textgenerationwebui':
+            return '/api/backends/text-completions/generate';
+        case 'novel':
+            return '/api/novelai/generate';
+        default:
+            throw new Error(`Unknown API: ${api}`);
     }
-    return generate_url;
 }
 
 function throwCircuitBreakerError() {
     callPopup(`Could not extract reply in ${MAX_GENERATION_LOOPS} attempts. Try generating again`, 'text');
     unblockGeneration();
-    throw new Error('Generate circuit breaker interruption');
 }
 
 function extractTitleFromData(data) {
@@ -7198,7 +7184,7 @@ function swipe_left() {      // when we swipe left..but no generation.
     }
 
     if (isStreamingEnabled() && streamingProcessor) {
-        streamingProcessor.isStopped = true;
+        streamingProcessor.onStopStreaming();
     }
 
     const swipe_duration = 120;
@@ -9327,8 +9313,6 @@ jQuery(async function () {
 
     $(document).on('click', '.mes_stop', function () {
         if (streamingProcessor) {
-            streamingProcessor.abortController.abort();
-            streamingProcessor.isStopped = true;
             streamingProcessor.onStopStreaming();
             streamingProcessor = null;
         }
@@ -9583,7 +9567,7 @@ jQuery(async function () {
         cancelTtsPlay();
         if (streamingProcessor) {
             console.log('Page reloaded. Aborting streaming...');
-            streamingProcessor.abortController.abort();
+            streamingProcessor.onStopStreaming();
         }
     });
 

public/style.css

@@ -637,6 +637,9 @@ hr {
     order: 2;
     padding-right: 2px;
     place-self: center;
+    cursor: pointer;
+    transition: 0.3s;
+    opacity: 0.7;
 }
 
 #options_button {