UNPKG

@ai-sdk/google-vertex

Version:

The **[Google Vertex provider](https://ai-sdk.dev/providers/ai-sdk-providers/google-vertex)** for the [AI SDK](https://ai-sdk.dev/docs) contains language model support for the [Google Vertex AI](https://cloud.google.com/vertex-ai) APIs.

162 lines (139 loc) 4.5 kB
import { loadOptionalSetting, loadSetting, withUserAgentSuffix, getRuntimeEnvironmentUserAgent, } from '@ai-sdk/provider-utils'; import { VERSION } from '../version'; export interface GoogleCredentials { /** * The client email for the Google Cloud service account. Defaults to the * value of the `GOOGLE_CLIENT_EMAIL` environment variable. */ clientEmail: string; /** * The private key for the Google Cloud service account. Defaults to the * value of the `GOOGLE_PRIVATE_KEY` environment variable. */ privateKey: string; /** * Optional. The private key ID for the Google Cloud service account. Defaults * to the value of the `GOOGLE_PRIVATE_KEY_ID` environment variable. */ privateKeyId?: string; } const loadCredentials = async (): Promise<GoogleCredentials> => { try { return { clientEmail: loadSetting({ settingValue: undefined, settingName: 'clientEmail', environmentVariableName: 'GOOGLE_CLIENT_EMAIL', description: 'Google client email', }), privateKey: loadSetting({ settingValue: undefined, settingName: 'privateKey', environmentVariableName: 'GOOGLE_PRIVATE_KEY', description: 'Google private key', }), privateKeyId: loadOptionalSetting({ settingValue: undefined, environmentVariableName: 'GOOGLE_PRIVATE_KEY_ID', }), }; } catch (error: any) { throw new Error(`Failed to load Google credentials: ${error.message}`); } }; // Convert a string to base64url const base64url = (str: string) => { return btoa(str).replace(/\+/g, '-').replace(/\//g, '_').replace(/=/g, ''); }; const importPrivateKey = async (pemKey: string) => { const pemHeader = '-----BEGIN PRIVATE KEY-----'; const pemFooter = '-----END PRIVATE KEY-----'; // Remove header, footer, and any whitespace/newlines const pemContents = pemKey .replace(pemHeader, '') .replace(pemFooter, '') .replace(/\s/g, ''); // Decode base64 to binary const binaryString = atob(pemContents); // Convert binary string to Uint8Array const binaryData = new Uint8Array(binaryString.length); for (let i = 0; i < binaryString.length; i++) { binaryData[i] = binaryString.charCodeAt(i); } return await crypto.subtle.importKey( 'pkcs8', binaryData, { name: 'RSASSA-PKCS1-v1_5', hash: 'SHA-256' }, true, ['sign'], ); }; const buildJwt = async (credentials: GoogleCredentials) => { const now = Math.floor(Date.now() / 1000); // Only include kid in header if privateKeyId is provided const header: { alg: string; typ: string; kid?: string } = { alg: 'RS256', typ: 'JWT', }; if (credentials.privateKeyId) { header.kid = credentials.privateKeyId; } const payload = { iss: credentials.clientEmail, scope: 'https://www.googleapis.com/auth/cloud-platform', aud: 'https://oauth2.googleapis.com/token', exp: now + 3600, iat: now, }; const privateKey = await importPrivateKey(credentials.privateKey); const signingInput = `${base64url(JSON.stringify(header))}.${base64url( JSON.stringify(payload), )}`; const encoder = new TextEncoder(); const data = encoder.encode(signingInput); const signature = await crypto.subtle.sign( 'RSASSA-PKCS1-v1_5', privateKey, data, ); const signatureBase64 = base64url( String.fromCharCode(...new Uint8Array(signature)), ); return `${base64url(JSON.stringify(header))}.${base64url( JSON.stringify(payload), )}.${signatureBase64}`; }; /** * Generate an authentication token for Google Vertex AI in a manner compatible * with the Edge runtime. */ export async function generateAuthToken(credentials?: GoogleCredentials) { try { const creds = credentials || (await loadCredentials()); const jwt = await buildJwt(creds); const response = await fetch('https://oauth2.googleapis.com/token', { method: 'POST', headers: withUserAgentSuffix( { 'Content-Type': 'application/x-www-form-urlencoded' }, `ai-sdk/google-vertex/${VERSION}`, getRuntimeEnvironmentUserAgent(), ), body: new URLSearchParams({ grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer', assertion: jwt, }), }); if (!response.ok) { throw new Error(`Token request failed: ${response.statusText}`); } const data = await response.json(); return data.access_token; } catch (error) { throw error; } }