Skip to content

Commit 1d16e84

Browse files
BobDuBobDu
authored
feat: add configurable max search result (#631)
* fix: try catch search timeout error Signed-off-by: Bob Du <i@bobdu.cc> * feat: add configurable max search result Signed-off-by: Bob Du <i@bobdu.cc> * feat: implement search test endpoint Signed-off-by: Bob Du <i@bobdu.cc> --------- Signed-off-by: Bob Du <i@bobdu.cc>
1 parent 4880b1f commit 1d16e84

File tree

12 files changed

+176
-70
lines changed

12 files changed

+176
-70
lines changed

‎README.en.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,6 +523,7 @@ Recommended for current IdP to use OIDC protocol, using [oauth2-proxy](https://o
523523
4. Fill in the following configurations:
524524
- **Enable Status**: Turn on/off global search functionality
525525
- **API Key**: Enter Tavily API Key
526+
- **Max Search Results**: Set the maximum number of search results returned per search (1-20, default 10)
526527
- **Search Query System Message**: Prompt template for extracting search keywords
527528
- **Search Result System Message**: Prompt template for processing search results
528529

@@ -571,6 +572,7 @@ Current time: {current_time}
571572
- **Result Format**: JSON format to store complete search results
572573
- **Data Storage**: MongoDB stores search queries and results
573574
- **Timeout Setting**: Search request timeout is 300 seconds
575+
- **Result Count Control**: Support configuration of maximum search results returned per search (1-20)
574576
575577
### Notes
576578
@@ -579,6 +581,7 @@ Current time: {current_time}
579581
- It is recommended to enable selectively based on actual needs
580582
- Administrators can control the global search functionality status
581583
- Each session can independently control whether to use search functionality
584+
- The maximum search results setting affects the detail level of search and API costs
582585
583586
584587
## Contributing

‎README.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,7 @@ pnpm build
404404
4. 填写以下配置：
405405
- **启用状态**: 开启/关闭全局搜索功能
406406
- **API Key**: 填入 Tavily API Key
407+
- **最大搜索结果数**: 设置每次搜索返回的最大结果数量（1-20，默认10）
407408
- **搜索查询系统消息**: 用于提取搜索关键词的提示模板
408409
- **搜索结果系统消息**: 用于处理搜索结果的提示模板
409410

@@ -452,6 +453,7 @@ Current time: {current_time}
452453
- **结果格式**: JSON 格式存储完整搜索结果
453454
- **数据存储**: MongoDB 存储搜索查询和结果
454455
- **超时设置**: 搜索请求超时时间为 300 秒
456+
- **结果数量控制**: 支持配置每次搜索返回的最大结果数量（1-20）
455457

456458
### 注意事项
457459

@@ -460,6 +462,7 @@ Current time: {current_time}
460462
- 建议根据实际需求选择性开启
461463
- 管理员可以控制全局搜索功能的开启状态
462464
- 每个会话可以独立控制是否使用搜索功能
465+
- 最大搜索结果数设置会影响搜索的详细程度和 API 费用
463466

464467

465468
## 上下文窗口控制

‎service/src/chatgpt/index.ts‎

Lines changed: 67 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -106,73 +106,78 @@ async function chatReplyProcess(options: RequestOptions) {
106106
let hasSearchResult = false
107107
const searchConfig = globalConfig.searchConfig
108108
if (searchConfig.enabled && searchConfig?.options?.apiKey && searchEnabled) {
109-
messages[0].content = renderSystemMessage(searchConfig.systemMessageGetSearchQuery, dayjs().format('YYYY-MM-DD HH:mm:ss'))
109+
try {
110+
messages[0].content = renderSystemMessage(searchConfig.systemMessageGetSearchQuery, dayjs().format('YYYY-MM-DD HH:mm:ss'))
110111

111-
const getSearchQueryChatCompletionCreateBody: OpenAI.ChatCompletionCreateParamsNonStreaming = {
112-
model,
113-
messages,
114-
}
115-
if (key.keyModel === 'VLLM') {
116-
// @ts-expect-error vLLM supports a set of parameters that are not part of the OpenAI API.
117-
getSearchQueryChatCompletionCreateBody.chat_template_kwargs = {
118-
enable_thinking: false,
112+
const getSearchQueryChatCompletionCreateBody: OpenAI.ChatCompletionCreateParamsNonStreaming = {
113+
model,
114+
messages,
119115
}
120-
}
121-
const completion = await openai.chat.completions.create(getSearchQueryChatCompletionCreateBody)
122-
let searchQuery: string = completion.choices[0].message.content
123-
const match = searchQuery.match(/<search_query>([\s\S]*)<\/search_query>/i)
124-
if (match)
125-
searchQuery = match[1].trim()
126-
else
127-
searchQuery = ''
128-
129-
if (searchQuery) {
130-
await updateChatSearchQuery(messageId, searchQuery)
131-
132-
process?.({
133-
searchQuery,
134-
})
135-
136-
const tvly = tavily({ apiKey: searchConfig.options?.apiKey })
137-
const response = await tvly.search(
138-
searchQuery,
139-
{
140-
// https://docs.tavily.com/documentation/best-practices/best-practices-search#search-depth%3Dadvanced-ideal-for-higher-relevance-in-search-results
141-
searchDepth: 'advanced',
142-
chunksPerSource: 3,
143-
includeRawContent: true,
144-
// 0 <= x <= 20 https://docs.tavily.com/documentation/api-reference/endpoint/search#body-max-results
145-
// https://docs.tavily.com/documentation/best-practices/best-practices-search#max-results-limiting-the-number-of-results
146-
maxResults: 10,
147-
// Max 120s, default to 60 https://github.com/tavily-ai/tavily-js/blob/de69e479c5d3f6c5d443465fa2c29407c0d3515d/src/search.ts#L118
148-
timeout: 120,
149-
},
150-
)
151-
152-
const searchResults = response.results as SearchResult[]
153-
const searchUsageTime = response.responseTime
154-
155-
await updateChatSearchResult(messageId, searchResults, searchUsageTime)
156-
157-
process?.({
158-
searchResults,
159-
searchUsageTime,
160-
})
161-
162-
let searchResultContent = JSON.stringify(searchResults)
163-
// remove image url
164-
const base64Pattern = /!\[([^\]]*)\]\([^)]*\)/g
165-
searchResultContent = searchResultContent.replace(base64Pattern, '$1')
166-
167-
messages.push({
168-
role: 'user',
169-
content: `Additional information from web searche engine.
116+
if (key.keyModel === 'VLLM') {
117+
// @ts-expect-error vLLM supports a set of parameters that are not part of the OpenAI API.
118+
getSearchQueryChatCompletionCreateBody.chat_template_kwargs = {
119+
enable_thinking: false,
120+
}
121+
}
122+
const completion = await openai.chat.completions.create(getSearchQueryChatCompletionCreateBody)
123+
let searchQuery: string = completion.choices[0].message.content
124+
const match = searchQuery.match(/<search_query>([\s\S]*)<\/search_query>/i)
125+
if (match)
126+
searchQuery = match[1].trim()
127+
else
128+
searchQuery = ''
129+
130+
if (searchQuery) {
131+
await updateChatSearchQuery(messageId, searchQuery)
132+
133+
process?.({
134+
searchQuery,
135+
})
136+
137+
const tvly = tavily({ apiKey: searchConfig.options?.apiKey })
138+
const response = await tvly.search(
139+
searchQuery,
140+
{
141+
// https://docs.tavily.com/documentation/best-practices/best-practices-search#search-depth%3Dadvanced-ideal-for-higher-relevance-in-search-results
142+
searchDepth: 'advanced',
143+
chunksPerSource: 3,
144+
includeRawContent: true,
145+
// 0 <= x <= 20 https://docs.tavily.com/documentation/api-reference/endpoint/search#body-max-results
146+
// https://docs.tavily.com/documentation/best-practices/best-practices-search#max-results-limiting-the-number-of-results
147+
maxResults: searchConfig.options?.maxResults || 10,
148+
// Max 120s, default to 60 https://github.com/tavily-ai/tavily-js/blob/de69e479c5d3f6c5d443465fa2c29407c0d3515d/src/search.ts#L118
149+
timeout: 120,
150+
},
151+
)
152+
153+
const searchResults = response.results as SearchResult[]
154+
const searchUsageTime = response.responseTime
155+
156+
await updateChatSearchResult(messageId, searchResults, searchUsageTime)
157+
158+
process?.({
159+
searchResults,
160+
searchUsageTime,
161+
})
162+
163+
let searchResultContent = JSON.stringify(searchResults)
164+
// remove image url
165+
const base64Pattern = /!\[([^\]]*)\]\([^)]*\)/g
166+
searchResultContent = searchResultContent.replace(base64Pattern, '$1')
167+
168+
messages.push({
169+
role: 'user',
170+
content: `Additional information from web searche engine.
170171
search query: <search_query>${searchQuery}</search_query>
171172
search result: <search_result>${searchResultContent}</search_result>`,
172-
})
173+
})
173174

174-
messages[0].content = renderSystemMessage(searchConfig.systemMessageWithSearchResult, dayjs().format('YYYY-MM-DD HH:mm:ss'))
175-
hasSearchResult = true
175+
messages[0].content = renderSystemMessage(searchConfig.systemMessageWithSearchResult, dayjs().format('YYYY-MM-DD HH:mm:ss'))
176+
hasSearchResult = true
177+
}
178+
}
179+
catch (e) {
180+
globalThis.console.error('search error from tavily, ', e)
176181
}
177182
}
178183

‎service/src/index.ts‎

Lines changed: 60 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { AnnounceConfig, AuditConfig, Config, GiftCard, KeyConfig, MailConfig, SiteConfig, UserInfo } from './storage/model'
1+
import type { AnnounceConfig, AuditConfig, Config, GiftCard, KeyConfig, MailConfig, SearchResult, SiteConfig, UserInfo } from './storage/model'
22
import type { AuthJwtPayload } from './types'
33
import * as path from 'node:path'
44
import * as process from 'node:process'
@@ -822,13 +822,66 @@ router.post('/setting-search', rootAuth, async (req, res) => {
822822

823823
router.post('/search-test', rootAuth, async (req, res) => {
824824
try {
825-
const { text } = req.body as { search: import('./storage/model').SearchConfig, text: string }
826-
// TODO: Implement actual search test logic with Tavily API
827-
// For now, just return a success response
828-
res.send({ status: 'Success', message: '搜索测试成功 | Search test successful', data: { query: text, results: [] } })
825+
const { search, text } = req.body as { search: import('./storage/model').SearchConfig, text: string }
826+
827+
// Validate search configuration
828+
if (!search.enabled) {
829+
res.send({ status: 'Fail', message: '搜索功能未启用 | Search functionality is not enabled', data: null })
830+
return
831+
}
832+
833+
if (!search.options?.apiKey) {
834+
res.send({ status: 'Fail', message: '搜索 API 密钥未配置 | Search API key is not configured', data: null })
835+
return
836+
}
837+
838+
if (!text || text.trim() === '') {
839+
res.send({ status: 'Fail', message: '搜索文本不能为空 | Search text cannot be empty', data: null })
840+
return
841+
}
842+
843+
// Validate maxResults range
844+
const maxResults = search.options?.maxResults || 10
845+
if (maxResults < 1 || maxResults > 20) {
846+
res.send({ status: 'Fail', message: '最大搜索结果数必须在 1-20 之间 | Max search results must be between 1-20', data: null })
847+
return
848+
}
849+
850+
// Import required modules
851+
const { tavily } = await import('@tavily/core')
852+
853+
// Execute search
854+
const tvly = tavily({ apiKey: search.options.apiKey })
855+
const response = await tvly.search(
856+
text.trim(),
857+
{
858+
searchDepth: 'advanced',
859+
chunksPerSource: 3,
860+
includeRawContent: true,
861+
maxResults,
862+
timeout: 120,
863+
},
864+
)
865+
866+
const searchResults = response.results as SearchResult[]
867+
const searchUsageTime = response.responseTime
868+
869+
// Return search results
870+
res.send({
871+
status: 'Success',
872+
message: `搜索测试成功 | Search test successful (用时 ${searchUsageTime}ms, 找到 ${searchResults.length} 个结果)`,
873+
data: {
874+
query: text.trim(),
875+
results: searchResults,
876+
usageTime: searchUsageTime,
877+
resultCount: searchResults.length,
878+
maxResults,
879+
},
880+
})
829881
}
830-
catch (error) {
831-
res.send({ status: 'Fail', message: error.message, data: null })
882+
catch (error: any) {
883+
console.error('Search test error:', error)
884+
res.send({ status: 'Fail', message: `搜索测试失败 | Search test failed: ${error.message}`, data: null })
832885
}
833886
})
834887

‎service/src/storage/config.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ export async function getOriginConfig() {
9595
if (!config.searchConfig) {
9696
config.searchConfig = new SearchConfig()
9797
config.searchConfig.enabled = false
98+
config.searchConfig.options = { apiKey: '', maxResults: 10 }
9899
}
99100

100101
if (!isNotEmptyString(config.siteConfig.chatModels))

‎service/src/storage/model.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,7 @@ export enum SearchServiceProvider {
199199

200200
export class SearchServiceOptions {
201201
public apiKey: string
202+
public maxResults?: number
202203
}
203204

204205
export class Config {

‎src/components/common/Setting/Search.vue‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ async function fetchConfig() {
2525
data.searchConfig = new SearchConfig(false, '', { apiKey: '' }, '', '')
2626
if (!data.searchConfig.options)
2727
data.searchConfig.options = { apiKey: '' }
28+
if (!data.searchConfig.options.maxResults)
29+
data.searchConfig.options.maxResults = 10
2830
config.value = data.searchConfig
2931
}
3032
finally {
@@ -96,6 +98,18 @@ onMounted(() => {
9698
/>
9799
</div>
98100
</div>
101+
<div v-if="config && config.enabled" class="flex items-center space-x-4">
102+
<span class="shrink-0 w-[100px]">{{ $t('setting.searchMaxResults') }}</span>
103+
<div class="flex-1">
104+
<NInputNumber
105+
v-model:value="config.options.maxResults"
106+
:min="1"
107+
:max="20"
108+
placeholder="1-20"
109+
style="width: 140px"
110+
/>
111+
</div>
112+
</div>
99113
<div v-if="config && config.enabled" class="flex items-center space-x-4">
100114
<span class="shrink-0 w-[100px]">{{ $t('setting.searchTest') }}</span>
101115
<div class="flex-1">

‎src/components/common/Setting/model.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,7 @@ export type SearchServiceProvider = 'tavily' | ''
185185

186186
export interface SearchServiceOptions {
187187
apiKey: string
188+
maxResults?: number
188189
}
189190

190191
export class SearchConfig {
@@ -199,5 +200,8 @@ export class SearchConfig {
199200
this.options = options
200201
this.systemMessageWithSearchResult = systemMessageWithSearchResult
201202
this.systemMessageGetSearchQuery = systemMessageGetSearchQuery
203+
if (!this.options.maxResults) {
204+
this.options.maxResults = 10
205+
}
202206
}
203207
}

‎src/locales/en-US.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,7 @@ export default {
183183
searchEnabled: 'Search Enabled',
184184
searchProvider: 'Search Provider',
185185
searchApiKey: 'Search API Key',
186+
searchMaxResults: 'Max Search Results',
186187
systemMessageWithSearchResult: 'System message for conversations with search results',
187188
systemMessageGetSearchQuery: 'System message for getting search query',
188189
systemMessageWithSearchResultPlaceholder: 'System message template when with search results. Use {\'{current_time}\'} as placeholder for current time.',

‎src/locales/ko-KR.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,16 @@ export default {
174174
auditBaiduLabelLink: '레이블 세부 정보로 이동',
175175
auditCustomizeEnabled: '맞춤화하다',
176176
auditCustomizeWords: '단어 맞춤설정',
177+
searchConfig: '검색 구성',
178+
searchEnabled: '검색 기능',
179+
searchProvider: '검색 공급자',
180+
searchApiKey: '검색 API 키',
181+
searchMaxResults: '최대 검색 결과 수',
182+
systemMessageWithSearchResult: '검색 결과가 포함된 대화 시스템 프롬프트',
183+
systemMessageGetSearchQuery: '검색 쿼리어를 가져오는 시스템 프롬프트',
184+
systemMessageWithSearchResultPlaceholder: '검색 결과가 포함된 시스템 메시지 템플릿, {\'{current_time}\'}를 현재 시간의 플레이스홀더로 사용',
185+
systemMessageGetSearchQueryPlaceholder: '검색 쿼리어를 생성하는 시스템 메시지 템플릿, {\'{current_time}\'}를 현재 시간의 플레이스홀더로 사용, LLM이 <search_query>예시 검색 쿼리어</search_query> 태그에서 쿼리어를 반환하도록 요구, 검색이 필요하지 않으면 빈 <search_query></search_query> 태그를 반환',
186+
searchTest: '검색 테스트',
177187
accessTokenExpiredTime: '만료된 시간',
178188
userConfig: 'Users',
179189
keysConfig: 'Keys Manager',

0 commit comments

Comments
 (0)