@lobehub/chat 1.96.12 → 1.96.14

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

@@ -330,4 +330,133 @@ describe('VertexAIStream', () => {
       ].map((i) => i + '\n'),
     );
   });
+
+  it('should return empty text chunk without candidates', 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' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 20,
+          candidatesTokenCount: 3,
+          totalTokenCount: 23,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }],
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+      {
+        usageMetadata: {
+          promptTokenCount: 20,
+          candidatesTokenCount: 3,
+          totalTokenCount: 23,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }],
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = VertexAIStream(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`].map(
+        (i) => i + '\n',
+      ),
+    );
+  });
+
+  it('should return stop chunk with empty content candidates', 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' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 20,
+          candidatesTokenCount: 3,
+          totalTokenCount: 23,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }],
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+      {
+        candidates: [{}],
+        usageMetadata: {
+          promptTokenCount: 20,
+          candidatesTokenCount: 3,
+          totalTokenCount: 23,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }],
+        },
+        modelVersion: 'gemini-2.5-flash-preview-04-17',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = VertexAIStream(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: stop', `data: ""\n`].map(
+        (i) => i + '\n',
+      ),
+    );
+  });
 });

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

@@ -1,4 +1,4 @@
-import { EnhancedGenerateContentResponse, GenerateContentResponse } from '@google/generative-ai';
+import { GenerateContentResponse } from '@google/genai';
 
 import { ModelTokensUsage } from '@/types/message';
 import { nanoid } from '@/utils/uuid';
@@ -22,8 +22,9 @@ const transformVertexAIStream = (
   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);
+    const outputReasoningTokens = usage.thoughtsTokenCount || undefined;
+    const outputTextTokens = usage.candidatesTokenCount ?? 0;
+    const totalOutputTokens = outputTextTokens + (outputReasoningTokens ?? 0);
 
     usageChunks.push(
       { data: candidate.finishReason, id: context?.id, type: 'stop' },
@@ -37,7 +38,7 @@ const transformVertexAIStream = (
             (i: any) => i.modality === 'TEXT',
           )?.tokenCount,
           outputReasoningTokens,
-          outputTextTokens: totalOutputTokens - (outputReasoningTokens ?? 0),
+          outputTextTokens,
           totalInputTokens: usage.promptTokenCount,
           totalOutputTokens,
           totalTokens: usage.totalTokenCount,
@@ -50,7 +51,7 @@ const transformVertexAIStream = (
 
   if (
     candidate && // 首先检查是否为 reasoning 内容 (thought: true)
-    Array.isArray(candidate.content.parts) &&
+    Array.isArray(candidate.content?.parts) &&
     candidate.content.parts.length > 0
   ) {
     for (const part of candidate.content.parts) {
@@ -60,19 +61,18 @@ const transformVertexAIStream = (
     }
   }
 
-  const candidates = chunk.candidates;
-  if (!candidates)
+  if (!candidate) {
     return {
       data: '',
       id: context?.id,
       type: 'text',
     };
+  }
 
-  const item = candidates[0];
-  if (item.content) {
-    const part = item.content.parts[0];
+  if (candidate.content) {
+    const part = candidate.content.parts?.[0];
 
-    if (part.functionCall) {
+    if (part?.functionCall) {
       const functionCall = part.functionCall;
 
       return [
@@ -95,18 +95,18 @@ const transformVertexAIStream = (
       ];
     }
 
-    if (item.finishReason) {
+    if (candidate.finishReason) {
       if (chunk.usageMetadata) {
         return [
-          !!part.text ? { data: part.text, id: context?.id, type: 'text' } : undefined,
+          !!part?.text ? { data: part.text, id: context?.id, type: 'text' } : undefined,
           ...usageChunks,
         ].filter(Boolean) as StreamProtocolChunk[];
       }
-      return { data: item.finishReason, id: context?.id, type: 'stop' };
+      return { data: candidate.finishReason, id: context?.id, type: 'stop' };
     }
 
     return {
-      data: part.text,
+      data: part?.text,
       id: context?.id,
       type: 'text',
     };
@@ -120,7 +120,7 @@ const transformVertexAIStream = (
 };
 
 export const VertexAIStream = (
-  rawStream: ReadableStream<EnhancedGenerateContentResponse>,
+  rawStream: ReadableStream<GenerateContentResponse>,
   { callbacks, inputStartAt }: GoogleAIStreamOptions = {},
 ) => {
   const streamStack: StreamContext = { id: 'chat_' + nanoid() };

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

@@ -1,13 +1,19 @@
-import { VertexAI, VertexInit } from '@google-cloud/vertexai';
+import { GoogleGenAI, GoogleGenAIOptions } from '@google/genai';
 
 import { AgentRuntimeErrorType } from '../error';
 import { LobeGoogleAI } from '../google';
 import { AgentRuntimeError } from '../utils/createError';
 
+const DEFAULT_VERTEXAI_LOCATION = 'global';
+
 export class LobeVertexAI extends LobeGoogleAI {
-  static initFromVertexAI(params?: VertexInit) {
+  static initFromVertexAI(params?: GoogleGenAIOptions) {
     try {
-      const client = new VertexAI({ ...params });
+      const client = new GoogleGenAI({
+        ...params,
+        location: params?.location ?? DEFAULT_VERTEXAI_LOCATION, // @google/genai 不传 location 会报错
+        vertexai: true,
+      });
 
       return new LobeGoogleAI({ apiKey: 'avoid-error', client, isVertexAi: true });
     } catch (e) {