From 97fc6517851d4ff5c6f0923b93bce6a0d9665243 Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Tue, 16 Apr 2024 12:32:46 +0800
Subject: [PATCH] improve fetch onmessage

---
 src/services/apis/custom-api.mjs | 26 ++++++++--------
 src/services/apis/openai-api.mjs | 52 +++++++++++++++++---------------
 2 files changed, 42 insertions(+), 36 deletions(-)

diff --git a/src/services/apis/custom-api.mjs b/src/services/apis/custom-api.mjs
index e8fbaaeb..4a6b8a42 100644
--- a/src/services/apis/custom-api.mjs
+++ b/src/services/apis/custom-api.mjs
@@ -32,6 +32,12 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
 
   let answer = ''
   let finished = false
+  const finish = () => {
+    finished = true
+    pushRecord(session, question, answer)
+    console.debug('conversation history', { content: session.conversationRecords })
+    port.postMessage({ answer: null, done: true, session: session })
+  }
   await fetchSSE(apiUrl, {
     method: 'POST',
     signal: controller.signal,
@@ -48,11 +54,9 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
     }),
     onMessage(message) {
       console.debug('sse message', message)
-      if (!finished && message.trim() === '[DONE]') {
-        finished = true
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: null, done: true, session: session })
+      if (finished) return
+      if (message.trim() === '[DONE]') {
+        finish()
         return
       }
       let data
@@ -62,13 +66,6 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
         console.debug('json error', error)
         return
       }
-      if (!finished && data.choices[0]?.finish_reason) {
-        finished = true
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: null, done: true, session: session })
-        return
-      }
 
       if (data.response) answer = data.response
       else {
@@ -84,6 +81,11 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
         }
       }
       port.postMessage({ answer: answer, done: false, session: null })
+
+      if (data.choices[0]?.finish_reason) {
+        finish()
+        return
+      }
     },
     async onStart() {},
     async onEnd() {
diff --git a/src/services/apis/openai-api.mjs b/src/services/apis/openai-api.mjs
index 6eaee816..9b4bb4bf 100644
--- a/src/services/apis/openai-api.mjs
+++ b/src/services/apis/openai-api.mjs
@@ -39,6 +39,12 @@ export async function generateAnswersWithGptCompletionApi(
 
   let answer = ''
   let finished = false
+  const finish = () => {
+    finished = true
+    pushRecord(session, question, answer)
+    console.debug('conversation history', { content: session.conversationRecords })
+    port.postMessage({ answer: null, done: true, session: session })
+  }
   await fetchSSE(`${apiUrl}/v1/completions`, {
     method: 'POST',
     signal: controller.signal,
@@ -56,11 +62,9 @@ export async function generateAnswersWithGptCompletionApi(
     }),
     onMessage(message) {
       console.debug('sse message', message)
-      if (!finished && message.trim() === '[DONE]') {
-        finished = true
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: null, done: true, session: session })
+      if (finished) return
+      if (message.trim() === '[DONE]') {
+        finish()
         return
       }
       let data
@@ -70,16 +74,14 @@ export async function generateAnswersWithGptCompletionApi(
         console.debug('json error', error)
         return
       }
-      if (!finished && data.choices[0]?.finish_reason) {
-        finished = true
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: null, done: true, session: session })
-        return
-      }
 
       answer += data.choices[0].text
       port.postMessage({ answer: answer, done: false, session: null })
+
+      if (data.choices[0]?.finish_reason) {
+        finish()
+        return
+      }
     },
     async onStart() {},
     async onEnd() {
@@ -136,6 +138,12 @@ export async function generateAnswersWithChatgptApiCompat(
 
   let answer = ''
   let finished = false
+  const finish = () => {
+    finished = true
+    pushRecord(session, question, answer)
+    console.debug('conversation history', { content: session.conversationRecords })
+    port.postMessage({ answer: null, done: true, session: session })
+  }
   await fetchSSE(`${baseUrl}/v1/chat/completions`, {
     method: 'POST',
     signal: controller.signal,
@@ -152,11 +160,9 @@ export async function generateAnswersWithChatgptApiCompat(
     }),
     onMessage(message) {
       console.debug('sse message', message)
-      if (!finished && message.trim() === '[DONE]') {
-        finished = true
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: null, done: true, session: session })
+      if (finished) return
+      if (message.trim() === '[DONE]') {
+        finish()
         return
       }
       let data
@@ -166,13 +172,6 @@ export async function generateAnswersWithChatgptApiCompat(
         console.debug('json error', error)
         return
       }
-      if (!finished && data.choices[0]?.finish_reason) {
-        finished = true
-        pushRecord(session, question, answer)
-        console.debug('conversation history', { content: session.conversationRecords })
-        port.postMessage({ answer: null, done: true, session: session })
-        return
-      }
 
       const delta = data.choices[0]?.delta?.content
       const content = data.choices[0]?.message?.content
@@ -185,6 +184,11 @@ export async function generateAnswersWithChatgptApiCompat(
         answer += text
       }
       port.postMessage({ answer: answer, done: false, session: null })
+
+      if (data.choices[0]?.finish_reason) {
+        finish()
+        return
+      }
     },
     async onStart() {},
     async onEnd() {