From 00badc3bc90198788269e538f5b0d7e19dbf9e1a Mon Sep 17 00:00:00 2001 From: josc146 Date: Wed, 22 Mar 2023 19:54:54 +0800 Subject: [PATCH] feat: offline/self-hosted large language model support (#11) --- src/background/apis/custom-api.mjs | 91 ++++++++++++++++++++++++++++++ src/background/index.mjs | 14 +++-- src/config/index.mjs | 7 +++ src/popup/Popup.jsx | 12 ++++ src/utils/init-session.mjs | 3 - 5 files changed, 120 insertions(+), 7 deletions(-) create mode 100644 src/background/apis/custom-api.mjs diff --git a/src/background/apis/custom-api.mjs b/src/background/apis/custom-api.mjs new file mode 100644 index 0000000..97fdb7f --- /dev/null +++ b/src/background/apis/custom-api.mjs @@ -0,0 +1,91 @@ +// custom api version + +// There is a lot of duplicated code here, but it is very easy to refactor. +// The current state is mainly convenient for making targeted changes at any time, +// and it has not yet had a negative impact on maintenance. +// If necessary, I will refactor. + +import { getUserConfig, maxResponseTokenLength, Models } from '../../config/index.mjs' +import { fetchSSE } from '../../utils/fetch-sse' +import { getConversationPairs } from '../../utils/get-conversation-pairs' +import { isEmpty } from 'lodash-es' + +const getCustomApiPromptBase = async () => { + return `I am a helpful, creative, clever, and very friendly assistant. I am familiar with various languages in the world.` +} + +/** + * @param {Browser.Runtime.Port} port + * @param {string} question + * @param {Session} session + * @param {string} apiKey + * @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 }) + controller.abort() + port.onMessage.removeListener(stopListener) + } + } + port.onMessage.addListener(stopListener) + port.onDisconnect.addListener(() => { + console.debug('port disconnected') + controller.abort() + }) + + const prompt = getConversationPairs(session.conversationRecords, true) + prompt.unshift({ role: 'system', content: await getCustomApiPromptBase() }) + prompt.push({ role: 'user', content: question }) + const apiUrl = (await getUserConfig()).customModelApiUrl + + let answer = '' + await fetchSSE(apiUrl, { + method: 'POST', + signal: controller.signal, + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify({ + messages: prompt, + model: Models[modelName].value, + stream: true, + max_tokens: maxResponseTokenLength, + }), + onMessage(message) { + console.debug('sse message', message) + if (message === '[DONE]') { + session.conversationRecords.push({ question: question, answer: answer }) + console.debug('conversation history', { content: session.conversationRecords }) + port.postMessage({ answer: null, done: true, session: session }) + return + } + let data + try { + data = JSON.parse(message) + } catch (error) { + console.debug('json error', error) + return + } + if (data.response) answer = data.response + port.postMessage({ answer: answer, done: false, session: null }) + }, + async onStart() {}, + async onEnd() { + port.onMessage.removeListener(stopListener) + }, + async onError(resp) { + if (resp instanceof Error) throw resp + port.onMessage.removeListener(stopListener) + 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/index.mjs b/src/background/index.mjs index 102c2ce..8e245b2 100644 --- a/src/background/index.mjs +++ b/src/background/index.mjs @@ -6,13 +6,14 @@ import { generateAnswersWithChatgptApi, generateAnswersWithGptCompletionApi, } from './apis/openai-api' +import { generateAnswersWithCustomApi } from './apis/custom-api.mjs' import { chatgptApiModelKeys, chatgptWebModelKeys, + customApiModelKeys, defaultConfig, getUserConfig, gptApiModelKeys, - isUsingApiKey, } from '../config/index.mjs' import { isSafari } from '../utils/is-safari' import { isFirefox } from '../utils/is-firefox' @@ -55,9 +56,6 @@ Browser.runtime.onConnect.addListener((port) => { const session = msg.session if (!session) return const config = await getUserConfig() - if (session.useApiKey == null) { - session.useApiKey = isUsingApiKey(config) - } try { if (chatgptWebModelKeys.includes(config.modelName)) { @@ -83,6 +81,14 @@ Browser.runtime.onConnect.addListener((port) => { config.apiKey, config.modelName, ) + } else if (customApiModelKeys.includes(config.modelName)) { + await generateAnswersWithCustomApi( + port, + session.question, + session, + config.apiKey, + config.modelName, + ) } } catch (err) { console.error(err) diff --git a/src/config/index.mjs b/src/config/index.mjs index 0c53bb0..aa64e7a 100644 --- a/src/config/index.mjs +++ b/src/config/index.mjs @@ -17,11 +17,13 @@ export const Models = { chatgptApi4_8k: { value: 'gpt-4', desc: 'ChatGPT (GPT-4-8k)' }, chatgptApi4_32k: { value: 'gpt-4-32k', desc: 'ChatGPT (GPT-4-32k)' }, gptApiDavinci: { value: 'text-davinci-003', desc: 'GPT-3.5' }, + chatglm6bInt4: { value: 'chatglm-6b-int4', desc: 'ChatGLM-6B-Int4' }, } export const chatgptWebModelKeys = ['chatgptFree35', 'chatgptPlus4'] export const gptApiModelKeys = ['gptApiDavinci'] export const chatgptApiModelKeys = ['chatgptApi35', 'chatgptApi4_8k', 'chatgptApi4_32k'] +export const customApiModelKeys = ['chatglm6bInt4'] export const TriggerMode = { always: 'Always', @@ -59,6 +61,7 @@ export const defaultConfig = { customChatGptWebApiUrl: 'https://chat.openai.com', customChatGptWebApiPath: '/backend-api/conversation', customOpenAiApiUrl: 'https://api.openai.com', + customModelApiUrl: 'http://localhost:8000/chat/completions', siteRegex: 'match nothing', userSiteRegexOnly: false, inputQuery: '', @@ -123,6 +126,10 @@ export function isUsingApiKey(config) { ) } +export function isUsingCustomModel(config) { + return customApiModelKeys.includes(config.modelName) +} + /** * get user config from local storage * @returns {Promise} diff --git a/src/popup/Popup.jsx b/src/popup/Popup.jsx index 2f3b995..c42c51b 100644 --- a/src/popup/Popup.jsx +++ b/src/popup/Popup.jsx @@ -8,6 +8,7 @@ import { defaultConfig, Models, isUsingApiKey, + isUsingCustomModel, } from '../config/index.mjs' import { Tab, Tabs, TabList, TabPanel } from 'react-tabs' import 'react-tabs/style/react-tabs.css' @@ -122,6 +123,17 @@ function GeneralPart({ config, updateConfig }) { )} + {isUsingCustomModel(config) && ( + { + const value = e.target.value + updateConfig({ customModelApiUrl: value }) + }} + /> + )}