@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
text/typescript
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;
}
}