From 09652b2e31a0de10ae5eef0c12751b85b2c99a2d Mon Sep 17 00:00:00 2001
From: Chen Tao <70054568+eeee0717@users.noreply.github.com>
Date: Tue, 25 Feb 2025 23:46:51 +0800
Subject: [PATCH] feat: add web search settings (#2314)

* fix: add time when using web search

* feat: add optional

* chore

* chore

* chore

* clean code

* feat: set search max results

* feat: add manual blacklist

* clean code

* chore

* chore

* clean
---
 src/renderer/src/i18n/locales/zh-cn.json      |  7 ++-
 .../src/pages/settings/WebSearchSettings.tsx  | 53 ++++++++++++++++++-
 src/renderer/src/services/WebSearchService.ts | 12 +++--
 src/renderer/src/store/migrate.ts             |  2 +
 src/renderer/src/store/websearch.ts           | 22 ++++++--
 src/renderer/src/utils/blacklist.ts           | 50 +++++++++++++++++
 6 files changed, 135 insertions(+), 11 deletions(-)
 create mode 100644 src/renderer/src/utils/blacklist.ts

diff --git a/src/renderer/src/i18n/locales/zh-cn.json b/src/renderer/src/i18n/locales/zh-cn.json
index c425621899..9973630d3c 100644
--- a/src/renderer/src/i18n/locales/zh-cn.json
+++ b/src/renderer/src/i18n/locales/zh-cn.json
@@ -828,7 +828,12 @@
           "api_key": "Tavily API 密钥",
           "api_key.placeholder": "请输入 Tavily API 密钥"
         },
