@lobehub/chat 1.96.12 → 1.96.14

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'),
+    );
+  });
 });

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

@@ -1,4 +1,4 @@
-import { EnhancedGenerateContentResponse } from '@google/generative-ai';
+import { GenerateContentResponse } from '@google/genai';
 
 import { ModelTokensUsage } from '@/types/message';
 import { GroundingSearch } from '@/types/search';
@@ -16,7 +16,7 @@ import {
 } from './protocol';
 
 const transformGoogleGenerativeAIStream = (
-  chunk: EnhancedGenerateContentResponse,
+  chunk: GenerateContentResponse,
   context: StreamContext,
 ): StreamProtocolChunk | StreamProtocolChunk[] => {
   // maybe need another structure to add support for multiple choices
@@ -24,22 +24,22 @@ const transformGoogleGenerativeAIStream = (
   const usage = chunk.usageMetadata;
   const usageChunks: StreamProtocolChunk[] = [];
   if (candidate?.finishReason && usage) {
-    const outputReasoningTokens = (usage as any).thoughtsTokenCount || undefined;
-    const totalOutputTokens = (usage.candidatesTokenCount ?? 0) + (outputReasoningTokens ?? 0);
+    // totalTokenCount = promptTokenCount + candidatesTokenCount + thoughtsTokenCount
+    const reasoningTokens = usage.thoughtsTokenCount;
+    const outputTextTokens = usage.candidatesTokenCount ?? 0;
+    const totalOutputTokens = outputTextTokens + (reasoningTokens ?? 0);
 
     usageChunks.push(
       { data: candidate.finishReason, id: context?.id, type: 'stop' },
       {
         data: {
           // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
-          inputImageTokens: (usage as any).promptTokensDetails?.find(
-            (i: any) => i.modality === 'IMAGE',
-          )?.tokenCount,
-          inputTextTokens: (usage as any).promptTokensDetails?.find(
-            (i: any) => i.modality === 'TEXT',
-          )?.tokenCount,
-          outputReasoningTokens,
-          outputTextTokens: totalOutputTokens - (outputReasoningTokens ?? 0),
+          inputImageTokens: usage.promptTokensDetails?.find((i: any) => i.modality === 'IMAGE')
+            ?.tokenCount,
+          inputTextTokens: usage.promptTokensDetails?.find((i: any) => i.modality === 'TEXT')
+            ?.tokenCount,
+          outputReasoningTokens: reasoningTokens,
+          outputTextTokens,
           totalInputTokens: usage.promptTokenCount,
           totalOutputTokens,
           totalTokens: usage.totalTokenCount,
@@ -50,7 +50,7 @@ const transformGoogleGenerativeAIStream = (
     );
   }
 
-  const functionCalls = chunk.functionCalls?.();
+  const functionCalls = chunk.functionCalls;
 
   if (functionCalls) {
     return [
@@ -73,11 +73,11 @@ const transformGoogleGenerativeAIStream = (
     ];
   }
 
-  const text = chunk.text?.();
+  const text = chunk.text;
 
   if (candidate) {
     // 首先检查是否为 reasoning 内容 (thought: true)
-    if (Array.isArray(candidate.content.parts) && candidate.content.parts.length > 0) {
+    if (Array.isArray(candidate.content?.parts) && candidate.content.parts.length > 0) {
       for (const part of candidate.content.parts) {
         if (part && part.text && (part as any).thought === true) {
           return { data: part.text, id: context.id, type: 'reasoning' };
@@ -122,7 +122,7 @@ const transformGoogleGenerativeAIStream = (
     if (!!text?.trim()) return { data: text, id: context?.id, type: 'text' };
 
     // streaming the image
-    if (Array.isArray(candidate.content.parts) && candidate.content.parts.length > 0) {
+    if (Array.isArray(candidate.content?.parts) && candidate.content.parts.length > 0) {
       const part = candidate.content.parts[0];
 
       if (part && part.inlineData && part.inlineData.data && part.inlineData.mimeType) {
@@ -148,7 +148,7 @@ export interface GoogleAIStreamOptions {
 }
 
 export const GoogleGenerativeAIStream = (
-  rawStream: ReadableStream<EnhancedGenerateContentResponse>,
+  rawStream: ReadableStream<GenerateContentResponse>,
   { callbacks, inputStartAt }: GoogleAIStreamOptions = {},
 ) => {
   const streamStack: StreamContext = { id: 'chat_' + nanoid() };