Skip to content

Commit 96dc04d

Browse files
arvinxxarvinxx
committed
✨ feat: add vertex ai model runtime
1 parent e75c71a commit 96dc04d

File tree

11 files changed

+303
-3
lines changed

11 files changed

+303
-3
lines changed

‎package.json

+1
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@
119119
"@codesandbox/sandpack-react": "^2.19.10",
120120
"@cyntler/react-doc-viewer": "^1.17.0",
121121
"@electric-sql/pglite": "0.2.13",
122+
"@google-cloud/vertexai": "^1.9.2",
122123
"@google/generative-ai": "^0.21.0",
123124
"@huggingface/inference": "^2.8.1",
124125
"@icons-pack/react-simple-icons": "9.6.0",

‎src/app/(backend)/webapi/chat/vertexai/route.ts

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import { readFileSync } from 'node:fs';
2+
import { resolve } from 'node:path';
3+
4+
import { AgentRuntime, ModelProvider } from '@/libs/agent-runtime';
5+
import { LobeVertexAI } from '@/libs/agent-runtime/vertexai';
6+
import { safeParseJSON } from '@/utils/safeParseJSON';
7+
8+
import { POST as UniverseRoute } from '../[provider]/route';
9+
10+
export const POST = async (req: Request) =>
11+
UniverseRoute(req, {
12+
createRuntime: () => {
13+
const credentialsContent =
14+
process.env.VERTEXAI_CREDENTIALS ??
15+
(process.env.VERTEXAI_CREDENTIALS_PATH
16+
? readFileSync(resolve(process.cwd(), process.env.VERTEXAI_CREDENTIALS_PATH), 'utf8')
17+
: undefined);
18+
19+
const googleAuthOptions = credentialsContent ? safeParseJSON(credentialsContent) : undefined;
20+
21+
const instance = LobeVertexAI.initFromVertexAI({
22+
googleAuthOptions: googleAuthOptions,
23+
location: process.env.VERTEXAI_LOCATION,
24+
project: process.env.VERTEXAI_PROJECT,
25+
});
26+
27+
return new AgentRuntime(instance);
28+
},
29+
params: Promise.resolve({ provider: ModelProvider.VertexAI }),
30+
});

‎src/config/aiModels/index.ts

