From 7216e9943c57df5a1f4a8842b3ad86ce5236a728 Mon Sep 17 00:00:00 2001 From: MyPrototypeWhat Date: Tue, 29 Jul 2025 19:26:29 +0800 Subject: [PATCH] refactor: streamline async function syntax and enhance plugin event handling - Simplified async function syntax in `RuntimeExecutor` and `PluginEngine` for improved readability. - Updated `AiSdkToChunkAdapter` to refine condition checks for Google metadata. - Enhanced `searchOrchestrationPlugin` to log conversation messages and improve memory storage logic. - Improved memory processing by ensuring fallback for existing memories. - Added new citation block handling in `toolCallbacks` for better integration with web search results. --- packages/aiCore/src/core/runtime/executor.ts | 8 ++----- .../aiCore/src/core/runtime/pluginEngine.ts | 2 +- .../src/aiCore/chunk/AiSdkToChunkAdapter.ts | 19 ++++++++------- .../plugins/searchOrchestrationPlugin.ts | 10 ++++---- src/renderer/src/services/MemoryProcessor.ts | 3 +-- .../messageStreaming/callbacks/index.ts | 3 ++- .../callbacks/textCallbacks.ts | 5 ++-- .../callbacks/toolCallbacks.ts | 24 +++++++++++++++++-- 8 files changed, 47 insertions(+), 27 deletions(-) diff --git a/packages/aiCore/src/core/runtime/executor.ts b/packages/aiCore/src/core/runtime/executor.ts index c36bb5a798..1f29baa747 100644 --- a/packages/aiCore/src/core/runtime/executor.ts +++ b/packages/aiCore/src/core/runtime/executor.ts @@ -165,9 +165,7 @@ export class RuntimeExecutor { 'generateObject', typeof modelOrId === 'string' ? modelOrId : modelOrId.modelId, params, - async (model, transformedParams) => { - return await generateObject({ model, ...transformedParams }) - } + async (model, transformedParams) => await generateObject({ model, ...transformedParams }) ) } @@ -201,9 +199,7 @@ export class RuntimeExecutor { 'streamObject', typeof modelOrId === 'string' ? modelOrId : modelOrId.modelId, params, - async (model, transformedParams) => { - return await streamObject({ model, ...transformedParams }) - } + async (model, transformedParams) => await streamObject({ model, ...transformedParams }) ) } diff --git a/packages/aiCore/src/core/runtime/pluginEngine.ts b/packages/aiCore/src/core/runtime/pluginEngine.ts index db14c3d5b1..e646501144 100644 --- a/packages/aiCore/src/core/runtime/pluginEngine.ts +++ b/packages/aiCore/src/core/runtime/pluginEngine.ts @@ -162,7 +162,7 @@ export class PluginEngine { const transformedResult = await this.pluginManager.executeSequential('transformResult', result, context) // 6. 触发完成事件（注意：对于流式调用，这里触发的是开始流式响应的事件） - await this.pluginManager.executeParallel('onRequestEnd', context, { stream: true }) + await this.pluginManager.executeParallel('onRequestEnd', context, transformedResult) return transformedResult } catch (error) { diff --git a/src/renderer/src/aiCore/chunk/AiSdkToChunkAdapter.ts b/src/renderer/src/aiCore/chunk/AiSdkToChunkAdapter.ts index e9b5430863..e5ea0aea86 100644 --- a/src/renderer/src/aiCore/chunk/AiSdkToChunkAdapter.ts +++ b/src/renderer/src/aiCore/chunk/AiSdkToChunkAdapter.ts @@ -174,7 +174,7 @@ export class AiSdkToChunkAdapter { case 'finish-step': { const { providerMetadata } = chunk // googel web search - if (providerMetadata?.google) { + if (providerMetadata?.google?.groundingMetadata) { this.onChunk({ type: ChunkType.LLM_WEB_SEARCH_COMPLETE, llm_web_search: { @@ -182,15 +182,16 @@ export class AiSdkToChunkAdapter { source: WebSearchSource.GEMINI } }) - } else { - this.onChunk({ - type: ChunkType.LLM_WEB_SEARCH_COMPLETE, - llm_web_search: { - results: final.webSearchResults, - source: WebSearchSource.AISDK - } - }) } + // else { + // this.onChunk({ + // type: ChunkType.LLM_WEB_SEARCH_COMPLETE, + // llm_web_search: { + // results: final.webSearchResults, + // source: WebSearchSource.AISDK + // } + // }) + // } final.webSearchResults = [] // final.reasoningId = '' break diff --git a/src/renderer/src/aiCore/plugins/searchOrchestrationPlugin.ts b/src/renderer/src/aiCore/plugins/searchOrchestrationPlugin.ts index 916ee074ec..cbf1f0baeb 100644 --- a/src/renderer/src/aiCore/plugins/searchOrchestrationPlugin.ts +++ b/src/renderer/src/aiCore/plugins/searchOrchestrationPlugin.ts @@ -178,7 +178,7 @@ async function storeConversationMemory( content: getMessageContent(msg) || '' })) .filter((msg) => msg.content.trim().length > 0) - + console.log('conversationMessages', conversationMessages) if (conversationMessages.length < 2) { console.log('Need at least a user message and assistant response for memory processing') return @@ -368,10 +368,12 @@ export const searchOrchestrationPlugin = (assistant: Assistant) => { /** * 💾 Step 3: 记忆存储阶段 */ - // eslint-disable-next-line @typescript-eslint/no-unused-vars - onRequestEnd: async (context: AiRequestContext, _result: any) => { - console.log('💾 [SearchOrchestration] Starting memory storage...', context.requestId) + onRequestEnd: async (context: AiRequestContext, result: any) => { + // context.isAnalyzing = false + console.log('context.isAnalyzing', context, result) + console.log('💾 [SearchOrchestration] Starting memory storage...', context.requestId) + if (context.isAnalyzing) return try { const messages = context.originalParams.messages diff --git a/src/renderer/src/services/MemoryProcessor.ts b/src/renderer/src/services/MemoryProcessor.ts index 11a5a9c528..21180c4825 100644 --- a/src/renderer/src/services/MemoryProcessor.ts +++ b/src/renderer/src/services/MemoryProcessor.ts @@ -100,8 +100,7 @@ export class MemoryProcessor { if (!memoryConfig.llmApiClient) { throw new Error('No LLM model configured for memory processing') } - - const existingMemoriesResult = window.keyv.get(`memory-search-${lastMessageId}`) as MemoryItem[] | [] + const existingMemoriesResult = (window.keyv.get(`memory-search-${lastMessageId}`) as MemoryItem[]) || [] const existingMemories = existingMemoriesResult.map((memory) => ({ id: memory.id, diff --git a/src/renderer/src/services/messageStreaming/callbacks/index.ts b/src/renderer/src/services/messageStreaming/callbacks/index.ts index 2b6fc5968a..866147be41 100644 --- a/src/renderer/src/services/messageStreaming/callbacks/index.ts +++ b/src/renderer/src/services/messageStreaming/callbacks/index.ts @@ -59,7 +59,8 @@ export const createCallbacks = (deps: CallbacksDependencies) => { blockManager, getState, assistantMsgId, - getCitationBlockId: citationCallbacks.getCitationBlockId + getCitationBlockId: citationCallbacks.getCitationBlockId, + getCitationBlockIdFromTool: toolCallbacks.getCitationBlockId }) // 组合所有回调 diff --git a/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts index d9e9e7f17c..b36f4f3fa9 100644 --- a/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts +++ b/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts @@ -9,10 +9,11 @@ interface TextCallbacksDependencies { getState: any assistantMsgId: string getCitationBlockId: () => string | null + getCitationBlockIdFromTool: () => string | null } export const createTextCallbacks = (deps: TextCallbacksDependencies) => { - const { blockManager, getState, assistantMsgId, getCitationBlockId } = deps + const { blockManager, getState, assistantMsgId, getCitationBlockId, getCitationBlockIdFromTool } = deps // 内部维护的状态 let mainTextBlockId: string | null = null @@ -37,7 +38,7 @@ export const createTextCallbacks = (deps: TextCallbacksDependencies) => { }, onTextChunk: async (text: string) => { - const citationBlockId = getCitationBlockId() + const citationBlockId = getCitationBlockId() || getCitationBlockIdFromTool() const citationBlockSource = citationBlockId ? (getState().messageBlocks.entities[citationBlockId] as CitationMessageBlock).response?.source : WebSearchSource.WEBSEARCH diff --git a/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts index 8d89bd05dd..4f66bbd2c0 100644 --- a/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts +++ b/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts @@ -1,6 +1,7 @@ import type { MCPToolResponse } from '@renderer/types' +import { WebSearchSource } from '@renderer/types' import { MessageBlockStatus, MessageBlockType, ToolMessageBlock } from '@renderer/types/newMessage' -import { createToolBlock } from '@renderer/utils/messageUtils/create' +import { createCitationBlock, createToolBlock } from '@renderer/utils/messageUtils/create' import { BlockManager } from '../BlockManager' @@ -15,6 +16,7 @@ export const createToolCallbacks = (deps: ToolCallbacksDependencies) => { // 内部维护的状态 const toolCallIdToBlockIdMap = new Map() let toolBlockId: string | null = null + let citationBlockId: string | null = null return { onToolCallPending: (toolResponse: MCPToolResponse) => { @@ -94,6 +96,21 @@ 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) { + const citationBlock = createCitationBlock( + assistantMsgId, + { + response: { results: toolResponse.response.rawResults, source: WebSearchSource.AISDK } + }, + { + status: MessageBlockStatus.SUCCESS + } + ) + citationBlockId = citationBlock.id + blockManager.handleBlockTransition(citationBlock, MessageBlockType.CITATION) + } } else { console.warn( `[onToolCallComplete] Received unhandled tool status: ${toolResponse.status} for ID: ${toolResponse.id}` @@ -101,6 +118,9 @@ export const createToolCallbacks = (deps: ToolCallbacksDependencies) => { } toolBlockId = null - } + }, + + // 暴露给 textCallbacks 使用的方法 + getCitationBlockId: () => citationBlockId } }