feat: offline/self-hosted large language model support (#11)

Author: josStorer

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}`)
+    },
+  })
+}

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)

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<UserConfig>}

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 }) {
             </span>
           )}
         </span>
+        {isUsingCustomModel(config) && (
+          <input
+            type="text"
+            value={config.customModelApiUrl}
+            placeholder="Custom Model API Url"
+            onChange={(e) => {
+              const value = e.target.value
+              updateConfig({ customModelApiUrl: value })
+            }}
+          />
+        )}
       </label>
       <label>
         <legend>Preferred Language</legend>

src/utils/init-session.mjs

@@ -5,7 +5,6 @@
  * @property {string|null} messageId - chatGPT web mode
  * @property {string|null} parentMessageId - chatGPT web mode
  * @property {Object[]|null} conversationRecords
- * @property {bool|null} useApiKey
  */
 /**
  * @param {Session} session
@@ -17,14 +16,12 @@ export function initSession({
   messageId = null,
   parentMessageId = null,
   conversationRecords = [],
-  useApiKey = null,
 } = {}) {
   return {
     question,
     conversationId,
     messageId,
     parentMessageId,
     conversationRecords,
-    useApiKey,
   }
 }