+3
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import { default as stepfun } from './stepfun';
3434
import { default as taichu } from './taichu';
3535
import { default as togetherai } from './togetherai';
3636
import { default as upstage } from './upstage';
37+
import { default as vertexai } from './vertexai';
3738
import { default as wenxin } from './wenxin';
3839
import { default as xai } from './xai';
3940
import { default as zeroone } from './zeroone';
@@ -93,6 +94,7 @@ export const LOBE_DEFAULT_MODEL_LIST = buildDefaultModelList({
9394
taichu,
9495
togetherai,
9596
upstage,
97+
vertexai,
9698
wenxin,
9799
xai,
98100
zeroone,
@@ -133,6 +135,7 @@ export { default as stepfun } from './stepfun';
133135
export { default as taichu } from './taichu';
134136
export { default as togetherai } from './togetherai';
135137
export { default as upstage } from './upstage';
138+
export { default as vertexai } from './vertexai';
136139
export { default as wenxin } from './wenxin';
137140
export { default as xai } from './xai';
138141
export { default as zeroone } from './zeroone';

‎src/config/aiModels/vertexai.ts

+207
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
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: { functionCall: true, vision: true },
7+
contextWindowTokens: 1_000_000 + 8192,
8+
description:
9+
'Gemini 1.5 Flash 是Google最新的多模态AI模型，具备快速处理能力，支持文本、图像和视频输入，适用于多种任务的高效扩展。',
10+
displayName: 'Gemini 1.5 Flash',
11+
enabled: true,
12+
id: 'gemini-1.5-flash-latest',
13+
maxOutput: 8192,
14+
pricing: {
15+
cachedInput: 0.018_75,
16+
input: 0.075,
17+
output: 0.3,
18+
},
19+
type: 'chat',
20+
},
21+
{
22+
abilities: { functionCall: true, vision: true },
23+
contextWindowTokens: 1_000_000 + 8192,
24+
description: 'Gemini 1.5 Flash 002 是一款高效的多模态模型，支持广泛应用的扩展。',
25+
displayName: 'Gemini 1.5 Flash 002',
26+
enabled: true,
27+
id: 'gemini-1.5-flash-002',
28+
maxOutput: 8192,
29+
pricing: {
30+
cachedInput: 0.018_75,
31+
input: 0.075,
32+
output: 0.3,
33+
},
34+
releasedAt: '2024-09-25',
35+
type: 'chat',
36+
},
37+
{
38+
abilities: { functionCall: true, vision: true },
39+
contextWindowTokens: 1_000_000 + 8192,
40+
description: 'Gemini 1.5 Flash 001 是一款高效的多模态模型，支持广泛应用的扩展。',
41+
displayName: 'Gemini 1.5 Flash 001',
42+
id: 'gemini-1.5-flash-001',
43+
maxOutput: 8192,
44+
pricing: {
45+
cachedInput: 0.018_75,
46+
input: 0.075,
47+
output: 0.3,
48+
},
49+
type: 'chat',
50+
},
51+
{
52+
abilities: { functionCall: true, vision: true },
53+
contextWindowTokens: 1_000_000 + 8192,
54+
description: 'Gemini 1.5 Flash 0827 提供了优化后的多模态处理能力，适用多种复杂任务场景。',
55+
displayName: 'Gemini 1.5 Flash 0827',
56+
id: 'gemini-1.5-flash-exp-0827',
57+
maxOutput: 8192,
58+
pricing: {
59+
cachedInput: 0.018_75,
60+
input: 0.075,
61+
output: 0.3,
62+
},
63+
releasedAt: '2024-08-27',
64+
type: 'chat',
65+
},
66+
67+
{
68+
abilities: { functionCall: true, vision: true },
69+
contextWindowTokens: 1_000_000 + 8192,
70+
description:
71+
'Gemini 1.5 Flash 8B 0924 是最新的实验性模型，在文本和多模态用例中都有显著的性能提升。',
72+
displayName: 'Gemini 1.5 Flash 8B 0924',
73+
id: 'gemini-1.5-flash-8b-exp-0924',
74+
maxOutput: 8192,
75+
pricing: {
76+
cachedInput: 0.018_75,
77+
input: 0.075,
78+
output: 0.3,
79+
},
80+
releasedAt: '2024-09-24',
81+
type: 'chat',
82+
},
83+
{
84+
abilities: { functionCall: true, vision: true },
85+
contextWindowTokens: 2_000_000 + 8192,
86+
description:
87+
'Gemini 1.5 Pro 支持高达200万个tokens，是中型多模态模型的理想选择，适用于复杂任务的多方面支持。',
88+
displayName: 'Gemini 1.5 Pro',
89+
enabled: true,
90+
id: 'gemini-1.5-pro-latest',
91+
maxOutput: 8192,
92+
pricing: {
93+
cachedInput: 0.875,
94+
input: 3.5,
95+
output: 10.5,
96+
},
97+
releasedAt: '2024-02-15',
98+
type: 'chat',
99+
},
100+
{
101+
abilities: { functionCall: true, vision: true },
102+
contextWindowTokens: 2_000_000 + 8192,
103+
description:
104+
'Gemini 1.5 Pro 002 是最新的生产就绪模型，提供更高质量的输出，特别在数学、长上下文和视觉任务方面有显著提升。',
105+
displayName: 'Gemini 1.5 Pro 002',
106+
enabled: true,
107+
id: 'gemini-1.5-pro-002',
108+
maxOutput: 8192,
109+
pricing: {
110+
cachedInput: 0.315,
111+
input: 1.25,
112+
output: 2.5,
113+
},
114+
releasedAt: '2024-09-24',
115+
type: 'chat',
116+
},
117+
{
118+
abilities: { functionCall: true, vision: true },
119+
contextWindowTokens: 2_000_000 + 8192,
120+
description: 'Gemini 1.5 Pro 001 是可扩展的多模态AI解决方案，支持广泛的复杂任务。',
121+
displayName: 'Gemini 1.5 Pro 001',
122+
id: 'gemini-1.5-pro-001',
123+
maxOutput: 8192,
124+
pricing: {
125+
cachedInput: 0.875,
126+
input: 3.5,
127+
output: 10.5,
128+
},
129+
releasedAt: '2024-02-15',
130+
type: 'chat',
131+
},
132+
{
133+
abilities: { functionCall: true, vision: true },
134+
contextWindowTokens: 2_000_000 + 8192,
135+
description: 'Gemini 1.5 Pro 0827 结合最新优化技术，带来更高效的多模态数据处理能力。',
136+
displayName: 'Gemini 1.5 Pro 0827',
137+
id: 'gemini-1.5-pro-exp-0827',
138+
maxOutput: 8192,
139+
pricing: {
140+
cachedInput: 0.875,
141+
input: 3.5,
142+
output: 10.5,
143+
},
144+
releasedAt: '2024-08-27',
145+
type: 'chat',
146+
},
147+
{
148+
abilities: { functionCall: true, vision: true },
149+
contextWindowTokens: 2_000_000 + 8192,
150+
description: 'Gemini 1.5 Pro 0801 提供出色的多模态处理能力，为应用开发带来更大灵活性。',
151+
displayName: 'Gemini 1.5 Pro 0801',
152+
id: 'gemini-1.5-pro-exp-0801',
153+
maxOutput: 8192,
154+
pricing: {
155+
cachedInput: 0.875,
156+
input: 3.5,
157+
output: 10.5,
158+
},
159+
releasedAt: '2024-08-01',
160+
type: 'chat',
161+
},
162+
{
163+
contextWindowTokens: 30_720 + 2048,
164+
description: 'Gemini 1.0 Pro 是Google的高性能AI模型，专为广泛任务扩展而设计。',
165+
displayName: 'Gemini 1.0 Pro',
166+
id: 'gemini-1.0-pro-latest',
167+
maxOutput: 2048,
168+
pricing: {
169+
input: 0.5,
170+
output: 1.5,
171+
},
172+
releasedAt: '2023-12-06',
173+
type: 'chat',
174+
},
175+
{
176+
abilities: { functionCall: true },
177+
contextWindowTokens: 30_720 + 2048,
178+
description:
179+
'Gemini 1.0 Pro 001 (Tuning) 提供稳定并可调优的性能，是复杂任务解决方案的理想选择。',
180+
displayName: 'Gemini 1.0 Pro 001 (Tuning)',
181+
id: 'gemini-1.0-pro-001',
182+
maxOutput: 2048,
183+
pricing: {
184+
input: 0.5,
185+
output: 1.5,
186+
},
187+
releasedAt: '2023-12-06',
188+
type: 'chat',
189+
},
190+
{
191+
contextWindowTokens: 30_720 + 2048,
192+
description: 'Gemini 1.0 Pro 002 (Tuning) 提供出色的多模态支持，专注于复杂任务的有效解决。',
193+
displayName: 'Gemini 1.0 Pro 002 (Tuning)',
194+
id: 'gemini-1.0-pro-002',
195+
maxOutput: 2048,
196+
pricing: {
197+
input: 0.5,
198+
output: 1.5,
199+
},
200+
releasedAt: '2023-12-06',
201+
type: 'chat',
202+
},
203+
];
204+
205+
export const allModels = [...vertexaiChatModels];
206+
207+
export default allModels;

