From 2fae6e4a3eb8fd7bb0677832473345217d06c239 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Thu, 12 Dec 2024 14:23:59 +0800
Subject: [PATCH] feat: add web search for google gemini modal
 gemini-2.0-flash-exp

---
 electron-builder.yml                          | 13 +++----
 package.json                                  |  2 +-
 .../src/components/Icons/WebSearchIcon.tsx    | 15 ++++++++
 .../components/Popups/SelectModelPopup.tsx    |  5 +--
 src/renderer/src/config/models.ts             |  4 +++
 src/renderer/src/i18n/locales/en-us.json      |  1 +
 src/renderer/src/i18n/locales/ru-ru.json      |  1 +
 src/renderer/src/i18n/locales/zh-cn.json      |  1 +
 src/renderer/src/i18n/locales/zh-tw.json      |  1 +
 .../src/pages/home/Inputbar/Inputbar.tsx      | 18 ++++++++--
 .../ProviderSettings/ProviderSetting.tsx      | 35 +++++++++++--------
 src/renderer/src/providers/GeminiProvider.ts  |  2 ++
 src/renderer/src/types/index.ts               |  1 +
 yarn.lock                                     | 10 +++---
 14 files changed, 74 insertions(+), 35 deletions(-)
 create mode 100644 src/renderer/src/components/Icons/WebSearchIcon.tsx

diff --git a/electron-builder.yml b/electron-builder.yml
index b0a66a89..0241cd77 100644
--- a/electron-builder.yml
+++ b/electron-builder.yml
@@ -62,12 +62,7 @@ electronDownload:
 afterSign: scripts/notarize.js
 releaseInfo:
   releaseNotes: |
-    增加快捷键切换助手和话题显示
-    历史消息懒加载 by @1355873789
-    更快的应用更新下载速度 by @1355873789
-    更加清晰的模型分组
-    修复部分代码块无法正常显示问题
-    增加应用更新内容显示
-    消息发送增加 Ctrl + Enter 快捷键
-    清除上下文消息点击可以撤销
-    增加 Top-P 设置选项
+    修复对话消息无法删除问题
+    支持 Gemini 模型联网搜索
+    发送消息增加 Command + Enter 快捷键 by @duanyongcheng
+    Windows 版本样式优化
diff --git a/package.json b/package.json
index 82cace0a..464e1629 100644
--- a/package.json
+++ b/package.json
@@ -57,7 +57,7 @@
     "@electron-toolkit/eslint-config-prettier": "^2.0.0",
     "@electron-toolkit/eslint-config-ts": "^1.0.1",
     "@electron-toolkit/tsconfig": "^1.0.1",
-    "@google/generative-ai": "^0.16.0",
+    "@google/generative-ai": "^0.21.0",
     "@hello-pangea/dnd": "^16.6.0",
     "@kangfenmao/keyv-storage": "^0.1.0",
     "@reduxjs/toolkit": "^2.2.5",
diff --git a/src/renderer/src/components/Icons/WebSearchIcon.tsx b/src/renderer/src/components/Icons/WebSearchIcon.tsx
new file mode 100644
index 00000000..42ab971e
--- /dev/null
+++ b/src/renderer/src/components/Icons/WebSearchIcon.tsx
@@ -0,0 +1,15 @@
+import { GlobalOutlined } from '@ant-design/icons'
+import React, { FC } from 'react'
+import styled from 'styled-components'
+
+const WebSearchIcon: FC<React.DetailedHTMLProps<React.HTMLAttributes<HTMLElement>, HTMLElement>> = (props) => {
+  return <Icon {...(props as any)} />
+}
+
+const Icon = styled(GlobalOutlined)`
+  color: var(--color-link);
+  font-size: 12px;
+  margin-left: 4px;
+`
+
+export default WebSearchIcon
diff --git a/src/renderer/src/components/Popups/SelectModelPopup.tsx b/src/renderer/src/components/Popups/SelectModelPopup.tsx
index 714fe276..beb2fbfa 100644
--- a/src/renderer/src/components/Popups/SelectModelPopup.tsx
+++ b/src/renderer/src/components/Popups/SelectModelPopup.tsx
@@ -1,7 +1,7 @@
 import { PushpinOutlined, SearchOutlined } from '@ant-design/icons'
 import VisionIcon from '@renderer/components/Icons/VisionIcon'
 import { TopView } from '@renderer/components/TopView'
