@lobehub/chat 1.96.13 → 1.96.14

package/src/libs/model-runtime/google/index.ts

@@ -1,14 +1,13 @@
-import type { VertexAI } from '@google-cloud/vertexai';
 import {
   Content,
-  FunctionCallPart,
   FunctionDeclaration,
+  GenerateContentConfig,
   Tool as GoogleFunctionCallTool,
-  GoogleGenerativeAI,
-  GoogleSearchRetrievalTool,
+  GoogleGenAI,
   Part,
-  SchemaType,
-} from '@google/generative-ai';
+  Type as SchemaType,
+  ThinkingConfig,
+} from '@google/genai';
 
 import { imageUrlToBase64 } from '@/utils/imageToBase64';
 import { safeParseJSON } from '@/utils/safeParseJSON';
@@ -77,16 +76,11 @@ const DEFAULT_BASE_URL = 'https://generativelanguage.googleapis.com';
 interface LobeGoogleAIParams {
   apiKey?: string;
   baseURL?: string;
-  client?: GoogleGenerativeAI | VertexAI;
+  client?: GoogleGenAI;
   id?: string;
   isVertexAi?: boolean;
 }
 
-interface GoogleAIThinkingConfig {
-  includeThoughts?: boolean;
-  thinkingBudget?: number;
-}
-
 const isAbortError = (error: Error): boolean => {
   const message = error.message.toLowerCase();
   return (
@@ -99,7 +93,7 @@ const isAbortError = (error: Error): boolean => {
 };
 
 export class LobeGoogleAI implements LobeRuntimeAI {
-  private client: GoogleGenerativeAI;
+  private client: GoogleGenAI;
   private isVertexAi: boolean;
   baseURL?: string;
   apiKey?: string;
@@ -108,9 +102,10 @@ export class LobeGoogleAI implements LobeRuntimeAI {
   constructor({ apiKey, baseURL, client, isVertexAi, id }: LobeGoogleAIParams = {}) {
     if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
 
-    this.client = new GoogleGenerativeAI(apiKey);
+    const httpOptions = baseURL ? { baseUrl: baseURL } : undefined;
+
     this.apiKey = apiKey;
-    this.client = client ? (client as GoogleGenerativeAI) : new GoogleGenerativeAI(apiKey);
+    this.client = client ? client : new GoogleGenAI({ apiKey, httpOptions });
     this.baseURL = client ? undefined : baseURL || DEFAULT_BASE_URL;
     this.isVertexAi = isVertexAi || false;
 
@@ -122,7 +117,7 @@ export class LobeGoogleAI implements LobeRuntimeAI {
       const payload = this.buildPayload(rawPayload);
       const { model, thinkingBudget } = payload;
 
-      const thinkingConfig: GoogleAIThinkingConfig = {
+      const thinkingConfig: ThinkingConfig = {
         includeThoughts:
           !!thinkingBudget ||
           (!thinkingBudget && model && (model.includes('-2.5-') || model.includes('thinking')))
@@ -155,8 +150,6 @@ export class LobeGoogleAI implements LobeRuntimeAI {
 
       const contents = await this.buildGoogleMessages(payload.messages);
 
-      const inputStartAt = Date.now();
-
       const controller = new AbortController();
       const originalSignal = options?.signal;
 
@@ -170,57 +163,50 @@ export class LobeGoogleAI implements LobeRuntimeAI {
         }
       }
 
-      const geminiStreamResult = await this.client
-        .getGenerativeModel(
+      const config: GenerateContentConfig = {
+        abortSignal: originalSignal,
+        maxOutputTokens: payload.max_tokens,
+        responseModalities: modelsWithModalities.has(model) ? ['Text', 'Image'] : undefined,
+        // avoid wide sensitive words
+        // refs: https://github.com/lobehub/lobe-chat/pull/1418
+        safetySettings: [
           {
-            generationConfig: {
-              maxOutputTokens: payload.max_tokens,
-              // @ts-expect-error - Google SDK 0.24.0 doesn't have this property for now with
-              response_modalities: modelsWithModalities.has(model) ? ['Text', 'Image'] : undefined,
-              temperature: payload.temperature,
-              topP: payload.top_p,
-              ...(modelsDisableInstuction.has(model) || model.toLowerCase().includes('learnlm')
-                ? {}
-                : { thinkingConfig }),
-            },
-            model,
-            // avoid wide sensitive words
-            // refs: https://github.com/lobehub/lobe-chat/pull/1418
-            safetySettings: [
-              {
-                category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-                threshold: getThreshold(model),
-              },
-              {
-                category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
-                threshold: getThreshold(model),
-              },
-              {
-                category: HarmCategory.HARM_CATEGORY_HARASSMENT,
-                threshold: getThreshold(model),
-              },
-              {
-                category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
-                threshold: getThreshold(model),
-              },
-            ],
+            category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+            threshold: getThreshold(model),
+          },
+          {
+            category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+            threshold: getThreshold(model),
           },
-          { apiVersion: 'v1beta', baseUrl: this.baseURL },
-        )
-        .generateContentStream(
           {
-            contents,
-            systemInstruction: modelsDisableInstuction.has(model)
-              ? undefined
-              : (payload.system as string),
-            tools: this.buildGoogleTools(payload.tools, payload),
+            category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+            threshold: getThreshold(model),
           },
           {
-            signal: controller.signal,
+            category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+            threshold: getThreshold(model),
           },
-        );
+        ],
+        systemInstruction: modelsDisableInstuction.has(model)
+          ? undefined
+          : (payload.system as string),
+        temperature: payload.temperature,
+        thinkingConfig:
+          modelsDisableInstuction.has(model) || model.toLowerCase().includes('learnlm')
+            ? undefined
+            : thinkingConfig,
+        tools: this.buildGoogleTools(payload.tools, payload),
+        topP: payload.top_p,
+      };
+
+      const inputStartAt = Date.now();
+      const geminiStreamResponse = await this.client.models.generateContentStream({
+        config,
+        contents,
+        model,
+      });
 
-      const googleStream = this.createEnhancedStream(geminiStreamResult.stream, controller.signal);
+      const googleStream = this.createEnhancedStream(geminiStreamResponse, controller.signal);
       const [prod, useForDebug] = googleStream.tee();
 
       const key = this.isVertexAi
@@ -405,7 +391,7 @@ export class LobeGoogleAI implements LobeRuntimeAI {
     const content = message.content as string | UserMessageContentPart[];
     if (!!message.tool_calls) {
       return {
-        parts: message.tool_calls.map<FunctionCallPart>((tool) => ({
+        parts: message.tool_calls.map<Part>((tool) => ({
           functionCall: {
             args: safeParseJSON(tool.function.arguments)!,
             name: tool.function.name,
@@ -527,7 +513,7 @@ export class LobeGoogleAI implements LobeRuntimeAI {
       return this.buildFunctionDeclarations(tools);
     }
     if (payload?.enabledSearch) {
-      return [{ googleSearch: {} } as GoogleSearchRetrievalTool];
+      return [{ googleSearch: {} }];
     }
 
     return this.buildFunctionDeclarations(tools);

package/src/libs/model-runtime/utils/streams/google-ai.test.ts

@@ -1,4 +1,4 @@
-import { EnhancedGenerateContentResponse } from '@google/generative-ai';
+import { GenerateContentResponse } from '@google/genai';
 import { describe, expect, it, vi } from 'vitest';
 
 import * as uuidModule from '@/utils/uuid';
@@ -11,10 +11,9 @@ describe('GoogleGenerativeAIStream', () => {
 
     const mockGenerateContentResponse = (text: string, functionCalls?: any[]) =>
       ({
-        text: () => text,
-        functionCall: () => functionCalls?.[0],
-        functionCalls: () => functionCalls,
-      }) as EnhancedGenerateContentResponse;
+        text: text,
+        functionCalls: functionCalls,
+      }) as unknown as GenerateContentResponse;
 
     const mockGoogleStream = new ReadableStream({
       start(controller) {
@@ -114,12 +113,6 @@ describe('GoogleGenerativeAIStream', () => {
       },
       modelVersion: 'gemini-2.0-flash-exp',
     };
-    const mockGenerateContentResponse = (text: string, functionCalls?: any[]) =>
-      ({
-        text: () => text,
-        functionCall: () => functionCalls?.[0],
-        functionCalls: () => functionCalls,
-      }) as EnhancedGenerateContentResponse;
 
     const mockGoogleStream = new ReadableStream({
       start(controller) {
@@ -209,7 +202,7 @@ describe('GoogleGenerativeAIStream', () => {
             ],
           },
         ],
-        text: () => '234',
+        text: '234',
         usageMetadata: {
           promptTokenCount: 20,
           totalTokenCount: 20,
@@ -218,7 +211,7 @@ describe('GoogleGenerativeAIStream', () => {
         modelVersion: 'gemini-2.0-flash-exp-image-generation',
       },
       {
-        text: () => '567890\n',
+        text: '567890\n',
         candidates: [
           {
             content: { parts: [{ text: '567890\n' }], role: 'model' },
@@ -299,7 +292,7 @@ describe('GoogleGenerativeAIStream', () => {
             ],
           },
         ],
-        text: () => '234',
+        text: '234',
         usageMetadata: {
           promptTokenCount: 19,
           candidatesTokenCount: 3,
@@ -307,10 +300,10 @@ describe('GoogleGenerativeAIStream', () => {
           promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
           thoughtsTokenCount: 100,
         },
-        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
       },
       {
-        text: () => '567890\n',
+        text: '567890\n',
         candidates: [
           {
             content: { parts: [{ text: '567890\n' }], role: 'model' },
@@ -331,7 +324,7 @@ describe('GoogleGenerativeAIStream', () => {
           candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }],
           thoughtsTokenCount: 100,
         },
-        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
       },
     ];
 
@@ -375,4 +368,410 @@ describe('GoogleGenerativeAIStream', () => {
       ].map((i) => i + '\n'),
     );
   });
+
+  it('should handle thought candidate part', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = [
+      {
+        candidates: [
+          {
+            content: {
+              parts: [{ text: '**Understanding the Conditional Logic**\n\n', thought: true }],
+              role: 'model',
+            },
+            index: 0,
+          },
+        ],
+        text: '**Understanding the Conditional Logic**\n\n',
+        usageMetadata: {
+          promptTokenCount: 38,
+          candidatesTokenCount: 7,
+          totalTokenCount: 301,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }],
+          thoughtsTokenCount: 256,
+        },
+        modelVersion: 'models/gemini-2.5-flash-preview-04-17',
+      },
+      {
+        candidates: [
+          {
+            content: {
+              parts: [{ text: '**Finalizing Interpretation**\n\n', thought: true }],
+              role: 'model',
+            },
+            index: 0,
+          },
+        ],
+        text: '**Finalizing Interpretation**\n\n',
+        usageMetadata: {
+          promptTokenCount: 38,
+          candidatesTokenCount: 13,
+          totalTokenCount: 355,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }],
+          thoughtsTokenCount: 304,
+        },
+        modelVersion: 'models/gemini-2.5-flash-preview-04-17',
+      },
+      {
+        candidates: [
+          {
+            content: {
+              parts: [{ text: '简单来说，' }],
+              role: 'model',
+            },
+            index: 0,
+          },
+        ],
+        text: '简单来说，',
+        usageMetadata: {
+          promptTokenCount: 38,
+          candidatesTokenCount: 16,
+          totalTokenCount: 358,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }],
+          thoughtsTokenCount: 304,
+        },
+        modelVersion: 'models/gemini-2.5-flash-preview-04-17',
+      },
+      {
+        candidates: [
+          {
+            content: { parts: [{ text: '文本内容。' }], role: 'model' },
+            finishReason: 'STOP',
+            index: 0,
+          },
+        ],
+        text: '文本内容。',
+        usageMetadata: {
+          promptTokenCount: 38,
+          candidatesTokenCount: 19,
+          totalTokenCount: 361,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }],
+          thoughtsTokenCount: 304,
+        },
+        modelVersion: 'models/gemini-2.5-flash-preview-04-17',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: chat_1',
+        'event: reasoning',
+        'data: "**Understanding the Conditional Logic**\\n\\n"\n',
+
+        'id: chat_1',
+        'event: reasoning',
+        `data: "**Finalizing Interpretation**\\n\\n"\n`,
+
+        'id: chat_1',
+        'event: text',
+        `data: "简单来说，"\n`,
+
+        'id: chat_1',
+        'event: text',
+        `data: "文本内容。"\n`,
+        // stop
+        'id: chat_1',
+        'event: stop',
+        `data: "STOP"\n`,
+        // usage
+        'id: chat_1',
+        'event: usage',
+        `data: {"inputTextTokens":38,"outputReasoningTokens":304,"outputTextTokens":19,"totalInputTokens":38,"totalOutputTokens":323,"totalTokens":361}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
+
+  it('should return undefined data without text', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = [
+      {
+        candidates: [
+          {
+            content: { parts: [{ text: '234' }], role: 'model' },
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        text: '234',
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 3,
+          totalTokenCount: 122,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          thoughtsTokenCount: 100,
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+      {
+        text: '',
+        candidates: [
+          {
+            content: { parts: [{ text: '' }], role: 'model' },
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 3,
+          totalTokenCount: 122,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }],
+          thoughtsTokenCount: 100,
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+      {
+        text: '567890\n',
+        candidates: [
+          {
+            content: { parts: [{ text: '567890\n' }], role: 'model' },
+            finishReason: 'STOP',
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 11,
+          totalTokenCount: 131,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }],
+          thoughtsTokenCount: 100,
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: chat_1',
+        'event: text',
+        'data: "234"\n',
+
+        'id: chat_1',
+        'event: text',
+        'data: ""\n',
+
+        'id: chat_1',
+        'event: text',
+        `data: "567890\\n"\n`,
+        // stop
+        'id: chat_1',
+        'event: stop',
+        `data: "STOP"\n`,
+        // usage
+        'id: chat_1',
+        'event: usage',
+        `data: {"inputTextTokens":19,"outputReasoningTokens":100,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":111,"totalTokens":131}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
+
+  it('should handle groundingMetadata', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = [
+      {
+        text: '123',
+        candidates: [
+          {
+            content: {
+              parts: [
+                {
+                  text: '123',
+                },
+              ],
+              role: 'model',
+            },
+            index: 0,
+            groundingMetadata: {},
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 9,
+          candidatesTokenCount: 18,
+          totalTokenCount: 27,
+          promptTokensDetails: [
+            {
+              modality: 'TEXT',
+              tokenCount: 9,
+            },
+          ],
+        },
+        modelVersion: 'models/gemini-2.5-flash-preview-04-17',
+      },
+      {
+        text: '45678',
+        candidates: [
+          {
+            content: {
+              parts: [
+                {
+                  text: '45678',
+                },
+              ],
+              role: 'model',
+            },
+            finishReason: 'STOP',
+            index: 0,
+            groundingMetadata: {
+              searchEntryPoint: {
+                renderedContent: 'content\n',
+              },
+              groundingChunks: [
+                {
+                  web: {
+                    uri: 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXG1234545',
+                    title: 'npmjs.com',
+                  },
+                },
+                {
+                  web: {
+                    uri: 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXE9288334',
+                    title: 'google.dev',
+                  },
+                },
+              ],
+              groundingSupports: [
+                {
+                  segment: {
+                    startIndex: 63,
+                    endIndex: 67,
+                    text: '1。',
+                  },
+                  groundingChunkIndices: [0],
+                  confidenceScores: [1],
+                },
+                {
+                  segment: {
+                    startIndex: 69,
+                    endIndex: 187,
+                    text: 'SDK。',
+                  },
+                  groundingChunkIndices: [1],
+                  confidenceScores: [1],
+                },
+              ],
+              webSearchQueries: ['sdk latest version'],
+            },
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 9,
+          candidatesTokenCount: 122,
+          totalTokenCount: 131,
+          promptTokensDetails: [
+            {
+              modality: 'TEXT',
+              tokenCount: 9,
+            },
+          ],
+        },
+        modelVersion: 'models/gemini-2.5-flash-preview-04-17',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: chat_1',
+        'event: text',
+        'data: "123"\n',
+
+        'id: chat_1',
+        'event: grounding',
+        'data: {}\n',
+
+        'id: chat_1',
+        'event: text',
+        'data: "45678"\n',
+
+        'id: chat_1',
+        'event: grounding',
+        `data: {\"citations\":[{\"favicon\":\"npmjs.com\",\"title\":\"npmjs.com\",\"url\":\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXG1234545\"},{\"favicon\":\"google.dev\",\"title\":\"google.dev\",\"url\":\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXE9288334\"}],\"searchQueries\":[\"sdk latest version\"]}\n`,
+        // stop
+        'id: chat_1',
+        'event: stop',
+        `data: "STOP"\n`,
+        // usage
+        'id: chat_1',
+        'event: usage',
+        `data: {"inputTextTokens":9,"outputTextTokens":122,"totalInputTokens":9,"totalOutputTokens":122,"totalTokens":131}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
 });