‎src/config/modelProviders/index.ts

+3
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import StepfunProvider from './stepfun';
3434
import TaichuProvider from './taichu';
3535
import TogetherAIProvider from './togetherai';
3636
import UpstageProvider from './upstage';
37+
import VertexAIProvider from './vertexai';
3738
import WenxinProvider from './wenxin';
3839
import XAIProvider from './xai';
3940
import ZeroOneProvider from './zeroone';
@@ -88,6 +89,7 @@ export const DEFAULT_MODEL_PROVIDER_LIST = [
8889
AnthropicProvider,
8990
BedrockProvider,
9091
GoogleProvider,
92+
VertexAIProvider,
9193
DeepSeekProvider,
9294
HuggingFaceProvider,
9395
OpenRouterProvider,
@@ -165,6 +167,7 @@ export { default as StepfunProviderCard } from './stepfun';
165167
export { default as TaichuProviderCard } from './taichu';
166168
export { default as TogetherAIProviderCard } from './togetherai';
167169
export { default as UpstageProviderCard } from './upstage';
170+
export { default as VertexAIProviderCard } from './vertexai';
168171
export { default as WenxinProviderCard } from './wenxin';
169172
export { default as XAIProviderCard } from './xai';
170173
export { default as ZeroOneProviderCard } from './zeroone';

‎src/config/modelProviders/vertexai.ts

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
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-latest',
7+
description:
8+
'Google 的 Gemini 系列是其最先进、通用的 AI模型，由 Google DeepMind 打造，专为多模态设计，支持文本、代码、图像、音频和视频的无缝理解与处理。适用于从数据中心到移动设备的多种环境，极大提升了AI模型的效率与应用广泛性。',
9+
id: 'vertexai',
10+
modelsUrl: 'https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning',
11+
name: 'VertexAI',
12+
settings: {
13+
showApiKey: false,
14+
// showChecker: false,
15+
smoothing: {
16+
speed: 2,
17+
text: true,
18+
},
19+
},
20+
url: 'https://cloud.google.com/vertex-ai',
21+
};
22+
23+
export default VertexAI;

