Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 05754da

Browse files
arvinxxarvinxx
committedNov 6, 2024·
✨ feat: add vertex ai route
1 parent ad3a154 commit 05754da

File tree

13 files changed

+332
-14
lines changed

13 files changed

+332
-14
lines changed
 

‎package.json

+1
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@
116116
"@clerk/themes": "^2.1.37",
117117
"@codesandbox/sandpack-react": "^2.19.9",
118118
"@cyntler/react-doc-viewer": "^1.17.0",
119+
"@google-cloud/vertexai": "^1.9.0",
119120
"@google/generative-ai": "^0.21.0",
120121
"@huggingface/inference": "^2.8.1",
121122
"@icons-pack/react-simple-icons": "9.6.0",

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

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

‎src/app/(main)/settings/llm/ProviderList/providers.tsx

+3-1
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import {
2323
TaichuProviderCard,
2424
TogetherAIProviderCard,
2525
UpstageProviderCard,
26+
VertexAIProviderCard,
2627
ZeroOneProviderCard,
2728
ZhiPuProviderCard,
2829
} from '@/config/modelProviders';
@@ -34,8 +35,8 @@ import { useGithubProvider } from './Github';
3435
import { useHuggingFaceProvider } from './HuggingFace';
3536
import { useOllamaProvider } from './Ollama';
3637
import { useOpenAIProvider } from './OpenAI';
37-
import { useWenxinProvider } from './Wenxin';
3838
import { useSenseNovaProvider } from './SenseNova';
39+
import { useWenxinProvider } from './Wenxin';
3940

