From 0c4e8228af04dc3dc77dbd4473c1dde935f56987 Mon Sep 17 00:00:00 2001
From: lizhixuan <daoquqiexing@gmail.com>
Date: Sat, 12 Jul 2025 21:11:17 +0800
Subject: [PATCH] feat: enhance AiSdkToChunkAdapter for web search results
 handling

- Updated `AiSdkToChunkAdapter` to include `webSearchResults` in the final output structure for improved web search integration.
- Modified `convertAndEmitChunk` method to handle `finish-step` events, differentiating between Google and other web search results.
- Adjusted the handling of `source` events to accumulate web search results for better processing.
- Enhanced citation formatting in `messageBlock.ts` to support new web search result structures.
---
 .../src/aiCore/AiSdkToChunkAdapter.ts         | 62 ++++++++++++++-----
 src/renderer/src/store/messageBlock.ts        |  9 +++
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/renderer/src/aiCore/AiSdkToChunkAdapter.ts b/src/renderer/src/aiCore/AiSdkToChunkAdapter.ts
index d800b17e73..9cc9960807 100644
--- a/src/renderer/src/aiCore/AiSdkToChunkAdapter.ts
+++ b/src/renderer/src/aiCore/AiSdkToChunkAdapter.ts
@@ -4,7 +4,7 @@
  */
 
 import { TextStreamPart, ToolSet } from '@cherrystudio/ai-core'
-import { MCPTool, WebSearchSource } from '@renderer/types'
+import { MCPTool, WebSearchResults, WebSearchSource } from '@renderer/types'
 import { Chunk, ChunkType } from '@renderer/types/chunk'
 
 import { ToolCallChunkHandler } from './chunk/handleTooCallChunk'
@@ -55,7 +55,8 @@ export class AiSdkToChunkAdapter {
     const reader = fullStream.getReader()
     const final = {
       text: '',
-      reasoning_content: ''
+      reasoningContent: '',
+      webSearchResults: []
     }
     try {
       while (true) {
@@ -77,7 +78,10 @@ export class AiSdkToChunkAdapter {
    * 转换 AI SDK chunk 为 Cherry Studio chunk 并调用回调
    * @param chunk AI SDK 的 chunk 数据
    */
-  private convertAndEmitChunk(chunk: TextStreamPart<any>, final: { text: string; reasoning_content: string }) {
+  private convertAndEmitChunk(
+    chunk: TextStreamPart<any>,
+    final: { text: string; reasoningContent: string; webSearchResults: any[] }
+  ) {
     console.log('AI SDK chunk type:', chunk.type, chunk)
     switch (chunk.type) {
       // === 文本相关事件 ===
@@ -147,16 +151,37 @@ export class AiSdkToChunkAdapter {
       //   final.text = ''
       //   break
 
-      // case 'finish-step': {
-      //   const { totalUsage, finishReason, providerMetadata } = chunk
-      // }
+      case 'finish-step': {
+        const { providerMetadata } = chunk
+        // googel web search
+        if (providerMetadata?.google) {
+          this.onChunk({
+            type: ChunkType.LLM_WEB_SEARCH_COMPLETE,
+            llm_web_search: {
+              results: providerMetadata.google?.groundingMetadata as WebSearchResults,
+              source: WebSearchSource.GEMINI
+            }
+          })
+        } else {
+          this.onChunk({
+            type: ChunkType.LLM_WEB_SEARCH_COMPLETE,
+            llm_web_search: {
+              results: final.webSearchResults,
+              source: WebSearchSource.AISDK
+            }
+          })
+        }
+        final.webSearchResults = []
+        break
+        // const { totalUsage, finishReason, providerMetadata } = chunk
+      }
 
       case 'finish':
         this.onChunk({
           type: ChunkType.BLOCK_COMPLETE,
           response: {
             text: final.text || '',
-            reasoning_content: final.reasoning_content || '',
+            reasoning_content: final.reasoningContent || '',
             usage: {
               completion_tokens: chunk.totalUsage.outputTokens || 0,
               prompt_tokens: chunk.totalUsage.inputTokens || 0,
@@ -174,7 +199,7 @@ export class AiSdkToChunkAdapter {
           type: ChunkType.LLM_RESPONSE_COMPLETE,
           response: {
             text: final.text || '',
-            reasoning_content: final.reasoning_content || '',
+            reasoning_content: final.reasoningContent || '',
             usage: {
               completion_tokens: chunk.totalUsage.outputTokens || 0,
               prompt_tokens: chunk.totalUsage.inputTokens || 0,
@@ -192,13 +217,20 @@ export class AiSdkToChunkAdapter {
 
       // === 源和文件相关事件 ===
       case 'source':
-        this.onChunk({
-          type: ChunkType.LLM_WEB_SEARCH_COMPLETE,
-          llm_web_search: {
-            source: WebSearchSource.AISDK,
-            results: [{}]
-          }
-        })
+        if (chunk.sourceType === 'url') {
+          // if (final.webSearchResults.length === 0) {
+          // eslint-disable-next-line @typescript-eslint/no-unused-vars
+          const { sourceType: _, ...rest } = chunk
+          final.webSearchResults.push(rest)
+          // }
+          // this.onChunk({
+          //   type: ChunkType.LLM_WEB_SEARCH_COMPLETE,
+          //   llm_web_search: {
+          //     source: WebSearchSource.AISDK,
+          //     results: final.webSearchResults
+          //   }
+          // })
+        }
         break
       // case 'file':
       //   // 文件相关事件，可能是图片生成
diff --git a/src/renderer/src/store/messageBlock.ts b/src/renderer/src/store/messageBlock.ts
index f9b8c34cd5..37714a80c1 100644
--- a/src/renderer/src/store/messageBlock.ts
+++ b/src/renderer/src/store/messageBlock.ts
@@ -207,6 +207,15 @@ export const formatCitationsFromBlock = (block: CitationMessageBlock | undefined
             type: 'websearch'
           })) || []
         break
+      case WebSearchSource.AISDK:
+        formattedCitations =
+          (block.response.results as any[])?.map((result, index) => ({
+            number: index + 1,
+            url: result.url,
+            title: result.title,
+            providerMetadata: result?.providerMetadata
+          })) || []
+        break
     }
   }
   // 3. Handle Knowledge Base References