‎src/libs/agent-runtime/error.ts

+1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ export const AgentRuntimeErrorType = {
1313
OllamaBizError: 'OllamaBizError',
1414

1515
InvalidBedrockCredentials: 'InvalidBedrockCredentials',
16+
InvalidVertexCredentials: 'InvalidVertexCredentials',
1617
StreamChunkError: 'StreamChunkError',
1718

1819
InvalidGithubToken: 'InvalidGithubToken',

‎src/libs/agent-runtime/google/index.ts

+10-3
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { VertexAI } from '@google-cloud/vertexai';
12
import {
23
Content,
34
FunctionCallPart,
@@ -46,15 +47,21 @@ function getThreshold(model: string): HarmBlockThreshold {
4647
return HarmBlockThreshold.BLOCK_NONE;
4748
}
4849

50+
interface LobeGoogleAIParams {
51+
apiKey?: string;
52+
baseURL?: string;
53+
client?: GoogleGenerativeAI | VertexAI;
54+
}
55+
4956
export class LobeGoogleAI implements LobeRuntimeAI {
5057
private client: GoogleGenerativeAI;
5158
baseURL?: string;
5259

53-
constructor({ apiKey, baseURL }: { apiKey?: string; baseURL?: string } = {}) {
60+
constructor({ apiKey, baseURL, client }: LobeGoogleAIParams = {}) {
5461
if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
5562

56-
this.client = new GoogleGenerativeAI(apiKey);
57-
this.baseURL = baseURL;
63+
this.client = client ? (client as GoogleGenerativeAI) : new GoogleGenerativeAI(apiKey);
64+
this.baseURL = client ? undefined : baseURL;
5865
}
5966

6067
async chat(rawPayload: ChatStreamPayload, options?: ChatCompetitionOptions) {

‎src/libs/agent-runtime/types/type.ts

+1
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ export enum ModelProvider {
5656
Taichu = 'taichu',
5757
TogetherAI = 'togetherai',
5858
Upstage = 'upstage',
59+
VertexAI = 'vertexai',
5960
Wenxin = 'wenxin',
6061
XAI = 'xai',
6162
ZeroOne = 'zeroone',

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

+23
Original file line numberDiff line numberDiff line change
@@ -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 });
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+
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ export interface UserKeyVaults {
6767
taichu?: OpenAICompatibleKeyVault;
6868
togetherai?: OpenAICompatibleKeyVault;
6969
upstage?: OpenAICompatibleKeyVault;
70+
vertexai?: undefined;
7071
wenxin?: WenxinKeyVault;
7172
xai?: OpenAICompatibleKeyVault;
7273
zeroone?: OpenAICompatibleKeyVault;

0 commit comments

Comments
 (0)