From 7221549c65f51501ac4e11fce4373b05b6c372b4 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Fri, 5 Apr 2024 01:25:48 +0300
Subject: [PATCH] #2013 Fix smooth stream event processing.

---
 public/scripts/sse-stream.js | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/public/scripts/sse-stream.js b/public/scripts/sse-stream.js
index 8e6a16f7c..9e335600d 100644
--- a/public/scripts/sse-stream.js
+++ b/public/scripts/sse-stream.js
@@ -106,7 +106,7 @@ function getDelay(s) {
 /**
  * Parses the stream data and returns the parsed data and the chunk to be sent.
  * @param {object} json The JSON data.
- * @returns {AsyncGenerator<{data: object, chunk: string} | null>} The parsed data and the chunk to be sent.
+ * @returns {AsyncGenerator<{data: object, chunk: string}>} The parsed data and the chunk to be sent.
  */
 async function* parseStreamData(json) {
     // Claude
@@ -120,6 +120,7 @@ async function* parseStreamData(json) {
                 };
             }
         }
+        return;
     }
     // MakerSuite
     else if (Array.isArray(json.candidates)) {
@@ -145,6 +146,7 @@ async function* parseStreamData(json) {
                 }
             }
         }
+        return;
     }
     // NovelAI / KoboldCpp Classic
     else if (typeof json.token === 'string' && json.token.length > 0) {
@@ -155,6 +157,7 @@ async function* parseStreamData(json) {
                 chunk: str,
             };
         }
+        return;
     }
     // llama.cpp?
     else if (typeof json.content === 'string' && json.content.length > 0) {
@@ -165,6 +168,7 @@ async function* parseStreamData(json) {
                 chunk: str,
             };
         }
+        return;
     }
     // OpenAI-likes
     else if (Array.isArray(json.choices)) {
@@ -184,6 +188,7 @@ async function* parseStreamData(json) {
                     chunk: str,
                 };
             }
+            return;
         }
         else if (typeof json.choices[0].delta === 'object') {
             if (typeof json.choices[0].delta.text === 'string' && json.choices[0].delta.text.length > 0) {
@@ -197,6 +202,7 @@ async function* parseStreamData(json) {
                         chunk: str,
                     };
                 }
+                return;
             }
             else if (typeof json.choices[0].delta.content === 'string' && json.choices[0].delta.content.length > 0) {
                 for (let j = 0; j < json.choices[0].delta.content.length; j++) {
@@ -209,6 +215,7 @@ async function* parseStreamData(json) {
                         chunk: str,
                     };
                 }
+                return;
             }
         }
         else if (typeof json.choices[0].message === 'object') {
@@ -223,11 +230,12 @@ async function* parseStreamData(json) {
                         chunk: str,
                     };
                 }
+                return;
             }
         }
     }
 
-    yield null;
+    throw new Error('Unknown event data format');
 }
 
 /**
@@ -257,11 +265,6 @@ export class SmoothEventSourceStream extends EventSourceStream {
                     }
 
                     for await (const parsed of parseStreamData(json)) {
-                        if (!parsed) {
-                            lastStr = '';
-                            return controller.enqueue(event);
-                        }
-
                         hasFocus && await delay(getDelay(lastStr));
                         controller.enqueue(new MessageEvent(event.type, { data: JSON.stringify(parsed.data) }));
                         lastStr = parsed.chunk;