vertex gemini consolidation

Author: icecrasher321

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/credential-selector/components/oauth-required-modal.tsx

@@ -43,6 +43,8 @@ const SCOPE_DESCRIPTIONS: Record<string, string> = {
   'https://www.googleapis.com/auth/admin.directory.group.readonly': 'View Google Workspace groups',
   'https://www.googleapis.com/auth/admin.directory.group.member.readonly':
     'View Google Workspace group memberships',
+  'https://www.googleapis.com/auth/cloud-platform':
+    'Full access to Google Cloud resources for Vertex AI',
   'read:confluence-content.all': 'Read all Confluence content',
   'read:confluence-space.summary': 'Read Confluence space information',
   'read:space:confluence': 'View Confluence spaces',

apps/sim/blocks/blocks/agent.ts

@@ -314,6 +314,19 @@ export const AgentBlock: BlockConfig<AgentResponse> = {
         value: providers.vertex.models,
       },
     },
+    {
+      id: 'vertexCredential',
+      title: 'Google Cloud Account',
+      type: 'oauth-input',
+      serviceId: 'vertex-ai',
+      requiredScopes: ['https://www.googleapis.com/auth/cloud-platform'],
+      placeholder: 'Select Google Cloud account',
+      required: true,
+      condition: {
+        field: 'model',
+        value: providers.vertex.models,
+      },
+    },
     {
       id: 'tools',
       title: 'Tools',
@@ -328,17 +341,21 @@ export const AgentBlock: BlockConfig<AgentResponse> = {
       password: true,
       connectionDroppable: false,
       required: true,
-      // Hide API key for hosted models, Ollama models, and vLLM models
+      // Hide API key for hosted models, Ollama models, vLLM models, and Vertex models (uses OAuth)
       condition: isHosted
         ? {
             field: 'model',
-            value: getHostedModels(),
+            value: [...getHostedModels(), ...providers.vertex.models],
             not: true, // Show for all models EXCEPT those listed
           }
         : () => ({
             field: 'model',
-            value: [...getCurrentOllamaModels(), ...getCurrentVLLMModels()],
-            not: true, // Show for all models EXCEPT Ollama and vLLM models
+            value: [
+              ...getCurrentOllamaModels(),
+              ...getCurrentVLLMModels(),
+              ...providers.vertex.models,
+            ],
+            not: true, // Show for all models EXCEPT Ollama, vLLM, and Vertex models
           }),
     },
     {

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

@@ -1,8 +1,9 @@
 import { db } from '@sim/db'
-import { mcpServers } from '@sim/db/schema'
+import { account, mcpServers } from '@sim/db/schema'
 import { and, eq, inArray, isNull } from 'drizzle-orm'
 import { createLogger } from '@/lib/logs/console/logger'
 import { createMcpToolId } from '@/lib/mcp/utils'
+import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
 import { getAllBlocks } from '@/blocks'
 import type { BlockOutput } from '@/blocks/types'
 import { AGENT, BlockType, DEFAULTS, HTTP } from '@/executor/constants'
@@ -919,6 +920,7 @@ export class AgentBlockHandler implements BlockHandler {
       azureApiVersion: inputs.azureApiVersion,
       vertexProject: inputs.vertexProject,
       vertexLocation: inputs.vertexLocation,
+      vertexCredential: inputs.vertexCredential,
       responseFormat,
       workflowId: ctx.workflowId,
       workspaceId: ctx.workspaceId,
@@ -997,7 +999,17 @@ export class AgentBlockHandler implements BlockHandler {
     responseFormat: any,
     providerStartTime: number
   ) {
-    const finalApiKey = this.getApiKey(providerId, model, providerRequest.apiKey)
+    let finalApiKey: string
+
+    // For Vertex AI, resolve OAuth credential to access token
+    if (providerId === 'vertex' && providerRequest.vertexCredential) {
+      finalApiKey = await this.resolveVertexCredential(
+        providerRequest.vertexCredential,
+        ctx.workflowId
+      )
+    } else {
+      finalApiKey = this.getApiKey(providerId, model, providerRequest.apiKey)
+    }
 
     const { blockData, blockNameMapping } = collectBlockData(ctx)
 
@@ -1024,7 +1036,6 @@ export class AgentBlockHandler implements BlockHandler {
       blockNameMapping,
     })
 
-    this.logExecutionSuccess(providerId, model, ctx, block, providerStartTime, response)
     return this.processProviderResponse(response, block, responseFormat)
   }
 
@@ -1049,15 +1060,6 @@ export class AgentBlockHandler implements BlockHandler {
       throw new Error(errorMessage)
     }
 
-    this.logExecutionSuccess(
-      providerRequest.provider,
-      providerRequest.model,
-      ctx,
-      block,
-      providerStartTime,
-      'HTTP response'
-    )
-
     const contentType = response.headers.get('Content-Type')
     if (contentType?.includes(HTTP.CONTENT_TYPE.EVENT_STREAM)) {
       return this.handleStreamingResponse(response, block, ctx, inputs)
@@ -1117,21 +1119,33 @@ export class AgentBlockHandler implements BlockHandler {
     }
   }
 
-  private logExecutionSuccess(
-    provider: string,
-    model: string,
-    ctx: ExecutionContext,
-    block: SerializedBlock,
-    startTime: number,
-    response: any
-  ) {
-    const executionTime = Date.now() - startTime
-    const responseType =
-      response instanceof ReadableStream
-        ? 'stream'
-        : response && typeof response === 'object' && 'stream' in response
-          ? 'streaming-execution'
-          : 'json'
+  /**
+   * Resolves a Vertex AI OAuth credential to an access token
+   */
+  private async resolveVertexCredential(credentialId: string, workflowId: string): Promise<string> {
+    const requestId = `vertex-${Date.now()}`
+
+    logger.info(`[${requestId}] Resolving Vertex AI credential: ${credentialId}`)
+
+    // Get the credential - we need to find the owner
+    // Since we're in a workflow context, we can query the credential directly
+    const credential = await db.query.account.findFirst({
+      where: eq(account.id, credentialId),
+    })
+
+    if (!credential) {
+      throw new Error(`Vertex AI credential not found: ${credentialId}`)
+    }
+
+    // Refresh the token if needed
+    const { accessToken } = await refreshTokenIfNeeded(requestId, credential, credentialId)
+
+    if (!accessToken) {
+      throw new Error('Failed to get Vertex AI access token')
+    }
+
+    logger.info(`[${requestId}] Successfully resolved Vertex AI credential`)
+    return accessToken
   }
 
   private handleExecutionError(

apps/sim/executor/handlers/agent/types.ts

@@ -21,6 +21,7 @@ export interface AgentInputs {
   azureApiVersion?: string
   vertexProject?: string
   vertexLocation?: string
+  vertexCredential?: string
   reasoningEffort?: string
   verbosity?: string
 }

apps/sim/lib/auth/auth.ts

@@ -579,6 +579,21 @@ export const auth = betterAuth({
           redirectURI: `${getBaseUrl()}/api/auth/oauth2/callback/google-groups`,
         },
 
+        {
+          providerId: 'vertex-ai',
+          clientId: env.GOOGLE_CLIENT_ID as string,
+          clientSecret: env.GOOGLE_CLIENT_SECRET as string,
+          discoveryUrl: 'https://accounts.google.com/.well-known/openid-configuration',
+          accessType: 'offline',
+          scopes: [
+            'https://www.googleapis.com/auth/userinfo.email',
+            'https://www.googleapis.com/auth/userinfo.profile',
+            'https://www.googleapis.com/auth/cloud-platform',
+          ],
+          prompt: 'consent',
+          redirectURI: `${getBaseUrl()}/api/auth/oauth2/callback/vertex-ai`,
+        },
+
         {
           providerId: 'microsoft-teams',
           clientId: env.MICROSOFT_CLIENT_ID as string,

apps/sim/lib/oauth/oauth.ts

@@ -32,6 +32,7 @@ import {
   SlackIcon,
   SpotifyIcon,
   TrelloIcon,
+  VertexIcon,
   WealthboxIcon,
   WebflowIcon,
   WordpressIcon,
@@ -80,6 +81,7 @@ export type OAuthService =
   | 'google-vault'
   | 'google-forms'
   | 'google-groups'
+  | 'vertex-ai'
   | 'github'
   | 'x'
   | 'confluence'
@@ -237,6 +239,16 @@ export const OAUTH_PROVIDERS: Record<string, OAuthProviderConfig> = {
         ],
         scopeHints: ['admin.directory.group'],
       },
+      'vertex-ai': {
+        id: 'vertex-ai',
+        name: 'Vertex AI',
+        description: 'Access Google Cloud Vertex AI for Gemini models with OAuth.',
+        providerId: 'vertex-ai',
+        icon: (props) => VertexIcon(props),
+        baseProviderIcon: (props) => VertexIcon(props),
+        scopes: ['https://www.googleapis.com/auth/cloud-platform'],
+        scopeHints: ['cloud-platform', 'vertex', 'aiplatform'],
+      },
     },
     defaultService: 'gmail',
   },
@@ -1099,6 +1111,12 @@ export function parseProvider(provider: OAuthProvider): ProviderConfig {
       featureType: 'microsoft-planner',
     }
   }
+  if (provider === 'vertex-ai') {
+    return {
+      baseProvider: 'google',
+      featureType: 'vertex-ai',
+    }
+  }
 
   // Handle compound providers (e.g., 'google-email' -> { baseProvider: 'google', featureType: 'email' })
   const [base, feature] = provider.split('-')

apps/sim/providers/gemini/client.ts

@@ -0,0 +1,58 @@
+import { GoogleGenAI } from '@google/genai'
+import { createLogger } from '@/lib/logs/console/logger'
+import type { GeminiClientConfig } from './types'
+
+const logger = createLogger('GeminiClient')
+
+/**
+ * Creates a GoogleGenAI client configured for either Google Gemini API or Vertex AI
+ *
+ * For Google Gemini API:
+ *   - Uses API key authentication
+ *
+ * For Vertex AI:
+ *   - Uses OAuth access token via HTTP Authorization header
+ *   - Requires project and location
+ */
+export function createGeminiClient(config: GeminiClientConfig): GoogleGenAI {
+  if (config.vertexai) {
+    if (!config.project) {
+      throw new Error('Vertex AI requires a project ID')
+    }
+    if (!config.accessToken) {
+      throw new Error('Vertex AI requires an access token')
+    }
+
+    const location = config.location ?? 'us-central1'
+
+    logger.info('Creating Vertex AI client', {
+      project: config.project,
+      location,
+      hasAccessToken: !!config.accessToken,
+    })
+
+    // Create client with Vertex AI configuration
+    // Use httpOptions.headers to pass the access token directly
+    return new GoogleGenAI({
+      vertexai: true,
+      project: config.project,
+      location,
+      httpOptions: {
+        headers: {
+          Authorization: `Bearer ${config.accessToken}`,
+        },
+      },
+    })
+  }
+
+  // Google Gemini API with API key
+  if (!config.apiKey) {
+    throw new Error('Google Gemini API requires an API key')
+  }
+
+  logger.info('Creating Google Gemini client')
+
+  return new GoogleGenAI({
+    apiKey: config.apiKey,
+  })
+}