@lobehub/chat 1.60.8 → 1.61.0

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

@@ -0,0 +1,236 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import * as uuidModule from '@/utils/uuid';
+
+import { VertexAIStream } from './vertex-ai';
+
+describe('VertexAIStream', () => {
+  it('should transform Vertex AI stream to protocol stream', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+    const rawChunks = [
+      {
+        candidates: [
+          {
+            content: { role: 'model', parts: [{ text: '你好' }] },
+            safetyRatings: [
+              {
+                category: 'HARM_CATEGORY_HATE_SPEECH',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.06298828,
+                severity: 'HARM_SEVERY_NEGLIGIBLE',
+                severityScore: 0.10986328,
+              },
+              {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.05029297,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.078125,
+              },
+              {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.19433594,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.16015625,
+              },
+              {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.059326172,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.064453125,
+              },
+            ],
+            index: 0,
+          },
+        ],
+        usageMetadata: {},
+        modelVersion: 'gemini-1.5-flash-001',
+      },
+      {
+        candidates: [
+          {
+            content: { role: 'model', parts: [{ text: '！ 😊' }] },
+            safetyRatings: [
+              {
+                category: 'HARM_CATEGORY_HATE_SPEECH',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.052734375,
+                severity: 'HARM_SEVRITY_NEGLIGIBLE',
+                severityScore: 0.08642578,
+              },
+              {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.071777344,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.095214844,
+              },
+              {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.1640625,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.10498047,
+              },
+              {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.075683594,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.053466797,
+              },
+            ],
+            index: 0,
+          },
+        ],
+        modelVersion: 'gemini-1.5-flash-001',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        rawChunks.forEach((chunk) => controller.enqueue(chunk));
+
+        controller.close();
+      },
+    });
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onToolCallMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = VertexAIStream(mockGoogleStream, {
+      onStart: onStartMock,
+      onText: onTextMock,
+      onToken: onTokenMock,
+      onToolCall: onToolCallMock,
+      onCompletion: onCompletionMock,
+    });
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      // text
+      'id: chat_1\n',
+      'event: text\n',
+      `data: "你好"\n\n`,
+
+      // text
+      'id: chat_1\n',
+      'event: text\n',
+      `data: "！ 😊"\n\n`,
+    ]);
+
+    expect(onStartMock).toHaveBeenCalledTimes(1);
+    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onCompletionMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('tool_calls', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+    const rawChunks = [
+      {
+        candidates: [
+          {
+            content: {
+              role: 'model',
+              parts: [
+                {
+                  functionCall: {
+                    name: 'realtime-weather____fetchCurrentWeather',
+                    args: { city: '杭州' },
+                  },
+                },
+              ],
+            },
+            finishReason: 'STOP',
+            safetyRatings: [
+              {
+                category: 'HARM_CATERY_HATE_SPEECH',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.09814453,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.07470703,
+              },
+              {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.1484375,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.15136719,
+              },
+              {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.11279297,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.10107422,
+              },
+              {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                probability: 'NEGLIGIBLE',
+                probabilityScore: 0.048828125,
+                severity: 'HARM_SEVERITY_NEGLIGIBLE',
+                severityScore: 0.05493164,
+              },
+            ],
+            index: 0,
+          },
+        ],
+        usageMetadata: { promptTokenCount: 95, candidatesTokenCount: 9, totalTokenCount: 104 },
+        modelVersion: 'gemini-1.5-flash-001',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        rawChunks.forEach((chunk) => controller.enqueue(chunk));
+
+        controller.close();
+      },
+    });
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onToolCallMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = VertexAIStream(mockGoogleStream, {
+      onStart: onStartMock,
+      onText: onTextMock,
+      onToken: onTokenMock,
+      onToolCall: onToolCallMock,
+      onCompletion: onCompletionMock,
+    });
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      // text
+      'id: chat_1\n',
+      'event: tool_calls\n',
+      `data: [{"function":{"arguments":"{\\"city\\":\\"杭州\\"}","name":"realtime-weather____fetchCurrentWeather"},"id":"realtime-weather____fetchCurrentWeather_0","index":0,"type":"function"}]\n\n`,
+    ]);
+
+    expect(onStartMock).toHaveBeenCalledTimes(1);
+    expect(onToolCallMock).toHaveBeenCalledTimes(1);
+    expect(onCompletionMock).toHaveBeenCalledTimes(1);
+  });
+});

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

