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

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)