新增字幕提问功能

2024-03-17 23:31:48 +08:00
parent d6d7e17f84
commit a86ba9e09f
8 changed files with 271 additions and 16 deletions
--- a/src/biz/Body.tsx
+++ b/src/biz/Body.tsx
@@ -1,5 +1,7 @@
-import React, {useCallback, useEffect, useRef} from 'react'
+import React, {useCallback, useEffect, useMemo, useRef} from 'react'
 import {
+  setAskFold,
+  setAskQuestion,
  setAutoScroll,
  setAutoTranslate,
  setCheckAutoScroll,
@@ -14,6 +16,9 @@ import {useAppDispatch, useAppSelector} from '../hooks/redux'
 import {
  AiOutlineAim,
  AiOutlineCloseCircle,
+  BsDashSquare,
+  BsPlusSquare,
+  FaQuestion,
  FaRegArrowAltCircleDown,
  IoWarning,
  MdExpand,
@@ -24,6 +29,7 @@ import classNames from 'classnames'
 import toast from 'react-hot-toast'
 import SegmentCard from './SegmentCard'
 import {
+  ASK_ENABLED_DEFAULT,
  HEADER_HEIGHT,
  PAGE_SETTINGS,
  SEARCH_BAR_HEIGHT,
@@ -35,6 +41,7 @@ import {FaClipboardList} from 'react-icons/fa'
 import useTranslate from '../hooks/useTranslate'
 import {getSummarize} from '../util/biz_util'
 import {openUrl} from '@kky002/kky-util'
+import Markdown from '../components/Markdown'

 const Body = () => {
  const dispatch = useAppDispatch()
@@ -48,7 +55,12 @@ 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 {addSummarizeTask} = useTranslate()
+  const {addSummarizeTask, addAskTask} = useTranslate()
+  const askFold = useAppSelector(state => state.env.askFold)
+  const askQuestion = useAppSelector(state => state.env.askQuestion)
+  const askContent = useAppSelector(state => state.env.askContent)
+  const askStatus = useAppSelector(state => state.env.askStatus)
+  const askError = useAppSelector(state => state.env.askError)
  const bodyRef = useRef<any>()
  const curOffsetTop = useAppSelector(state => state.env.curOffsetTop)
  const checkAutoScroll = useAppSelector(state => state.env.checkAutoScroll)
@@ -56,7 +68,23 @@ const Body = () => {
  const totalHeight = useAppSelector(state => state.env.totalHeight)
  const curSummaryType = useAppSelector(state => state.env.tempData.curSummaryType)
  const title = useAppSelector(state => state.env.title)
+  const fontSize = useAppSelector(state => state.env.envData.fontSize)
  const searchText = useAppSelector(state => state.env.searchText)
+  const searchPlaceholder = useMemo(() => {
+    let placeholder = ''
+    if (envData.searchEnabled) {
+      if (envData.askEnabled??ASK_ENABLED_DEFAULT) {
+        placeholder = '搜索或提问字幕内容'
+      } else {
+        placeholder = '搜索字幕内容'
+      }
+    } else {
+      if (envData.askEnabled??ASK_ENABLED_DEFAULT) {
+        placeholder = '提问字幕内容'
+      }
+    }
+    return placeholder
+  }, [envData.askEnabled, envData.searchEnabled])

  const normalCallback = useCallback(() => {
    dispatch(setTempData({
@@ -102,6 +130,7 @@ const Body = () => {

  const onFoldAll = useCallback(() => {
    dispatch(setFoldAll(!foldAll))
+    dispatch(setAskFold(!foldAll))
    for (const segment of segments ?? []) {
      dispatch(setSegmentFold({
        segmentStartIdx: segment.startIdx,
@@ -149,6 +178,29 @@ const Body = () => {
    dispatch(setSearchText(''))
  }, [dispatch])

+  const onAsk = useCallback(() => {
+    if ((envData.askEnabled??ASK_ENABLED_DEFAULT) && searchText) {
+      const apiKey = envData.aiType === 'gemini'?envData.geminiApiKey:envData.apiKey
+      if (apiKey) {
+        if (segments != null && segments.length > 0) {
+          dispatch(setAskQuestion(searchText))
+          addAskTask(segments[0], searchText).catch(console.error)
+        }
+      } else {
+        dispatch(setPage(PAGE_SETTINGS))
+        toast.error('需要先设置ApiKey!')
+      }
+    }
+  }, [addAskTask, dispatch, envData.aiType, envData.apiKey, envData.askEnabled, envData.geminiApiKey, searchText, segments])
+
+  const onSetAsk = useCallback(() => {
+    dispatch(setSearchText(askQuestion??''))
+  }, [askQuestion, dispatch])
+
+  const onAskFold = useCallback(() => {
+    dispatch(setAskFold(!askFold))
+  }, [askFold, dispatch])
+
  // 自动滚动
  useEffect(() => {
    if (checkAutoScroll && curOffsetTop && autoScroll && !needScroll) {
@@ -193,8 +245,8 @@ const Body = () => {
    </div>

    {/* search */}
-    {envData.searchEnabled && <div className='px-2 py-1 flex flex-col relative'>
-      <input type='text' className='input input-xs bg-base-200' placeholder='搜索字幕内容' value={searchText} onChange={onSearchTextChange}/>
+    {(envData.searchEnabled ? envData.searchEnabled : (envData.askEnabled ?? ASK_ENABLED_DEFAULT)) && <div className='px-2 py-1 flex flex-col relative'>
+      <input type='text' className='input input-xs bg-base-200' placeholder={searchPlaceholder} value={searchText} onChange={onSearchTextChange}/>
      {searchText && <button className='absolute top-1 right-2 btn btn-ghost btn-xs btn-circle text-base-content/75' onClick={onClearSearchText}><AiOutlineCloseCircle/></button>}
    </div>}

@@ -213,6 +265,35 @@ const Body = () => {
           height: `${totalHeight - HEADER_HEIGHT - TITLE_HEIGHT - (envData.searchEnabled ? SEARCH_BAR_HEIGHT : 0)}px`
         }}
    >
+      {/* ask */}
+      {(envData.askEnabled??ASK_ENABLED_DEFAULT) && (searchText || askQuestion) &&
+        <div className='shadow bg-base-200 my-0.5 mx-1.5 p-1.5 rounded flex flex-col justify-center items-center'>
+          <div className='w-full relative flex justify-center min-h-[20px]'>
+            <div className='absolute left-0 top-0 bottom-0 text-xs select-none flex-center desc'>
+              {askFold
+                ? <BsPlusSquare className='cursor-pointer' onClick={onAskFold}/> :
+                <BsDashSquare className='cursor-pointer' onClick={onAskFold}/>}
+            </div>
+            <div className="tabs">
+              <a className="tab tab-lifted tab-xs tab-disabled cursor-default"></a>
+              <a className='tab tab-lifted tab-xs tab-active'><FaQuestion/>提问</a>
+              <a className="tab tab-lifted tab-xs tab-disabled cursor-default"></a>
+            </div>
+          </div>
+          {!askFold && askQuestion &&
+            <div className='link link-hover text-sm font-medium max-w-[90%]' onClick={onSetAsk}>{askQuestion}</div>}
+          {!askFold && askContent &&
+            <div className={classNames('font-medium max-w-[90%] mt-1', fontSize === 'large' ? 'text-sm' : 'text-xs')}>
+              <Markdown content={askContent}/>
+            </div>}
+          {!askFold && <button disabled={askStatus === 'pending'}
+                  className={classNames('btn btn-link btn-xs', askStatus === 'pending' && 'loading')}
+                  onClick={onAsk}>{askStatus === 'init' ? '点击提问' : (askStatus === 'pending' ? '生成中' : '重新生成')}</button>}
+          {!askFold && askStatus === 'init' && <div className='desc-lighter text-xs'>提问举例：这个视频说了什么</div>}
+          {!askFold && askError && <div className='text-xs text-error'>{askError}</div>}
+        </div>}
+
+      {/* segments */}
      {segments?.map((segment, segmentIdx) => <SegmentCard key={segment.startIdx} segment={segment}
                                                           segmentIdx={segmentIdx} bodyRef={bodyRef}/>)}

--- a/src/biz/Settings.tsx
+++ b/src/biz/Settings.tsx
@@ -398,7 +398,6 @@ const Settings = () => {
          <input id='askEnabled' type='checkbox' className='toggle toggle-primary' checked={askEnabledValue}
                 onChange={setAskEnabledValue}/>
        </FormItem>
-        <div className='desc text-xs'>在搜索框输入提问内容，然后按<span className='font-semibold font-mono'>Enter</span>即可提问。</div>
      </Section>
      <div className='flex justify-center gap-5'>
        <button className='btn btn-primary btn-sm' onClick={onSave}>保存</button>
--- a/src/chrome/openaiService.ts
+++ b/src/chrome/openaiService.ts
@@ -1,4 +1,12 @@
-import {getServerUrl} from '../util/biz_util'
+const getServerUrl = (serverUrl?: string) => {
+  if (!serverUrl) {
+    return 'https://api.openai.com'
+  }
+  if (serverUrl.endsWith('/')) {
+    serverUrl = serverUrl.slice(0, -1)
+  }
+  return serverUrl
+}

 export const handleChatCompleteTask = async (task: Task) => {
  const data = task.def.data
--- a/src/components/Markdown.tsx
+++ b/src/components/Markdown.tsx
@@ -0,0 +1,46 @@
+import classNames from 'classnames'
+import ReactMarkdown from 'react-markdown'
+import toast from 'react-hot-toast'
+
+function CopyBtn(props: {
+  content: string
+}) {
+  const {content} = props
+  return <div className='flex justify-center mt-1'>
+    <button className="btn btn-xs px-10 btn-primary normal-case" onClick={() => {
+      navigator.clipboard.writeText(content).then(() => {
+        toast.success('Copied!')
+      }).catch(console.error)
+    }}>Copy
+    </button>
+  </div>
+}
+
+function Markdown(props: {
+  content: string
+  codeBlockClass?: string
+}) {
+  const {content, codeBlockClass} = props
+
+  return <ReactMarkdown
+    className='markdown prose prose-sm dark:prose-invert prose-h1:text-center prose-h1:font-bold prose-h1:underline-offset-4 overflow-y-auto scrollbar-hide'
+    linkTarget={'_blank'}
+    components={{
+      code({node, inline, className, children, ...props}) {
+        if (inline) {
+          return <code
+            className={classNames(className, 'md-inline-block kbd kbd-xs rounded text-base-content/80')} {...props}>
+            {children}
+          </code>
+        } else {
+          return <code className={classNames(className, 'relative', codeBlockClass)} {...props}>
+            {children}
+            {className?.includes('language-copy') && <CopyBtn content={children[0] as string}/>}
+          </code>
+        }
+      }
+    }}
+  >{content}</ReactMarkdown>
+}
+
+export default Markdown
--- a/src/const.tsx
+++ b/src/const.tsx
@@ -9,6 +9,7 @@ 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_TYPE_ASK = 'ask'
 export const PROMPT_TYPES = [{
  name: '翻译',
  type: PROMPT_TYPE_TRANSLATE,
@@ -21,6 +22,9 @@ export const PROMPT_TYPES = [{
 }, {
  name: '总结',
  type: PROMPT_TYPE_SUMMARIZE_BRIEF,
+}, {
+  name: '提问',
+  type: PROMPT_TYPE_ASK,
 }]

 export const SUMMARIZE_TYPES = {
@@ -119,7 +123,20 @@ The video's subtitles:

 '''
 {{segment}}
-'''`
+'''`,
+  [PROMPT_TYPE_ASK]: `You are a helpful assistant who answers question related to video subtitles.
+Answer in language '{{language}}'.
+
+The video's title: '''{{title}}'''.
+The video's subtitles:
+
+'''
+{{segment}}
+'''
+
+Question: '''{{question}}'''
+Answer:
+`,
 }

 export const EVENT_EXPAND = 'expand'
@@ -157,7 +174,7 @@ export const SERVER_URL_THIRD = 'https://op.kongkongye.com'
 export const MODELS = [{
  code: 'gpt-3.5-turbo',
  name: 'gpt-3.5-turbo',
-  tokens: 16385,
+  tokens: 4096,
 }, {
  code: 'gpt-3.5-turbo-0125',
  name: 'gpt-3.5-turbo-0125',
@@ -167,7 +184,8 @@ export const MODELS = [{
  name: 'gpt-3.5-turbo-1106',
  tokens: 16385,
 }]
-export const MODEL_DEFAULT = MODELS[0].code
+export const GEMINI_TOKENS = 32768
+export const MODEL_DEFAULT = MODELS[1].code
 export const MODEL_MAP: {[key: string]: typeof MODELS[number]} = {}
 for (const model of MODELS) {
  MODEL_MAP[model.code] = model
--- a/src/hooks/useSearchService.ts
+++ b/src/hooks/useSearchService.ts
@@ -21,6 +21,9 @@ const useSearchService = () => {

  // reset search
  useEffect(() => {
+    if (!envData.searchEnabled) {
+      return
+    }
    const startTime = Date.now()
    const docs: Document[] = []
    for (const item of data?.body??[]) {
@@ -35,13 +38,13 @@ const useSearchService = () => {
    // 日志
    const endTime = Date.now()
    console.debug(`[Search]reset ${docs.length} docs, cost ${endTime-startTime}ms`)
-  }, [data?.body, dispatch, reset])
+  }, [data?.body, dispatch, envData.searchEnabled, reset])

  // search text
  useEffect(() => {
    const searchResult: Set<number> = new Set()

-    if (searchText) {
+    if (envData.searchEnabled && searchText) {
      // @ts-expect-error
      const documents: Document[] | undefined = search(searchText)
      if (documents != null) {
@@ -52,7 +55,7 @@ const useSearchService = () => {
    }

    dispatch(setSearchResult(searchResult))
-  }, [dispatch, search, searchText])
+  }, [dispatch, envData.searchEnabled, search, searchText])
 }

 export default useSearchService
--- a/src/hooks/useTranslate.ts
+++ b/src/hooks/useTranslate.ts
@@ -4,6 +4,9 @@ import {
  addTaskId,
  addTransResults,
  delTaskId,
+  setAskContent,
+  setAskError,
+  setAskStatus,
  setLastSummarizeTime,
  setLastTransTime,
  setSummaryContent,
@@ -15,6 +18,7 @@ import {
  LANGUAGES_MAP,
  MODEL_DEFAULT,
  PROMPT_DEFAULTS,
+  PROMPT_TYPE_ASK,
  PROMPT_TYPE_TRANSLATE,
  SUMMARIZE_LANGUAGE_DEFAULT,
  SUMMARIZE_THRESHOLD,
@@ -107,7 +111,7 @@ const useTranslate = () => {
                    content: prompt,
                  }
                ],
-                temperature: 0,
+                temperature: 0.25,
                n: 1,
                stream: false,
              },
@@ -176,7 +180,7 @@ const useTranslate = () => {
                  content: prompt,
                }
              ],
-              temperature: 0,
+              temperature: 0.5,
              n: 1,
              stream: false,
            },
@@ -196,6 +200,59 @@ const useTranslate = () => {
    }
  }, [dispatch, envData.aiType, envData.apiKey, envData.geminiApiKey, envData.model, envData.prompts, envData.serverUrl, summarizeLanguage.name, title])

+  const addAskTask = useCallback(async (segment: Segment, question: string) => {
+    if (segment.text.length >= SUMMARIZE_THRESHOLD) {
+      let prompt: string = envData.prompts?.[PROMPT_TYPE_ASK]??PROMPT_DEFAULTS[PROMPT_TYPE_ASK]
+      // replace params
+      prompt = prompt.replaceAll('{{language}}', summarizeLanguage.name)
+      prompt = prompt.replaceAll('{{title}}', title??'')
+      prompt = prompt.replaceAll('{{segment}}', segment.text)
+      prompt = prompt.replaceAll('{{question}}', question)
+
+      const taskDef: TaskDef = {
+        type: envData.aiType === 'gemini'?'geminiChatComplete':'chatComplete',
+        serverUrl: envData.serverUrl,
+        data: envData.aiType === 'gemini'
+          ?{
+              contents: [
+                {
+                  parts: [
+                    {
+                      text: prompt
+                    }
+                  ]
+                }
+              ],
+              generationConfig: {
+                maxOutputTokens: 2048
+              }
+            }
+          :{
+              model: envData.model??MODEL_DEFAULT,
+              messages: [
+                {
+                  role: 'user',
+                  content: prompt,
+                }
+              ],
+              temperature: 0.5,
+              n: 1,
+              stream: false,
+            },
+        extra: {
+          type: 'ask',
+          // startIdx: segment.startIdx,
+          apiKey: envData.apiKey,
+          geminiApiKey: envData.geminiApiKey,
+        }
+      }
+      console.debug('addAskTask', taskDef)
+      dispatch(setAskStatus({status: 'pending'}))
+      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])
+
  const handleTranslate = useMemoizedFn((task: Task, content: string) => {
    let map: {[key: string]: string} = {}
    try {
@@ -247,6 +304,13 @@ const useTranslate = () => {
    console.debug('setSummary', task.def.extra.startIdx, summaryType, obj, task.error)
  })

+  const handleAsk = useMemoizedFn((task: Task, content?: string) => {
+    dispatch(setAskContent({content}))
+    dispatch(setAskStatus({status: 'done'}))
+    dispatch(setAskError({error: task.error}))
+    console.debug('setAsk', content, task.error)
+  })
+
  const getTask = useCallback(async (taskId: string) => {
    const taskResp = await chrome.runtime.sendMessage({type: 'getTask', taskId})
    if (taskResp.code === 'ok') {
@@ -266,14 +330,16 @@ const useTranslate = () => {
          handleTranslate(task, content)
        } else if (taskType === 'summarize') { // 总结
          handleSummarize(task, content)
+        } else if (taskType === 'ask') { // 总结
+          handleAsk(task, content)
        }
      }
    } else {
      dispatch(delTaskId(taskId))
    }
-  }, [dispatch, envData.aiType, handleSummarize, handleTranslate])
+  }, [dispatch, envData.aiType, handleAsk, handleSummarize, handleTranslate])

-  return {getFetch, getTask, addTask, addSummarizeTask}
+  return {getFetch, getTask, addTask, addSummarizeTask, addAskTask}
 }

 export default useTranslate
--- a/src/redux/envReducer.ts
+++ b/src/redux/envReducer.ts
@@ -37,6 +37,13 @@ interface EnvState {
  lastTransTime?: number
  lastSummarizeTime?: number

+  // ask
+  askFold?: boolean
+  askQuestion?: string
+  askStatus: SummaryStatus
+  askError?: string
+  askContent?: string
+
  searchText: string
  searchResult: Set<number>
 }
@@ -53,6 +60,7 @@ const initialState: EnvState = {
  tempData: {
    curSummaryType: 'overview',
  },
+  askStatus: 'init',
  totalHeight: TOTAL_HEIGHT_DEF,
  autoScroll: true,
  currentTime: import.meta.env.VITE_ENV === 'web-dev' ? 30 : undefined,
@@ -195,6 +203,27 @@ export const slice = createSlice({
        }
      }
    },
+    setAskFold: (state, action: PayloadAction<boolean>) => {
+      state.askFold = action.payload
+    },
+    setAskQuestion: (state, action: PayloadAction<string | undefined>) => {
+      state.askQuestion = action.payload
+    },
+    setAskContent: (state, action: PayloadAction<{
+      content?: any
+    }>) => {
+      state.askContent = action.payload.content
+    },
+    setAskStatus: (state, action: PayloadAction<{
+      status: SummaryStatus
+    }>) => {
+      state.askStatus = action.payload.status
+    },
+    setAskError: (state, action: PayloadAction<{
+      error?: string
+    }>) => {
+      state.askError = action.payload.error
+    },
    setSegmentFold: (state, action: PayloadAction<{
      segmentStartIdx: number
      fold: boolean
@@ -259,6 +288,11 @@ export const slice = createSlice({
 })

 export const {
+  setAskFold,
+  setAskQuestion,
+  setAskStatus,
+  setAskError,
+  setAskContent,
  setTempReady,
  setTempData,
  setUploadedTranscript,