From 5078baca1b91637d96dfb0fe4df642c45c9637eb Mon Sep 17 00:00:00 2001
From: IndieKKY <indiekky@gmail.com>
Date: Sat, 20 Jan 2024 12:20:09 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81gemini?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/biz/Body.tsx            |  9 ++--
 src/biz/SegmentCard.tsx     |  5 ++-
 src/biz/Settings.tsx        | 29 +++++++-----
 src/chrome/openaiService.ts | 13 ++++++
 src/chrome/taskService.ts   |  5 ++-
 src/hooks/useTranslate.ts   | 90 +++++++++++++++++++++++++------------
 src/typings.d.ts            |  2 +-
 7 files changed, 104 insertions(+), 49 deletions(-)

diff --git a/src/biz/Body.tsx b/src/biz/Body.tsx
index 04006c4..cf8d3a9 100644
--- a/src/biz/Body.tsx
+++ b/src/biz/Body.tsx
@@ -45,7 +45,6 @@ const Body = () => {
   const foldAll = useAppSelector(state => state.env.foldAll)
   const envData = useAppSelector(state => state.env.envData)
   const compact = useAppSelector(state => state.env.tempData.compact)
-  const apiKey = useAppSelector(state => state.env.envData.apiKey)
   const floatKeyPointsSegIdx = useAppSelector(state => state.env.floatKeyPointsSegIdx)
   const translateEnable = useAppSelector(state => state.env.envData.translateEnable)
   const summarizeEnable = useAppSelector(state => state.env.envData.summarizeEnable)
@@ -76,6 +75,7 @@ const Body = () => {
   }, [dispatch])
 
   const onSummarizeAll = useCallback(() => {
+    const apiKey = envData.aiType === 'gemini'?envData.geminiApiKey:envData.apiKey
     if (!apiKey) {
       dispatch(setPage(PAGE_SETTINGS))
       toast.error('需要先设置ApiKey!')
@@ -98,7 +98,7 @@ const Body = () => {
       }
       toast.success(`已添加${segments_.length}个总结任务!`)
     }
-  }, [addSummarizeTask, apiKey, curSummaryType, dispatch, segments])
+  }, [addSummarizeTask, curSummaryType, dispatch, envData.aiType, envData.apiKey, envData.geminiApiKey, segments])
 
   const onFoldAll = useCallback(() => {
     dispatch(setFoldAll(!foldAll))
@@ -111,13 +111,14 @@ const Body = () => {
   }, [dispatch, foldAll, segments])
 
   const toggleAutoTranslateCallback = useCallback(() => {
-    if (envData.apiKey) {
+    const apiKey = envData.aiType === 'gemini'?envData.geminiApiKey:envData.apiKey
+    if (apiKey) {
       dispatch(setAutoTranslate(!autoTranslate))
     } else {
       dispatch(setPage(PAGE_SETTINGS))
       toast.error('需要先设置ApiKey!')
     }
-  }, [autoTranslate, dispatch, envData.apiKey])
+  }, [autoTranslate, dispatch, envData.aiType, envData.apiKey, envData.geminiApiKey])
 
   const onEnableAutoScroll = useCallback(() => {
     dispatch(setAutoScroll(true))
diff --git a/src/biz/SegmentCard.tsx b/src/biz/SegmentCard.tsx
index 8eed898..66fa963 100644
--- a/src/biz/SegmentCard.tsx
+++ b/src/biz/SegmentCard.tsx
@@ -66,19 +66,20 @@ const Summarize = (props: {
   const {segment, segmentIdx, summary, float} = props
 
   const dispatch = useAppDispatch()
-  const apiKey = useAppSelector(state => state.env.envData.apiKey)
+  const envData = useAppSelector(state => state.env.envData)
   const fontSize = useAppSelector(state => state.env.envData.fontSize)
   const curSummaryType = useAppSelector(state => state.env.tempData.curSummaryType)
   const {addSummarizeTask} = useTranslate()
 
   const onGenerate = useCallback(() => {
+    const apiKey = envData.aiType === 'gemini'?envData.geminiApiKey:envData.apiKey
     if (apiKey) {
       addSummarizeTask(curSummaryType, segment).catch(console.error)
     } else {
       dispatch(setPage(PAGE_SETTINGS))
       toast.error('需要先设置ApiKey!')
     }
-  }, [addSummarizeTask, apiKey, curSummaryType, dispatch, segment])
+  }, [addSummarizeTask, curSummaryType, dispatch, envData.aiType, envData.apiKey, envData.geminiApiKey, segment])
 
   const onCopy = useCallback(() => {
     if (summary != null) {
diff --git a/src/biz/Settings.tsx b/src/biz/Settings.tsx
index e3c481b..6047420 100644
--- a/src/biz/Settings.tsx
+++ b/src/biz/Settings.tsx
@@ -262,18 +262,21 @@ const Settings = () => {
 
       {aiTypeValue === 'gemini' && <Section title='gemini配置'>
         <FormItem title='ApiKey' htmlFor='geminiApiKey'>
-          <input id='geminiApiKey' type='text' className='input input-sm input-bordered w-full' placeholder='xxx' value={geminiApiKeyValue} onChange={onChangeGeminiApiKeyValue}/>
+          <input id='geminiApiKey' type='text' className='input input-sm input-bordered w-full' placeholder='xxx'
+                 value={geminiApiKeyValue} onChange={onChangeGeminiApiKeyValue}/>
         </FormItem>
         <div className='flex justify-center'>
-          <a className='link text-xs' onClick={toggleMoreFold}>{moreFold?'点击查看说明':'点击折叠说明'}</a>
+          <a className='link text-xs' onClick={toggleMoreFold}>{moreFold ? '点击查看说明' : '点击折叠说明'}</a>
         </div>
         {!moreFold && <div>
           <ul className='pl-3 list-decimal desc text-xs'>
-            <li>官方网址：<a className='link' href='https://makersuite.google.com/app/apikey' target='_blank' rel="noreferrer">Google AI Studio</a></li>
+            <li>官方网址：<a className='link' href='https://makersuite.google.com/app/apikey' target='_blank'
+                            rel="noreferrer">Google AI Studio</a> (目前免费)</li>
           </ul>
         </div>}
         <div className='flex justify-center'>
-          <a className='link text-xs' onClick={togglePromptsFold}>{promptsFold?'点击查看提示词':'点击折叠提示词'}</a>
+          <a className='link text-xs'
+             onClick={togglePromptsFold}>{promptsFold ? '点击查看提示词' : '点击折叠提示词'}</a>
         </div>
         {!promptsFold && <div>
           {PROMPT_TYPES.map((item, idx) => <FormItem key={item.type} title={<div>
@@ -282,16 +285,18 @@ const Settings = () => {
               setPromptsValue({
                 ...promptsValue,
                 // @ts-expect-error
-                [item.type]: PROMPT_DEFAULTS[item.type]??''
+                [item.type]: PROMPT_DEFAULTS[item.type] ?? ''
               })
-            }}>点击填充默认</div>
+            }}>点击填充默认
+            </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
-              })
-            }}/>
+            <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>}
diff --git a/src/chrome/openaiService.ts b/src/chrome/openaiService.ts
index efda32a..d6b0af1 100644
--- a/src/chrome/openaiService.ts
+++ b/src/chrome/openaiService.ts
@@ -18,3 +18,16 @@ export const handleChatCompleteTask = async (task: Task) => {
     throw new Error(`${task.resp.error.code as string??''} ${task.resp.error.message as string ??''}`)
   }
 }
+
+export const handleGeminiChatCompleteTask = async (task: Task) => {
+  const data = task.def.data
+  const resp = await fetch('https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent', {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'x-goog-api-key': task.def.extra.geminiApiKey,
+    },
+    body: JSON.stringify(data),
+  })
+  task.resp = await resp.json()
+}
diff --git a/src/chrome/taskService.ts b/src/chrome/taskService.ts
index d4c85e4..beb14f0 100644
--- a/src/chrome/taskService.ts
+++ b/src/chrome/taskService.ts
@@ -1,5 +1,5 @@
 import {TASK_EXPIRE_TIME} from '../const'
-import {handleChatCompleteTask} from './openaiService'
+import {handleChatCompleteTask, handleGeminiChatCompleteTask} from './openaiService'
 
 export const tasksMap = new Map<string, Task>()
 
@@ -11,6 +11,9 @@ export const handleTask = async (task: Task) => {
       case 'chatComplete':
         await handleChatCompleteTask(task)
         break
+      case 'geminiChatComplete':
+        await handleGeminiChatCompleteTask(task)
+        break
       default:
         // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
         throw new Error(`任务类型不支持: ${task.def.type}`)
diff --git a/src/hooks/useTranslate.ts b/src/hooks/useTranslate.ts
index a6cb1c2..fe5203c 100644
--- a/src/hooks/useTranslate.ts
+++ b/src/hooks/useTranslate.ts
@@ -82,23 +82,39 @@ const useTranslate = () => {
         prompt = prompt.replaceAll('{{subtitles}}', lineStr)
 
         const taskDef: TaskDef = {
-          type: 'chatComplete',
+          type: envData.aiType === 'gemini'?'geminiChatComplete':'chatComplete',
           serverUrl: envData.serverUrl,
-          data: {
-            model: envData.model??MODEL_DEFAULT,
-            messages: [
-              {
-                role: 'user',
-                content: prompt,
+          data: envData.aiType === 'gemini'
+            ?{
+                contents: [
+                  {
+                    parts: [
+                      {
+                        text: prompt
+                      }
+                    ]
+                  }
+                ],
+                generationConfig: {
+                  maxOutputTokens: 2048
+                }
               }
-            ],
-            temperature: 0,
-            n: 1,
-            stream: false,
-          },
+            :{
+                model: envData.model??MODEL_DEFAULT,
+                messages: [
+                  {
+                    role: 'user',
+                    content: prompt,
+                  }
+                ],
+                temperature: 0,
+                n: 1,
+                stream: false,
+              },
           extra: {
             type: 'translate',
             apiKey: envData.apiKey,
+            geminiApiKey: envData.geminiApiKey,
             startIdx,
             size: lines.length,
           }
@@ -117,10 +133,10 @@ const useTranslate = () => {
         dispatch(addTaskId(task.id))
       }
     }
-  }, [data?.body, dispatch, envData.apiKey, envData.fetchAmount, envData.serverUrl, envData.prompts, title, language.name])
+  }, [data?.body, envData.fetchAmount, envData.prompts, envData.aiType, envData.serverUrl, envData.model, envData.apiKey, envData.geminiApiKey, language.name, title, dispatch])
 
   const addSummarizeTask = useCallback(async (type: SummaryType, segment: Segment) => {
-    if (segment.text.length >= SUMMARIZE_THRESHOLD && envData.apiKey) {
+    if (segment.text.length >= SUMMARIZE_THRESHOLD) {
       let subtitles = ''
       for (const item of segment.items) {
         subtitles += formatTime(item.from) + ' ' + item.content + '\n'
@@ -135,25 +151,41 @@ const useTranslate = () => {
       prompt = prompt.replaceAll('{{segment}}', segment.text)
 
       const taskDef: TaskDef = {
-        type: 'chatComplete',
+        type: envData.aiType === 'gemini'?'geminiChatComplete':'chatComplete',
         serverUrl: envData.serverUrl,
-        data: {
-          model: envData.model??MODEL_DEFAULT,
-          messages: [
-            {
-              role: 'user',
-              content: prompt,
+        data: envData.aiType === 'gemini'
+          ?{
+              contents: [
+                {
+                  parts: [
+                    {
+                      text: prompt
+                    }
+                  ]
+                }
+              ],
+              generationConfig: {
+                maxOutputTokens: 2048
+              }
             }
-          ],
-          temperature: 0,
-          n: 1,
-          stream: false,
-        },
+          :{
+              model: envData.model??MODEL_DEFAULT,
+              messages: [
+                {
+                  role: 'user',
+                  content: prompt,
+                }
+              ],
+              temperature: 0,
+              n: 1,
+              stream: false,
+            },
         extra: {
           type: 'summarize',
           summaryType: type,
           startIdx: segment.startIdx,
           apiKey: envData.apiKey,
+          geminiApiKey: envData.geminiApiKey,
         }
       }
       console.debug('addSummarizeTask', taskDef)
@@ -162,7 +194,7 @@ const useTranslate = () => {
       const task = await chrome.runtime.sendMessage({type: 'addTask', taskDef})
       dispatch(addTaskId(task.id))
     }
-  }, [dispatch, envData.apiKey, envData.prompts, envData.serverUrl, summarizeLanguage.name, title])
+  }, [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} = {}
@@ -221,7 +253,7 @@ const useTranslate = () => {
       console.debug('getTask', taskResp.task)
       const task: Task = taskResp.task
       const taskType: string | undefined = task.def.extra?.type
-      const content = task.resp?.choices?.[0]?.message?.content?.trim()
+      const content = envData.aiType === 'gemini'?task.resp?.candidates[0]?.content?.parts[0]?.text?.trim():task.resp?.choices?.[0]?.message?.content?.trim()
       if (task.status === 'done') {
         // 异常提示
         if (task.error) {
@@ -239,7 +271,7 @@ const useTranslate = () => {
     } else {
       dispatch(delTaskId(taskId))
     }
-  }, [dispatch, handleSummarize, handleTranslate])
+  }, [dispatch, envData.aiType, handleSummarize, handleTranslate])
 
   return {getFetch, getTask, addTask, addSummarizeTask}
 }
diff --git a/src/typings.d.ts b/src/typings.d.ts
index 8245ce9..af1cd27 100644
--- a/src/typings.d.ts
+++ b/src/typings.d.ts
@@ -38,7 +38,7 @@ interface TempData {
 }
 
 interface TaskDef {
-  type: 'chatComplete'
+  type: 'chatComplete' | 'geminiChatComplete'
   serverUrl?: string
   data: any
   extra?: any