add additional finish conditions for OpenAI API and Custom API (both …

…can be customized, thus requiring more condition checks, now the api has better support for ollama and LM Studio) (#631, #648)
josStorer · Mar 23, 2024 · 680900b · 680900b
1 parent d49280c
commit 680900b
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 4 deletions.
diff --git a/src/services/apis/azure-openai-api.mjs b/src/services/apis/azure-openai-api.mjs
@@ -51,7 +51,7 @@ export async function generateAnswersWithAzureOpenaiApi(port, question, session)
           answer += data.choices[0].delta.content
           port.postMessage({ answer: answer, done: false, session: null })
         }
-        if (data.choices[0].finish_reason === 'stop') {
+        if (data.choices[0]?.finish_reason) {
           pushRecord(session, question, answer)
           console.debug('conversation history', { content: session.conversationRecords })
           port.postMessage({ answer: null, done: true, session: session })

diff --git a/src/services/apis/custom-api.mjs b/src/services/apis/custom-api.mjs
@@ -31,6 +31,7 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
   const apiUrl = config.customModelApiUrl
 
   let answer = ''
+  let finished = false
   await fetchSSE(apiUrl, {
     method: 'POST',
     signal: controller.signal,
@@ -47,7 +48,8 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
     }),
     onMessage(message) {
       console.debug('sse message', message)
-      if (message.trim() === '[DONE]') {
+      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 })
@@ -60,6 +62,12 @@ 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 })
+      }
 
       if (data.response) answer = data.response
       else {

diff --git a/src/services/apis/openai-api.mjs b/src/services/apis/openai-api.mjs
@@ -38,6 +38,7 @@ export async function generateAnswersWithGptCompletionApi(
   const apiUrl = config.customOpenAiApiUrl
 
   let answer = ''
+  let finished = false
   await fetchSSE(`${apiUrl}/v1/completions`, {
     method: 'POST',
     signal: controller.signal,
@@ -55,7 +56,8 @@ export async function generateAnswersWithGptCompletionApi(
     }),
     onMessage(message) {
       console.debug('sse message', message)
-      if (message.trim() === '[DONE]') {
+      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 })
@@ -68,6 +70,13 @@ 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 })
+      }
+
       answer += data.choices[0].text
       port.postMessage({ answer: answer, done: false, session: null })
     },
@@ -125,6 +134,7 @@ export async function generateAnswersWithChatgptApiCompat(
   prompt.push({ role: 'user', content: question })
 
   let answer = ''
+  let finished = false
   await fetchSSE(`${baseUrl}/v1/chat/completions`, {
     method: 'POST',
     signal: controller.signal,
@@ -141,7 +151,8 @@ export async function generateAnswersWithChatgptApiCompat(
     }),
     onMessage(message) {
       console.debug('sse message', message)
-      if (message.trim() === '[DONE]') {
+      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 })
@@ -154,6 +165,13 @@ 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 })
+      }
+
       const delta = data.choices[0]?.delta?.content
       const content = data.choices[0]?.message?.content
       const text = data.choices[0]?.text