feat: 修复工具调用和流式响应处理

Dango233 · aider-chat-bot · Dango233 · commit e18dafd95175 · 2025-07-02T18:28:55.000+08:00
Co-authored-by: aider (vertex_ai/gemini-2.5-pro) &lt;aider@aider.chat&gt;
diff --git a/packages/mcp-server/src/bridge/stream-transformer.ts b/packages/mcp-server/src/bridge/stream-transformer.ts
@@ -16,38 +16,66 @@ interface OpenAIDelta {
   }[];
 }
 
-interface OpenAIChoice {
-  index: number;
-  delta: OpenAIDelta;
-  finish_reason: string | null;
-}
-
 interface OpenAIChunk {
   id: string;
   object: 'chat.completion.chunk';
   created: number;
   model: string;
-  choices: OpenAIChoice[];
+  choices: {
+    index: number;
+    delta: OpenAIDelta;
+    finish_reason: string | null;
+  }[];
 }
 
-// --- 更新的转换器 ---
+type ToolCallState = {
+  id: string;
+  name: string;
+  arguments: string;
+};
+
+// --- 更新的、有状态的转换器 ---
 export function createOpenAIStreamTransformer(
   model: string,
 ): TransformStream<GenerateContentResponse, Uint8Array> {
   const chatID = `chatcmpl-${randomUUID()}`;
   const creationTime = Math.floor(Date.now() / 1000);
   const encoder = new TextEncoder();
   let isFirstChunk = true;
+  const toolCallStates: ToolCallState[] = [];
 
-  return new TransformStream({
-    transform(chunk, controller) {
-      const parts = chunk.candidates?.[0]?.content?.parts || [];
-      const finishReason = chunk.candidates?.[0]?.finishReason;
+  const createChunk = (
+    delta: OpenAIDelta,
+    finish_reason: string | null = null,
+  ): OpenAIChunk => ({
+    id: chatID,
+    object: 'chat.completion.chunk',
+    created: creationTime,
+    model: model,
+    choices: [
+      {
+        index: 0,
+        delta,
+        finish_reason,
+      },
+    ],
+  });
 
-      let hasContent = false;
+  const enqueueChunk = (
+    controller: TransformStreamDefaultController<Uint8Array>,
+    chunk: OpenAIChunk,
+  ) => {
+    const sseString = `data: ${JSON.stringify(chunk)}\n\n`;
+    controller.enqueue(encoder.encode(sseString));
+  };
+
+  return new TransformStream({
+    transform(geminiChunk, controller) {
+      const parts = geminiChunk.candidates?.[0]?.content?.parts || [];
+      const finishReason = geminiChunk.candidates?.[0]?.finishReason;
 
       for (const part of parts) {
-        const delta: OpenAIDelta = {};
+        let delta: OpenAIDelta = {};
 
         if (isFirstChunk) {
           delta.role = 'assistant';
@@ -56,73 +84,67 @@ export function createOpenAIStreamTransformer(
 
         if (part.text) {
           delta.content = part.text;
-          hasContent = true;
+          enqueueChunk(controller, createChunk(delta));
         }
 
         if (part.functionCall) {
           const fc = part.functionCall;
-          const callId = `call_${randomUUID()}`; // 为每个调用生成一个唯一的ID
+          const callId = `call_${randomUUID()}`;
 
-          // OpenAI的流式工具调用是分块的，我们这里简化为一次性发送
-          // Gemini通常也是一次性返回一个完整的functionCall
-          delta.tool_calls = [
-            {
-              index: 0, // 假设只有一个工具调用
-              id: callId,
-              type: 'function',
-              function: {
-                name: fc.name,
-                arguments: JSON.stringify(fc.args), // 参数必须是字符串
+          // 模拟分块发送 tool_calls
+          // 1. 发送带有 name 的块
+          const nameDelta: OpenAIDelta = {
+            tool_calls: [
+              {
+                index: toolCallStates.length,
+                id: callId,
+                type: 'function',
+                function: { name: fc.name, arguments: '' },
               },
-            },
-          ];
-          hasContent = true;
-        }
+            ],
+          };
+          if (isFirstChunk) {
+            nameDelta.role = 'assistant';
+            isFirstChunk = false;
+          }
+          enqueueChunk(controller, createChunk(nameDelta));
 
-        if (hasContent) {
-          const openAIChunk: OpenAIChunk = {
-            id: chatID,
-            object: 'chat.completion.chunk',
-            created: creationTime,
-            model: model,
-            choices: [
+          // 2. 发送带有 arguments 的块
+          const argsDelta: OpenAIDelta = {
+            tool_calls: [
               {
-                index: 0,
-                delta: delta,
-                finish_reason: null,
+                index: toolCallStates.length,
+                id: callId,
+                type: 'function',
+                function: { arguments: JSON.stringify(fc.args) },
               },
             ],
           };
-          const sseString = `data: ${JSON.stringify(openAIChunk)}\n\n`;
-          controller.enqueue(encoder.encode(sseString));
+          enqueueChunk(controller, createChunk(argsDelta));
+
+          toolCallStates.push({
+            id: callId,
+            name: fc.name,
+            arguments: JSON.stringify(fc.args),
+          });
         }
       }
 
-      // 如果有 finishReason，发送一个带有 finish_reason 的块
       if (
         finishReason &&
-        finishReason !== 'FINISH_REASON_UNSPECIFIED'
+        finishReason !== 'FINISH_REASON_UNSPECIFIED' &&
+        finishReason !== 'NOT_SET'
       ) {
-        const finishDelta: OpenAIDelta = {};
-        const openAIChunk: OpenAIChunk = {
-          id: chatID,
-          object: 'chat.completion.chunk',
-          created: creationTime,
-          model: model,
-          choices: [
-            {
-              index: 0,
-              delta: finishDelta,
-              finish_reason: finishReason === 'STOP' ? 'stop' : 'tool_calls',
-            },
-          ],
-        };
-        const sseString = `data: ${JSON.stringify(openAIChunk)}\n\n`;
-        controller.enqueue(encoder.encode(sseString));
+        const reason =
+          finishReason === 'STOP'
+            ? 'stop'
+            : finishReason === 'TOOL_CALL'
+              ? 'tool_calls'
+              : finishReason.toLowerCase();
+        enqueueChunk(controller, createChunk({}, reason));
       }
     },
     flush(controller) {
-      // 流结束时，发送 [DONE] 消息
       const doneString = `data: [DONE]\n\n`;
       controller.enqueue(encoder.encode(doneString));
     },
diff --git a/packages/mcp-server/src/gemini-client.ts b/packages/mcp-server/src/gemini-client.ts
@@ -5,8 +5,18 @@
  */
 
 import { type Config, GeminiChat } from '@google/gemini-cli-core';
-import { type Content, type Part, type Tool } from '@google/genai';
-import { type OpenAIMessage, type MessageContentPart } from './types.js';
+import {
+  type Content,
+  type Part,
+  type Tool,
+  type FunctionDeclaration,
+  type GenerateContentConfig,
+} from '@google/genai';
+import {
+  type OpenAIMessage,
+  type MessageContentPart,
+  type OpenAIChatCompletionRequest,
+} from './types.js';
 
 export class GeminiApiClient {
   private readonly config: Config;
@@ -17,13 +27,51 @@ export class GeminiApiClient {
     this.contentGenerator = this.config.getGeminiClient().getContentGenerator();
   }
 
+  /**
+   * 将 OpenAI 的工具定义转换为 Gemini 的工具定义。
+   */
+  private convertOpenAIToolsToGemini(
+    openAITools?: OpenAIChatCompletionRequest['tools'],
+  ): Tool[] | undefined {
+    if (!openAITools || openAITools.length === 0) {
+      return undefined;
+    }
+
+    const functionDeclarations: FunctionDeclaration[] = openAITools
+      .filter((tool) => tool.type === 'function' && tool.function)
+      .map((tool) => ({
+        name: tool.function.name,
+        description: tool.function.description,
+        parameters: tool.function.parameters,
+      }));
+
+    if (functionDeclarations.length === 0) {
+      return undefined;
+    }
+
+    return [{ functionDeclarations }];
+  }
+
   /**
    * 将 OpenAI 格式的消息转换为 Gemini 格式的 Content 对象。
-   * 这个函数现在能处理文本和图片（多模态）输入。
    */
   private openAIMessageToGemini(msg: OpenAIMessage): Content {
     const role = msg.role === 'assistant' ? 'model' : 'user';
 
+    if (msg.role === 'tool') {
+      return {
+        role: 'user', // Gemini 使用 'user' role 来承载 functionResponse
+        parts: [
+          {
+            functionResponse: {
+              name: msg.tool_call_id || 'unknown_tool', // 需要一个工具名
+              response: { content: msg.content },
+            },
+          },
+        ],
+      };
+    }
+
     if (typeof msg.content === 'string') {
       return { role, parts: [{ text: msg.content }] };
     }
@@ -41,32 +89,15 @@ export class GeminiApiClient {
               const mimeType = mimePart.split(':')[1].split(';')[0];
               return { inlineData: { mimeType, data: dataPart } };
             }
-            // Gemini API 可能不支持直接传递 URL，但我们先按协议转换
             return { fileData: { mimeType: 'image/jpeg', fileUri: imageUrl } };
           }
-          // 对于不支持的 part 类型，返回一个空文本 part
           return { text: '' };
         })
-        .filter((p) => p.text !== '' || p.inlineData || p.fileData); // 过滤掉完全空的 part
+        .filter((p) => p.text !== '' || p.inlineData || p.fileData);
 
       return { role, parts };
     }
 
-    // 针对 tool role 的转换
-    if (msg.role === 'tool' && msg.tool_call_id && msg.content) {
-      return {
-        role: 'user', // Gemini 使用 'user' role 来承载 functionResponse
-        parts: [
-          {
-            functionResponse: {
-              name: msg.tool_call_id, // 这里的映射关系需要确认，通常是工具名
-              response: { content: msg.content },
-            },
-          },
-        ],
-      };
-    }
-
     return { role, parts: [{ text: '' }] };
   }
 
@@ -81,32 +112,41 @@ export class GeminiApiClient {
   }: {
     model: string;
     messages: OpenAIMessage[];
-    tools?: Tool[];
+    tools?: OpenAIChatCompletionRequest['tools'];
     tool_choice?: any;
   }) {
-    // 1. 转换消息格式
-    const history = messages.map((msg) => this.openAIMessageToGemini(msg));
+    const history = messages.map(this.openAIMessageToGemini);
     const lastMessage = history.pop();
     if (!lastMessage) {
       throw new Error('No message to send.');
     }
 
-    // 2. 创建一个一次性的 GeminiChat 实例
     const oneShotChat = new GeminiChat(
       this.config,
       this.contentGenerator,
-      {}, // generationConfig
-      history, // 传入历史记录
+      {},
+      history,
     );
 
-    // 3. 构造请求，包含工具定义
+    const geminiTools = this.convertOpenAIToolsToGemini(tools);
+
+    const generationConfig: GenerateContentConfig = {};
+    if (tool_choice && tool_choice !== 'auto') {
+      generationConfig.toolConfig = {
+        functionCallingConfig: {
+          mode: tool_choice.type === 'function' ? 'ANY' : 'AUTO',
+          allowedFunctionNames: tool_choice.function
+            ? [tool_choice.function.name]
+            : undefined,
+        },
+      };
+    }
+
     const geminiStream = await oneShotChat.sendMessageStream({
       message: lastMessage.parts || [],
       config: {
-        tools: tools,
-        toolConfig: tool_choice
-          ? { functionCallingConfig: { mode: tool_choice } }
-          : undefined,
+        tools: geminiTools,
+        ...generationConfig,
       },
     });
 