-        "search_with_time": "搜索包含日期"
+        "search_with_time": "搜索包含日期",
+        "search_max_result": "搜索结果个数",
+        "search_result_default": "默认",
+        "blacklist": "黑名单",
+        "blacklist_description": "在搜索结果中不会出现以下网站的结果",
+        "blacklist_tooltip": "请使用以下格式(换行分隔)\nexample.com\nhttps://www.example.com\nhttps://example.com\n*://*.example.com"
       }
     },
     "translate": {
diff --git a/src/renderer/src/pages/settings/WebSearchSettings.tsx b/src/renderer/src/pages/settings/WebSearchSettings.tsx
index c95044b77b..307be41101 100644
--- a/src/renderer/src/pages/settings/WebSearchSettings.tsx
+++ b/src/renderer/src/pages/settings/WebSearchSettings.tsx
@@ -4,8 +4,10 @@ import { HStack } from '@renderer/components/Layout'
 import { useTheme } from '@renderer/context/ThemeProvider'
 import { useWebSearchProvider } from '@renderer/hooks/useWebSearchProviders'
 import { useAppDispatch, useAppSelector } from '@renderer/store'
-import { setSearchWithTime } from '@renderer/store/websearch'
-import { Input, Switch, Typography } from 'antd'
+import { setExcludeDomains, setMaxResult, setSearchWithTime } from '@renderer/store/websearch'
+import { formatDomains } from '@renderer/utils/blacklist'
+import { Alert, Input, Slider, Switch, Typography } from 'antd'
+import TextArea from 'antd/es/input/TextArea'
 import { FC, useEffect, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
@@ -29,6 +31,11 @@ const WebSearchSettings: FC = () => {
   const [apiKey, setApiKey] = useState(provider.apiKey)
   const logo = theme === 'dark' ? tavilyLogoDark : tavilyLogo
   const searchWithTime = useAppSelector((state) => state.websearch.searchWithTime)
+  const maxResults = useAppSelector((state) => state.websearch.maxResults)
+  const excludeDomains = useAppSelector((state) => state.websearch.excludeDomains)
+  const [errFormat, setErrFormat] = useState(false)
+  const [blacklistInput, setBlacklistInput] = useState('')
+
   const dispatch = useAppDispatch()
 
   useEffect(() => {
@@ -39,6 +46,20 @@ const WebSearchSettings: FC = () => {
     }
   }, [apiKey, provider, updateProvider])
 
+  useEffect(() => {
+    if (excludeDomains) {
+      setBlacklistInput(excludeDomains.join('\n'))
+    }
+  }, [excludeDomains])
+
+  function updateManualBlacklist(blacklist: string) {
+    const blacklistDomains = blacklist.split('\n').filter((url) => url.trim() !== '')
+    const { formattedDomains, hasError } = formatDomains(blacklistDomains)
+    setErrFormat(hasError)
+    if (hasError) return
+    dispatch(setExcludeDomains(formattedDomains))
+  }
+
   return (
     <SettingContainer theme={theme}>
       <SettingGroup theme={theme}>
@@ -70,6 +91,34 @@ const WebSearchSettings: FC = () => {
           <SettingRowTitle>{t('settings.websearch.search_with_time')}</SettingRowTitle>
           <Switch checked={searchWithTime} onChange={(checked) => dispatch(setSearchWithTime(checked))} />
         </SettingRow>
+        <SettingRow>
+          <SettingRowTitle>{t('settings.websearch.search_max_result')}</SettingRowTitle>
+          <Slider
+            defaultValue={maxResults}
+            style={{ width: '200px' }}
+            min={1}
+            max={20}
+            step={1}
+            marks={{ 1: '1', 5: t('settings.websearch.search_result_default'), 20: '20' }}
+            onChangeComplete={(value) => dispatch(setMaxResult(value))}
+          />
+        </SettingRow>
+      </SettingGroup>
+      <SettingGroup theme={theme}>
+        <SettingTitle>{t('settings.websearch.blacklist')}</SettingTitle>
+        <SettingDivider />
+        <SettingRow>
+          <SettingRowTitle>{t('settings.websearch.blacklist_description')}</SettingRowTitle>
+        </SettingRow>
+        <TextArea
+          value={blacklistInput}
+          onChange={(e) => setBlacklistInput(e.target.value)}
+          onBlur={() => updateManualBlacklist(blacklistInput)}
+          placeholder={t('settings.websearch.blacklist_tooltip')}
+          autoSize={{ minRows: 2, maxRows: 6 }}
+          rows={4}
+        />
+        {errFormat && <Alert message={t('settings.websearch.blacklist_tooltip')} type="error" />}
       </SettingGroup>
     </SettingContainer>
   )
diff --git a/src/renderer/src/services/WebSearchService.ts b/src/renderer/src/services/WebSearchService.ts
index 932fb3cf17..000195b7ba 100644
--- a/src/renderer/src/services/WebSearchService.ts
+++ b/src/renderer/src/services/WebSearchService.ts
@@ -25,18 +25,20 @@ class WebSearchService {
 
   public async search(query: string) {
     const searchWithTime = store.getState().websearch.searchWithTime
+    const maxResults = store.getState().websearch.maxResults
+    const excludeDomains = store.getState().websearch.excludeDomains
     let formatted_query = query
-
     if (searchWithTime) {
       formatted_query = `today is ${dayjs().format('YYYY-MM-DD')} \r\n ${query}`
     }
-
     const provider = this.getWebSearchProvider()
     const tvly = tavily({ apiKey: provider.apiKey })
-
-    return await tvly.search(formatted_query, {
-      maxResults: 5
+    const result = await tvly.search(formatted_query, {
+      maxResults: maxResults,
+      excludeDomains: excludeDomains
     })
+
+    return result
   }
 }
 
diff --git a/src/renderer/src/store/migrate.ts b/src/renderer/src/store/migrate.ts
index 2acc0f08d0..34ffd2f74f 100644
--- a/src/renderer/src/store/migrate.ts
+++ b/src/renderer/src/store/migrate.ts
@@ -1124,6 +1124,8 @@ const migrateConfig = {
   '73': (state: RootState) => {
     if (state.websearch) {
       state.websearch.searchWithTime = true
+      state.websearch.maxResults = 5
+      state.websearch.excludeDomains = []
     }
     if (!state.llm.providers.find((provider) => provider.id === 'lmstudio')) {
       state.llm.providers.push({
diff --git a/src/renderer/src/store/websearch.ts b/src/renderer/src/store/websearch.ts
index a986e0d019..1c82496f17 100644
--- a/src/renderer/src/store/websearch.ts
+++ b/src/renderer/src/store/websearch.ts
@@ -4,6 +4,8 @@ export interface WebSearchState {
   defaultProvider: string
   providers: WebSearchProvider[]
   searchWithTime: boolean
+  maxResults: number
+  excludeDomains: string[]
 }
 
 const initialState: WebSearchState = {
@@ -15,7 +17,9 @@ const initialState: WebSearchState = {
       apiKey: ''
     }
   ],
-  searchWithTime: true
+  searchWithTime: true,
+  maxResults: 5,
+  excludeDomains: []
 }
 
 const websearchSlice = createSlice({
@@ -36,11 +40,23 @@ const websearchSlice = createSlice({
     },
     setSearchWithTime: (state, action: PayloadAction<boolean>) => {
       state.searchWithTime = action.payload
+    },
+    setMaxResult: (state, action: PayloadAction<number>) => {
+      state.maxResults = action.payload
+    },
+    setExcludeDomains: (state, action: PayloadAction<string[]>) => {
+      state.excludeDomains = action.payload
     }
   }
 })
 
-export const { setWebSearchProviders, updateWebSearchProvider, setDefaultProvider, setSearchWithTime } =
-  websearchSlice.actions
+export const {
+  setWebSearchProviders,
+  updateWebSearchProvider,
+  setDefaultProvider,
+  setSearchWithTime,
+  setExcludeDomains,
+  setMaxResult
+} = websearchSlice.actions
 
 export default websearchSlice.reducer
diff --git a/src/renderer/src/utils/blacklist.ts b/src/renderer/src/utils/blacklist.ts
new file mode 100644
index 0000000000..b2ff241b24
--- /dev/null
+++ b/src/renderer/src/utils/blacklist.ts
@@ -0,0 +1,50 @@
+interface FormatDomainsResult {
+  formattedDomains: string[]
+  hasError: boolean
+}
+
+export function formatDomains(urls: string[]): FormatDomainsResult {
+  let hasError = false
+  const formattedDomains: string[] = []
+
+  for (const urlString of urls) {
+    try {
+      let modifiedUrlString = urlString
+
+      // 1. 处理通配符协议 (*://)
+      if (modifiedUrlString.startsWith('*://')) {
+        modifiedUrlString = modifiedUrlString.substring(4)
+      }
+
+      // 2. 检查并添加协议前缀
+      if (!modifiedUrlString.match(/^[a-zA-Z]+:\/\//)) {
+        modifiedUrlString = 'https://' + modifiedUrlString
+      }
+
+      // 3. URL 解析和验证
+      const url = new URL(modifiedUrlString)
+      if (url.protocol !== 'https:') {
+        if (url.protocol !== 'http:') {
+          hasError = true
+        } else {
+          url.protocol = 'https:'
+        }
+      }
+
+      // 4. 通配符处理
+      let domain = url.hostname
+      if (domain.startsWith('*.')) {
+        domain = domain.substring(2)
+      }
+
+      // 5. 格式化
+      const formattedDomain = `https://${domain}`
+      formattedDomains.push(formattedDomain)
+    } catch (error) {
+      hasError = true
+      console.error('Error formatting URL:', urlString, error)
+    }
+  }
+
+  return { formattedDomains, hasError }
+}