4041
export const useProviderList = (): ProviderItem[] => {
4142
const AzureProvider = useAzureProvider();
@@ -55,6 +56,7 @@ export const useProviderList = (): ProviderItem[] => {
5556
AnthropicProviderCard,
5657
BedrockProvider,
5758
GoogleProviderCard,
59+
VertexAIProviderCard,
5860
DeepSeekProviderCard,
5961
HuggingFaceProvider,
6062
OpenRouterProviderCard,

‎src/config/llm.ts

+17-1
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,13 @@ export const getLLMConfig = () => {
3131
GOOGLE_PROXY_URL: z.string().optional(),
3232
GOOGLE_MODEL_LIST: z.string().optional(),
3333

34+
ENABLED_VERTEXAI: z.boolean(),
35+
VERTEXAI_CREDENTIALS: z.string().optional(),
36+
VERTEXAI_CREDENTIALS_PATH: z.string().optional(),
37+
VERTEXAI_PROJECT: z.string().optional(),
38+
VERTEXAI_LOCATION: z.string().optional(),
39+
VERTEXAI_MODEL_LIST: z.string().optional(),
40+
3441
ENABLED_MOONSHOT: z.boolean(),
3542
MOONSHOT_API_KEY: z.string().optional(),
3643
MOONSHOT_MODEL_LIST: z.string().optional(),
@@ -177,6 +184,14 @@ export const getLLMConfig = () => {
177184
GOOGLE_PROXY_URL: process.env.GOOGLE_PROXY_URL,
178185
GOOGLE_MODEL_LIST: process.env.GOOGLE_MODEL_LIST,
179186

187+
ENABLED_VERTEXAI:
188+
!!process.env.VERTEXAI_CREDENTIALS_PATH || !!process.env.VERTEXAI_CREDENTIALS,
189+
VERTEXAI_CREDENTIALS_PATH: process.env.VERTEXAI_CREDENTIALS_PATH,
190+
VERTEXAI_CREDENTIALS: process.env.VERTEXAI_CREDENTIALS,
191+
VERTEXAI_LOCATION: process.env.VERTEXAI_LOCATION,
192+
VERTEXAI_PROJECT: process.env.VERTEXAI_PROJECT,
193+
VERTEXAI_MODEL_LIST: process.env.VERTEXAI_MODEL_LIST,
194+
180195
ENABLED_PERPLEXITY: !!process.env.PERPLEXITY_API_KEY,
181196
PERPLEXITY_API_KEY: process.env.PERPLEXITY_API_KEY,
182197
PERPLEXITY_MODEL_LIST: process.env.PERPLEXITY_MODEL_LIST,
@@ -291,7 +306,8 @@ export const getLLMConfig = () => {
291306
HUGGINGFACE_PROXY_URL: process.env.HUGGINGFACE_PROXY_URL,
292307
HUGGINGFACE_MODEL_LIST: process.env.HUGGINGFACE_MODEL_LIST,
293308

294-
ENABLED_SENSENOVA: !!process.env.SENSENOVA_ACCESS_KEY_ID && !!process.env.SENSENOVA_ACCESS_KEY_SECRET,
309+
ENABLED_SENSENOVA:
310+
!!process.env.SENSENOVA_ACCESS_KEY_ID && !!process.env.SENSENOVA_ACCESS_KEY_SECRET,
295311
SENSENOVA_ACCESS_KEY_ID: process.env.SENSENOVA_ACCESS_KEY_ID,
296312
SENSENOVA_ACCESS_KEY_SECRET: process.env.SENSENOVA_ACCESS_KEY_SECRET,
297313
SENSENOVA_MODEL_LIST: process.env.SENSENOVA_MODEL_LIST,

‎src/config/modelProviders/index.ts

+4
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import StepfunProvider from './stepfun';
2929
import TaichuProvider from './taichu';
3030
import TogetherAIProvider from './togetherai';
3131
import UpstageProvider from './upstage';
32+
import VertexAIProvider from './vertexai';
3233
import WenxinProvider from './wenxin';
3334
import ZeroOneProvider from './zeroone';
3435
import ZhiPuProvider from './zhipu';
@@ -61,6 +62,7 @@ export const LOBE_DEFAULT_MODEL_LIST: ChatModelCard[] = [
6162
SiliconCloudProvider.chatModels,
6263
UpstageProvider.chatModels,
6364
SparkProvider.chatModels,
65+
VertexAIProvider.chatModels,
6466
Ai21Provider.chatModels,
6567
HunyuanProvider.chatModels,
6668
WenxinProvider.chatModels,
@@ -100,6 +102,7 @@ export const DEFAULT_MODEL_PROVIDER_LIST = [
100102
Ai360Provider,
101103
TaichuProvider,
102104
SiliconCloudProvider,
105+
VertexAIProvider,
103106
];
104107

105108
export const filterEnabledModels = (provider: ModelProviderCard) => {
@@ -140,6 +143,7 @@ export { default as StepfunProviderCard } from './stepfun';
140143
export { default as TaichuProviderCard } from './taichu';
141144
export { default as TogetherAIProviderCard } from './togetherai';
142145
export { default as UpstageProviderCard } from './upstage';
146+
export { default as VertexAIProviderCard } from './vertexai';
143147
export { default as WenxinProviderCard } from './wenxin';
144148
export { default as ZeroOneProviderCard } from './zeroone';
145149
export { default as ZhiPuProviderCard } from './zhipu';

‎src/config/modelProviders/vertexai.ts

+217
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
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+
{
7+
description:
8+
'Gemini 1.5 Flash 是Google最新的多模态AI模型，具备快速处理能力，支持文本、图像和视频输入，适用于多种任务的高效扩展。',
9+
displayName: 'Gemini 1.5 Flash',
10+
enabled: true,
11+
functionCall: 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+
tokens: 1_000_000 + 8192,
20+
vision: true,
21+
},
22+
{
23+
description: 'Gemini 1.5 Flash 002 是一款高效的多模态模型，支持广泛应用的扩展。',
24+
displayName: 'Gemini 1.5 Flash 002',
25+
enabled: true,
26+
functionCall: 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+
tokens: 1_000_000 + 8192,
36+
vision: true,
37+
},
38+
{
39+
description: 'Gemini 1.5 Flash 001 是一款高效的多模态模型，支持广泛应用的扩展。',
40+
displayName: 'Gemini 1.5 Flash 001',
41+
functionCall: true,
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+
tokens: 1_000_000 + 8192,
50+
vision: true,
51+
},
52+
{
53+
description: 'Gemini 1.5 Flash 0827 提供了优化后的多模态处理能力，适用多种复杂任务场景。',
54+
displayName: 'Gemini 1.5 Flash 0827',
55+
functionCall: true,
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+
tokens: 1_000_000 + 8192,
65+
vision: true,
66+
},
67+
68+
{
69+
description:
70+
'Gemini 1.5 Flash 8B 0924 是最新的实验性模型，在文本和多模态用例中都有显著的性能提升。',
71+
displayName: 'Gemini 1.5 Flash 8B 0924',
72+
functionCall: true,
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+
tokens: 1_000_000 + 8192,
82+
vision: true,
83+
},
84+
{
85+
description:
86+
'Gemini 1.5 Pro 支持高达200万个tokens，是中型多模态模型的理想选择，适用于复杂任务的多方面支持。',
87+
displayName: 'Gemini 1.5 Pro',
88+
enabled: true,
89+
functionCall: 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+
tokens: 2_000_000 + 8192,
99+
vision: true,
100+
},
101+
{
102+
description:
103+
'Gemini 1.5 Pro 002 是最新的生产就绪模型，提供更高质量的输出，特别在数学、长上下文和视觉任务方面有显著提升。',
104+
displayName: 'Gemini 1.5 Pro 002',
105+
enabled: true,
106+
functionCall: 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+
tokens: 2_000_000 + 8192,
116+
vision: true,
117+
},
118+
{
119+
description: 'Gemini 1.5 Pro 001 是可扩展的多模态AI解决方案，支持广泛的复杂任务。',
120+
displayName: 'Gemini 1.5 Pro 001',
121+
functionCall: true,
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+
tokens: 2_000_000 + 8192,
131+
vision: true,
132+
},
133+
{
134+
description: 'Gemini 1.5 Pro 0827 结合最新优化技术，带来更高效的多模态数据处理能力。',
135+
displayName: 'Gemini 1.5 Pro 0827',
136+
functionCall: true,
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+
tokens: 2_000_000 + 8192,
146+
vision: true,
147+
},
148+
{
149+
description: 'Gemini 1.5 Pro 0801 提供出色的多模态处理能力，为应用开发带来更大灵活性。',
150+
displayName: 'Gemini 1.5 Pro 0801',
151+
functionCall: true,
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+
tokens: 2_000_000 + 8192,
161+
vision: true,
162+
},
163+
{
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+
tokens: 30_720 + 2048,
174+
},
175+
{
176+
description:
177+
'Gemini 1.0 Pro 001 (Tuning) 提供稳定并可调优的性能，是复杂任务解决方案的理想选择。',
178+
displayName: 'Gemini 1.0 Pro 001 (Tuning)',
179+
functionCall: true,
180+
id: 'gemini-1.0-pro-001',
181+
maxOutput: 2048,
182+
pricing: {
183+
input: 0.5,
184+
output: 1.5,
185+
},
186+
releasedAt: '2023-12-06',
187+
tokens: 30_720 + 2048,
188+
},
189+
{
190+
description: 'Gemini 1.0 Pro 002 (Tuning) 提供出色的多模态支持，专注于复杂任务的有效解决。',
191+
displayName: 'Gemini 1.0 Pro 002 (Tuning)',
192+
id: 'gemini-1.0-pro-002',
193+
maxOutput: 2048,
194+
pricing: {
195+
input: 0.5,
196+
output: 1.5,
197+
},
198+
releasedAt: '2023-12-06',
199+
tokens: 30_720 + 2048,
200+
},
201+
],
202+
checkModel: 'gemini-1.5-flash-latest',
203+
description:
204+
'Google 的 Gemini 系列是其最先进、通用的 AI模型，由 Google DeepMind 打造，专为多模态设计，支持文本、代码、图像、音频和视频的无缝理解与处理。适用于从数据中心到移动设备的多种环境，极大提升了AI模型的效率与应用广泛性。',
205+
id: 'vertexai',
206+
modelsUrl: 'https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning',
207+
name: 'VertexAI',
208+
showApiKey: false,
209+
showChecker: false,
210+
smoothing: {
211+
speed: 2,
212+
text: true,
213+
},
214+
url: 'https://cloud.google.com/vertex-ai',
215+
};
216+
217+
export default VertexAI;

‎src/const/settings/llm.ts

+4
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,10 @@ export const DEFAULT_LLM_CONFIG: UserModelProviderConfig = {
152152
enabled: false,
153153
enabledModels: filterEnabledModels(UpstageProviderCard),
154154
},
155+
vertexai: {
156+
enabled: false,
157+
enabledModels: filterEnabledModels(GoogleProviderCard),
158+
},
155159
wenxin: {
156160
enabled: false,
157161
enabledModels: filterEnabledModels(WenxinProviderCard),

‎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-12
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
1+
import type { VertexAI } from '@google-cloud/vertexai';
12
import {
23
Content,
34
FunctionCallPart,
45
FunctionDeclaration,
56
Tool as GoogleFunctionCallTool,
67
GoogleGenerativeAI,
8+
HarmBlockThreshold,
9+
HarmCategory,
710
Part,
811
SchemaType,
912
} from '@google/generative-ai';
@@ -27,26 +30,21 @@ import { StreamingResponse } from '../utils/response';
2730
import { GoogleGenerativeAIStream, convertIterableToStream } from '../utils/streams';
2831
import { parseDataUri } from '../utils/uriParser';
2932

30-
enum HarmCategory {
31-
HARM_CATEGORY_DANGEROUS_CONTENT = 'HARM_CATEGORY_DANGEROUS_CONTENT',
32-
HARM_CATEGORY_HARASSMENT = 'HARM_CATEGORY_HARASSMENT',
33-
HARM_CATEGORY_HATE_SPEECH = 'HARM_CATEGORY_HATE_SPEECH',
34-
HARM_CATEGORY_SEXUALLY_EXPLICIT = 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
35-
}
36-
37-
enum HarmBlockThreshold {
38-
BLOCK_NONE = 'BLOCK_NONE',
33+
interface LobeGoogleAIParams {
34+
apiKey?: string;
35+
baseURL?: string;
36+
client?: GoogleGenerativeAI | VertexAI;
3937
}
4038

4139
export class LobeGoogleAI implements LobeRuntimeAI {
4240
private client: GoogleGenerativeAI;
4341
baseURL?: string;
4442

45-
constructor({ apiKey, baseURL }: { apiKey?: string; baseURL?: string } = {}) {
43+
constructor({ apiKey, baseURL, client }: LobeGoogleAIParams = {}) {
4644
if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
4745

48-
this.client = new GoogleGenerativeAI(apiKey);
49-
this.baseURL = baseURL;
46+
this.client = client ? (client as GoogleGenerativeAI) : new GoogleGenerativeAI(apiKey);
47+
this.baseURL = client ? undefined : baseURL;
5048
}
5149

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

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

+1
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ export enum ModelProvider {
5151
Taichu = 'taichu',
5252
TogetherAI = 'togetherai',
5353
Upstage = 'upstage',
54+
VertexAI = 'vertexai',
5455
Wenxin = 'wenxin',
5556
ZeroOne = 'zeroone',
5657
ZhiPu = 'zhipu',

‎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/server/globalConfig/index.ts

+12
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import {
3232
TaichuProviderCard,
3333
TogetherAIProviderCard,
3434
UpstageProviderCard,
35+
VertexAIProviderCard,
3536
WenxinProviderCard,
3637
ZeroOneProviderCard,
3738
ZhiPuProviderCard,
@@ -143,6 +144,9 @@ export const getServerGlobalConfig = () => {
143144

144145
ENABLED_HUGGINGFACE,
145146
HUGGINGFACE_MODEL_LIST,
147+
148+
ENABLED_VERTEXAI,
149+
VERTEXAI_MODEL_LIST,
146150
} = getLLMConfig();
147151

148152
const config: GlobalServerConfig = {
@@ -387,6 +391,14 @@ export const getServerGlobalConfig = () => {
387391
modelString: UPSTAGE_MODEL_LIST,
388392
}),
389393
},
394+
vertexai: {
395+
enabled: ENABLED_VERTEXAI,
396+
enabledModels: extractEnabledModels(VERTEXAI_MODEL_LIST),
397+
serverModelCards: transformToChatModelCards({
398+
defaultChatModels: VertexAIProviderCard.chatModels,
399+
modelString: VERTEXAI_MODEL_LIST,
400+
}),
401+
},
390402
wenxin: {
391403
enabled: ENABLED_WENXIN,
392404
enabledModels: extractEnabledModels(WENXIN_MODEL_LIST),

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

+1
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ export interface UserKeyVaults {
5858
taichu?: OpenAICompatibleKeyVault;
5959
togetherai?: OpenAICompatibleKeyVault;
6060
upstage?: OpenAICompatibleKeyVault;
61+
vertexai?: undefined;
6162
wenxin?: WenxinKeyVault;
6263
zeroone?: OpenAICompatibleKeyVault;
6364
zhipu?: OpenAICompatibleKeyVault;

0 commit comments

Comments
 (0)
Please sign in to comment.