Add Ollama native API to support keep alive parameters (#748)

lzskyline · web-flow · commit 488170068c0e · 2024-07-30T16:53:08.000+08:00
* Add Chinese translation.

* Optimize style.

* Add Ollama native API to support keep alive parameters.

* Optimized popup page style.

* fix: Fixed data type for Ollama keep_alive parameter forever
diff --git a/src/_locales/zh-hans/main.json b/src/_locales/zh-hans/main.json
@@ -77,6 +77,7 @@
   "ChatGPT (GPT-4-8k)": "ChatGPT (GPT-4-8k)",
   "ChatGPT (GPT-4-32k)": "ChatGPT (GPT-4-32k)",
   "GPT-3.5": "GPT-3.5",
+  "Ollama API": "Ollama API",
   "Custom Model": "自定义模型",
   "Balanced": "平衡",
   "Creative": "有创造力",
@@ -142,5 +143,13 @@
   "Icon": "图标",
   "Prompt Template": "提示模板",
   "Explain this: {{selection}}": "解释这个: {{selection}}",
-  "New": "新建"
+  "New": "新建",
+  "DisplayMode": "显示方式",
+  "Display in sidebar": "在侧边栏显示",
+  "Display in floating toolbar": "在浮动工具栏显示",
+  "Temperature": "温度",
+  "keep-alive Time": "保活时间",
+  "5m": "5分钟",
+  "30m": "半小时",
+  "Forever": "永久"
 }
diff --git a/src/background/index.mjs b/src/background/index.mjs
@@ -10,6 +10,7 @@ import {
   generateAnswersWithGptCompletionApi,
 } from '../services/apis/openai-api'
 import { generateAnswersWithCustomApi } from '../services/apis/custom-api.mjs'
+import { generateAnswersWithOllamaApi } from '../services/apis/ollama-api.mjs'
 import { generateAnswersWithAzureOpenaiApi } from '../services/apis/azure-openai-api.mjs'
 import { generateAnswersWithClaudeApi } from '../services/apis/claude-api.mjs'
 import { generateAnswersWithChatGLMApi } from '../services/apis/chatglm-api.mjs'
