Skip to content

Commit ab95a3e

Browse files
josStorerjosStorer
committed
refactor: setAbortController
1 parent a01f677 commit ab95a3e

File tree

5 files changed

+44
-87
lines changed

5 files changed

+44
-87
lines changed

‎src/background/apis/bing-web.mjs

Lines changed: 4 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import BingAIClient from '../clients/BingAIClient'
22
import { getUserConfig } from '../../config/index.mjs'
3-
import { pushRecord } from './shared.mjs'
3+
import { pushRecord, setAbortController } from './shared.mjs'
44

55
/**
66
* @param {Runtime.Port} port
@@ -17,20 +17,7 @@ export async function generateAnswersWithBingWebApi(
1717
// eslint-disable-next-line
1818
modelName,
1919
) {
20-
const controller = new AbortController()
21-
const stopListener = (msg) => {
22-
if (msg.stop) {
23-
console.debug('stop generating')
24-
port.postMessage({ done: true })
25-
port.onMessage.removeListener(stopListener)
26-
controller.abort()
27-
}
28-
}
29-
port.onMessage.addListener(stopListener)
30-
port.onDisconnect.addListener(() => {
31-
console.debug('port disconnected')
32-
controller.abort()
33-
})
20+
const { controller, messageListener } = setAbortController(port)
3421

3522
const bingAIClient = new BingAIClient({ userToken: accessToken })
3623

@@ -55,7 +42,7 @@ export async function generateAnswersWithBingWebApi(
5542
: {}),
5643
})
5744
.catch((err) => {
58-
port.onMessage.removeListener(stopListener)
45+
port.onMessage.removeListener(messageListener)
5946
throw err
6047
})
6148

@@ -66,6 +53,6 @@ export async function generateAnswersWithBingWebApi(
6653

6754
pushRecord(session, question, answer)
6855
console.debug('conversation history', { content: session.conversationRecords })
69-
port.onMessage.removeListener(stopListener)
56+
port.onMessage.removeListener(messageListener)
7057
port.postMessage({ answer: answer, done: true, session: session })
7158
}

‎src/background/apis/chatgpt-web.mjs

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import { fetchSSE } from '../../utils/fetch-sse'
44
import { isEmpty } from 'lodash-es'
55
import { chatgptWebModelKeys, getUserConfig, Models } from '../../config/index.mjs'
6-
import { pushRecord } from './shared.mjs'
6+
import { pushRecord, setAbortController } from './shared.mjs'
77

88
async function request(token, method, path, data) {
99
const apiUrl = (await getUserConfig()).customChatGptWebApiUrl
@@ -53,24 +53,12 @@ export async function getModels(token) {
5353
* @param {string} accessToken
5454
*/
5555
export async function generateAnswersWithChatgptWebApi(port, question, session, accessToken) {
56-
const controller = new AbortController()
57-
const stopListener = (msg) => {
58-
if (msg.stop) {
59-
console.debug('stop generating')
60-
port.postMessage({ done: true })
61-
port.onMessage.removeListener(stopListener)
62-
controller.abort()
63-
}
64-
}
65-
port.onMessage.addListener(stopListener)
66-
port.onDisconnect.addListener(() => {
67-
console.debug('port disconnected')
68-
controller.abort()
56+
const { controller, messageListener } = setAbortController(port, null, () => {
6957
if (session.autoClean) deleteConversation(accessToken, session.conversationId)
7058
})
7159

7260
const models = await getModels(accessToken).catch(() => {
73-
port.onMessage.removeListener(stopListener)
61+
port.onMessage.removeListener(messageListener)
7462
})
7563
console.debug('models', models)
7664
const config = await getUserConfig()
@@ -130,10 +118,10 @@ export async function generateAnswersWithChatgptWebApi(port, question, session,
130118
// sendModerations(accessToken, question, session.conversationId, session.messageId)
131119
},
132120
async onEnd() {
133-
port.onMessage.removeListener(stopListener)
121+
port.onMessage.removeListener(messageListener)
134122
},
135123
async onError(resp) {
136-
port.onMessage.removeListener(stopListener)
124+
port.onMessage.removeListener(messageListener)
137125
if (resp instanceof Error) throw resp
138126
if (resp.status === 403) {
139127
throw new Error('CLOUDFLARE')

‎src/background/apis/custom-api.mjs

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { getUserConfig, maxResponseTokenLength } from '../../config/index.mjs'
99
import { fetchSSE } from '../../utils/fetch-sse'
1010
import { getConversationPairs } from '../../utils/get-conversation-pairs'
1111
import { isEmpty } from 'lodash-es'
12-
import { pushRecord } from './shared.mjs'
12+
import { pushRecord, setAbortController } from './shared.mjs'
1313

1414
const getCustomApiPromptBase = async () => {
1515
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 () => {
2323
* @param {string} modelName
2424
*/
2525
export async function generateAnswersWithCustomApi(port, question, session, apiKey, modelName) {
26-
const controller = new AbortController()
27-
const stopListener = (msg) => {
28-
if (msg.stop) {
29-
console.debug('stop generating')
30-
port.postMessage({ done: true })
31-
port.onMessage.removeListener(stopListener)
32-
controller.abort()
33-
}
34-
}
35-
port.onMessage.addListener(stopListener)
36-
port.onDisconnect.addListener(() => {
37-
console.debug('port disconnected')
38-
controller.abort()
39-
})
26+
const { controller, messageListener } = setAbortController(port)
4027

4128
const prompt = getConversationPairs(session.conversationRecords, true)
4229
prompt.unshift({ role: 'system', content: await getCustomApiPromptBase() })
@@ -77,14 +64,11 @@ export async function generateAnswersWithCustomApi(port, question, session, apiK
7764
},
7865
async onStart() {},
7966
async onEnd() {
80-
port.onMessage.removeListener(stopListener)
67+
port.onMessage.removeListener(messageListener)
8168
},
8269
async onError(resp) {
83-
port.onMessage.removeListener(stopListener)
70+
port.onMessage.removeListener(messageListener)
8471
if (resp instanceof Error) throw resp
85-
if (resp.status === 403) {
86-
throw new Error('CLOUDFLARE')
87-
}
8872
const error = await resp.json().catch(() => ({}))
8973
throw new Error(!isEmpty(error) ? JSON.stringify(error) : `${resp.status} ${resp.statusText}`)
9074
},

‎src/background/apis/openai-api.mjs

Lines changed: 7 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { maxResponseTokenLength, Models, getUserConfig } from '../../config/inde
44
import { fetchSSE } from '../../utils/fetch-sse'
55
import { getConversationPairs } from '../../utils/get-conversation-pairs'
66
import { isEmpty } from 'lodash-es'
7-
import { pushRecord } from './shared.mjs'
7+
import { pushRecord, setAbortController } from './shared.mjs'
88

99
const getChatgptPromptBase = async () => {
1010
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(
3535
apiKey,
3636
modelName,
3737
) {
38-
const controller = new AbortController()
39-
const stopListener = (msg) => {
40-
if (msg.stop) {
41-
console.debug('stop generating')
42-
port.postMessage({ done: true })
43-
port.onMessage.removeListener(stopListener)
44-
controller.abort()
45-
}
46-
}
47-
port.onMessage.addListener(stopListener)
48-
port.onDisconnect.addListener(() => {
49-
console.debug('port disconnected')
50-
controller.abort()
51-
})
38+
const { controller, messageListener } = setAbortController(port)
5239

5340
const prompt =
5441
(await getGptPromptBase()) +
@@ -90,10 +77,10 @@ export async function generateAnswersWithGptCompletionApi(
9077
},
9178
async onStart() {},
9279
async onEnd() {
93-
port.onMessage.removeListener(stopListener)
80+
port.onMessage.removeListener(messageListener)
9481
},
9582
async onError(resp) {
96-
port.onMessage.removeListener(stopListener)
83+
port.onMessage.removeListener(messageListener)
9784
if (resp instanceof Error) throw resp
9885
if (resp.status === 403) {
9986
throw new Error('CLOUDFLARE')
@@ -112,20 +99,7 @@ export async function generateAnswersWithGptCompletionApi(
11299
* @param {string} modelName
113100
*/
114101
export async function generateAnswersWithChatgptApi(port, question, session, apiKey, modelName) {
115-
const controller = new AbortController()
116-
const stopListener = (msg) => {
117-
if (msg.stop) {
118-
console.debug('stop generating')
119-
port.postMessage({ done: true })
120-
port.onMessage.removeListener(stopListener)
121-
controller.abort()
122-
}
123-
}
124-
port.onMessage.addListener(stopListener)
125-
port.onDisconnect.addListener(() => {
126-
console.debug('port disconnected')
127-
controller.abort()
128-
})
102+
const { controller, messageListener } = setAbortController(port)
129103

130104
const prompt = getConversationPairs(session.conversationRecords, true)
131105
prompt.unshift({ role: 'system', content: await getChatgptPromptBase() })
@@ -166,10 +140,10 @@ export async function generateAnswersWithChatgptApi(port, question, session, api
166140
},
167141
async onStart() {},
168142
async onEnd() {
169-
port.onMessage.removeListener(stopListener)
143+
port.onMessage.removeListener(messageListener)
170144
},
171145
async onError(resp) {
172-
port.onMessage.removeListener(stopListener)
146+
port.onMessage.removeListener(messageListener)
173147
if (resp instanceof Error) throw resp
174148
if (resp.status === 403) {
175149
throw new Error('CLOUDFLARE')

‎src/background/apis/shared.mjs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,27 @@
1+
export function setAbortController(port, onStop, onDisconnect) {
2+
const controller = new AbortController()
3+
const messageListener = (msg) => {
4+
if (msg.stop) {
5+
port.onMessage.removeListener(messageListener)
6+
console.debug('stop generating')
7+
port.postMessage({ done: true })
8+
controller.abort()
9+
if (onStop) onStop()
10+
}
11+
}
12+
port.onMessage.addListener(messageListener)
13+
14+
const disconnectListener = () => {
15+
port.onDisconnect.removeListener(disconnectListener)
16+
console.debug('port disconnected')
17+
controller.abort()
18+
if (onDisconnect) onDisconnect()
19+
}
20+
port.onDisconnect.addListener(disconnectListener)
21+
22+
return { controller, messageListener }
23+
}
24+
125
export function pushRecord(session, question, answer) {
226
const recordLength = session.conversationRecords.length
327
let lastRecord

0 commit comments

Comments
 (0)