ollama支持

2024-04-12 10:18:44 +08:00
parent c40338a5f5
commit 60697769bb
9 changed files with 85 additions and 80 deletions
--- a/manifest.json
+++ b/manifest.json
@@ -6,6 +6,10 @@
  "permissions": [
    "storage"
  ],
+  "host_permissions": [
+    "http://localhost/*",
+    "http://127.0.0.1/*"
+  ],
  "background": {
    "service_worker": "src/chrome/background.ts"
  },
--- a/src/biz/Settings.tsx
+++ b/src/biz/Settings.tsx
@@ -3,6 +3,8 @@ import {setEnvData, setPage} from '../redux/envReducer'
 import {useAppDispatch, useAppSelector} from '../hooks/redux'
 import {
  ASK_ENABLED_DEFAULT,
+  CUSTOM_MODEL_TOKENS,
+  DEFAULT_SERVER_URL_OPENAI,
  GEMINI_TOKENS,
  HEADER_HEIGHT,
  LANGUAGE_DEFAULT,
@@ -69,6 +71,8 @@ const Settings = () => {
  const [geminiApiKeyValue, { onChange: onChangeGeminiApiKeyValue }] = useEventTarget({initialValue: envData.geminiApiKey??''})
  const [languageValue, { onChange: onChangeLanguageValue }] = useEventTarget({initialValue: envData.language??LANGUAGE_DEFAULT})
  const [modelValue, { onChange: onChangeModelValue }] = useEventTarget({initialValue: envData.model??MODEL_DEFAULT})
+  const [customModelValue, { onChange: onChangeCustomModelValue }] = useEventTarget({initialValue: envData.customModel})
+  const [customModelTokensValue, setCustomModelTokensValue] = useState(envData.customModelTokens)
  const [summarizeLanguageValue, { onChange: onChangeSummarizeLanguageValue }] = useEventTarget({initialValue: envData.summarizeLanguage??SUMMARIZE_LANGUAGE_DEFAULT})
  const [hideOnDisableAutoTranslateValue, setHideOnDisableAutoTranslateValue] = useState(envData.hideOnDisableAutoTranslate)
  const [themeValue, setThemeValue] = useState(envData.theme)
@@ -114,6 +118,8 @@ const Settings = () => {
      apiKey: apiKeyValue,
      serverUrl: serverUrlValue,
      model: modelValue,
+      customModel: customModelValue,
+      customModelTokens: customModelTokensValue,
      geminiApiKey: geminiApiKeyValue,
      translateEnable: translateEnableValue,
      language: languageValue,
@@ -133,7 +139,7 @@ const Settings = () => {
    }))
    dispatch(setPage(PAGE_MAIN))
    toast.success('保存成功')
-  }, [dispatch, autoExpandValue, aiTypeValue, apiKeyValue, serverUrlValue, modelValue, geminiApiKeyValue, translateEnableValue, languageValue, hideOnDisableAutoTranslateValue, themeValue, transDisplayValue, summarizeEnableValue, summarizeFloatValue, summarizeLanguageValue, wordsValue, fetchAmountValue, fontSizeValue, promptsValue, searchEnabledValue, cnSearchEnabledValue, askEnabledValue])
+  }, [dispatch, autoExpandValue, aiTypeValue, apiKeyValue, serverUrlValue, modelValue, customModelValue, customModelTokensValue, geminiApiKeyValue, translateEnableValue, languageValue, hideOnDisableAutoTranslateValue, themeValue, transDisplayValue, summarizeEnableValue, summarizeFloatValue, summarizeLanguageValue, wordsValue, fetchAmountValue, fontSizeValue, promptsValue, searchEnabledValue, cnSearchEnabledValue, askEnabledValue])

  const onCancel = useCallback(() => {
    dispatch(setPage(PAGE_MAIN))
@@ -224,7 +230,7 @@ const Settings = () => {
        </FormItem>
        <FormItem title='服务器' htmlFor='serverUrl'>
          <input id='serverUrl' type='text' className='input input-sm input-bordered w-full'
-                 placeholder='默认使用官方地址' value={serverUrlValue}
+                 placeholder={DEFAULT_SERVER_URL_OPENAI} value={serverUrlValue}
                 onChange={e => setServerUrlValue(e.target.value)}/>
        </FormItem>
        <div>
@@ -233,7 +239,7 @@ const Settings = () => {
            <div>官方网址：<a className='link link-primary' href='https://platform.openai.com/' target='_blank'
                             rel="noreferrer">点击访问</a></div>
            <div>服务器地址：<a className='link link-primary'
-                               onClick={() => setServerUrlValue('https://api.openai.com')}
+                               onClick={() => setServerUrlValue(DEFAULT_SERVER_URL_OPENAI)}
                               rel='noreferrer'>点击设置</a></div>
            <div className='flex justify-center font-semibold'>【第三方代理】</div>
            <div>代理网址：<a className='link link-primary' href='https://api.openai-up.com/register?aff=varM'
@@ -251,31 +257,14 @@ const Settings = () => {
            {MODELS.map(model => <option key={model.code} value={model.code}>{model.name}</option>)}
          </select>
        </FormItem>
-        <div className='flex justify-center'>
-          <a className='link text-xs'
-             onClick={togglePromptsFold}>{promptsFold ? '点击查看提示词' : '点击折叠提示词'}</a>
-        </div>
-        {!promptsFold && <div>
-          {PROMPT_TYPES.map((item, idx) => <FormItem key={item.type} title={<div>
-            <div>{item.name}</div>
-            <div className='link text-xs' onClick={() => {
-              setPromptsValue({
-                ...promptsValue,
-                // @ts-expect-error
-                [item.type]: PROMPT_DEFAULTS[item.type] ?? ''
-              })
-            }}>点击填充默认
-            </div>
-          </div>} htmlFor={`prompt-${item.type}`}>
-            <textarea id={`prompt-${item.type}`} className='mt-2 textarea input-bordered w-full'
-                      placeholder='留空使用默认提示词' value={promptsValue[item.type] ?? ''} onChange={(e) => {
-                        setPromptsValue({
-                          ...promptsValue,
-                          [item.type]: e.target.value
-                        })
-                      }}/>
-          </FormItem>)}
-        </div>}
+        {modelValue === 'custom' && <FormItem title='模型名' htmlFor='customModel'>
+          <input id='customModel' type='text' className='input input-sm input-bordered w-full' placeholder='llama2'
+                 value={customModelValue} onChange={onChangeCustomModelValue}/>
+        </FormItem>}
+        {modelValue === 'custom' && <FormItem title='Token上限' htmlFor='customModelTokens'>
+          <input id='customModelTokens' type='number' className='input input-sm input-bordered w-full' placeholder={''+CUSTOM_MODEL_TOKENS}
+                 value={customModelTokensValue} onChange={e => setCustomModelTokensValue(e.target.value?parseInt(e.target.value):undefined)}/>
+        </FormItem>}
      </Section>}

      {aiTypeValue === 'gemini' && <Section title='gemini配置'>
@@ -291,11 +280,9 @@ const Settings = () => {
            <div className='text-xs text-error flex items-center'><IoWarning className='text-sm text-warning'/>谷歌模型安全要求比较高，有些视频可能无法生成总结!</div>
          </div>
        </div>
-        <div className='flex justify-center'>
-          <a className='link text-xs'
-             onClick={togglePromptsFold}>{promptsFold ? '点击查看提示词' : '点击折叠提示词'}</a>
-        </div>
-        {!promptsFold && <div>
+      </Section>}
+
+      <Section title='提示词配置'>
        {PROMPT_TYPES.map((item, idx) => <FormItem key={item.type} title={<div>
          <div>{item.name}</div>
          <div className='link text-xs' onClick={() => {
@@ -315,8 +302,7 @@ const Settings = () => {
                      })
                    }}/>
        </FormItem>)}
-        </div>}
-      </Section>}
+      </Section>

      <Section title={<div className='flex items-center'>
        翻译配置
@@ -329,7 +315,8 @@ const Settings = () => {
                 onChange={setTranslateEnableValue}/>
        </FormItem>
        <FormItem title='目标语言' htmlFor='language'>
-          <select id='language' className="select select-sm select-bordered" value={languageValue} onChange={onChangeLanguageValue}>
+          <select id='language' className="select select-sm select-bordered" value={languageValue}
+                  onChange={onChangeLanguageValue}>
            {LANGUAGES.map(language => <option key={language.code} value={language.code}>{language.name}</option>)}
          </select>
        </FormItem>
--- a/src/chrome/openaiService.ts
+++ b/src/chrome/openaiService.ts
@@ -1,6 +1,8 @@
+import {DEFAULT_SERVER_URL_OPENAI} from '../const'
+
 const getServerUrl = (serverUrl?: string) => {
  if (!serverUrl) {
-    return 'https://api.openai.com'
+    return DEFAULT_SERVER_URL_OPENAI
  }
  if (serverUrl.endsWith('/')) {
    serverUrl = serverUrl.slice(0, -1)
--- a/src/const.tsx
+++ b/src/const.tsx
@@ -169,7 +169,8 @@ export const SUMMARIZE_THRESHOLD = 100
 export const SUMMARIZE_LANGUAGE_DEFAULT = 'cn'
 export const SUMMARIZE_ALL_THRESHOLD = 5
 export const ASK_ENABLED_DEFAULT = true
-export const SERVER_URL_OPENAI = 'https://api.openai.com'
+export const DEFAULT_SERVER_URL_OPENAI = 'https://api.openai.com'
+export const CUSTOM_MODEL_TOKENS = 16385

 export const MODELS = [{
  code: 'gpt-3.5-turbo',
@@ -183,6 +184,9 @@ export const MODELS = [{
  code: 'gpt-3.5-turbo-1106',
  name: 'gpt-3.5-turbo-1106',
  tokens: 16385,
+}, {
+  code: 'custom',
+  name: '自定义',
 }]
 export const GEMINI_TOKENS = 32768
 export const MODEL_DEFAULT = MODELS[1].code
--- a/src/hooks/useSubtitleService.ts
+++ b/src/hooks/useSubtitleService.ts
@@ -15,18 +15,9 @@ import {
  setUrl,
 } from '../redux/envReducer'
 import {EventBusContext} from '../Router'
-import {
-  EVENT_EXPAND,
-  GEMINI_TOKENS,
-  MODEL_DEFAULT,
-  MODEL_MAP,
-  TOTAL_HEIGHT_MAX,
-  TOTAL_HEIGHT_MIN,
-  WORDS_MIN,
-  WORDS_RATE
-} from '../const'
+import {EVENT_EXPAND, GEMINI_TOKENS, TOTAL_HEIGHT_MAX, TOTAL_HEIGHT_MIN, WORDS_MIN, WORDS_RATE} from '../const'
 import {useInterval} from 'ahooks'
-import {getWholeText} from '../util/biz_util'
+import {getModelMaxTokens, getWholeText} from '../util/biz_util'

 /**
 * Service是单例，类似后端的服务概念
@@ -172,7 +163,7 @@ const useSubtitleService = () => {
          if (envData.aiType === 'gemini') {
            size = GEMINI_TOKENS*WORDS_RATE
          } else {
-            size = (MODEL_MAP[envData.model??MODEL_DEFAULT]?.tokens??4000)*WORDS_RATE
+            size = getModelMaxTokens(envData)*WORDS_RATE
          }
        }
        size = Math.max(size, WORDS_MIN)
@@ -209,7 +200,7 @@ const useSubtitleService = () => {
      }
    }
    dispatch(setSegments(segments))
-  }, [data?.body, dispatch, envData.aiType, envData.model, envData.summarizeEnable, envData.words])
+  }, [data?.body, dispatch, envData])

  // 每秒更新当前视频时间
  useInterval(() => {
--- a/src/hooks/useTranslate.ts
+++ b/src/hooks/useTranslate.ts
@@ -16,7 +16,6 @@ import {
 import {
  LANGUAGE_DEFAULT,
  LANGUAGES_MAP,
-  MODEL_DEFAULT,
  PROMPT_DEFAULTS,
  PROMPT_TYPE_ASK,
  PROMPT_TYPE_TRANSLATE,
@@ -28,7 +27,7 @@ import {
 } from '../const'
 import toast from 'react-hot-toast'
 import {useMemoizedFn} from 'ahooks/es'
-import {extractJsonArray, extractJsonObject} from '../util/biz_util'
+import {extractJsonArray, extractJsonObject, getModel} from '../util/biz_util'
 import {formatTime} from '../util/util'

 const useTranslate = () => {
@@ -104,7 +103,7 @@ const useTranslate = () => {
                }
              }
            :{
-                model: envData.model??MODEL_DEFAULT,
+                model: getModel(envData),
                messages: [
                  {
                    role: 'user',
@@ -137,7 +136,7 @@ const useTranslate = () => {
        dispatch(addTaskId(task.id))
      }
    }
-  }, [data?.body, envData.fetchAmount, envData.prompts, envData.aiType, envData.serverUrl, envData.model, envData.apiKey, envData.geminiApiKey, language.name, title, dispatch])
+  }, [data?.body, envData, language.name, title, dispatch])

  const addSummarizeTask = useCallback(async (type: SummaryType, segment: Segment) => {
    if (segment.text.length >= SUMMARIZE_THRESHOLD) {
@@ -173,7 +172,7 @@ const useTranslate = () => {
              }
            }
          :{
-              model: envData.model??MODEL_DEFAULT,
+              model: getModel(envData),
              messages: [
                {
                  role: 'user',
@@ -198,7 +197,7 @@ const useTranslate = () => {
      const task = await chrome.runtime.sendMessage({type: 'addTask', taskDef})
      dispatch(addTaskId(task.id))
    }
-  }, [dispatch, envData.aiType, envData.apiKey, envData.geminiApiKey, envData.model, envData.prompts, envData.serverUrl, summarizeLanguage.name, title])
+  }, [dispatch, envData, summarizeLanguage.name, title])

  const addAskTask = useCallback(async (segment: Segment, question: string) => {
    if (segment.text.length >= SUMMARIZE_THRESHOLD) {
@@ -228,7 +227,7 @@ const useTranslate = () => {
              }
            }
          :{
-              model: envData.model??MODEL_DEFAULT,
+              model: getModel(envData),
              messages: [
                {
                  role: 'user',
@@ -251,7 +250,7 @@ const useTranslate = () => {
      const task = await chrome.runtime.sendMessage({type: 'addTask', taskDef})
      dispatch(addTaskId(task.id))
    }
-  }, [dispatch, envData.aiType, envData.apiKey, envData.geminiApiKey, envData.model, envData.prompts, envData.serverUrl, summarizeLanguage.name, title])
+  }, [dispatch, envData, summarizeLanguage.name, title])

  const handleTranslate = useMemoizedFn((task: Task, content: string) => {
    let map: {[key: string]: string} = {}
--- a/src/redux/envReducer.ts
+++ b/src/redux/envReducer.ts
@@ -1,7 +1,7 @@
 import {createSlice, PayloadAction} from '@reduxjs/toolkit'
 import {find} from 'lodash-es'
 import {getDevData} from '../util/biz_util'
-import {SERVER_URL_OPENAI, TOTAL_HEIGHT_DEF} from '../const'
+import {DEFAULT_SERVER_URL_OPENAI, TOTAL_HEIGHT_DEF} from '../const'

 interface EnvState {
  envData: EnvData
@@ -51,7 +51,7 @@ interface EnvState {

 const initialState: EnvState = {
  envData: {
-    serverUrl: SERVER_URL_OPENAI,
+    serverUrl: DEFAULT_SERVER_URL_OPENAI,
    translateEnable: true,
    summarizeEnable: true,
    autoExpand: true,
--- a/src/typings.d.ts
+++ b/src/typings.d.ts
@@ -7,6 +7,8 @@ interface EnvData {
  apiKey?: string
  serverUrl?: string
  model?: string
+  customModel?: string
+  customModelTokens?: number
  // gemini
  geminiApiKey?: string

--- a/src/util/biz_util.ts
+++ b/src/util/biz_util.ts
@@ -1,5 +1,5 @@
 import devData from '../data/data.json'
-import {APP_DOM_ID, SUMMARIZE_TYPES} from '../const'
+import {APP_DOM_ID, CUSTOM_MODEL_TOKENS, MODEL_DEFAULT, MODEL_MAP, SUMMARIZE_TYPES} from '../const'
 import {isDarkMode} from '@kky002/kky-util'
 import toast from 'react-hot-toast'
 import {findIndex} from 'lodash-es'
@@ -123,6 +123,22 @@ export const getServerUrl = (serverUrl?: string) => {
  return serverUrl
 }

+export const getModel = (envData: EnvData) => {
+  if (envData.model === 'custom') {
+    return envData.customModel
+  } else {
+    return envData.model
+  }
+}
+
+export const getModelMaxTokens = (envData: EnvData) => {
+  if (envData.model === 'custom') {
+    return envData.customModelTokens??CUSTOM_MODEL_TOKENS
+  } else {
+    return MODEL_MAP[envData.model??MODEL_DEFAULT]?.tokens??4000
+  }
+}
+
 export const setTheme = (theme: EnvData['theme']) => {
  const appRoot = document.getElementById(APP_DOM_ID)
  if (appRoot != null) {