-import { getModelLogo, isVisionModel } from '@renderer/config/models'
+import { getModelLogo, isVisionModel, isWebSearchModel } from '@renderer/config/models'
 import db from '@renderer/databases'
 import { useProviders } from '@renderer/hooks/useProvider'
 import { getModelUniqId } from '@renderer/services/ModelService'
@@ -12,6 +12,7 @@ import { useEffect, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
 
+import WebSearchIcon from '../Icons/WebSearchIcon'
 import { HStack } from '../Layout'
 import Scrollbar from '../Scrollbar'
 
@@ -73,7 +74,7 @@ const PopupContainer: React.FC<PopupContainerProps> = ({ model, resolve }) => {
           label: (
             <ModelItem>
               <span>
-                {m?.name} {isVisionModel(m) && <VisionIcon />}
+                {m?.name} {isVisionModel(m) && <VisionIcon />} {isWebSearchModel(m) && <WebSearchIcon />}
               </span>
               <PinIcon
                 onClick={(e) => {
diff --git a/src/renderer/src/config/models.ts b/src/renderer/src/config/models.ts
index a3b2be42..eee2af0f 100644
--- a/src/renderer/src/config/models.ts
+++ b/src/renderer/src/config/models.ts
@@ -1056,3 +1056,7 @@ export function isVisionModel(model: Model): boolean {
 export function isSupportedModel(model: OpenAI.Models.Model): boolean {
   return !NOT_SUPPORTED_REGEX.test(model.id)
 }
+
+export function isWebSearchModel(model: Model): boolean {
+  return model?.provider === 'gemini' && model?.id === 'gemini-2.0-flash-exp'
+}
diff --git a/src/renderer/src/i18n/locales/en-us.json b/src/renderer/src/i18n/locales/en-us.json
index 792eb633..9d6c22b1 100644
--- a/src/renderer/src/i18n/locales/en-us.json
+++ b/src/renderer/src/i18n/locales/en-us.json
@@ -80,6 +80,7 @@
       "input.topics": " Topics ",
       "input.translate": "Translate to English",
       "input.upload": "Upload image or document file",
+      "input.web_search": "Enable web search",
       "message.new.branch": "New Branch",
       "message.new.branch.created": "New Branch Created",
       "message.new.context": "New Context",
diff --git a/src/renderer/src/i18n/locales/ru-ru.json b/src/renderer/src/i18n/locales/ru-ru.json
index 96e2f777..1fff8ae2 100644
--- a/src/renderer/src/i18n/locales/ru-ru.json
+++ b/src/renderer/src/i18n/locales/ru-ru.json
@@ -80,6 +80,7 @@
       "input.topics": " Топики ",
       "input.translate": "Перевести на английский",
       "input.upload": "Загрузить изображение или документ",
+      "input.web_search": "Включить веб-поиск",
       "message.new.branch": "Новая ветка",
       "message.new.branch.created": "Новая ветка создана",
       "message.new.context": "Новый контекст",
diff --git a/src/renderer/src/i18n/locales/zh-cn.json b/src/renderer/src/i18n/locales/zh-cn.json
index 8a5b9dfa..92c26120 100644
--- a/src/renderer/src/i18n/locales/zh-cn.json
+++ b/src/renderer/src/i18n/locales/zh-cn.json
@@ -80,6 +80,7 @@
       "input.topics": " 话题 ",
       "input.translate": "翻译成英文",
       "input.upload": "上传图片或文档",
+      "input.web_search": "开启网络搜索",
       "message.new.branch": "新分支",
       "message.new.branch.created": "新分支已创建",
       "message.new.context": "清除上下文",
diff --git a/src/renderer/src/i18n/locales/zh-tw.json b/src/renderer/src/i18n/locales/zh-tw.json
index 43f0e708..c230edda 100644
--- a/src/renderer/src/i18n/locales/zh-tw.json
+++ b/src/renderer/src/i18n/locales/zh-tw.json
@@ -80,6 +80,7 @@
       "input.topics": " 話題 ",
       "input.translate": "翻譯成英文",
       "input.upload": "上傳圖片或文檔",
+      "input.web_search": "開啟網路搜索",
       "message.new.branch": "新分支",
       "message.new.branch.created": "新分支已建立",
       "message.new.context": "新上下文",
diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
index 91a868de..df5309a1 100644
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@@ -4,13 +4,14 @@ import {
   FormOutlined,
   FullscreenExitOutlined,
   FullscreenOutlined,
+  GlobalOutlined,
   PauseCircleOutlined,
   QuestionCircleOutlined
 } from '@ant-design/icons'
 import { PicCenterOutlined } from '@ant-design/icons'
 import TranslateButton from '@renderer/components/TranslateButton'
 import { isMac } from '@renderer/config/constant'
-import { isVisionModel } from '@renderer/config/models'
+import { isVisionModel, isWebSearchModel } from '@renderer/config/models'
 import db from '@renderer/databases'
 import { useAssistant } from '@renderer/hooks/useAssistant'
 import { useRuntime } from '@renderer/hooks/useRuntime'
@@ -48,10 +49,10 @@ interface Props {
 let _text = ''
 let _files: FileType[] = []
 
-const Inputbar: FC<Props> = ({ assistant, setActiveTopic }) => {
+const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
   const [text, setText] = useState(_text)
   const [inputFocus, setInputFocus] = useState(false)
-  const { addTopic, model, setModel } = useAssistant(assistant.id)
+  const { assistant, addTopic, model, setModel, updateAssistant } = useAssistant(_assistant.id)
   const {
     sendMessageShortcut,
     fontSize,
@@ -400,6 +401,17 @@ const Inputbar: FC<Props> = ({ assistant, setActiveTopic }) => {
                 <FormOutlined />
               </ToolbarButton>
             </Tooltip>
+            {isWebSearchModel(model) && (
+              <Tooltip placement="top" title={t('chat.input.web_search')} arrow>
+                <ToolbarButton
+                  type="text"
+                  onClick={() => updateAssistant({ ...assistant, enableWebSearch: !assistant.enableWebSearch })}>
+                  <GlobalOutlined
+                    style={{ color: assistant.enableWebSearch ? 'var(--color-link)' : 'var(--color-icon)' }}
+                  />
+                </ToolbarButton>
+              </Tooltip>
+            )}
             <Tooltip placement="top" title={t('chat.input.clear')} arrow>
               <Popconfirm
                 title={t('chat.input.clear.content')}
diff --git a/src/renderer/src/pages/settings/ProviderSettings/ProviderSetting.tsx b/src/renderer/src/pages/settings/ProviderSettings/ProviderSetting.tsx
index 19a79386..e82f95cb 100644
--- a/src/renderer/src/pages/settings/ProviderSettings/ProviderSetting.tsx
+++ b/src/renderer/src/pages/settings/ProviderSettings/ProviderSetting.tsx
@@ -8,7 +8,8 @@ import {
   SettingOutlined
 } from '@ant-design/icons'
 import VisionIcon from '@renderer/components/Icons/VisionIcon'
-import { getModelLogo, isVisionModel, VISION_REGEX } from '@renderer/config/models'
+import WebSearchIcon from '@renderer/components/Icons/WebSearchIcon'
+import { getModelLogo, isVisionModel, isWebSearchModel, VISION_REGEX } from '@renderer/config/models'
 import { PROVIDER_CONFIG } from '@renderer/config/providers'
 import { useTheme } from '@renderer/context/ThemeProvider'
 import { useAssistants, useDefaultModel } from '@renderer/hooks/useAssistant'
@@ -261,20 +262,24 @@ const ProviderSetting: FC<Props> = ({ provider: _provider }) => {
           title={group}
           style={{ marginBottom: '10px', border: '0.5px solid var(--color-border)' }}
           size="small">
-          {modelGroups[group].map((model) => (
-            <ModelListItem key={model.id}>
-              <ModelListHeader>
-                <Avatar src={getModelLogo(model.id)} size={22} style={{ marginRight: '8px' }}>
-                  {model.name[0].toUpperCase()}
-                </Avatar>
-                {model.name} {isVisionModel(model) && <VisionIcon />}
-                <Popover content={modelTypeContent(model)} title={t('model.type.select')} trigger="click">
-                  <SettingIcon />
-                </Popover>
-              </ModelListHeader>
-              <RemoveIcon onClick={() => removeModel(model)} />
-            </ModelListItem>
-          ))}
+          {modelGroups[group].map((model) => {
+            console.debug(model)
+            return (
+              <ModelListItem key={model.id}>
+                <ModelListHeader>
+                  <Avatar src={getModelLogo(model.id)} size={22} style={{ marginRight: '8px' }}>
+                    {model.name[0].toUpperCase()}
+                  </Avatar>
+                  {model.name} {isVisionModel(model) && <VisionIcon />}
+                  {isWebSearchModel(model) && <WebSearchIcon />}
+                  <Popover content={modelTypeContent(model)} title={t('model.type.select')} trigger="click">
+                    <SettingIcon />
+                  </Popover>
+                </ModelListHeader>
+                <RemoveIcon onClick={() => removeModel(model)} />
+              </ModelListItem>
+            )
+          })}
         </Card>
       ))}
       {docsWebsite && (
diff --git a/src/renderer/src/providers/GeminiProvider.ts b/src/renderer/src/providers/GeminiProvider.ts
index 0ed68219..d00c8776 100644
--- a/src/renderer/src/providers/GeminiProvider.ts
+++ b/src/renderer/src/providers/GeminiProvider.ts
@@ -84,6 +84,8 @@ export default class GeminiProvider extends BaseProvider {
       {
         model: model.id,
         systemInstruction: assistant.prompt,
+        // @ts-ignore googleSearch is not a valid tool for Gemini
+        tools: assistant.enableWebSearch ? [{ googleSearch: {} }] : [],
         generationConfig: {
           maxOutputTokens: maxTokens,
           temperature: assistant?.settings?.temperature,
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index 4f30546c..d086540e 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -13,6 +13,7 @@ export type Assistant = {
   defaultModel?: Model
   settings?: Partial<AssistantSettings>
   messages?: AssistantMessage[]
+  enableWebSearch?: boolean
 }
 
 export type AssistantMessage = {
diff --git a/yarn.lock b/yarn.lock
index 50c38fa2..b98b0b24 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -798,10 +798,10 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@google/generative-ai@npm:^0.16.0":
-  version: 0.16.1
-  resolution: "@google/generative-ai@npm:0.16.1"
-  checksum: 10c0/bad13a040e210a48a44272b81ef78213d23c17491ff3cb0a937952518c79e96005b78d8c3324342b5b8613b6b13f56fe9bc53447c5d2c7da4a9cd26c6a8d7de4
+"@google/generative-ai@npm:^0.21.0":
+  version: 0.21.0
+  resolution: "@google/generative-ai@npm:0.21.0"
+  checksum: 10c0/cff5946c5964f2380e5097d82bd563d79be27a1a5ac604aaaad3f9ba3382992e4f0a371bd255baabfba4e5bdf296d8ce1410cbd65424afa98e64b2590fe49f3b
   languageName: node
   linkType: hard
 
@@ -2341,7 +2341,7 @@ __metadata:
     "@electron-toolkit/preload": "npm:^3.0.0"
     "@electron-toolkit/tsconfig": "npm:^1.0.1"
     "@electron-toolkit/utils": "npm:^3.0.0"
-    "@google/generative-ai": "npm:^0.16.0"
+    "@google/generative-ai": "npm:^0.21.0"
     "@hello-pangea/dnd": "npm:^16.6.0"
     "@kangfenmao/keyv-storage": "npm:^0.1.0"
     "@reduxjs/toolkit": "npm:^2.2.5"