refactor: setAbortController

josStorer · Apr 5, 2023 · ab95a3e · ab95a3e
1 parent a01f677
commit ab95a3e
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 87 deletions.
diff --git a/src/background/apis/bing-web.mjs b/src/background/apis/bing-web.mjs
@@ -1,6 +1,6 @@
 import BingAIClient from '../clients/BingAIClient'
 import { getUserConfig } from '../../config/index.mjs'
-import { pushRecord } from './shared.mjs'
+import { pushRecord, setAbortController } from './shared.mjs'
 
 /**
  * @param {Runtime.Port} port
@@ -17,20 +17,7 @@ export async function generateAnswersWithBingWebApi(
   // eslint-disable-next-line
   modelName,
 ) {
-  const controller = new AbortController()
-  const stopListener = (msg) => {
-    if (msg.stop) {
-      console.debug('stop generating')
-      port.postMessage({ done: true })
-      port.onMessage.removeListener(stopListener)
-      controller.abort()
-    }
-  }
-  port.onMessage.addListener(stopListener)
-  port.onDisconnect.addListener(() => {
-    console.debug('port disconnected')
-    controller.abort()
-  })
+  const { controller, messageListener } = setAbortController(port)
 
   const bingAIClient = new BingAIClient({ userToken: accessToken })
 
@@ -55,7 +42,7 @@ export async function generateAnswersWithBingWebApi(
         : {}),
     })
     .catch((err) => {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
       throw err
     })
 
@@ -66,6 +53,6 @@ export async function generateAnswersWithBingWebApi(
 
   pushRecord(session, question, answer)
   console.debug('conversation history', { content: session.conversationRecords })
-  port.onMessage.removeListener(stopListener)
+  port.onMessage.removeListener(messageListener)
   port.postMessage({ answer: answer, done: true, session: session })
 }
diff --git a/src/background/apis/chatgpt-web.mjs b/src/background/apis/chatgpt-web.mjs
@@ -3,7 +3,7 @@
 import { fetchSSE } from '../../utils/fetch-sse'
 import { isEmpty } from 'lodash-es'
 import { chatgptWebModelKeys, getUserConfig, Models } from '../../config/index.mjs'
-import { pushRecord } from './shared.mjs'
+import { pushRecord, setAbortController } from './shared.mjs'
 
 async function request(token, method, path, data) {
   const apiUrl = (await getUserConfig()).customChatGptWebApiUrl
@@ -53,24 +53,12 @@ export async function getModels(token) {
  * @param {string} accessToken
  */
 export async function generateAnswersWithChatgptWebApi(port, question, session, accessToken) {
-  const controller = new AbortController()
-  const stopListener = (msg) => {
-    if (msg.stop) {
-      console.debug('stop generating')
-      port.postMessage({ done: true })
-      port.onMessage.removeListener(stopListener)
-      controller.abort()
-    }
-  }
-  port.onMessage.addListener(stopListener)
-  port.onDisconnect.addListener(() => {
-    console.debug('port disconnected')
-    controller.abort()
+  const { controller, messageListener } = setAbortController(port, null, () => {
     if (session.autoClean) deleteConversation(accessToken, session.conversationId)
   })
 
   const models = await getModels(accessToken).catch(() => {
-    port.onMessage.removeListener(stopListener)
+    port.onMessage.removeListener(messageListener)
   })
   console.debug('models', models)
   const config = await getUserConfig()
@@ -130,10 +118,10 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
       // sendModerations(accessToken, question, session.conversationId, session.messageId)
     },
     async onEnd() {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
     },
     async onError(resp) {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
       if (resp instanceof Error) throw resp
       if (resp.status === 403) {
         throw new Error('CLOUDFLARE')

diff --git a/src/background/apis/custom-api.mjs b/src/background/apis/custom-api.mjs
@@ -9,7 +9,7 @@ import { getUserConfig, maxResponseTokenLength } from '../../config/index.mjs'
 import { fetchSSE } from '../../utils/fetch-sse'
 import { getConversationPairs } from '../../utils/get-conversation-pairs'
 import { isEmpty } from 'lodash-es'
-import { pushRecord } from './shared.mjs'
+import { pushRecord, setAbortController } from './shared.mjs'
 
 const getCustomApiPromptBase = async () => {
   return `I am a helpful, creative, clever, and very friendly assistant. I am familiar with various languages in the world.`
@@ -23,20 +23,7 @@ const getCustomApiPromptBase = async () => {
  * @param {string} modelName
  */
 export async function generateAnswersWithCustomApi(port, question, session, apiKey, modelName) {
-  const controller = new AbortController()
-  const stopListener = (msg) => {
-    if (msg.stop) {
-      console.debug('stop generating')
-      port.postMessage({ done: true })
-      port.onMessage.removeListener(stopListener)
-      controller.abort()
-    }
-  }
-  port.onMessage.addListener(stopListener)
-  port.onDisconnect.addListener(() => {
-    console.debug('port disconnected')
-    controller.abort()
-  })
+  const { controller, messageListener } = setAbortController(port)
 
   const prompt = getConversationPairs(session.conversationRecords, true)
   prompt.unshift({ role: 'system', content: await getCustomApiPromptBase() })
@@ -77,14 +64,11 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
     },
     async onStart() {},
     async onEnd() {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
     },
     async onError(resp) {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
       if (resp instanceof Error) throw resp
-      if (resp.status === 403) {
-        throw new Error('CLOUDFLARE')
-      }
       const error = await resp.json().catch(() => ({}))
       throw new Error(!isEmpty(error) ? JSON.stringify(error) : `${resp.status} ${resp.statusText}`)
     },

diff --git a/src/background/apis/openai-api.mjs b/src/background/apis/openai-api.mjs
@@ -4,7 +4,7 @@ import { maxResponseTokenLength, Models, getUserConfig } from '../../config/inde
 import { fetchSSE } from '../../utils/fetch-sse'
 import { getConversationPairs } from '../../utils/get-conversation-pairs'
 import { isEmpty } from 'lodash-es'
-import { pushRecord } from './shared.mjs'
+import { pushRecord, setAbortController } from './shared.mjs'
 
 const getChatgptPromptBase = async () => {
   return `You are a helpful, creative, clever, and very friendly assistant. You are familiar with various languages in the world.`
@@ -35,20 +35,7 @@ export async function generateAnswersWithGptCompletionApi(
   apiKey,
   modelName,
 ) {
-  const controller = new AbortController()
-  const stopListener = (msg) => {
-    if (msg.stop) {
-      console.debug('stop generating')
-      port.postMessage({ done: true })
-      port.onMessage.removeListener(stopListener)
-      controller.abort()
-    }
-  }
-  port.onMessage.addListener(stopListener)
-  port.onDisconnect.addListener(() => {
-    console.debug('port disconnected')
-    controller.abort()
-  })
+  const { controller, messageListener } = setAbortController(port)
 
   const prompt =
     (await getGptPromptBase()) +
@@ -90,10 +77,10 @@ export async function generateAnswersWithGptCompletionApi(
     },
     async onStart() {},
     async onEnd() {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
     },
     async onError(resp) {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
       if (resp instanceof Error) throw resp
       if (resp.status === 403) {
         throw new Error('CLOUDFLARE')
@@ -112,20 +99,7 @@ export async function generateAnswersWithGptCompletionApi(
  * @param {string} modelName
  */
 export async function generateAnswersWithChatgptApi(port, question, session, apiKey, modelName) {
-  const controller = new AbortController()
-  const stopListener = (msg) => {
-    if (msg.stop) {
-      console.debug('stop generating')
-      port.postMessage({ done: true })
-      port.onMessage.removeListener(stopListener)
-      controller.abort()
-    }
-  }
-  port.onMessage.addListener(stopListener)
-  port.onDisconnect.addListener(() => {
-    console.debug('port disconnected')
-    controller.abort()
-  })
+  const { controller, messageListener } = setAbortController(port)
 
   const prompt = getConversationPairs(session.conversationRecords, true)
   prompt.unshift({ role: 'system', content: await getChatgptPromptBase() })
@@ -166,10 +140,10 @@ export async function generateAnswersWithChatgptApi(port, question, session, api
     },
     async onStart() {},
     async onEnd() {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
     },
     async onError(resp) {
-      port.onMessage.removeListener(stopListener)
+      port.onMessage.removeListener(messageListener)
       if (resp instanceof Error) throw resp
       if (resp.status === 403) {
         throw new Error('CLOUDFLARE')

diff --git a/src/background/apis/shared.mjs b/src/background/apis/shared.mjs
@@ -1,3 +1,27 @@
+export function setAbortController(port, onStop, onDisconnect) {
+  const controller = new AbortController()
+  const messageListener = (msg) => {
+    if (msg.stop) {
+      port.onMessage.removeListener(messageListener)
+      console.debug('stop generating')
+      port.postMessage({ done: true })
+      controller.abort()
+      if (onStop) onStop()
+    }
+  }
+  port.onMessage.addListener(messageListener)
+
+  const disconnectListener = () => {
+    port.onDisconnect.removeListener(disconnectListener)
+    console.debug('port disconnected')
+    controller.abort()
+    if (onDisconnect) onDisconnect()
+  }
+  port.onDisconnect.addListener(disconnectListener)
+
+  return { controller, messageListener }
+}
+
 export function pushRecord(session, question, answer) {
   const recordLength = session.conversationRecords.length
   let lastRecord