From ca4e7e3d2b4107d08c121825cfe222cd4c610aff Mon Sep 17 00:00:00 2001 From: MyPrototypeWhat Date: Fri, 22 Aug 2025 19:35:09 +0800 Subject: [PATCH] feat(tools): refactor MemorySearchTool and WebSearchTool for improved response handling - Updated MemorySearchTool to utilize aiSdk for better integration and removed unused imports. - Refactored WebSearchTool to streamline search results handling, changing from an array to a structured object for clarity. - Adjusted MessageTool and MessageWebSearchTool components to reflect changes in tool response structure. - Enhanced error handling and logging in tool callbacks for improved debugging and user feedback. --- .../src/aiCore/tools/MemorySearchTool.ts | 4 +- .../src/aiCore/tools/WebSearchTool.ts | 44 +++++++------------ .../pages/home/Messages/Tools/MessageTool.tsx | 14 +++--- .../Messages/Tools/MessageWebSearchTool.tsx | 2 +- src/renderer/src/services/ApiService.ts | 3 +- .../callbacks/toolCallbacks.ts | 4 +- 6 files changed, 33 insertions(+), 38 deletions(-) diff --git a/src/renderer/src/aiCore/tools/MemorySearchTool.ts b/src/renderer/src/aiCore/tools/MemorySearchTool.ts index cd62f7de8b..67ece6d3ac 100644 --- a/src/renderer/src/aiCore/tools/MemorySearchTool.ts +++ b/src/renderer/src/aiCore/tools/MemorySearchTool.ts @@ -1,11 +1,13 @@ +import { aiSdk, type InferToolOutput } from '@cherrystudio/ai-core' import store from '@renderer/store' import { selectCurrentUserId, selectGlobalMemoryEnabled, selectMemoryConfig } from '@renderer/store/memory' import type { Assistant } from '@renderer/types' -import { type InferToolOutput, tool } from 'ai' +// import { type InferToolOutput, tool } from 'ai' import { z } from 'zod' import { MemoryProcessor } from '../../services/MemoryProcessor' +const { tool } = aiSdk /** * 🧠 基础记忆搜索工具 * AI 可以主动调用的简单记忆搜索 diff --git a/src/renderer/src/aiCore/tools/WebSearchTool.ts b/src/renderer/src/aiCore/tools/WebSearchTool.ts index 199b7a08d2..f998f91d37 100644 --- a/src/renderer/src/aiCore/tools/WebSearchTool.ts +++ b/src/renderer/src/aiCore/tools/WebSearchTool.ts @@ -54,15 +54,17 @@ Call this tool to execute the search. You can optionally provide additional cont } } - const searchResults: WebSearchProviderResponse[] = [] + let searchResults: WebSearchProviderResponse = { + query: '', + results: [] + } // 检查是否需要搜索 if (finalQueries[0] === 'not_needed') { return { summary: 'No search needed based on the query analysis.', - searchResults: [], + searchResults, sources: '', - instructions: '', - rawResults: [] + instructions: '' } } @@ -74,29 +76,24 @@ Call this tool to execute the search. You can optionally provide additional cont links: extractedKeywords.links } } - const response = await WebSearchService.processWebsearch(webSearchProvider!, extractResults, requestId) - searchResults.push(response) + searchResults = await WebSearchService.processWebsearch(webSearchProvider!, extractResults, requestId) } catch (error) { return { summary: `Search failed: ${error instanceof Error ? error.message : 'Unknown error'}`, - searchResults: [], - sources: '', - instructions: '', - rawResults: [] + sources: [], + instructions: '' } } - - if (searchResults.length === 0 || !searchResults[0].results) { + console.log('searchResults', searchResults) + if (searchResults.results.length === 0) { return { summary: 'No search results found for the given query.', - searchResults: [], - sources: '', - instructions: '', - rawResults: [] + sources: [], + instructions: '' } } - const results = searchResults[0].results + const results = searchResults.results const citationData = results.map((result, index) => ({ number: index + 1, title: result.title, @@ -105,25 +102,18 @@ Call this tool to execute the search. You can optionally provide additional cont })) // 🔑 返回引用友好的格式,复用 REFERENCE_PROMPT 逻辑 - const referenceContent = `\`\`\`json\n${JSON.stringify(citationData, null, 2)}\n\`\`\`` + // const referenceContent = `\`\`\`json\n${JSON.stringify(citationData, null, 2)}\n\`\`\`` // 构建完整的引用指导文本 const fullInstructions = REFERENCE_PROMPT.replace( '{question}', "Based on the search results, please answer the user's question with proper citations." - ).replace('{references}', referenceContent) + ).replace('{references}', 'searchResults:') return { summary: `Found ${citationData.length} relevant sources. Use [number] format to cite specific information.`, searchResults, - sources: citationData - .map((source) => `[${source.number}] ${source.title}\n${source.content}\nURL: ${source.url}`) - .join('\n\n'), - - instructions: fullInstructions, - - // 原始数据,便于后续处理 - rawResults: citationData + instructions: fullInstructions } } }) diff --git a/src/renderer/src/pages/home/Messages/Tools/MessageTool.tsx b/src/renderer/src/pages/home/Messages/Tools/MessageTool.tsx index afd9450a65..38b2b0846e 100644 --- a/src/renderer/src/pages/home/Messages/Tools/MessageTool.tsx +++ b/src/renderer/src/pages/home/Messages/Tools/MessageTool.tsx @@ -2,8 +2,8 @@ import { MCPToolResponse } from '@renderer/types' import type { ToolMessageBlock } from '@renderer/types/newMessage' import { Collapse } from 'antd' -import { MessageKnowledgeSearchToolBody, MessageKnowledgeSearchToolTitle } from './MessageKnowledgeSearch' -import { MessageWebSearchToolBody, MessageWebSearchToolTitle } from './MessageWebSearchTool' +import { MessageKnowledgeSearchToolTitle } from './MessageKnowledgeSearch' +import { MessageWebSearchToolTitle } from './MessageWebSearchTool' interface Props { block: ToolMessageBlock @@ -69,12 +69,12 @@ const ChooseTool = (toolResponse: MCPToolResponse): { label: React.ReactNode; bo case 'web_search_preview': return { label: , - body: + body: null } case 'knowledge_search': return { label: , - body: + body: null } default: return null @@ -82,7 +82,7 @@ const ChooseTool = (toolResponse: MCPToolResponse): { label: React.ReactNode; bo } export default function MessageTool({ block }: Props) { - // FIXME: 语义错误,这里已经不是 MCP tool 了 + // FIXME: 语义错误,这里已经不是 MCP tool 了,更改rawMcpToolResponse需要改用户数据, 所以暂时保留 const toolResponse = block.metadata?.rawMcpToolResponse if (!toolResponse) return null @@ -91,7 +91,7 @@ export default function MessageTool({ block }: Props) { if (!toolRenderer) return null - return ( + return toolRenderer.body ? ( + ) : ( + toolRenderer.label ) } // const PrepareToolWrapper = styled.span` diff --git a/src/renderer/src/pages/home/Messages/Tools/MessageWebSearchTool.tsx b/src/renderer/src/pages/home/Messages/Tools/MessageWebSearchTool.tsx index 6ccb3c4c2b..9f3321bf4e 100644 --- a/src/renderer/src/pages/home/Messages/Tools/MessageWebSearchTool.tsx +++ b/src/renderer/src/pages/home/Messages/Tools/MessageWebSearchTool.tsx @@ -25,7 +25,7 @@ export const MessageWebSearchToolTitle = ({ toolResponse }: { toolResponse: MCPT {i18n.t('message.websearch.fetch_complete', { - count: toolOutput?.searchResults?.reduce((acc, result) => acc + result.results.length, 0) ?? 0 + count: toolOutput?.searchResults?.results?.length ?? 0 })} ) diff --git a/src/renderer/src/services/ApiService.ts b/src/renderer/src/services/ApiService.ts index 55d1c0aae6..25ba538d3d 100644 --- a/src/renderer/src/services/ApiService.ts +++ b/src/renderer/src/services/ApiService.ts @@ -380,7 +380,7 @@ export async function fetchGenerate({ // } const middlewareConfig: AiSdkMiddlewareConfig = { - streamOutput: assistant.settings?.streamOutput ?? true, + streamOutput: assistant.settings?.streamOutput ?? false, enableReasoning: false, isPromptToolUse: false, isSupportedToolUse: false, @@ -393,6 +393,7 @@ export async function fetchGenerate({ const result = await AI.completions( model.id, { + system: prompt, prompt: content }, { diff --git a/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts index 9cc14a937f..2a970d8afa 100644 --- a/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts +++ b/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts @@ -103,11 +103,11 @@ export const createToolCallbacks = (deps: ToolCallbacksDependencies) => { blockManager.smartBlockUpdate(existingBlockId, changes, MessageBlockType.TOOL, true) // Handle citation block creation for web search results - if (toolResponse.tool.name === 'builtin_web_search' && toolResponse.response?.rawResults) { + if (toolResponse.tool.name === 'builtin_web_search' && toolResponse.response?.searchResults) { const citationBlock = createCitationBlock( assistantMsgId, { - response: { results: toolResponse.response.rawResults, source: WebSearchSource.AISDK } + response: { results: toolResponse.response.searchResults, source: WebSearchSource.WEBSEARCH } }, { status: MessageBlockStatus.SUCCESS