@lobehub/chat 1.60.9 → 1.61.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (93) hide show
  1. package/CHANGELOG.md +33 -0
  2. package/changelog/v1.json +12 -0
  3. package/locales/ar/error.json +1 -0
  4. package/locales/ar/modelProvider.json +7 -0
  5. package/locales/ar/models.json +3 -12
  6. package/locales/ar/providers.json +3 -0
  7. package/locales/bg-BG/error.json +1 -0
  8. package/locales/bg-BG/modelProvider.json +7 -0
  9. package/locales/bg-BG/models.json +3 -12
  10. package/locales/bg-BG/providers.json +3 -0
  11. package/locales/de-DE/error.json +1 -0
  12. package/locales/de-DE/modelProvider.json +7 -0
  13. package/locales/de-DE/models.json +3 -12
  14. package/locales/de-DE/providers.json +3 -0
  15. package/locales/en-US/error.json +1 -0
  16. package/locales/en-US/modelProvider.json +7 -0
  17. package/locales/en-US/models.json +3 -12
  18. package/locales/en-US/providers.json +3 -0
  19. package/locales/es-ES/error.json +1 -0
  20. package/locales/es-ES/modelProvider.json +7 -0
  21. package/locales/es-ES/models.json +3 -12
  22. package/locales/es-ES/providers.json +3 -0
  23. package/locales/fa-IR/error.json +1 -0
  24. package/locales/fa-IR/modelProvider.json +7 -0
  25. package/locales/fa-IR/models.json +3 -12
  26. package/locales/fa-IR/providers.json +3 -0
  27. package/locales/fr-FR/error.json +1 -0
  28. package/locales/fr-FR/modelProvider.json +7 -0
  29. package/locales/fr-FR/models.json +3 -12
  30. package/locales/fr-FR/providers.json +3 -0
  31. package/locales/it-IT/error.json +1 -0
  32. package/locales/it-IT/modelProvider.json +7 -0
  33. package/locales/it-IT/models.json +3 -12
  34. package/locales/it-IT/providers.json +3 -0
  35. package/locales/ja-JP/error.json +1 -0
  36. package/locales/ja-JP/modelProvider.json +7 -0
  37. package/locales/ja-JP/models.json +3 -12
  38. package/locales/ja-JP/providers.json +3 -0
  39. package/locales/ko-KR/error.json +1 -0
  40. package/locales/ko-KR/modelProvider.json +7 -0
  41. package/locales/ko-KR/models.json +3 -12
  42. package/locales/ko-KR/providers.json +3 -0
  43. package/locales/nl-NL/error.json +1 -0
  44. package/locales/nl-NL/modelProvider.json +7 -0
  45. package/locales/nl-NL/models.json +3 -12
  46. package/locales/nl-NL/providers.json +3 -0
  47. package/locales/pl-PL/error.json +1 -0
  48. package/locales/pl-PL/modelProvider.json +7 -0
  49. package/locales/pl-PL/models.json +3 -12
  50. package/locales/pl-PL/providers.json +3 -0
  51. package/locales/pt-BR/error.json +1 -0
  52. package/locales/pt-BR/modelProvider.json +7 -0
  53. package/locales/pt-BR/models.json +3 -12
  54. package/locales/pt-BR/providers.json +3 -0
  55. package/locales/ru-RU/error.json +1 -0
  56. package/locales/ru-RU/modelProvider.json +7 -0
  57. package/locales/ru-RU/models.json +3 -12
  58. package/locales/ru-RU/providers.json +3 -0
  59. package/locales/tr-TR/error.json +1 -0
  60. package/locales/tr-TR/modelProvider.json +7 -0
  61. package/locales/tr-TR/models.json +3 -12
  62. package/locales/tr-TR/providers.json +3 -0
  63. package/locales/vi-VN/error.json +1 -0
  64. package/locales/vi-VN/modelProvider.json +7 -0
  65. package/locales/vi-VN/models.json +3 -12
  66. package/locales/vi-VN/providers.json +3 -0
  67. package/locales/zh-CN/error.json +1 -0
  68. package/locales/zh-CN/modelProvider.json +7 -0
  69. package/locales/zh-CN/models.json +3 -12
  70. package/locales/zh-CN/providers.json +3 -0
  71. package/locales/zh-TW/error.json +1 -0
  72. package/locales/zh-TW/modelProvider.json +7 -0
  73. package/locales/zh-TW/models.json +3 -12
  74. package/locales/zh-TW/providers.json +3 -0
  75. package/package.json +2 -1
  76. package/src/app/(backend)/webapi/chat/vertexai/route.ts +35 -0
  77. package/src/app/[variants]/(main)/settings/provider/(detail)/huggingface/page.tsx +3 -3
  78. package/src/app/[variants]/(main)/settings/provider/(detail)/vertexai/page.tsx +67 -0
  79. package/src/config/aiModels/index.ts +3 -0
  80. package/src/config/aiModels/vertexai.ts +200 -0
  81. package/src/config/modelProviders/index.ts +3 -0
  82. package/src/config/modelProviders/vertexai.ts +22 -0
  83. package/src/database/client/db.ts +2 -1
  84. package/src/libs/agent-runtime/error.ts +1 -0
  85. package/src/libs/agent-runtime/google/index.ts +22 -4
  86. package/src/libs/agent-runtime/types/type.ts +1 -0
  87. package/src/libs/agent-runtime/utils/streams/vertex-ai.test.ts +236 -0
  88. package/src/libs/agent-runtime/utils/streams/vertex-ai.ts +75 -0
  89. package/src/libs/agent-runtime/vertexai/index.ts +23 -0
  90. package/src/locales/default/error.ts +1 -0
  91. package/src/locales/default/modelProvider.ts +7 -0
  92. package/src/types/user/settings/keyVaults.ts +1 -0
  93. package/src/utils/safeParseJSON.ts +1 -1
