diff --git a/src/background/apis/bing-web.mjs b/src/background/apis/bing-web.mjs index a8239d0..b456dfb 100644 --- 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 index 17f1631..79e0126 100644 --- 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 index 18df6e4..70fa3bf 100644 --- 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 index 3079d31..7d1a1bc 100644 --- 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 index 8519851..dee0966 100644 --- 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