@@ -0,0 +1,75 @@
+import { EnhancedGenerateContentResponse, GenerateContentResponse } from '@google/generative-ai';
+
+import { nanoid } from '@/utils/uuid';
+
+import { ChatStreamCallbacks } from '../../types';
+import {
+  StreamProtocolChunk,
+  StreamStack,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+  generateToolCallId,
+} from './protocol';
+
+const transformVertexAIStream = (
+  chunk: GenerateContentResponse,
+  stack: StreamStack,
+): StreamProtocolChunk => {
+  // maybe need another structure to add support for multiple choices
+  const candidates = chunk.candidates;
+
+  if (!candidates)
+    return {
+      data: '',
+      id: stack?.id,
+      type: 'text',
+    };
+
+  const item = candidates[0];
+  if (item.content) {
+    const part = item.content.parts[0];
+
+    if (part.functionCall) {
+      const functionCall = part.functionCall;
+
+      return {
+        data: [
+          {
+            function: {
+              arguments: JSON.stringify(functionCall.args),
+              name: functionCall.name,
+            },
+            id: generateToolCallId(0, functionCall.name),
+            index: 0,
+            type: 'function',
+          },
+        ],
+        id: stack?.id,
+        type: 'tool_calls',
+      };
+    }
+
+    return {
+      data: part.text,
+      id: stack?.id,
+      type: 'text',
+    };
+  }
+
+  return {
+    data: '',
+    id: stack?.id,
+    type: 'stop',
+  };
+};
+
+export const VertexAIStream = (
+  rawStream: ReadableStream<EnhancedGenerateContentResponse>,
+  callbacks?: ChatStreamCallbacks,
+) => {
+  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+
+  return rawStream
+    .pipeThrough(createSSEProtocolTransformer(transformVertexAIStream, streamStack))
+    .pipeThrough(createCallbacksTransformer(callbacks));
+};

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

@@ -0,0 +1,23 @@
+import { VertexAI, VertexInit } from '@google-cloud/vertexai';
+
+import { AgentRuntimeError, AgentRuntimeErrorType, LobeGoogleAI } from '@/libs/agent-runtime';
+
+export class LobeVertexAI extends LobeGoogleAI {
+  static initFromVertexAI(params?: VertexInit) {
+    try {
+      const client = new VertexAI({ ...params });
+
+      return new LobeGoogleAI({ apiKey: 'avoid-error', client, isVertexAi: true });
+    } catch (e) {
+      const err = e as Error;
+
+      if (err.name === 'IllegalArgumentError') {
+        throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidVertexCredentials, {
+          message: err.message,
+        });
+      }
+
+      throw e;
+    }
+  }
+}

package/src/locales/default/error.ts

@@ -106,6 +106,7 @@ export default {
      */
     OpenAIBizError: '请求 OpenAI 服务出错，请根据以下信息排查或重试',
 
+    InvalidVertexCredentials: 'Vertex 鉴权未通过，请检查鉴权凭证后重试',
     InvalidBedrockCredentials: 'Bedrock 鉴权未通过，请检查 AccessKeyId/SecretAccessKey 后重试',
     StreamChunkError:
       '流式请求的消息块解析错误，请检查当前 API 接口是否符合标准规范，或联系你的 API 供应商咨询',

package/src/locales/default/modelProvider.ts

@@ -325,6 +325,13 @@ export default {
     tooltip: '更新服务商基础配置',
     updateSuccess: '更新成功',
   },
+  vertexai: {
+    apiKey: {
+      desc: '填入你的 Vertex Ai Keys',
+      placeholder: `{ "type": "service_account", "project_id": "xxx", "private_key_id": ... }`,
+      title: 'Vertex AI Keys',
+    },
+  },
   zeroone: {
     title: '01.AI 零一万物',
   },

package/src/types/user/settings/keyVaults.ts

@@ -68,6 +68,7 @@ export interface UserKeyVaults {
   tencentcloud?: OpenAICompatibleKeyVault;
   togetherai?: OpenAICompatibleKeyVault;
   upstage?: OpenAICompatibleKeyVault;
+  vertexai?: OpenAICompatibleKeyVault;
   vllm?: OpenAICompatibleKeyVault;
   volcengine?: OpenAICompatibleKeyVault;
   wenxin?: OpenAICompatibleKeyVault;

package/src/utils/safeParseJSON.ts

@@ -1,4 +1,4 @@
-export const safeParseJSON = <T = Record<string, any>>(text: string) => {
+export const safeParseJSON = <T = Record<string, any>>(text?: string) => {
   if (typeof text !== 'string') return undefined;
 
   let json: T;