From a86ba9e09f53b48263290bc5f6ddd58dcbf365e5 Mon Sep 17 00:00:00 2001
From: IndieKKY <indiekky@gmail.com>
Date: Sun, 17 Mar 2024 23:31:48 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E5=AD=97=E5=B9=95=E6=8F=90?=
 =?UTF-8?q?=E9=97=AE=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/biz/Body.tsx              | 89 +++++++++++++++++++++++++++++++++--
 src/biz/Settings.tsx          |  1 -
 src/chrome/openaiService.ts   | 10 +++-
 src/components/Markdown.tsx   | 46 ++++++++++++++++++
 src/const.tsx                 | 24 ++++++++--
 src/hooks/useSearchService.ts |  9 ++--
 src/hooks/useTranslate.ts     | 74 +++++++++++++++++++++++++++--
 src/redux/envReducer.ts       | 34 +++++++++++++
 8 files changed, 271 insertions(+), 16 deletions(-)
 create mode 100644 src/components/Markdown.tsx

diff --git a/src/biz/Body.tsx b/src/biz/Body.tsx
index a4f6d28..2384e72 100644
--- 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}/>)}
 
diff --git a/src/biz/Settings.tsx b/src/biz/Settings.tsx
index ff88a7a..9bbced1 100644
--- 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>
diff --git a/src/chrome/openaiService.ts b/src/chrome/openaiService.ts
index d6b0af1..58bca3b 100644
--- 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
diff --git a/src/components/Markdown.tsx b/src/components/Markdown.tsx
new file mode 100644
index 0000000..ac22bbd
--- /dev/null
+++ 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
diff --git a/src/const.tsx b/src/const.tsx
index 179f887..c704b6a 100644
--- 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
diff --git a/src/hooks/useSearchService.ts b/src/hooks/useSearchService.ts
index b549ef8..cd16d64 100644
--- 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
diff --git a/src/hooks/useTranslate.ts b/src/hooks/useTranslate.ts
index fe5203c..12ae77a 100644
--- 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
diff --git a/src/redux/envReducer.ts b/src/redux/envReducer.ts
index ba2570b..b88f5b2 100644
--- 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,