Skip to content

Commit d2792fc

Browse files
waleedlatif1waleedlatif1
committed
improvement(vertex): added vertex to all LLM-based blocks, fixed refresh
1 parent 641ac58 commit d2792fc

File tree

10 files changed

+244
-18
lines changed

10 files changed

+244
-18
lines changed

‎apps/sim/app/api/providers/route.ts‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
1+
import { db } from '@sim/db'
2+
import { account } from '@sim/db/schema'
3+
import { eq } from 'drizzle-orm'
14
import { type NextRequest, NextResponse } from 'next/server'
25
import { generateRequestId } from '@/lib/core/utils/request'
36
import { createLogger } from '@/lib/logs/console/logger'
7+
import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
48
import type { StreamingExecution } from '@/executor/types'
59
import { executeProviderRequest } from '@/providers'
610
import { getApiKey } from '@/providers/utils'
@@ -37,6 +41,7 @@ export async function POST(request: NextRequest) {
3741
azureApiVersion,
3842
vertexProject,
3943
vertexLocation,
44+
vertexCredential,
4045
responseFormat,
4146
workflowId,
4247
workspaceId,
@@ -62,6 +67,7 @@ export async function POST(request: NextRequest) {
6267
hasAzureApiVersion: !!azureApiVersion,
6368
hasVertexProject: !!vertexProject,
6469
hasVertexLocation: !!vertexLocation,
70+
hasVertexCredential: !!vertexCredential,
6571
hasResponseFormat: !!responseFormat,
6672
workflowId,
6773
stream: !!stream,
@@ -76,13 +82,18 @@ export async function POST(request: NextRequest) {
7682

7783
let finalApiKey: string
7884
try {
79-
finalApiKey = getApiKey(provider, model, apiKey)
85+
if (provider === 'vertex' && vertexCredential) {
86+
finalApiKey = await resolveVertexCredential(requestId, vertexCredential)
87+
} else {
88+
finalApiKey = getApiKey(provider, model, apiKey)
89+
}
8090
} catch (error) {
8191
logger.error(`[${requestId}] Failed to get API key:`, {
8292
provider,
8393
model,
8494
error: error instanceof Error ? error.message : String(error),
8595
hasProvidedApiKey: !!apiKey,
96+
hasVertexCredential: !!vertexCredential,
8697
})
8798
return NextResponse.json(
8899
{ error: error instanceof Error ? error.message : 'API key error' },
@@ -324,3 +335,27 @@ function sanitizeObject(obj: any): any {
324335

325336
return result
326337
}
338+
339+
/**
340+
* Resolves a Vertex AI OAuth credential to an access token
341+
*/
342+
async function resolveVertexCredential(requestId: string, credentialId: string): Promise<string> {
343+
logger.info(`[${requestId}] Resolving Vertex AI credential: ${credentialId}`)
344+
345+
const credential = await db.query.account.findFirst({
346+
where: eq(account.id, credentialId),
347+
})
348+
349+
if (!credential) {
350+
throw new Error(`Vertex AI credential not found: ${credentialId}`)
351+
}
352+
353+
const { accessToken } = await refreshTokenIfNeeded(requestId, credential, credentialId)
354+
355+
if (!accessToken) {
356+
throw new Error('Failed to get Vertex AI access token')
357+
}
358+
359+
logger.info(`[${requestId}] Successfully resolved Vertex AI credential`)
360+
return accessToken
361+
}

‎apps/sim/blocks/blocks/agent.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -310,6 +310,7 @@ export const AgentBlock: BlockConfig<AgentResponse> = {
310310
type: 'short-input',
311311
placeholder: 'your-gcp-project-id',
312312
connectionDroppable: false,
313+
required: true,
313314
condition: {
314315
field: 'model',
315316
value: providers.vertex.models,
@@ -321,6 +322,7 @@ export const AgentBlock: BlockConfig<AgentResponse> = {
321322
type: 'short-input',
322323
placeholder: 'us-central1',
323324
connectionDroppable: false,
325+
required: true,
324326
condition: {
325327
field: 'model',
326328
value: providers.vertex.models,

‎apps/sim/blocks/blocks/evaluator.ts‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ const getCurrentOllamaModels = () => {
1818
return useProvidersStore.getState().providers.ollama.models
1919
}
2020

21+
const getCurrentVLLMModels = () => {
22+
return useProvidersStore.getState().providers.vllm.models
23+
}
24+
2125
interface Metric {
2226
name: string
2327
description: string
@@ -196,6 +200,19 @@ export const EvaluatorBlock: BlockConfig<EvaluatorResponse> = {
196200
})
197201
},
198202
},
203+
{
204+
id: 'vertexCredential',
205+
title: 'Google Cloud Account',
206+
type: 'oauth-input',
207+
serviceId: 'vertex-ai',
208+
requiredScopes: ['https://www.googleapis.com/auth/cloud-platform'],
209+
placeholder: 'Select Google Cloud account',
210+
required: true,
211+
condition: {
212+
field: 'model',
213+
value: providers.vertex.models,
214+
},
215+
},
199216
{
200217
id: 'apiKey',
201218
title: 'API Key',
@@ -204,16 +221,21 @@ export const EvaluatorBlock: BlockConfig<EvaluatorResponse> = {
204221
password: true,
205222
connectionDroppable: false,
206223
required: true,
224+
// Hide API key for hosted models, Ollama models, vLLM models, and Vertex models (uses OAuth)
207225
condition: isHosted
208226
? {
209227
field: 'model',
210-
value: getHostedModels(),
228+
value: [...getHostedModels(), ...providers.vertex.models],
211229
not: true, // Show for all models EXCEPT those listed
212230
}
213231
: () => ({
214232
field: 'model',
215-
value: getCurrentOllamaModels(),
216-
not: true, // Show for all models EXCEPT Ollama models
233+
value: [
234+
...getCurrentOllamaModels(),
235+
...getCurrentVLLMModels(),
236+
...providers.vertex.models,
237+
],
238+
not: true, // Show for all models EXCEPT Ollama, vLLM, and Vertex models
217239
}),
218240
},
219241
{
@@ -245,6 +267,7 @@ export const EvaluatorBlock: BlockConfig<EvaluatorResponse> = {
245267
type: 'short-input',
246268
placeholder: 'your-gcp-project-id',
247269
connectionDroppable: false,
270+
required: true,
248271
condition: {
249272
field: 'model',
250273
value: providers.vertex.models,
@@ -256,6 +279,7 @@ export const EvaluatorBlock: BlockConfig<EvaluatorResponse> = {
256279
type: 'short-input',
257280
placeholder: 'us-central1',
258281
connectionDroppable: false,
282+
required: true,
259283
condition: {
260284
field: 'model',
261285
value: providers.vertex.models,
@@ -386,6 +410,10 @@ export const EvaluatorBlock: BlockConfig<EvaluatorResponse> = {
386410
type: 'string' as ParamType,
387411
description: 'Google Cloud location for Vertex AI',
388412
},
413+
vertexCredential: {
414+
type: 'string' as ParamType,
415+
description: 'Google Cloud OAuth credential ID for Vertex AI',
416+
},
389417
temperature: {
390418
type: 'number' as ParamType,
391419
description: 'Response randomness level (low for consistent evaluation)',

‎apps/sim/blocks/blocks/router.ts‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@ const getCurrentOllamaModels = () => {
1515
return useProvidersStore.getState().providers.ollama.models
1616
}
1717

18+
const getCurrentVLLMModels = () => {
19+
return useProvidersStore.getState().providers.vllm.models
20+
}
21+
1822
interface RouterResponse extends ToolResponse {
1923
output: {
2024
prompt: string
@@ -144,6 +148,19 @@ export const RouterBlock: BlockConfig<RouterResponse> = {
144148
})
145149
},
146150
},
151+
{
152+
id: 'vertexCredential',
153+
title: 'Google Cloud Account',
154+
type: 'oauth-input',
155+
serviceId: 'vertex-ai',
156+
requiredScopes: ['https://www.googleapis.com/auth/cloud-platform'],
157+
placeholder: 'Select Google Cloud account',
158+
required: true,
159+
condition: {
160+
field: 'model',
161+
value: providers.vertex.models,
162+
},
163+
},
147164
{
148165
id: 'apiKey',
149166
title: 'API Key',
@@ -152,17 +169,21 @@ export const RouterBlock: BlockConfig<RouterResponse> = {
152169
password: true,
153170
connectionDroppable: false,
154171
required: true,
155-
// Hide API key for hosted models and Ollama models
172+
// Hide API key for hosted models, Ollama models, vLLM models, and Vertex models (uses OAuth)
156173
condition: isHosted
157174
? {
158175
field: 'model',
159-
value: getHostedModels(),
176+
value: [...getHostedModels(), ...providers.vertex.models],
160177
not: true, // Show for all models EXCEPT those listed
161178
}
162179
: () => ({
163180
field: 'model',
164-
value: getCurrentOllamaModels(),
165-
not: true, // Show for all models EXCEPT Ollama models
181+
value: [
182+
...getCurrentOllamaModels(),
183+
...getCurrentVLLMModels(),
184+
...providers.vertex.models,
185+
],
186+
not: true, // Show for all models EXCEPT Ollama, vLLM, and Vertex models
166187
}),
167188
},
168189
{
@@ -194,6 +215,7 @@ export const RouterBlock: BlockConfig<RouterResponse> = {
194215
type: 'short-input',
195216
placeholder: 'your-gcp-project-id',
196217
connectionDroppable: false,
218+
required: true,
197219
condition: {
198220
field: 'model',
199221
value: providers.vertex.models,
@@ -205,6 +227,7 @@ export const RouterBlock: BlockConfig<RouterResponse> = {
205227
type: 'short-input',
206228
placeholder: 'us-central1',
207229
connectionDroppable: false,
230+
required: true,
208231
condition: {
209232
field: 'model',
210233
value: providers.vertex.models,
@@ -259,6 +282,10 @@ export const RouterBlock: BlockConfig<RouterResponse> = {
259282
azureApiVersion: { type: 'string', description: 'Azure API version' },
260283
vertexProject: { type: 'string', description: 'Google Cloud project ID for Vertex AI' },
261284
vertexLocation: { type: 'string', description: 'Google Cloud location for Vertex AI' },
285+
vertexCredential: {
286+
type: 'string',
287+
description: 'Google Cloud OAuth credential ID for Vertex AI',
288+
},
262289
temperature: {
263290
type: 'number',
264291
description: 'Response randomness level (low for consistent routing)',

‎apps/sim/blocks/blocks/translate.ts‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@ const getCurrentOllamaModels = () => {
88
return useProvidersStore.getState().providers.ollama.models
99
}
1010

11+
const getCurrentVLLMModels = () => {
12+
return useProvidersStore.getState().providers.vllm.models
13+
}
14+
1115
const getTranslationPrompt = (targetLanguage: string) =>
1216
`Translate the following text into ${targetLanguage || 'English'}. Output ONLY the translated text with no additional commentary, explanations, or notes.`
1317

@@ -55,6 +59,19 @@ export const TranslateBlock: BlockConfig = {
5559
})
5660
},
5761
},
62+
{
63+
id: 'vertexCredential',
64+
title: 'Google Cloud Account',
65+
type: 'oauth-input',
66+
serviceId: 'vertex-ai',
67+
requiredScopes: ['https://www.googleapis.com/auth/cloud-platform'],
68+
placeholder: 'Select Google Cloud account',
69+
required: true,
70+
condition: {
71+
field: 'model',
72+
value: providers.vertex.models,
73+
},
74+
},
5875
{
5976
id: 'apiKey',
6077
title: 'API Key',
@@ -63,17 +80,21 @@ export const TranslateBlock: BlockConfig = {
6380
password: true,
6481
connectionDroppable: false,
6582
required: true,
66-
// Hide API key for hosted models and Ollama models
83+
// Hide API key for hosted models, Ollama models, vLLM models, and Vertex models (uses OAuth)
6784
condition: isHosted
6885
? {
6986
field: 'model',
70-
value: getHostedModels(),
87+
value: [...getHostedModels(), ...providers.vertex.models],
7188
not: true, // Show for all models EXCEPT those listed
7289
}
7390
: () => ({
7491
field: 'model',
75-
value: getCurrentOllamaModels(),
76-
not: true, // Show for all models EXCEPT Ollama models
92+
value: [
93+
...getCurrentOllamaModels(),
94+
...getCurrentVLLMModels(),
95+
...providers.vertex.models,
96+
],
97+
not: true, // Show for all models EXCEPT Ollama, vLLM, and Vertex models
7798
}),
7899
},
79100
{
@@ -105,6 +126,7 @@ export const TranslateBlock: BlockConfig = {
105126
type: 'short-input',
106127
placeholder: 'your-gcp-project-id',
107128
connectionDroppable: false,
129+
required: true,
108130
condition: {
109131
field: 'model',
110132
value: providers.vertex.models,
@@ -116,6 +138,7 @@ export const TranslateBlock: BlockConfig = {
116138
type: 'short-input',
117139
placeholder: 'us-central1',
118140
connectionDroppable: false,
141+
required: true,
119142
condition: {
120143
field: 'model',
121144
value: providers.vertex.models,
@@ -144,6 +167,7 @@ export const TranslateBlock: BlockConfig = {
144167
azureApiVersion: params.azureApiVersion,
145168
vertexProject: params.vertexProject,
146169
vertexLocation: params.vertexLocation,
170+
vertexCredential: params.vertexCredential,
147171
}),
148172
},
149173
},
@@ -155,6 +179,10 @@ export const TranslateBlock: BlockConfig = {
155179
azureApiVersion: { type: 'string', description: 'Azure API version' },
156180
vertexProject: { type: 'string', description: 'Google Cloud project ID for Vertex AI' },
157181
vertexLocation: { type: 'string', description: 'Google Cloud location for Vertex AI' },
182+
vertexCredential: {
183+
type: 'string',
184+
description: 'Google Cloud OAuth credential ID for Vertex AI',
185+
},
158186
systemPrompt: { type: 'string', description: 'Translation instructions' },
159187
},
160188
outputs: {

‎apps/sim/executor/handlers/agent/agent-handler.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1001,7 +1001,6 @@ export class AgentBlockHandler implements BlockHandler {
10011001
) {
10021002
let finalApiKey: string
10031003

1004-
// For Vertex AI, resolve OAuth credential to access token
10051004
if (providerId === 'vertex' && providerRequest.vertexCredential) {
10061005
finalApiKey = await this.resolveVertexCredential(
10071006
providerRequest.vertexCredential,

0 commit comments

Comments
 (0)