@@ -0,0 +1,200 @@
1
+ import { AIChatModelCard } from '@/types/aiModel';
2
+
3
+ // ref: https://ai.google.dev/gemini-api/docs/models/gemini
4
+ const vertexaiChatModels: AIChatModelCard[] = [
5
+ {
6
+ abilities: {
7
+ functionCall: true,
8
+ vision: true,
9
+ },
10
+ contextWindowTokens: 2_097_152 + 8192,
11
+ description:
12
+ 'Gemini 2.0 Pro Experimental 是 Google 最新的实验性多模态AI模型,与历史版本相比有一定的质量提升,特别是对于世界知识、代码和长上下文。',
13
+ displayName: 'Gemini 2.0 Pro Experimental 02-05',
14
+ enabled: true,
15
+ id: 'gemini-2.0-pro-exp-02-05',
16
+ maxOutput: 8192,
17
+ pricing: {
18
+ cachedInput: 0,
19
+ input: 0,
20
+ output: 0,
21
+ },
22
+ releasedAt: '2025-02-05',
23
+ type: 'chat',
24
+ },
25
+ {
26
+ abilities: {
27
+ functionCall: true,
28
+ vision: true,
29
+ },
30
+ contextWindowTokens: 1_048_576 + 8192,
31
+ description:
32
+ 'Gemini 2.0 Flash 提供下一代功能和改进,包括卓越的速度、原生工具使用、多模态生成和1M令牌上下文窗口。',
33
+ displayName: 'Gemini 2.0 Flash',
34
+ enabled: true,
35
+ id: 'gemini-2.0-flash',
36
+ maxOutput: 8192,
37
+ pricing: {
38
+ cachedInput: 0.025,
39
+ input: 0.1,
40
+ output: 0.4,
41
+ },
42
+ releasedAt: '2025-02-05',
43
+ type: 'chat',
44
+ },
45
+ {
46
+ abilities: {
47
+ functionCall: true,
48
+ vision: true,
49
+ },
50
+ contextWindowTokens: 1_048_576 + 8192,
51
+ description:
52
+ 'Gemini 2.0 Flash 提供下一代功能和改进,包括卓越的速度、原生工具使用、多模态生成和1M令牌上下文窗口。',
53
+ displayName: 'Gemini 2.0 Flash 001',
54
+ id: 'gemini-2.0-flash-001',
55
+ maxOutput: 8192,
56
+ pricing: {
57
+ cachedInput: 0.025,
58
+ input: 0.1,
59
+ output: 0.4,
60
+ },
61
+ releasedAt: '2025-02-05',
62
+ type: 'chat',
63
+ },
64
+ {
65
+ abilities: {
66
+ vision: true,
67
+ },
68
+ contextWindowTokens: 1_048_576 + 8192,
69
+ description: '一个 Gemini 2.0 Flash 模型,针对成本效益和低延迟等目标进行了优化。',
70
+ displayName: 'Gemini 2.0 Flash-Lite Preview 02-05',
71
+ id: 'gemini-2.0-flash-lite-preview-02-05',
72
+ maxOutput: 8192,
73
+ pricing: {
74
+ cachedInput: 0.018_75,
75
+ input: 0.075,
76
+ output: 0.3,
77
+ },
78
+ releasedAt: '2025-02-05',
79
+ type: 'chat',
80
+ },
81
+ {
82
+ abilities: {
83
+ reasoning: true,
84
+ vision: true,
85
+ },
86
+ contextWindowTokens: 1_048_576 + 65_536,
87
+ description:
88
+ 'Gemini 2.0 Flash Thinking Exp 是 Google 的实验性多模态推理AI模型,能对复杂问题进行推理,拥有新的思维能力。',
89
+ displayName: 'Gemini 2.0 Flash Thinking Experimental 01-21',
90
+ enabled: true,
91
+ id: 'gemini-2.0-flash-thinking-exp-01-21',
92
+ maxOutput: 65_536,
93
+ pricing: {
94
+ cachedInput: 0,
95
+ input: 0,
96
+ output: 0,
97
+ },
98
+ releasedAt: '2025-01-21',
99
+ type: 'chat',
100
+ },
101
+ {
102
+ abilities: { functionCall: true, vision: true },
103
+ contextWindowTokens: 1_000_000 + 8192,
104
+ description:
105
+ 'Gemini 1.5 Flash 是Google最新的多模态AI模型,具备快速处理能力,支持文本、图像和视频输入,适用于多种任务的高效扩展。',
106
+ displayName: 'Gemini 1.5 Flash',
107
+ enabled: true,
108
+ id: 'gemini-1.5-flash',
109
+ maxOutput: 8192,
110
+ pricing: {
111
+ cachedInput: 0.018_75,
112
+ input: 0.075,
113
+ output: 0.3,
114
+ },
115
+ type: 'chat',
116
+ },
117
+ {
118
+ abilities: { functionCall: true, vision: true },
119
+ contextWindowTokens: 1_000_000 + 8192,
120
+ description: 'Gemini 1.5 Flash 002 是一款高效的多模态模型,支持广泛应用的扩展。',
121
+ displayName: 'Gemini 1.5 Flash 002',
122
+ enabled: true,
123
+ id: 'gemini-1.5-flash-002',
124
+ maxOutput: 8192,
125
+ pricing: {
126
+ cachedInput: 0.018_75,
127
+ input: 0.075,
128
+ output: 0.3,
129
+ },
130
+ releasedAt: '2024-09-25',
131
+ type: 'chat',
132
+ },
133
+ {
134
+ abilities: { functionCall: true, vision: true },
135
+ contextWindowTokens: 1_000_000 + 8192,
136
+ description: 'Gemini 1.5 Flash 001 是一款高效的多模态模型,支持广泛应用的扩展。',
137
+ displayName: 'Gemini 1.5 Flash 001',
138
+ id: 'gemini-1.5-flash-001',
139
+ maxOutput: 8192,
140
+ pricing: {
141
+ cachedInput: 0.018_75,
142
+ input: 0.075,
143
+ output: 0.3,
144
+ },
145
+ type: 'chat',
146
+ },
147
+ {
148
+ abilities: { functionCall: true, vision: true },
149
+ contextWindowTokens: 2_000_000 + 8192,
150
+ description:
151
+ 'Gemini 1.5 Pro 支持高达200万个tokens,是中型多模态模型的理想选择,适用于复杂任务的多方面支持。',
152
+ displayName: 'Gemini 1.5 Pro',
153
+ enabled: true,
154
+ id: 'gemini-1.5-pro-latest',
155
+ maxOutput: 8192,
156
+ pricing: {
157
+ cachedInput: 0.875,
158
+ input: 3.5,
159
+ output: 10.5,
160
+ },
161
+ releasedAt: '2024-02-15',
162
+ type: 'chat',
163
+ },
164
+ {
165
+ abilities: { functionCall: true, vision: true },
166
+ contextWindowTokens: 2_000_000 + 8192,
167
+ description:
168
+ 'Gemini 1.5 Pro 002 是最新的生产就绪模型,提供更高质量的输出,特别在数学、长上下文和视觉任务方面有显著提升。',
169
+ displayName: 'Gemini 1.5 Pro 002',
170
+ enabled: true,
171
+ id: 'gemini-1.5-pro-002',
172
+ maxOutput: 8192,
173
+ pricing: {
174
+ cachedInput: 0.315,
175
+ input: 1.25,
176
+ output: 2.5,
177
+ },
178
+ releasedAt: '2024-09-24',
179
+ type: 'chat',
180
+ },
181
+ {
182
+ abilities: { functionCall: true, vision: true },
183
+ contextWindowTokens: 2_000_000 + 8192,
184
+ description: 'Gemini 1.5 Pro 001 是可扩展的多模态AI解决方案,支持广泛的复杂任务。',
185
+ displayName: 'Gemini 1.5 Pro 001',
186
+ id: 'gemini-1.5-pro-001',
187
+ maxOutput: 8192,
188
+ pricing: {
189
+ cachedInput: 0.875,
190
+ input: 3.5,
191
+ output: 10.5,
192
+ },
193
+ releasedAt: '2024-02-15',
194
+ type: 'chat',
195
+ },
196
+ ];
197
+
198
+ export const allModels = [...vertexaiChatModels];
199
+
200
+ export default allModels;
@@ -40,6 +40,7 @@ import TaichuProvider from './taichu';
40
40
  import TencentcloudProvider from './tencentcloud';
