Add proper processing of streaming aborting

This commit is contained in:
SillyLossy
2023-04-21 20:29:18 +03:00
parent 9af7c63d9c
commit f25ecbd95c
6 changed files with 62 additions and 10 deletions

View File

@ -1273,6 +1273,7 @@ class StreamingProcessor {
this.isStopped = false;
this.isFinished = false;
this.generator = this.nullStreamingGeneration;
this.abortController = new AbortController();
}
async generate() {
@ -1927,7 +1928,7 @@ async function Generate(type, automatic_trigger, force_name2) {
let prompt = await prepareOpenAIMessages(name2, storyString, worldInfoBefore, worldInfoAfter, afterScenarioAnchor, promptBias, type);
if (isStreamingEnabled()) {
streamingProcessor.generator = await sendOpenAIRequest(prompt);
streamingProcessor.generator = await sendOpenAIRequest(prompt, streamingProcessor.abortController.signal);
}
else {
sendOpenAIRequest(prompt).then(onSuccess).catch(onError);
@ -1938,14 +1939,14 @@ async function Generate(type, automatic_trigger, force_name2) {
}
else if (main_api == 'poe') {
if (isStreamingEnabled()) {
streamingProcessor.generator = await generatePoe(type, finalPromt);
streamingProcessor.generator = await generatePoe(type, finalPromt, streamingProcessor.abortController.signal);
}
else {
generatePoe(type, finalPromt).then(onSuccess).catch(onError);
}
}
else if (main_api == 'textgenerationwebui' && textgenerationwebui_settings.streaming) {
streamingProcessor.generator = await generateTextGenWithStreaming(generate_data);
streamingProcessor.generator = await generateTextGenWithStreaming(generate_data, streamingProcessor.abortController.signal);
}
else {
jQuery.ajax({
@ -5013,6 +5014,7 @@ $(document).ready(function () {
$(document).on("click", ".mes_stop", function () {
if (streamingProcessor) {
streamingProcessor.abortController.abort();
streamingProcessor.isStopped = true;
streamingProcessor.onStopStreaming();
streamingProcessor = null;
@ -5106,4 +5108,11 @@ $(document).ready(function () {
}
});
});
$(document).on('beforeunload', () => {
if (streamingProcessor) {
console.log('Page reloaded. Aborting streaming...');
streamingProcessor.abortController.abort();
}
});
})