From 18f84979f2cd7576d4fdd51cbbb5f2a09ab7fff5 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sat, 10 Feb 2024 02:43:50 +0200
Subject: [PATCH] Use SSE streaming for MakerSuite

---
 public/scripts/openai.js                   | 19 ++++--------
 src/endpoints/backends/chat-completions.js | 34 ++--------------------
 2 files changed, 9 insertions(+), 44 deletions(-)

diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index c9e82ee76..aa100e695 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -1692,24 +1692,17 @@ async function sendOpenAIRequest(type, messages, signal) {
         throw new Error(`Got response status ${response.status}`);
     }
     if (stream) {
-        let reader;
-        let isSSEStream = oai_settings.chat_completion_source !== chat_completion_sources.MAKERSUITE;
-        if (isSSEStream) {
-            const eventStream = new EventSourceStream();
-            response.body.pipeThrough(eventStream);
-            reader = eventStream.readable.getReader();
-        } else {
-            reader = response.body.getReader();
-        }
+        const eventStream = new EventSourceStream();
+        response.body.pipeThrough(eventStream);
+        const reader = eventStream.readable.getReader();
         return async function* streamData() {
             let text = '';
-            let utf8Decoder = new TextDecoder();
             const swipes = [];
             while (true) {
                 const { done, value } = await reader.read();
                 if (done) return;
-                const rawData = isSSEStream ? value.data : utf8Decoder.decode(value, { stream: true });
-                if (isSSEStream && rawData === '[DONE]') return;
+                const rawData = value.data;
+                if (rawData === '[DONE]') return;
                 tryParseStreamingError(response, rawData);
                 const parsed = JSON.parse(rawData);
 
@@ -1750,7 +1743,7 @@ function getStreamingReply(data) {
     if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
         return data?.completion || '';
     } else if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {
-        return data?.candidates[0].content.parts[0].text || '';
+        return data?.candidates?.[0]?.content?.parts?.[0]?.text || '';
     } else {
         return data.choices[0]?.delta?.content || data.choices[0]?.message?.content || data.choices[0]?.text || '';
     }
diff --git a/src/endpoints/backends/chat-completions.js b/src/endpoints/backends/chat-completions.js
index 0f3bb1335..30413aab5 100644
--- a/src/endpoints/backends/chat-completions.js
+++ b/src/endpoints/backends/chat-completions.js
@@ -267,7 +267,7 @@ async function sendMakerSuiteRequest(request, response) {
             ? (stream ? 'streamGenerateContent' : 'generateContent')
             : (isText ? 'generateText' : 'generateMessage');
 
-        const generateResponse = await fetch(`https://generativelanguage.googleapis.com/${apiVersion}/models/${model}:${responseType}?key=${apiKey}`, {
+        const generateResponse = await fetch(`https://generativelanguage.googleapis.com/${apiVersion}/models/${model}:${responseType}?key=${apiKey}${stream ? '&alt=sse' : ''}`, {
             body: JSON.stringify(body),
             method: 'POST',
             headers: {
@@ -279,36 +279,8 @@ async function sendMakerSuiteRequest(request, response) {
         // have to do this because of their busted ass streaming endpoint
         if (stream) {
             try {
-                let partialData = '';
-                generateResponse.body.on('data', (data) => {
-                    const chunk = data.toString();
-                    if (chunk.startsWith(',') || chunk.endsWith(',') || chunk.startsWith('[') || chunk.endsWith(']')) {
-                        partialData = chunk.slice(1);
-                    } else {
-                        partialData += chunk;
-                    }
-                    while (true) {
-                        let json;
-                        try {
-                            json = JSON.parse(partialData);
-                        } catch (e) {
-                            break;
-                        }
-                        response.write(JSON.stringify(json));
-                        partialData = '';
-                    }
-                });
-
-                request.socket.on('close', function () {
-                    if (generateResponse.body instanceof Readable) generateResponse.body.destroy();
-                    response.end();
-                });
-
-                generateResponse.body.on('end', () => {
-                    console.log('Streaming request finished');
-                    response.end();
-                });
-
+                // Pipe remote SSE stream to Express response
+                forwardFetchResponse(generateResponse, response);
             } catch (error) {
                 console.log('Error forwarding streaming response:', error);
                 if (!response.headersSent) {