feat: LLM可以根据需求自行选择使用知识库或者网络搜索

src/renderer/src/config/prompts.ts

@@ -51,59 +51,74 @@ export const SUMMARIZE_PROMPT =
 
 // https://github.com/ItzCrazyKns/Perplexica/blob/master/src/lib/prompts/webSearch.ts
 export const SEARCH_SUMMARY_PROMPT = `
-  You are an AI question rephraser. You will be given a conversation and a follow-up question,  you will have to rephrase the follow up question so it is a standalone question and can be used by another LLM to search the web for information to answer it.
-  If it is a simple writing task or a greeting (unless the greeting contains a question after it) like Hi, Hello, How are you, etc. than a question then you need to return \`not_needed\` as the response (This is because the LLM won't need to search the web for finding information on this topic).
-  If the user asks some question from some URL or wants you to summarize a PDF or a webpage (via URL) you need to return the links inside the \`links\` XML block and the question inside the \`question\` XML block. If the user wants to you to summarize the webpage or the PDF you need to return \`summarize\` inside the \`question\` XML block in place of a question and the link to summarize in the \`links\` XML block.
-  You must always return the rephrased question inside the \`question\` XML block, if there are no links in the follow-up question then don't insert a \`links\` XML block in your response.
+  You are an AI question rephraser. Your role is to rephrase follow-up questions from a conversation into standalone questions that can be used by another LLM to retrieve information, either through web search or from a knowledge base.
+  Follow these guidelines:
+  1. If the question is a simple writing task, greeting (e.g., Hi, Hello, How are you), or does not require searching for information (unless the greeting contains a follow-up question), return 'not_needed' in the 'question' XML block. This indicates that no search is required.
+  2. If the user asks a question related to a specific URL, PDF, or webpage, include the links in the 'links' XML block and the question in the 'question' XML block. If the request is to summarize content from a URL or PDF, return 'summarize' in the 'question' XML block and include the relevant link in the 'links' XML block.
+  3. Always return the rephrased question inside the 'question' XML block. If there are no links in the follow-up question, do not insert a 'links' XML block in your response.
+  4. Add a 'tools' XML block to specify the tool(s) to be used for answering the question. Use 'websearch' for queries requiring real-time or external information, 'knowledge' for queries that can be answered from a pre-existing knowledge base, or both ('websearch, knowledge') if either tool could be applicable.
 
-  There are several examples attached for your reference inside the below \`examples\` XML block
+  There are several examples attached for your reference inside the below 'examples' XML block.
 
   <examples>
   1. Follow up question: What is the capital of France
   Rephrased question:\`
   <question>
-  Capital of france
+  What is the capital of France?
   </question>
+  <tools>
+  websearch, knowledge
+  </tools>
   \`
 
-  2. Hi, how are you?
-  Rephrased question\`
+  2. Follow up question: Hi, how are you?
+  Rephrased question:\`
   <question>
   not_needed
   </question>
+  <tools>
+  none
+  </tools>
   \`
 
   3. Follow up question: What is Docker?
   Rephrased question: \`
   <question>
-  What is Docker
+  What is Docker?
   </question>
+  <tools>
+  websearch, knowledge
+  </tools>
   \`
 
   4. Follow up question: Can you tell me what is X from https://example.com
   Rephrased question: \`
   <question>
-  Can you tell me what is X?
+  What is X?
   </question>
-
   <links>
   https://example.com
   </links>
+  <tools>
+  websearch
+  </tools>
   \`
 
   5. Follow up question: Summarize the content from https://example.com
   Rephrased question: \`
   <question>
   summarize
   </question>
-
   <links>
   https://example.com
   </links>
+  <tools>
+  websearch
+  </tools>
   \`
   </examples>
 
-  Anything below is the part of the actual conversation and you need to use conversation and the follow-up question to rephrase the follow-up question as a standalone question based on the guidelines shared above.
+  Anything below is part of the actual conversation. Use the conversation history and the follow-up question to rephrase the follow-up question as a standalone question based on the guidelines shared above.
 
   <conversation>
   {chat_history}

src/renderer/src/providers/AiProvider/BaseProvider.ts

@@ -1,7 +1,6 @@
 import { FOOTNOTE_PROMPT, REFERENCE_PROMPT } from '@renderer/config/prompts'
 import { getLMStudioKeepAliveTime } from '@renderer/hooks/useLMStudio'
 import { getOllamaKeepAliveTime } from '@renderer/hooks/useOllama'
-import { getKnowledgeBaseReferences } from '@renderer/services/KnowledgeService'
 import type {
   Assistant,
   GenerateImageParams,
@@ -98,14 +97,15 @@ export default abstract class BaseProvider {
       return message.content
     }
 
-    const webSearchReferences = await this.getWebSearchReferences(message)
+    const webSearchReferences = await this.getWebSearchReferencesFromCache(message)
 
     if (!isEmpty(webSearchReferences)) {
       const referenceContent = `\`\`\`json\n${JSON.stringify(webSearchReferences, null, 2)}\n\`\`\``
       return REFERENCE_PROMPT.replace('{question}', message.content).replace('{references}', referenceContent)
     }
 