41
41
  import TogetherAIProvider from './togetherai';
42
42
  import UpstageProvider from './upstage';
43
+ import VertexAIProvider from './vertexai';
43
44
  import VLLMProvider from './vllm';
44
45
  import VolcengineProvider from './volcengine';
45
46
  import WenxinProvider from './wenxin';
@@ -102,6 +103,7 @@ export const DEFAULT_MODEL_PROVIDER_LIST = [
102
103
  AnthropicProvider,
103
104
  BedrockProvider,
104
105
  GoogleProvider,
106
+ VertexAIProvider,
105
107
  DeepSeekProvider,
106
108
  HuggingFaceProvider,
107
109
  OpenRouterProvider,
@@ -191,6 +193,7 @@ export { default as TaichuProviderCard } from './taichu';
191
193
  export { default as TencentCloudProviderCard } from './tencentcloud';
192
194
  export { default as TogetherAIProviderCard } from './togetherai';
193
195
  export { default as UpstageProviderCard } from './upstage';
196
+ export { default as VertexAIProviderCard } from './vertexai';
194
197
  export { default as VLLMProviderCard } from './vllm';
195
198
  export { default as VolcengineProviderCard } from './volcengine';
196
199
  export { default as WenxinProviderCard } from './wenxin';
@@ -0,0 +1,22 @@
1
+ import { ModelProviderCard } from '@/types/llm';
2
+
3
+ // ref: https://ai.google.dev/gemini-api/docs/models/gemini
4
+ const VertexAI: ModelProviderCard = {
5
+ chatModels: [],
6
+ checkModel: 'gemini-1.5-flash-001',
7
+ description:
8
+ 'Google 的 Gemini 系列是其最先进、通用的 AI模型,由 Google DeepMind 打造,专为多模态设计,支持文本、代码、图像、音频和视频的无缝理解与处理。适用于从数据中心到移动设备的多种环境,极大提升了AI模型的效率与应用广泛性。',
9
+ id: 'vertexai',
10
+ modelsUrl: 'https://console.cloud.google.com/vertex-ai/model-garden',
11
+ name: 'VertexAI',
12
+ settings: {
13
+ disableBrowserRequest: true,
14
+ smoothing: {
15
+ speed: 2,
16
+ text: true,
17
+ },
18
+ },
19
+ url: 'https://cloud.google.com/vertex-ai',
20
+ };
21
+
22
+ export default VertexAI;
@@ -201,7 +201,8 @@ export class DatabaseManager {
201
201
  const dbName = 'lobechat';
202
202
 
203
203
  // make db as web worker if worker is available
204
- if (typeof Worker !== 'undefined') {
204
+ // https://github.com/lobehub/lobe-chat/issues/5785
205
+ if (typeof Worker !== 'undefined' && typeof navigator.locks !== 'undefined') {
205
206
  db = await initPgliteWorker({
206
207
  dbName,
207
208
  fsBundle: fsBundle as Blob,
@@ -13,6 +13,7 @@ export const AgentRuntimeErrorType = {
13
13
  OllamaBizError: 'OllamaBizError',
14
14
 
15
15
  InvalidBedrockCredentials: 'InvalidBedrockCredentials',
16
+ InvalidVertexCredentials: 'InvalidVertexCredentials',
16
17
  StreamChunkError: 'StreamChunkError',
17
18
 
18
19
  InvalidGithubToken: 'InvalidGithubToken',
@@ -1,3 +1,4 @@
1
+ import type { VertexAI } from '@google-cloud/vertexai';
1
2
  import {
2
3
  Content,
3
4
  FunctionCallPart,
@@ -9,6 +10,7 @@ import {
9
10
  } from '@google/generative-ai';
10
11
 
11
12
  import type { ChatModelCard } from '@/types/llm';
13
+ import { VertexAIStream } from '@/libs/agent-runtime/utils/streams/vertex-ai';
12
14
  import { imageUrlToBase64 } from '@/utils/imageToBase64';
13
15
  import { safeParseJSON } from '@/utils/safeParseJSON';
14
16
 
@@ -56,17 +58,27 @@ function getThreshold(model: string): HarmBlockThreshold {
56
58
 
57
59
  const DEFAULT_BASE_URL = 'https://generativelanguage.googleapis.com';
58
60
 
61
+ interface LobeGoogleAIParams {
62
+ apiKey?: string;
63
+ baseURL?: string;
64
+ client?: GoogleGenerativeAI | VertexAI;
65
+ isVertexAi?: boolean;
66
+ }
67
+
59
68
  export class LobeGoogleAI implements LobeRuntimeAI {
60
69
  private client: GoogleGenerativeAI;
70
+ private isVertexAi: boolean;
61
71
  baseURL?: string;
62
72
  apiKey?: string;
63
73
 
64
- constructor({ apiKey, baseURL }: { apiKey?: string; baseURL?: string } = {}) {
74
+ constructor({ apiKey, baseURL, client, isVertexAi }: LobeGoogleAIParams = {}) {
65
75
  if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
66
76
 
67
77
  this.client = new GoogleGenerativeAI(apiKey);
68
- this.baseURL = baseURL || DEFAULT_BASE_URL;
69
78
  this.apiKey = apiKey;
79
+ this.client = client ? (client as GoogleGenerativeAI) : new GoogleGenerativeAI(apiKey);
80
+ this.baseURL = client ? undefined : baseURL || DEFAULT_BASE_URL;
81
+ this.isVertexAi = isVertexAi || false;
70
82
  }
71
83
 
72
84
  async chat(rawPayload: ChatStreamPayload, options?: ChatCompetitionOptions) {
@@ -117,18 +129,24 @@ export class LobeGoogleAI implements LobeRuntimeAI {
117
129
  const googleStream = convertIterableToStream(geminiStreamResult.stream);
118
130
  const [prod, useForDebug] = googleStream.tee();
119
131
 
120
- if (process.env.DEBUG_GOOGLE_CHAT_COMPLETION === '1') {
132
+ const key = this.isVertexAi
133
+ ? 'DEBUG_VERTEX_AI_CHAT_COMPLETION'
134
+ : 'DEBUG_GOOGLE_CHAT_COMPLETION';
135
+
136
+ if (process.env[key] === '1') {
121
137
  debugStream(useForDebug).catch();
122
138
  }
123
139
 
124
140
  // Convert the response into a friendly text-stream
125
- const stream = GoogleGenerativeAIStream(prod, options?.callback);
141
+ const Stream = this.isVertexAi ? VertexAIStream : GoogleGenerativeAIStream;
142
+ const stream = Stream(prod, options?.callback);
126
143
 
127
144
  // Respond with the stream
128
145
  return StreamingResponse(stream, { headers: options?.headers });
129
146
  } catch (e) {
130
147
  const err = e as Error;
131
148
 
149
+ console.log(err);
132
150
  const { errorType, error } = this.parseErrorMessage(err.message);
133
151
 
134
152
  throw AgentRuntimeError.chat({ error, errorType, provider: ModelProvider.Google });
@@ -66,6 +66,7 @@ export enum ModelProvider {
66
66
  TogetherAI = 'togetherai',
67
67
  Upstage = 'upstage',
68
68
  VLLM = 'vllm',
69
+ VertexAI = 'vertexai',
69
70
  Volcengine = 'volcengine',
70
71
  Wenxin = 'wenxin',
71
72
  XAI = 'xai',
@@ -0,0 +1,236 @@
1
+ import { describe, expect, it, vi } from 'vitest';
2
+
3
+ import * as uuidModule from '@/utils/uuid';
4
+
5
+ import { VertexAIStream } from './vertex-ai';
6
+
7
+ describe('VertexAIStream', () => {
8
+ it('should transform Vertex AI stream to protocol stream', async () => {
9
+ vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
10
+ const rawChunks = [
11
+ {
12
+ candidates: [
13
+ {
14
+ content: { role: 'model', parts: [{ text: '你好' }] },
15
+ safetyRatings: [
16
+ {
17
+ category: 'HARM_CATEGORY_HATE_SPEECH',
18
+ probability: 'NEGLIGIBLE',
19
+ probabilityScore: 0.06298828,
20
+ severity: 'HARM_SEVERY_NEGLIGIBLE',
21
+ severityScore: 0.10986328,
22
+ },
23
+ {
24
+ category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
25
+ probability: 'NEGLIGIBLE',
26
+ probabilityScore: 0.05029297,
27
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
28
+ severityScore: 0.078125,
29
+ },
30
+ {
31
+ category: 'HARM_CATEGORY_HARASSMENT',
32
+ probability: 'NEGLIGIBLE',
33
+ probabilityScore: 0.19433594,
34
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
35
+ severityScore: 0.16015625,
36
+ },
37
+ {
38
+ category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
39
+ probability: 'NEGLIGIBLE',
40
+ probabilityScore: 0.059326172,
41
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
42
+ severityScore: 0.064453125,
43
+ },
44
+ ],
45
+ index: 0,
46
+ },
47
+ ],
48
+ usageMetadata: {},
49
+ modelVersion: 'gemini-1.5-flash-001',
50
+ },
51
+ {
52
+ candidates: [
53
+ {
54
+ content: { role: 'model', parts: [{ text: '! 😊' }] },
55
+ safetyRatings: [
56
+ {
57
+ category: 'HARM_CATEGORY_HATE_SPEECH',
58
+ probability: 'NEGLIGIBLE',
59
+ probabilityScore: 0.052734375,
60
+ severity: 'HARM_SEVRITY_NEGLIGIBLE',
61
+ severityScore: 0.08642578,
62
+ },
63
+ {
64
+ category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
65
+ probability: 'NEGLIGIBLE',
66
+ probabilityScore: 0.071777344,
67
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
68
+ severityScore: 0.095214844,
69
+ },
70
+ {
71
+ category: 'HARM_CATEGORY_HARASSMENT',
72
+ probability: 'NEGLIGIBLE',
73
+ probabilityScore: 0.1640625,
74
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
75
+ severityScore: 0.10498047,
76
+ },
77
+ {
78
+ category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
79
+ probability: 'NEGLIGIBLE',
80
+ probabilityScore: 0.075683594,
81
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
82
+ severityScore: 0.053466797,
83
+ },
84
+ ],
85
+ index: 0,
86
+ },
87
+ ],
88
+ modelVersion: 'gemini-1.5-flash-001',
89
+ },
90
+ ];
91
+
92
+ const mockGoogleStream = new ReadableStream({
93
+ start(controller) {
94
+ rawChunks.forEach((chunk) => controller.enqueue(chunk));
95
+
96
+ controller.close();
97
+ },
98
+ });
99
+
100
+ const onStartMock = vi.fn();
101
+ const onTextMock = vi.fn();
102
+ const onTokenMock = vi.fn();
103
+ const onToolCallMock = vi.fn();
104
+ const onCompletionMock = vi.fn();
105
+
106
+ const protocolStream = VertexAIStream(mockGoogleStream, {
107
+ onStart: onStartMock,
108
+ onText: onTextMock,
109
+ onToken: onTokenMock,
110
+ onToolCall: onToolCallMock,
111
+ onCompletion: onCompletionMock,
112
+ });
113
+
114
+ const decoder = new TextDecoder();
115
+ const chunks = [];
116
+
117
+ // @ts-ignore
118
+ for await (const chunk of protocolStream) {
119
+ chunks.push(decoder.decode(chunk, { stream: true }));
120
+ }
121
+
122
+ expect(chunks).toEqual([
123
+ // text
124
+ 'id: chat_1\n',
125
+ 'event: text\n',
126
+ `data: "你好"\n\n`,
127
+
128
+ // text
129
+ 'id: chat_1\n',
130
+ 'event: text\n',
131
+ `data: "! 😊"\n\n`,
132
+ ]);
133
+
134
+ expect(onStartMock).toHaveBeenCalledTimes(1);
135
+ expect(onTokenMock).toHaveBeenCalledTimes(2);
136
+ expect(onCompletionMock).toHaveBeenCalledTimes(1);
137
+ });
138
+
139
+ it('tool_calls', async () => {
140
+ vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
141
+ const rawChunks = [
142
+ {
143
+ candidates: [
144
+ {
145
+ content: {
146
+ role: 'model',
147
+ parts: [
148
+ {
149
+ functionCall: {
150
+ name: 'realtime-weather____fetchCurrentWeather',
151
+ args: { city: '杭州' },
152
+ },
153
+ },
154
+ ],
155
+ },
156
+ finishReason: 'STOP',
157
+ safetyRatings: [
158
+ {
159
+ category: 'HARM_CATERY_HATE_SPEECH',
160
+ probability: 'NEGLIGIBLE',
161
+ probabilityScore: 0.09814453,
162
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
163
+ severityScore: 0.07470703,
164
+ },
165
+ {
166
+ category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
167
+ probability: 'NEGLIGIBLE',
168
+ probabilityScore: 0.1484375,
169
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
170
+ severityScore: 0.15136719,
171
+ },
172
+ {
173
+ category: 'HARM_CATEGORY_HARASSMENT',
174
+ probability: 'NEGLIGIBLE',
175
+ probabilityScore: 0.11279297,
176
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
177
+ severityScore: 0.10107422,
178
+ },
179
+ {
180
+ category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
181
+ probability: 'NEGLIGIBLE',
182
+ probabilityScore: 0.048828125,
183
+ severity: 'HARM_SEVERITY_NEGLIGIBLE',
184
+ severityScore: 0.05493164,
185
+ },
186
+ ],
187
+ index: 0,
188
+ },
189
+ ],
190
+ usageMetadata: { promptTokenCount: 95, candidatesTokenCount: 9, totalTokenCount: 104 },
191
+ modelVersion: 'gemini-1.5-flash-001',
192
+ },
193
+ ];
194
+
195
+ const mockGoogleStream = new ReadableStream({
196
+ start(controller) {
197
+ rawChunks.forEach((chunk) => controller.enqueue(chunk));
198
+
199
+ controller.close();
200
+ },
201
+ });
202
+
203
+ const onStartMock = vi.fn();
204
+ const onTextMock = vi.fn();
205
+ const onTokenMock = vi.fn();
206
+ const onToolCallMock = vi.fn();
207
+ const onCompletionMock = vi.fn();
208
+
209
+ const protocolStream = VertexAIStream(mockGoogleStream, {
210
+ onStart: onStartMock,
211
+ onText: onTextMock,
212
+ onToken: onTokenMock,
213
+ onToolCall: onToolCallMock,
214
+ onCompletion: onCompletionMock,
215
+ });
216
+
217
+ const decoder = new TextDecoder();
218
+ const chunks = [];
219
+
220
+ // @ts-ignore
221
+ for await (const chunk of protocolStream) {
222
+ chunks.push(decoder.decode(chunk, { stream: true }));
223
+ }
224
+
225
+ expect(chunks).toEqual([
226
+ // text
227
+ 'id: chat_1\n',
228
+ 'event: tool_calls\n',
229
+ `data: [{"function":{"arguments":"{\\"city\\":\\"杭州\\"}","name":"realtime-weather____fetchCurrentWeather"},"id":"realtime-weather____fetchCurrentWeather_0","index":0,"type":"function"}]\n\n`,
230
+ ]);
231
+
232
+ expect(onStartMock).toHaveBeenCalledTimes(1);
233
+ expect(onToolCallMock).toHaveBeenCalledTimes(1);
234
+ expect(onCompletionMock).toHaveBeenCalledTimes(1);
235
+ });
236
+ });
@@ -0,0 +1,75 @@
1
+ import { EnhancedGenerateContentResponse, GenerateContentResponse } from '@google/generative-ai';
2
+
3
+ import { nanoid } from '@/utils/uuid';
4
+
5
+ import { ChatStreamCallbacks } from '../../types';
6
+ import {
7
+ StreamProtocolChunk,
8
+ StreamStack,
9
+ createCallbacksTransformer,
10
+ createSSEProtocolTransformer,
11
+ generateToolCallId,
12
+ } from './protocol';
13
+
14
+ const transformVertexAIStream = (
15
+ chunk: GenerateContentResponse,
16
+ stack: StreamStack,
17
+ ): StreamProtocolChunk => {
18
+ // maybe need another structure to add support for multiple choices
19
+ const candidates = chunk.candidates;
20
+
21
+ if (!candidates)
22
+ return {
23
+ data: '',
24
+ id: stack?.id,
25
+ type: 'text',
26
+ };
27
+
28
+ const item = candidates[0];
29
+ if (item.content) {
30
+ const part = item.content.parts[0];
31
+
32
+ if (part.functionCall) {
33
+ const functionCall = part.functionCall;
34
+
35
+ return {
36
+ data: [
37
+ {
38
+ function: {
39
+ arguments: JSON.stringify(functionCall.args),
40
+ name: functionCall.name,
41
+ },
42
+ id: generateToolCallId(0, functionCall.name),
43
+ index: 0,
44
+ type: 'function',
45
+ },
46
+ ],
47
+ id: stack?.id,
48
+ type: 'tool_calls',
49
+ };
50
+ }
51
+
52
+ return {
53
+ data: part.text,
54
+ id: stack?.id,
55
+ type: 'text',
56
+ };
57
+ }
58
+
59
+ return {
60
+ data: '',
61
+ id: stack?.id,
62
+ type: 'stop',
63
+ };
64
+ };
65
+
66
+ export const VertexAIStream = (
67
+ rawStream: ReadableStream<EnhancedGenerateContentResponse>,
68
+ callbacks?: ChatStreamCallbacks,
69
+ ) => {
70
+ const streamStack: StreamStack = { id: 'chat_' + nanoid() };
71
+
72
+ return rawStream
73
+ .pipeThrough(createSSEProtocolTransformer(transformVertexAIStream, streamStack))
74
+ .pipeThrough(createCallbacksTransformer(callbacks));
75
+ };
@@ -0,0 +1,23 @@
1
+ import { VertexAI, VertexInit } from '@google-cloud/vertexai';
2
+
3
+ import { AgentRuntimeError, AgentRuntimeErrorType, LobeGoogleAI } from '@/libs/agent-runtime';
4
+
5
+ export class LobeVertexAI extends LobeGoogleAI {
6
+ static initFromVertexAI(params?: VertexInit) {
7
+ try {
8
+ const client = new VertexAI({ ...params });
9
+
10
+ return new LobeGoogleAI({ apiKey: 'avoid-error', client, isVertexAi: true });
11
+ } catch (e) {
12
+ const err = e as Error;
13
+
14
+ if (err.name === 'IllegalArgumentError') {
15
+ throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidVertexCredentials, {
16
+ message: err.message,
17
+ });
18
+ }
19
+
20
+ throw e;
21
+ }
22
+ }
23
+ }
@@ -106,6 +106,7 @@ export default {
106
106
  */
107
107
  OpenAIBizError: '请求 OpenAI 服务出错,请根据以下信息排查或重试',
108
108
 
109
+ InvalidVertexCredentials: 'Vertex 鉴权未通过,请检查鉴权凭证后重试',
109
110
  InvalidBedrockCredentials: 'Bedrock 鉴权未通过,请检查 AccessKeyId/SecretAccessKey 后重试',
110
111
  StreamChunkError:
111
112
  '流式请求的消息块解析错误,请检查当前 API 接口是否符合标准规范,或联系你的 API 供应商咨询',