Zenn
🐙

Workers環境でもVertex AI APIを使う

2024/09/29に公開
Google Cloud
Cloudflare Workers
Vertex AI
workerd
tech

問題

Vertex AIはGoogle共通のOAuthベースの認証システムを使っており、APIクライアントSDKの内部実装に使われているgoogle-auth-clientがworkerd環境に対応していない。

内部のOAuthクライアントをスキップして、自前でアクセストークンを取得する処理を書くことで回避する。

動作したコード

1. google-auth-clientを使わずにトークンを取得するコード

※大部分はLLMに書いてもらった

export async function createGoogleAccessToken(credentialString: string) {
  const credentials = JSON.parse(credentialString)
  const {private_key, client_email} = credentials
  const key = await createJwt({
    iss: client_email,
    scope: "https://www.googleapis.com/auth/cloud-platform",
    aud: "https://oauth2.googleapis.com/token",
    iat: Math.floor(Date.now() / 1000),
    exp: Math.floor(Date.now() / 1000) + 60 * 60,
  }, private_key)

  return await getAccessToken(key)
}

async function importPrivateKey(pem: string): Promise<CryptoKey> {
  const binaryDer = convertPemToBinary(pem);
  return crypto.subtle.importKey(
    "pkcs8",
    binaryDer,
    {
      name: "RSASSA-PKCS1-v1_5",
      hash: { name: "SHA-256" },
    },
    false,
    ["sign"]
  );
}

function convertPemToBinary(pem: string): ArrayBuffer {
  const base64 = pem
    .replace(/-----BEGIN PRIVATE KEY-----/, "")
    .replace(/-----END PRIVATE KEY-----/, "")
    .replace(/\n/g, "");
  const binaryString = atob(base64);
  const binaryLen = binaryString.length;
  const bytes = new Uint8Array(binaryLen);
  for (let i = 0; i < binaryLen; i++) {
    bytes[i] = binaryString.charCodeAt(i);
  }
  // @ts-ignore
  return bytes.buffer;
}

async function createJwt(payload: Record<string, any>, privateKeyPem: string): Promise<string> {
  const header = {
    alg: "RS256",
    typ: "JWT",
  };

  const encoder = new TextEncoder();

  const base64Header = btoa(JSON.stringify(header));
  const base64Payload = btoa(JSON.stringify(payload));

  const data = `${base64Header}.${base64Payload}`;

  const privateKey = await importPrivateKey(privateKeyPem);

  const signature = await crypto.subtle.sign(
    "RSASSA-PKCS1-v1_5",
    privateKey,
    encoder.encode(data)
  );

  const base64Signature = btoa(String.fromCharCode(...new Uint8Array(signature)));

  return `${data}.${base64Signature}`;
}

async function getAccessToken(key: string) {
  const tokenResponse = await fetch("https://oauth2.googleapis.com/token", {
    method: "POST",
    headers: {
      "Content-Type": "application/x-www-form-urlencoded"
    },
    body: new URLSearchParams({
      grant_type: "urn:ietf:params:oauth:grant-type:jwt-bearer",
      assertion: key
    })
  });
  
  const data = await tokenResponse.json();
  return data.access_token;
}

2. サービスアカウントのJSONからアクセストークンを取得

// サービスアカウントのJSONを1行にしたものをenvに入れておく
const credentials = getEnv().GOOGLE_VERTEX_AI_CREDENTIAL_JSON;
const accessToken = await createGoogleAccessToken(credentials)

// -> 1時間有効なので、kvなどでキャッシュしておく

3. アクセストークンを使う

ここまで来たら普通のAPIアクセストークンになっているのでAuthorization: Bearer XXXとして使うだけ。

Vercel AI SDKから使う場合は、ラップされてるので下記のように工夫する。

const authClient = new OAuth2Client();
authClient.setCredentials({
  access_token: accessToken, // 先ほど取得したトークン
});
const google = createVertex({
  project: getEnv().GOOGLE_VERTEX_AI_PROJECT,
  location: getEnv().GOOGLE_VERTEX_AI_LOCATION,
  googleAuthOptions: {
    // @ts-ignore
    authClient,
  },
});

感想

Workersに漏れず、Google Cloud APIは外部からめちゃくちゃ呼びにくいので普通のアクセストークンかリクエスト署名での認証に対応して欲しい！
(Vertex AIを通さないGoogle AI Studio APIは慣例を破ってて素晴らしい。)

なお、2024/9/29現在はVertex AI経由のasia-northeast1でgemini-1.5-flash-002のみ使えない不具合があるっぽい？
他のリージョンからは使えるし、日本リージョンから1.5-pro-002も使えるのでおそらく設定不備っぽい。

Discussion