自定义prompt功能

src/biz/Body.tsx

@@ -31,7 +31,6 @@ const Body = () => {
   const floatKeyPointsSegIdx = useAppSelector(state => state.env.floatKeyPointsSegIdx)
   const translateEnable = useAppSelector(state => state.env.envData.translateEnable)
   const summarizeEnable = useAppSelector(state => state.env.envData.summarizeEnable)
-  const title = useAppSelector(state => state.env.title)
   const {addSummarizeTask} = useTranslate()
   const bodyRef = useRef<any>()
   const curOffsetTop = useAppSelector(state => state.env.curOffsetTop)
@@ -71,11 +70,11 @@ const Body = () => {
     }
     if (segments_.length < SUMMARIZE_ALL_THRESHOLD || confirm(`确定总结${segments_.length}个段落?`)) {
       for (const segment of segments_) {
-        addSummarizeTask(title, curSummaryType, segment).catch(console.error)
+        addSummarizeTask(curSummaryType, segment).catch(console.error)
       }
       toast.success(`已添加${segments_.length}个总结任务!`)
     }
-  }, [addSummarizeTask, apiKey, curSummaryType, dispatch, segments, title])
+  }, [addSummarizeTask, apiKey, curSummaryType, dispatch, segments])
 
   const onFoldAll = useCallback(() => {
     dispatch(setFoldAll(!foldAll))

src/biz/SegmentCard.tsx

@@ -68,18 +68,17 @@ const Summarize = (props: {
   const dispatch = useAppDispatch()
   const apiKey = useAppSelector(state => state.env.envData.apiKey)
   const fontSize = useAppSelector(state => state.env.envData.fontSize)
-  const title = useAppSelector(state => state.env.title)
   const curSummaryType = useAppSelector(state => state.env.tempData.curSummaryType)
   const {addSummarizeTask} = useTranslate()
 
   const onGenerate = useCallback(() => {
     if (apiKey) {
-      addSummarizeTask(title, curSummaryType, segment).catch(console.error)
+      addSummarizeTask(curSummaryType, segment).catch(console.error)
     } else {
       dispatch(setPage(PAGE_SETTINGS))
       toast.error('需要先设置ApiKey!')
     }
-  }, [addSummarizeTask, apiKey, curSummaryType, dispatch, segment, title])
+  }, [addSummarizeTask, apiKey, curSummaryType, dispatch, segment])
 
   const onCopy = useCallback(() => {
     if (summary != null) {

src/biz/Settings.tsx

@@ -6,6 +6,8 @@ import {
   LANGUAGE_DEFAULT,
   LANGUAGES,
   PAGE_MAIN,
+  PROMPT_DEFAULTS,
+  PROMPT_TYPES,
   SERVER_URL_THIRD,
   SUMMARIZE_LANGUAGE_DEFAULT,
   TRANSLATE_FETCH_DEFAULT,
@@ -69,8 +71,10 @@ const Settings = () => {
   const [wordsValue, setWordsValue] = useState(envData.words??WORDS_DEFAULT)
   const [fetchAmountValue, setFetchAmountValue] = useState(envData.fetchAmount??TRANSLATE_FETCH_DEFAULT)
   const [moreFold, {toggle: toggleMoreFold}] = useBoolean(true)
+  const [promptsFold, {toggle: togglePromptsFold}] = useBoolean(true)
   const fold = useAppSelector(state => state.env.fold)
   const totalHeight = useAppSelector(state => state.env.totalHeight)
+  const [promptsValue, setPromptsValue] = useState<{[key: string]: string}>(envData.prompts??{})
   const wordsList = useMemo(() => {
     const list = []
     for (let i = WORDS_MIN; i <= WORDS_MAX; i += WORDS_STEP) {
@@ -106,10 +110,11 @@ const Settings = () => {
       words: wordsValue,
       fetchAmount: fetchAmountValue,
       fontSize: fontSizeValue,
+      prompts: promptsValue,
     }))
     dispatch(setPage(PAGE_MAIN))
     toast.success('保存成功')
-  }, [fontSizeValue, apiKeyValue, autoExpandValue, dispatch, fetchAmountValue, hideOnDisableAutoTranslateValue, languageValue, serverUrlValue, summarizeEnableValue, summarizeFloatValue, summarizeLanguageValue, themeValue, transDisplayValue, translateEnableValue, wordsValue])
+  }, [promptsValue, fontSizeValue, apiKeyValue, autoExpandValue, dispatch, fetchAmountValue, hideOnDisableAutoTranslateValue, languageValue, serverUrlValue, summarizeEnableValue, summarizeFloatValue, summarizeLanguageValue, themeValue, transDisplayValue, translateEnableValue, wordsValue])
 
   const onCancel = useCallback(() => {
     dispatch(setPage(PAGE_MAIN))
@@ -202,6 +207,28 @@ const Settings = () => {
             <li>支持其他第三方代理，有问题可加群交流</li>
           </ul>
         </div>}
+        <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>}
       </Section>
       <Section title={<div className='flex items-center'>
         翻译配置

src/const.tsx

@@ -5,6 +5,108 @@ export const IFRAME_ID = 'bilibili-subtitle-iframe'
 export const STORAGE_ENV = 'bilibili-subtitle_env'
 export const STORAGE_TEMP = 'bilibili-subtitle_temp'
 
+export const PROMPT_TYPE_TRANSLATE = 'translate'
+export const PROMPT_TYPE_SUMMARIZE_OVERVIEW = 'summarize_overview'
+export const PROMPT_TYPE_SUMMARIZE_KEYPOINT = 'summarize_keypoint'
+export const PROMPT_TYPE_SUMMARIZE_BRIEF = 'summarize_brief'
+export const PROMPT_TYPES = [{
+  name: '翻译',
+  type: PROMPT_TYPE_TRANSLATE,
+}, {
+  name: '概览',
+  type: PROMPT_TYPE_SUMMARIZE_OVERVIEW,
+}, {
+  name: '要点',
+  type: PROMPT_TYPE_SUMMARIZE_KEYPOINT,
+}, {
+  name: '总结',
+  type: PROMPT_TYPE_SUMMARIZE_BRIEF,
+}]
+
+export const SUMMARIZE_TYPE_TO_PROMPT_TYPE = {
+  overview: PROMPT_TYPE_SUMMARIZE_OVERVIEW,
+  keypoint: PROMPT_TYPE_SUMMARIZE_KEYPOINT,
+  brief: PROMPT_TYPE_SUMMARIZE_BRIEF,
+}
+
+export const PROMPT_DEFAULTS = {
+  [PROMPT_TYPE_TRANSLATE]: `You are a professional translator. Translate following video subtitles to language '{{language}}'.
+Preserve incomplete sentence.
+Translate in the same json format.
+Answer in markdown json format.
+
+video subtitles:
+
+\`\`\`
+{{subtitles}}
+\`\`\``,
+  [PROMPT_TYPE_SUMMARIZE_OVERVIEW]: `You are a helpful assistant that summarize key points of video subtitle.
+Summarize 3 to 8 brief key points in language '{{language}}'.
+Answer in markdown json format.
+The emoji should be related to the key point and 1 char length.
+
+example output format:
+
+\`\`\`json
+[
+  {
+    "time": "03:00",
+    "emoji": "👍",
+    "key": "key point 1"
+  },
+  {
+    "time": "10:05",
+    "emoji": "😊",
+    "key": "key point 2"
+  }
+]
+\`\`\`
+
+The video's title: '''{{title}}'''.
+The video's subtitles:
+
+'''
+{{subtitles}}
+'''`,
+  [PROMPT_TYPE_SUMMARIZE_KEYPOINT]: `You are a helpful assistant that summarize key points of video subtitle.
+Summarize brief key points in language '{{language}}'.
+Answer in markdown json format.
+
+example output format:
+
+\`\`\`json
+[
+  "key point 1",
+  "key point 2"
+]
+\`\`\`
+
+The video's title: '''{{title}}'''.
+The video's subtitles:
+
+'''
+{{segment}}
+'''`,
+  [PROMPT_TYPE_SUMMARIZE_BRIEF]: `You are a helpful assistant that summarize video subtitle.
+Summarize in language '{{language}}'.
+Answer in markdown json format.
+
+example output format:
+
+\`\`\`json
+{
+  "summary": "brief summary"
+}
+\`\`\`
+
+The video's title: '''{{title}}'''.
+The video's subtitles:
+
+'''
+{{segment}}
+'''`
+}
+
 export const EVENT_EXPAND = 'expand'
 
 export const TASK_EXPIRE_TIME = 15*60*1000

src/hooks/useTranslate.ts

@@ -13,8 +13,11 @@ import {
 import {
   LANGUAGE_DEFAULT,
   LANGUAGES_MAP,
+  PROMPT_DEFAULTS,
+  PROMPT_TYPE_TRANSLATE,
   SUMMARIZE_LANGUAGE_DEFAULT,
   SUMMARIZE_THRESHOLD,
+  SUMMARIZE_TYPE_TO_PROMPT_TYPE,
   TRANSLATE_COOLDOWN,
   TRANSLATE_FETCH_DEFAULT,
 } from '../const'
@@ -32,6 +35,7 @@ const useTranslate = () => {
   const envData = useAppSelector(state => state.env.envData)
   const language = LANGUAGES_MAP[envData.language??LANGUAGE_DEFAULT]
   const summarizeLanguage = LANGUAGES_MAP[envData.summarizeLanguage??SUMMARIZE_LANGUAGE_DEFAULT]
+  const title = useAppSelector(state => state.env.title)
 
   /**
    * 获取下一个需要翻译的行
@@ -69,28 +73,22 @@ const useTranslate = () => {
         })
         let lineStr = JSON.stringify(linesMap).replaceAll('\n', '')
         lineStr = '```' + lineStr + '```'
+
+        let prompt: string = envData.prompts?.[PROMPT_TYPE_TRANSLATE]??PROMPT_DEFAULTS[PROMPT_TYPE_TRANSLATE]
+        // replace params
+        prompt = prompt.replaceAll('{{language}}', language.name)
+        prompt = prompt.replaceAll('{{title}}', title??'')
+        prompt = prompt.replaceAll('{{subtitles}}', lineStr)
+
         const taskDef: TaskDef = {
           type: 'chatComplete',
           serverUrl: envData.serverUrl,
           data: {
             model: 'gpt-3.5-turbo',
             messages: [
-              {
-                role: 'system',
-                content: 'You are a professional translator.'
-              },
               {
                 role: 'user',
-                content: `Translate following video subtitles to language '${language.name}'.
+                content: prompt,
-Preserve incomplete sentence.
-Translate in the same json format.
-Answer in markdown json format.
-
-video subtitles:
-
-\`\`\`
-${lineStr}
-\`\`\``
               }
             ],
             temperature: 0,
@@ -118,85 +116,23 @@ ${lineStr}
         dispatch(addTaskId(task.id))
       }
     }
-  }, [data?.body, dispatch, envData.apiKey, envData.fetchAmount, envData.serverUrl, language.name])
+  }, [data?.body, dispatch, envData.apiKey, envData.fetchAmount, envData.serverUrl, envData.prompts, title, language.name])
 
-  const addSummarizeTask = useCallback(async (title: string | undefined, type: SummaryType, segment: Segment) => {
+  const addSummarizeTask = useCallback(async (type: SummaryType, segment: Segment) => {
     if (segment.text.length >= SUMMARIZE_THRESHOLD && envData.apiKey) {
-      const title_ = title?`The video's title is '${title}'.`:''
       let subtitles = ''
       for (const item of segment.items) {
         subtitles += formatTime(item.from) + ' ' + item.content + '\n'
       }
-      let content
+      // @ts-expect-error
-      if (type === 'overview') {
+      const promptType: keyof typeof PROMPT_DEFAULTS = SUMMARIZE_TYPE_TO_PROMPT_TYPE[type]
-        content = `You are a helpful assistant that summarize key points of video subtitle.
+      let prompt: string = envData.prompts?.[promptType]??PROMPT_DEFAULTS[promptType]
-Summarize 3 to 8 brief key points in language '${summarizeLanguage.name}'.
+      // replace params
-Answer in markdown json format.
+      prompt = prompt.replaceAll('{{language}}', summarizeLanguage.name)
-The emoji should be related to the key point and 1 char length.
+      prompt = prompt.replaceAll('{{title}}', title??'')
+      prompt = prompt.replaceAll('{{subtitles}}', subtitles)
+      prompt = prompt.replaceAll('{{segment}}', segment.text)
 
-example output format:
-
-\`\`\`json
-[
-  {
-    "time": "03:00",
-    "emoji": "👍",
-    "key": "key point 1"
-  },
-  {
-    "time": "10:05",
-    "emoji": "😊",
-    "key": "key point 2"
-  }
-]
-\`\`\`
-
-The video's title: '''${title_}'''.
-The video's subtitles:
-
-'''
-${subtitles}
-'''`
-      } else if (type === 'keypoint') {
-        content = `You are a helpful assistant that summarize key points of video subtitle.
-Summarize brief key points in language '${summarizeLanguage.name}'.
-Answer in markdown json format.
-
-example output format:
-
-\`\`\`json
-[
-  "key point 1",
-  "key point 2"
-]
-\`\`\`
-
-The video's title: '''${title_}'''.
-The video's subtitles:
-
-'''
-${segment.text}
-'''`
-      } else if (type === 'brief') {
-        content = `You are a helpful assistant that summarize video subtitle.
-Summarize in language '${summarizeLanguage.name}'.
-Answer in markdown json format.
-
-example output format:
-
-\`\`\`json
-{
-  "summary": "brief summary"
-}
-\`\`\`
-
-The video's title: '''${title_}'''.
-The video's subtitles:
-
-'''
-${segment.text}
-'''`
-      }
       const taskDef: TaskDef = {
         type: 'chatComplete',
         serverUrl: envData.serverUrl,
@@ -205,7 +141,7 @@ ${segment.text}
           messages: [
             {
               role: 'user',
-              content,
+              content: prompt,
             }
           ],
           temperature: 0,
@@ -225,7 +161,7 @@ ${segment.text}
       const task = await chrome.runtime.sendMessage({type: 'addTask', taskDef})
       dispatch(addTaskId(task.id))
     }
-  }, [dispatch, envData.apiKey, envData.serverUrl, summarizeLanguage.name])
+  }, [dispatch, envData.apiKey, envData.prompts, envData.serverUrl, summarizeLanguage.name, title])
 
   const handleTranslate = useMemoizedFn((task: Task, content: string) => {
     let map: {[key: string]: string} = {}

src/typings.d.ts

@@ -14,6 +14,10 @@ interface EnvData {
   summarizeFloat?: boolean
   theme?: 'system' | 'light' | 'dark'
   fontSize?: 'normal' | 'large'
+
+  prompts?: {
+    [key: string]: string
+  }
 }
 
 interface TempData {