Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow vertex to work on Vercel-Edge #3096

Open
cosbgn opened this issue Sep 23, 2024 · 1 comment
Open

Allow vertex to work on Vercel-Edge #3096

cosbgn opened this issue Sep 23, 2024 · 1 comment
Labels
ai/provider enhancement New feature or request

Comments

@cosbgn
Copy link

cosbgn commented Sep 23, 2024

Feature Description

The vertex model doesn't work on vercel edge. This is cause by the auth-library not being compatible with the edge. This could be easily replaced by a vanilla js script:

import { webcrypto } from 'crypto'
const tokenUrl = 'https://oauth2.googleapis.com/token';
const base64url = (source) => {
	let encodedSource = Buffer.from(source).toString('base64');
	encodedSource = encodedSource.replace(/=+$/, '');
	encodedSource = encodedSource.replace(/\+/g, '-');
	encodedSource = encodedSource.replace(/\//g, '_');
	return encodedSource;
  }
  
  const createJwtSegment = (segment) => base64url(JSON.stringify(segment));
  
  const importPrivateKey = async (pemKey) => {
	const pemHeader = "-----BEGIN PRIVATE KEY-----";
	const pemFooter = "-----END PRIVATE KEY-----";
	const pemContents = pemKey.substring(pemHeader.length, pemKey.length - pemFooter.length);
	const binaryDerString = Buffer.from(pemContents, 'base64').toString('binary');
	const binaryDer = new Uint8Array(binaryDerString.length);
  
	for (let i = 0; i < binaryDerString.length; i++) {
	  binaryDer[i] = binaryDerString.charCodeAt(i);
	}
  
	return webcrypto.subtle.importKey("pkcs8", binaryDer.buffer, { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" }, true, ["sign"] );
  }
  
  const signJwt = async (header, payload, privateKey) => {
	const data = `${header}.${payload}`;
	const encoder = new TextEncoder();
	const signature = await webcrypto.subtle.sign("RSASSA-PKCS1-v1_5", privateKey, encoder.encode(data) );
	return `${data}.${base64url(Buffer.from(signature))}`;
  };
  
  const buildJwt = async ({ client_email, private_key, private_key_id, scopes }) => {
	const iat = Math.floor(Date.now() / 1000);
	const exp = iat + 3600;
  
	const payload = { iss: client_email, sub: client_email, aud: tokenUrl, iat: iat, exp: exp, scope: scopes.join(" ") }
	const jwtHeaders = { kid: private_key_id, alg: 'RS256', typ: 'JWT' };
  
	const encodedHeader = createJwtSegment(jwtHeaders);
	const encodedPayload = createJwtSegment(payload);
  
	// Import the private key
	const privateKey = await importPrivateKey(private_key);
  
	// Sign the JWT
	const signedJwt = await signJwt(encodedHeader, encodedPayload, privateKey);
  
	return signedJwt;
}


// Function to get access token
const generateToken = async(credentials) => {
	// 1) Build and sign JWT
    const signedJwt = await buildJwt(credentials)
    // 2) Get access token
    const response = await fetch(tokenUrl, {
      method: 'POST',
      headers: {  'Content-Type': 'application/x-www-form-urlencoded' },
      body: new URLSearchParams({
        'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
        'assertion': signedJwt
      })
    })

    if (!response.ok) {
      throw new Error(`Token request failed: ${response.statusText}`);
    }
    return await response.json()
}

It can easily be used like this:

	const { expires_in, access_token } = await generateToken({
			client_email:process.env.GOOGLE_CLOUD_CLIENT_EMAIL, 
			private_key: process.env.GOOGLE_CLOUD_PRIVATE_KEY?.replace(/\\n/gm, "\n"), // Fix issue with new line
			private_key_id: process.env.GOOGLE_CLOUD_PRIVATE_KEY_ID, 
			scopes: ["https://www.googleapis.com/auth/cloud-platform"]
		})

Use Case

Would be great to decouple vertex and auto so that we can use our custom auth script like the one above or simply add the script above directly in the ai-sdk model so that it would work on vercel edge.

Additional context

No response

@cosbgn
Copy link
Author

cosbgn commented Sep 23, 2024

Also vertex has a very simply rest API, so most likely the whole ai-sdk/vertex could only depend on built in fetch and work perfectly everywhere.

@lgrammel lgrammel added enhancement New feature or request ai/provider labels Sep 23, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ai/provider enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

2 participants