Support Gemini API

chathub-dev · Dec 14, 2023 · e620074 · e620074
1 parent 81e25f7
commit e620074
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 0 deletions.
diff --git a/package.json b/package.json
@@ -45,6 +45,7 @@
   },
   "dependencies": {
     "@floating-ui/react": "^0.26.1",
+    "@google/generative-ai": "^0.1.1",
     "@headlessui/react": "^1.7.17",
     "@heroicons/react": "^2.0.18",
     "@radix-ui/react-tooltip": "^1.0.7",

diff --git a/src/app/bots/gemini-api/index.ts b/src/app/bots/gemini-api/index.ts
@@ -0,0 +1,58 @@
+import { GoogleGenerativeAI, ChatSession } from '@google/generative-ai'
+import { AbstractBot, AsyncAbstractBot, SendMessageParams } from '../abstract-bot'
+import { getUserConfig } from '~services/user-config'
+
+interface ConversationContext {
+  chatSession: ChatSession
+}
+
+export class GeminiApiBot extends AbstractBot {
+  private conversationContext?: ConversationContext
+  sdk: GoogleGenerativeAI
+
+  constructor(public apiKey: string) {
+    super()
+    this.sdk = new GoogleGenerativeAI(apiKey)
+  }
+
+  async doSendMessage(params: SendMessageParams) {
+    if (!this.conversationContext) {
+      const model = this.sdk.getGenerativeModel({ model: 'gemini-pro' })
+      const chatSession = model.startChat()
+      this.conversationContext = { chatSession }
+    }
+
+    const result = await this.conversationContext.chatSession.sendMessageStream(params.prompt)
+
+    let text = ''
+    for await (const chunk of result.stream) {
+      const chunkText = chunk.text()
+      console.debug('gemini stream', chunkText)
+      text += chunkText
+      params.onEvent({ type: 'UPDATE_ANSWER', data: { text } })
+    }
+
+    if (!text) {
+      params.onEvent({ type: 'UPDATE_ANSWER', data: { text: 'Empty response' } })
+    }
+    params.onEvent({ type: 'DONE' })
+  }
+
+  resetConversation() {
+    this.conversationContext = undefined
+  }
+
+  get name() {
+    return 'Gemini Pro'
+  }
+}
+
+export class GeminiBot extends AsyncAbstractBot {
+  async initializeBot() {
+    const { geminiApiKey } = await getUserConfig()
+    if (!geminiApiKey) {
+      throw new Error('Gemini API key missing')
+    }
+    return new GeminiApiBot(geminiApiKey)
+  }
+}
diff --git a/src/app/bots/index.ts b/src/app/bots/index.ts
@@ -3,6 +3,7 @@ import { BardBot } from './bard'
 import { BingWebBot } from './bing'
 import { ChatGPTBot } from './chatgpt'
 import { ClaudeBot } from './claude'
+import { GeminiBot } from './gemini-api'
 import { GrokWebBot } from './grok'
 import { LMSYSBot } from './lmsys'
 import { PerplexityBot } from './perplexity'
@@ -28,6 +29,7 @@ export type BotId =
   | 'baichuan'
   | 'yi'
   | 'grok'
+  | 'gemini'
 
 export function createBotInstance(botId: BotId) {
   switch (botId) {
@@ -65,6 +67,8 @@ export function createBotInstance(botId: BotId) {
       return new PerplexityBot()
     case 'grok':
       return new GrokWebBot()
+    case 'gemini':
+      return new GeminiBot()
   }
 }
 

diff --git a/src/app/consts.ts b/src/app/consts.ts
@@ -15,6 +15,7 @@ import xunfeiLogo from '~/assets/xunfei-logo.png'
 import pplxLogo from '~/assets/pplx-logo.jpg'
 import yiLogo from '~/assets/yi-logo.svg'
 import grokLogo from '~/assets/grok-logo.png'
+import geminiLogo from '~/assets/gemini-logo.png'
 import { BotId } from './bots'
 
 export const CHATBOTS: Record<BotId, { name: string; avatar: string }> = {
@@ -50,6 +51,10 @@ export const CHATBOTS: Record<BotId, { name: string; avatar: string }> = {
     name: 'Falcon',
     avatar: falconLogo,
   },
+  gemini: {
+    name: 'Gemini Pro',
+    avatar: geminiLogo,
+  },
   grok: {
     name: 'Grok',
     avatar: grokLogo,

diff --git a/src/app/pages/SettingPage.tsx b/src/app/pages/SettingPage.tsx
@@ -4,8 +4,10 @@ import { useTranslation } from 'react-i18next'
 import { BiExport, BiImport } from 'react-icons/bi'
 import Browser from 'webextension-polyfill'
 import Button, { MotionButton } from '~app/components/Button'
+import { Input } from '~app/components/Input'
 import RadioGroup from '~app/components/RadioGroup'
 import Select from '~app/components/Select'
+import Blockquote from '~app/components/Settings/Blockquote'
 import ChatGPTAPISettings from '~app/components/Settings/ChatGPTAPISettings'
 import ChatGPTAzureSettings from '~app/components/Settings/ChatGPTAzureSettings'
 import ChatGPTOpenRouterSettings from '~app/components/Settings/ChatGPTOpenRouterSettings'
@@ -196,6 +198,30 @@ function SettingPage() {
                 <ClaudePoeSettings userConfig={userConfig} updateConfigValue={updateConfigValue} />
               )}
             </ChatBotSettingPanel>
+            <ChatBotSettingPanel title="Gemini Pro">
+              <div className="flex flex-col gap-1">
+                <p className="font-medium text-sm">
+                  API Key (
+                  <a
+                    href="https://makersuite.google.com/app/apikey"
+                    target="_blank"
+                    rel="noreferrer"
+                    className="underline"
+                  >
+                    how to create key
+                  </a>
+                  )
+                </p>
+                <Input
+                  className="w-[300px]"
+                  placeholder="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+                  value={userConfig.geminiApiKey}
+                  onChange={(e) => updateConfigValue({ geminiApiKey: e.currentTarget.value })}
+                  type="password"
+                />
+                <Blockquote className="mt-1">{t('Your keys are stored locally')}</Blockquote>
+              </div>
+            </ChatBotSettingPanel>
             <ChatBotSettingPanel title="Bing">
               <div className="flex flex-row gap-5 items-center">
                 <p className="font-medium text-base">{t('Chat style')}</p>

diff --git a/src/assets/gemini-logo.png b/src/assets/gemini-logo.png
diff --git a/src/services/user-config.ts b/src/services/user-config.ts
@@ -81,6 +81,7 @@ const userConfigWithDefaultValue = {
   openrouterApiKey: '',
   perplexityMode: PerplexityMode.Webapp,
   perplexityApiKey: '',
+  geminiApiKey: '',
 }
 
 export type UserConfig = typeof userConfigWithDefaultValue

diff --git a/yarn.lock b/yarn.lock
@@ -861,6 +861,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@google/generative-ai@npm:^0.1.1":
+  version: 0.1.1
+  resolution: "@google/generative-ai@npm:0.1.1"
+  checksum: 1b42c54490670769eeb9eaa620fd9bc947a80849eeee14998c4fa91c37abbd2072cd746dd97d302d2e09a94d2b8011e7e3b3e492932a9a3c54dbbdb03cb8c44f
+  languageName: node
+  linkType: hard
+
 "@headlessui/react@npm:^1.7.17":
   version: 1.7.17
   resolution: "@headlessui/react@npm:1.7.17"
@@ -2461,6 +2468,7 @@ __metadata:
   dependencies:
     "@crxjs/vite-plugin": "npm:^2.0.0-beta.18"
     "@floating-ui/react": "npm:^0.26.1"
+    "@google/generative-ai": "npm:^0.1.1"
     "@headlessui/react": "npm:^1.7.17"
     "@headlessui/tailwindcss": "npm:^0.2.0"
     "@heroicons/react": "npm:^2.0.18"