-    const knowledgeReferences = await getKnowledgeBaseReferences(message)
+    // const knowledgeReferences = await getKnowledgeBaseReferences(message)
+    const knowledgeReferences = await this.getKnowledgeBaseReferencesFromCache(message)
 
     if (!isEmpty(message.knowledgeBaseIds) && isEmpty(knowledgeReferences)) {
       window.message.info({ content: t('knowledge.no_match'), key: 'knowledge-base-no-match-info' })
@@ -119,7 +119,7 @@ export default abstract class BaseProvider {
     return message.content
   }
 
-  private async getWebSearchReferences(message: Message) {
+  private async getWebSearchReferencesFromCache(message: Message) {
     if (isEmpty(message.content)) {
       return []
     }
@@ -140,6 +140,23 @@ export default abstract class BaseProvider {
     return []
   }
 
+  /**
+   * 从缓存中获取知识库引用
+   */
+  private async getKnowledgeBaseReferencesFromCache(message: Message): Promise<KnowledgeReference[]> {
+    if (isEmpty(message.content)) {
+      return []
+    }
+    const knowledgeReferences: KnowledgeReference[] = window.keyv.get(`knowledge-search-${message.id}`)
+
+    if (!isEmpty(knowledgeReferences)) {
+      console.log(`Found ${knowledgeReferences.length} knowledge base references in cache for ID: ${message.id}`)
+      return knowledgeReferences
+    }
+    console.log(`No knowledge base references found in cache for ID: ${message.id}`)
+    return []
+  }
+
   protected getCustomParameters(assistant: Assistant) {
     return (
       assistant?.settings?.customParameters?.reduce((acc, param) => {

src/renderer/src/services/ApiService.ts

@@ -8,9 +8,18 @@ import { SEARCH_SUMMARY_PROMPT } from '@renderer/config/prompts'
 import i18n from '@renderer/i18n'
 import store from '@renderer/store'
 import { setGenerating } from '@renderer/store/runtime'
-import { Assistant, MCPTool, Message, Model, Provider, Suggestion, WebSearchResponse } from '@renderer/types'
+import {
+  Assistant,
+  KnowledgeReference,
+  MCPTool,
+  Message,
+  Model,
+  Provider,
+  Suggestion,
+  WebSearchResponse
+} from '@renderer/types'
 import { formatMessageError, isAbortError } from '@renderer/utils/error'
-import { fetchWebContents } from '@renderer/utils/fetch'
+import { extractInfoFromXML, ExtractResults } from '@renderer/utils/extract'
 import { withGenerateImage } from '@renderer/utils/formats'
 import {
   cleanLinkCommas,
@@ -26,13 +35,13 @@ import { cloneDeep, findLast, isEmpty } from 'lodash'
 import AiProvider from '../providers/AiProvider'
 import {
   getAssistantProvider,
-  getDefaultAssistant,
   getDefaultModel,
   getProviderByModel,
   getTopNamingModel,
   getTranslateModel
 } from './AssistantService'
 import { EVENT_NAMES, EventEmitter } from './EventService'
+import { processKnowledgeSearch } from './KnowledgeService'
 import { filterContextMessages, filterMessages, filterUsefulMessages } from './MessagesService'
 import { estimateMessagesUsage } from './TokenService'
 import WebSearchService from './WebSearchService'
@@ -52,88 +61,100 @@ export async function fetchChatCompletion({
   const webSearchProvider = WebSearchService.getWebSearchProvider()
   const AI = new AiProvider(provider)
 
+  const lastUserMessage = findLast(messages, (m) => m.role === 'user')
+  const lastAnswer = findLast(messages, (m) => m.role === 'assistant')
+  const hasKnowledgeBase = !isEmpty(lastUserMessage?.knowledgeBaseIds)
+  if (!lastUserMessage) {
+    return
+  }
+
+  // 网络搜索/知识库 关键词提取
+  const extract = async () => {
+    // 更新消息状态为搜索中
+    onResponse({ ...message, status: 'searching' })
+    const summaryAssistant = {
+      ...assistant,
+      prompt: SEARCH_SUMMARY_PROMPT
+    }
+    const keywords = await fetchSearchSummary({
+      messages: lastAnswer ? [lastAnswer, lastUserMessage] : [lastUserMessage],
+      assistant: summaryAssistant
+    })
+    try {
+      return extractInfoFromXML(keywords || '')
+    } catch (e: any) {
+      console.error('extract error', e)
+      throw new Error('extract error')
+    }
+  }
+  let extractResults: ExtractResults
+  if (assistant.enableWebSearch || hasKnowledgeBase) {
+    extractResults = await extract()
+  }
+
   const searchTheWeb = async () => {
-    if (WebSearchService.isWebSearchEnabled() && assistant.enableWebSearch && assistant.model) {
-      let query = ''
-      let webSearchResponse: WebSearchResponse = {
-        results: []
+    // 检查是否需要进行网络搜索
+    const shouldSearch =
+      extractResults?.tools?.includes('websearch') &&
+      WebSearchService.isWebSearchEnabled() &&
+      assistant.enableWebSearch &&
+      assistant.model
+
+    if (!shouldSearch) return
+
+    // 检查是否使用OpenAI的网络搜索
+    const webSearchParams = getOpenAIWebSearchParams(assistant, assistant.model!)
+    if (!isEmpty(webSearchParams) || isOpenAIWebSearch(assistant.model!)) return
+
+    // 更新消息状态为搜索中
+    onResponse({ ...message, status: 'searching' })
+
+    try {
+      const webSearchResponse: WebSearchResponse = await WebSearchService.processWebsearch(
+        webSearchProvider,
+        extractResults
+      )
+      console.log('webSearchResponse', webSearchResponse)
+      // 处理搜索结果
+      message.metadata = {
+        ...message.metadata,
+        webSearch: webSearchResponse
       }
-      const webSearchParams = getOpenAIWebSearchParams(assistant, assistant.model)
-      if (isEmpty(webSearchParams) && !isOpenAIWebSearch(assistant.model)) {
-        const lastMessage = findLast(messages, (m) => m.role === 'user')
-        const lastAnswer = findLast(messages, (m) => m.role === 'assistant')
-        const hasKnowledgeBase = !isEmpty(lastMessage?.knowledgeBaseIds)
-
-        if (lastMessage) {
-          if (hasKnowledgeBase) {
-            window.message.info({
-              content: i18n.t('message.ignore.knowledge.base'),
-              key: 'knowledge-base-no-match-info'
-            })
-          }
 
-          // 更新消息状态为搜索中
-          onResponse({ ...message, status: 'searching' })
-
-          try {
-            // 等待关键词生成完成
-            const searchSummaryAssistant = getDefaultAssistant()
-            searchSummaryAssistant.model = assistant.model || getDefaultModel()
-            searchSummaryAssistant.prompt = SEARCH_SUMMARY_PROMPT
-
-            // 如果启用搜索增强模式，则使用搜索增强模式
-            if (WebSearchService.isEnhanceModeEnabled()) {
-              const keywords = await fetchSearchSummary({
-                messages: lastAnswer ? [lastAnswer, lastMessage] : [lastMessage],
-                assistant: searchSummaryAssistant
-              })
-
-              try {
-                const result = WebSearchService.extractInfoFromXML(keywords || '')
-                if (result.question === 'not_needed') {
-                  // 如果不需要搜索，则直接返回
-                  console.log('No need to search')
-                  return
-                } else if (result.question === 'summarize' && result.links && result.links.length > 0) {
-                  const contents = await fetchWebContents(result.links)
-                  webSearchResponse = {
-                    query: 'summaries',
-                    results: contents
-                  }
-                } else {
-                  query = result.question
-                  webSearchResponse = await WebSearchService.search(webSearchProvider, query)
-                }
-              } catch (error) {
-                console.error('Failed to extract info from XML:', error)
-              }
-            } else {
-              query = lastMessage.content
-            }
-
-            // 处理搜索结果
-            message.metadata = {
-              ...message.metadata,
-              webSearch: webSearchResponse
-            }
+      window.keyv.set(`web-search-${lastUserMessage?.id}`, webSearchResponse)
+    } catch (error) {
+      console.error('Web search failed:', error)
+    }
+  }
 
-            window.keyv.set(`web-search-${lastMessage?.id}`, webSearchResponse)
-          } catch (error) {
-            console.error('Web search failed:', error)
-          }
-        }
-      }
+  // --- 知识库搜索 ---
+  const searchKnowledgeBase = async () => {
+    const shouldSearch = hasKnowledgeBase && extractResults.tools?.includes('knowledge')
+
+    if (!shouldSearch) return
+
+    try {
+      const knowledgeReferences: KnowledgeReference[] = await processKnowledgeSearch(
+        extractResults,
+        lastUserMessage.knowledgeBaseIds
+      )
+      console.log('knowledgeReferences', knowledgeReferences)
+      window.keyv.set(`knowledge-search-${lastUserMessage?.id}`, knowledgeReferences)
+    } catch (error) {
+      console.error('Knowledge base search failed:', error)
+      window.keyv.set(`knowledge-search-${lastUserMessage?.id}`, [])
     }
   }
 
   try {
     let _messages: Message[] = []
     let isFirstChunk = true
 
-    // Search web
-    await searchTheWeb()
+    // --- 并行执行搜索 ---
+    // 更新消息状态为搜索中
+    onResponse({ ...message, status: 'searching' })
+    await Promise.all([searchTheWeb(), searchKnowledgeBase()])
 
-    const lastUserMessage = findLast(messages, (m) => m.role === 'user')
     // Get MCP tools
     const mcpTools: MCPTool[] = []
     const enabledMCPs = lastUserMessage?.enabledMCPs