@@ -25,6 +26,7 @@ import {
   claudeWebModelKeys,
   moonshotWebModelKeys,
   customApiModelKeys,
+  ollamaApiModelKeys,
   defaultConfig,
   getUserConfig,
   githubThirdPartyApiModelKeys,
@@ -124,6 +126,14 @@ async function executeApi(session, port, config) {
       config.customApiKey,
       config.customModelName,
     )
+  } else if (ollamaApiModelKeys.includes(session.modelName)) {
+    await generateAnswersWithOllamaApi(
+      port,
+      session.question,
+      session,
+      config.ollamaApiKey,
+      config.ollamaModelName,
+    )
   } else if (azureOpenAiApiModelKeys.includes(session.modelName)) {
     await generateAnswersWithAzureOpenaiApi(port, session.question, session)
   } else if (claudeApiModelKeys.includes(session.modelName)) {
diff --git a/src/components/ConversationItem/index.jsx b/src/components/ConversationItem/index.jsx
@@ -5,7 +5,7 @@ import ReadButton from '../ReadButton'
 import PropTypes from 'prop-types'
 import MarkdownRender from '../MarkdownRender/markdown.jsx'
 import { useTranslation } from 'react-i18next'
-import { isUsingCustomModel } from '../../config/index.mjs'
+import { isUsingCustomModel, isUsingOllamaModel } from '../../config/index.mjs'
 import { useConfig } from '../../hooks/use-config.mjs'
 
 function AnswerTitle({ descName, modelName }) {
diff --git a/src/config/index.mjs b/src/config/index.mjs
@@ -57,6 +57,7 @@ export const chatgptApiModelKeys = [
   'chatgptApi4_128k_0125_preview',
 ]
 export const customApiModelKeys = ['customModel']
+export const ollamaApiModelKeys = ['ollamaModel']
 export const azureOpenAiApiModelKeys = ['azureOpenAi']
 export const claudeApiModelKeys = [
   'claude12Api',
@@ -163,6 +164,7 @@ export const Models = {
   gptApiDavinci: { value: 'text-davinci-003', desc: 'GPT-3.5' },
 
   customModel: { value: '', desc: 'Custom Model' },
+  ollamaModel: { value: '', desc: 'Ollama API' },
   azureOpenAi: { value: '', desc: 'ChatGPT (Azure)' },
   waylaidwandererApi: { value: '', desc: 'Waylaidwanderer API (Github)' },
 
@@ -249,6 +251,10 @@ export const defaultConfig = {
   customModelName: 'gpt-3.5-turbo',
   githubThirdPartyUrl: 'http://127.0.0.1:3000/conversation',
 
+  ollamaEndpoint: 'http://127.0.0.1:11434',
+  ollamaModelName: 'gemma2',
+  keepAliveTime: '5m',
+
   // advanced
 
   maxResponseTokenLength: 1000,
@@ -281,6 +287,7 @@ export const defaultConfig = {
     'moonshotWebFree',
     'chatglmTurbo',
     'customModel',
+    'ollamaModel',
     'azureOpenAi',
   ],
   activeSelectionTools: ['translate', 'summary', 'polish', 'code', 'ask'],
@@ -381,6 +388,10 @@ export function isUsingCustomModel(configOrSession) {
   return customApiModelKeys.includes(configOrSession.modelName)
 }
 
+export function isUsingOllamaModel(configOrSession) {
+  return ollamaApiModelKeys.includes(configOrSession.modelName)
+}
+
 export function isUsingChatGLMApi(configOrSession) {
   return chatglmApiModelKeys.includes(configOrSession.modelName)
 }
diff --git a/src/content-script/index.jsx b/src/content-script/index.jsx
@@ -72,7 +72,7 @@ async function mountComponent(siteConfig, userConfig) {
   })
 
   const position = {
-    x: window.innerWidth - 300 - (Math.floor((20 / 100) * window.innerWidth)),
+    x: window.innerWidth - 300 - Math.floor((20 / 100) * window.innerWidth),
     y: window.innerHeight / 2 - 200,
   }
   const toolbarContainer = createElementAtPosition(position.x, position.y)
diff --git a/src/manifest.json b/src/manifest.json
@@ -1,7 +1,7 @@
 {
   "name": "ChatGPTBox",
   "description": "Integrating ChatGPT into your browser deeply, everything you need is here",
-  "version": "2.5.6",
+  "version": "2.5.7",
   "manifest_version": 3,
   "icons": {
     "16": "logo.png",
diff --git a/src/manifest.v2.json b/src/manifest.v2.json
@@ -1,7 +1,7 @@
 {
   "name": "ChatGPTBox",
   "description": "Integrating ChatGPT into your browser deeply, everything you need is here",
-  "version": "2.5.6",
+  "version": "2.5.7",
   "manifest_version": 2,
   "icons": {
     "16": "logo.png",
diff --git a/src/popup/sections/AdvancedPart.jsx b/src/popup/sections/AdvancedPart.jsx
@@ -1,5 +1,7 @@
+import '../styles.scss'
 import { useTranslation } from 'react-i18next'
 import { parseFloatWithClamp, parseIntWithClamp } from '../../utils/index.mjs'
+import { isUsingOllamaModel } from '../../config/index.mjs'
 import PropTypes from 'prop-types'
 import { Tab, TabList, TabPanel, Tabs } from 'react-tabs'
 import Browser from 'webextension-polyfill'
@@ -56,6 +58,49 @@ function ApiParams({ config, updateConfig }) {
           }}
         />
       </label>
+      {isUsingOllamaModel(config) && (
+        <label>
+          {`${t('keep-alive Time')}: `}
+          <div className="label-group">
+            <label>
+              <input
+                type="radio"
+                name="keepAliveTime"
+                value="5m"
+                checked={config.keepAliveTime === '5m'}
+                onChange={(e) => {
+                  updateConfig({ keepAliveTime: e.target.value })
+                }}
+              />
+              {t('5m')}
+            </label>
+            <label>
+              <input
+                type="radio"
+                name="keepAliveTime"
+                value="30m"
+                checked={config.keepAliveTime === '30m'}
+                onChange={(e) => {
+                  updateConfig({ keepAliveTime: e.target.value })
+                }}
+              />
+              {t('30m')}
+            </label>
+            <label>
+              <input
+                type="radio"
+                name="keepAliveTime"
+                value="-1"
+                checked={config.keepAliveTime === '-1'}
+                onChange={(e) => {
+                  updateConfig({ keepAliveTime: e.target.value })
+                }}
+              />
+              {t('Forever')}
+            </label>
+          </div>
+        </label>
+      )}
     </>
   )
 }
diff --git a/src/popup/sections/GeneralPart.jsx b/src/popup/sections/GeneralPart.jsx
@@ -8,6 +8,7 @@ import {
   isUsingClaudeApi,
   isUsingCustomModel,
   isUsingCustomNameOnlyModel,
+  isUsingOllamaModel,
   isUsingGithubThirdPartyApi,
   isUsingMultiModeModel,
   ModelMode,
@@ -163,6 +164,7 @@ export function GeneralPart({ config, updateConfig }) {
               isUsingOpenAiApiKey(config) ||
               isUsingMultiModeModel(config) ||
               isUsingCustomModel(config) ||
+              isUsingOllamaModel(config) ||
               isUsingAzureOpenAi(config) ||
               isUsingClaudeApi(config) ||
               isUsingCustomNameOnlyModel(config) ||
@@ -271,6 +273,18 @@ export function GeneralPart({ config, updateConfig }) {
               }}
             />
           )}
+          {isUsingOllamaModel(config) && (
+            <input
+              style="width: 50%;"
+              type="text"
+              value={config.ollamaModelName}
+              placeholder={t('Model Name')}
+              onChange={(e) => {
+                const ollamaModelName = e.target.value
+                updateConfig({ ollamaModelName: ollamaModelName })
+              }}
+            />
+          )}
           {isUsingAzureOpenAi(config) && (
             <input
               type="password"
@@ -354,6 +368,28 @@ export function GeneralPart({ config, updateConfig }) {
             }}
           />
         )}
+        {isUsingOllamaModel(config) && (
+          <input
+            type="text"
+            value={config.ollamaEndpoint}
+            placeholder={t('Ollama Endpoint')}
+            onChange={(e) => {
+              const value = e.target.value
+              updateConfig({ ollamaEndpoint: value })
+            }}
+          />
+        )}
+        {isUsingOllamaModel(config) && (
+          <input
+            type="password"
+            value={config.ollamaApiKey}
+            placeholder={t('API Key')}
+            onChange={(e) => {
+              const apiKey = e.target.value
+              updateConfig({ ollamaApiKey: apiKey })
+            }}
+          />
+        )}
         {isUsingAzureOpenAi(config) && (
           <input
             type="password"
diff --git a/src/popup/styles.scss b/src/popup/styles.scss
@@ -23,6 +23,11 @@
   --active-color: #eaecf0;
 }
 
+#app {
+  max-width: 600px;
+  margin: 0 auto;
+}
+
 .container-page-mode {
   display: flex;
   flex-direction: column;
@@ -58,6 +63,7 @@
 }
 
 .footer {
+  max-width: 580px;
   width: 90%;
   position: fixed;
   bottom: 10px;
@@ -88,3 +94,12 @@
     background: var(--active-color);
   }
 }
+
+.label-group {
+  display: flex;
+  align-items: center;
+}
+
+.label-group label {
+  margin-right: 10px;
+}
diff --git a/src/services/apis/ollama-api.mjs b/src/services/apis/ollama-api.mjs
@@ -0,0 +1,82 @@
+// ollama 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 } from '../../config/index.mjs'
+import { fetchSSE } from '../../utils/fetch-ollama.mjs'
+import { getConversationPairs } from '../../utils/get-conversation-pairs.mjs'
+import { isEmpty } from 'lodash-es'
+import { pushRecord, setAbortController } from './shared.mjs'
+
+/**
+ * @param {Browser.Runtime.Port} port
+ * @param {string} question
+ * @param {Session} session
+ * @param {string} apiKey
+ * @param {string} modelName
+ */
+export async function generateAnswersWithOllamaApi(port, question, session, apiKey, modelName) {
+  const { controller, messageListener, disconnectListener } = setAbortController(port)
+
+  const config = await getUserConfig()
+  const prompt = getConversationPairs(
+    session.conversationRecords.slice(-config.maxConversationContextLength),
+    false,
+  )
+  // prompt.unshift({ role: 'system', content: await getOllamaApiPromptBase() })
+  prompt.push({ role: 'user', content: question })
+  const apiUrl = config.ollamaEndpoint
+
+  let answer = ''
+  let finished = false
+  const finish = () => {
+    finished = true
+    pushRecord(session, question, answer)
+    console.debug('conversation history', { content: session.conversationRecords })
+    port.postMessage({ answer: null, done: true, session: session })
+  }
+  await fetchSSE(`${apiUrl}/api/chat`, {
+    method: 'POST',
+    signal: controller.signal,
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${apiKey}`,
+    },
+    body: JSON.stringify({
+      messages: prompt,
+      model: modelName,
+      stream: true,
+      keep_alive: config.keepAliveTime === '-1' ? -1 : config.keepAliveTime,
+    }),
+    onMessage(message) {
+      console.debug('sse message', message)
+      if (finished) return
+      let data = message
+      const delta = data.message?.content
+      if (delta) {
+        answer += delta
+        port.postMessage({ answer: answer, done: false, session: null })
+      }
+      if (data.done_reason) {
+        finish()
+        return
+      }
+    },
+    async onStart() {},
+    async onEnd() {
+      port.postMessage({ done: true })
+      port.onMessage.removeListener(messageListener)
+      port.onDisconnect.removeListener(disconnectListener)
+    },
+    async onError(resp) {
+      port.onMessage.removeListener(messageListener)
+      port.onDisconnect.removeListener(disconnectListener)
+      if (resp instanceof Error) throw resp
+      const error = await resp.json().catch(() => ({}))
+      throw new Error(!isEmpty(error) ? JSON.stringify(error) : `${resp.status} ${resp.statusText}`)
+    },
+  })
+}
diff --git a/src/utils/fetch-ollama.mjs b/src/utils/fetch-ollama.mjs
diff --git a/src/utils/fetch-sse.mjs b/src/utils/fetch-sse.mjs

Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ async function mountComponent(siteConfig, userConfig) {`
`72`	`72`	`})`
`73`	`73`
`74`	`74`	`const position = {`
`75`		`- x: window.innerWidth - 300 - (Math.floor((20 / 100) * window.innerWidth)),`
	`75`	`+ x: window.innerWidth - 300 - Math.floor((20 / 100) * window.innerWidth),`
`76`	`76`	`y: window.innerHeight / 2 - 200,`
`77`	`77`	`}`
`78`	`78`	`const toolbarContainer = createElementAtPosition(position.x, position.y)`
Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,7 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "ChatGPTBox",`
`3`	`3`	`"description": "Integrating ChatGPT into your browser deeply, everything you need is here",`
`4`		`- "version": "2.5.6",`
	`4`	`+ "version": "2.5.7",`
`5`	`5`	`"manifest_version": 3,`
`6`	`6`	`"icons": {`
`7`	`7`	`"16": "logo.png",`