feat: Implement Vertex AI support for Anthropic and Google models

app/utils/chat.ts

@@ -3,7 +3,7 @@ import {
   UPLOAD_URL,
   REQUEST_TIMEOUT_MS,
 } from "@/app/constant";
-import { RequestMessage } from "@/app/client/api";
+import { ChatOptions, RequestMessage } from "@/app/client/api";
 import Locale from "@/app/locales";
 import {
   EventStreamContentType,
@@ -167,7 +167,7 @@ export function stream(
     toolCallMessage: any,
     toolCallResult: any[],
   ) => void,
-  options: any,
+  options: ChatOptions,
 ) {
   let responseText = "";
   let remainText = "";

app/utils/gtoken.ts

@@ -0,0 +1,285 @@
+/**
+ * npm:gtoken patched version for nextjs edge runtime, by ryanhex53
+ */
+// import { default as axios } from "axios";
+import { SignJWT, importPKCS8 } from "jose";
+
+const GOOGLE_TOKEN_URL = "https://www.googleapis.com/oauth2/v4/token";
+const GOOGLE_REVOKE_TOKEN_URL =
+  "https://accounts.google.com/o/oauth2/revoke?token=";
+
+export type GetTokenCallback = (err: Error | null, token?: TokenData) => void;
+
+export interface Credentials {
+  privateKey: string;
+  clientEmail?: string;
+}
+
+export interface TokenData {
+  refresh_token?: string;
+  expires_in?: number;
+  access_token?: string;
+  token_type?: string;
+  id_token?: string;
+}
+
+export interface TokenOptions {
+  key: string;
+  email?: string;
+  iss?: string;
+  sub?: string;
+  scope?: string | string[];
+  additionalClaims?: Record<string | number | symbol, never>;
+  // Eagerly refresh unexpired tokens when they are within this many
+  // milliseconds from expiring".
+  // Defaults to 5 minutes (300,000 milliseconds).
+  eagerRefreshThresholdMillis?: number;
+}
+
+export interface GetTokenOptions {
+  forceRefresh?: boolean;
+}
+
+class ErrorWithCode extends Error {
+  constructor(
+    message: string,
+    public code: string,
+  ) {
+    super(message);
+  }
+}
+
+export class GoogleToken {
+  get accessToken() {
+    return this.rawToken ? this.rawToken.access_token : undefined;
+  }
+  get idToken() {
+    return this.rawToken ? this.rawToken.id_token : undefined;
+  }
+  get tokenType() {
+    return this.rawToken ? this.rawToken.token_type : undefined;
+  }
+  get refreshToken() {
+    return this.rawToken ? this.rawToken.refresh_token : undefined;
+  }
+  expiresAt?: number;
+  key: string = "";
+  iss?: string;
+  sub?: string;
+  scope?: string;
+  rawToken?: TokenData;
+  tokenExpires?: number;
+  email?: string;
+  additionalClaims?: Record<string | number | symbol, never>;
+  eagerRefreshThresholdMillis: number = 0;
+
+  #inFlightRequest?: undefined | Promise<TokenData | undefined>;
+
+  /**
+   * Create a GoogleToken.
+   *
+   * @param options  Configuration object.
+   */
+  constructor(options?: TokenOptions) {
+    this.#configure(options);
+  }
+
+  /**
+   * Returns whether the token has expired.
+   *
+   * @return true if the token has expired, false otherwise.
+   */
+  hasExpired() {
+    const now = new Date().getTime();
+    if (this.rawToken && this.expiresAt) {
+      return now >= this.expiresAt;
+    } else {
+      return true;
+    }
+  }
+
+  /**
+   * Returns whether the token will expire within eagerRefreshThresholdMillis
+   *
+   * @return true if the token will be expired within eagerRefreshThresholdMillis, false otherwise.
+   */
+  isTokenExpiring() {
+    const now = new Date().getTime();
+    const eagerRefreshThresholdMillis = this.eagerRefreshThresholdMillis ?? 0;
+    if (this.rawToken && this.expiresAt) {
+      return this.expiresAt <= now + eagerRefreshThresholdMillis;
+    } else {
+      return true;
+    }
+  }
+
+  /**
+   * Returns a cached token or retrieves a new one from Google.
+   *
+   * @param callback The callback function.
+   */
+  getToken(opts?: GetTokenOptions): Promise<TokenData | undefined>;
+  getToken(callback: GetTokenCallback, opts?: GetTokenOptions): void;
+  getToken(
+    callback?: GetTokenCallback | GetTokenOptions,
+    opts = {} as GetTokenOptions,
+  ): void | Promise<TokenData | undefined> {
+    if (typeof callback === "object") {
+      opts = callback as GetTokenOptions;
+      callback = undefined;
+    }
+    opts = Object.assign(
+      {
+        forceRefresh: false,
+      },
+      opts,
+    );
+
+    if (callback) {
+      const cb = callback as GetTokenCallback;
+      this.#getTokenAsync(opts).then((t) => cb(null, t), callback);
+      return;
+    }
+
+    return this.#getTokenAsync(opts);
+  }
+
+  async #getTokenAsync(opts: GetTokenOptions): Promise<TokenData | undefined> {
+    if (this.#inFlightRequest && !opts.forceRefresh) {
+      return this.#inFlightRequest;
+    }
+
+    try {
+      return await (this.#inFlightRequest = this.#getTokenAsyncInner(opts));
+    } finally {
+      this.#inFlightRequest = undefined;
+    }
+  }
+
+  async #getTokenAsyncInner(
+    opts: GetTokenOptions,
+  ): Promise<TokenData | undefined> {
+    if (this.isTokenExpiring() === false && opts.forceRefresh === false) {
+      return Promise.resolve(this.rawToken!);
+    }
+    if (!this.key) {
+      throw new Error("No key or keyFile set.");
+    }
+    if (!this.iss) {
+      throw new ErrorWithCode("email is required.", "MISSING_CREDENTIALS");
+    }
+    const token = await this.#requestToken();
+    return token;
+  }
+
+  /**
+   * Revoke the token if one is set.
+   *
+   * @param callback The callback function.
+   */
+  revokeToken(): Promise<void>;
+  revokeToken(callback: (err?: Error) => void): void;
+  revokeToken(callback?: (err?: Error) => void): void | Promise<void> {
+    if (callback) {
+      this.#revokeTokenAsync().then(() => callback(), callback);
+      return;
+    }
+    return this.#revokeTokenAsync();
+  }
+
+  async #revokeTokenAsync() {
+    if (!this.accessToken) {
+      throw new Error("No token to revoke.");
+    }
+    const url = GOOGLE_REVOKE_TOKEN_URL + this.accessToken;
+    // await axios.get(url, { timeout: 10000 });
+    // uncomment below if prefer using fetch, but fetch will not follow HTTPS_PROXY
+    await fetch(url, { method: "GET" });
+
+    this.#configure({
+      email: this.iss,
+      sub: this.sub,
+      key: this.key,
+      scope: this.scope,
+      additionalClaims: this.additionalClaims,
+    });
+  }
+
+  /**
+   * Configure the GoogleToken for re-use.
+   * @param  {object} options Configuration object.
+   */
+  #configure(options: TokenOptions = { key: "" }) {
+    this.key = options.key;
+    this.rawToken = undefined;
+    this.iss = options.email || options.iss;
+    this.sub = options.sub;
+    this.additionalClaims = options.additionalClaims;
+    if (typeof options.scope === "object") {
+      this.scope = options.scope.join(" ");
+    } else {
+      this.scope = options.scope;
+    }
+    this.eagerRefreshThresholdMillis =
+      options.eagerRefreshThresholdMillis || 5 * 60 * 1000;
+  }
+
+  /**
+   * Request the token from Google.
+   */
+  async #requestToken(): Promise<TokenData | undefined> {
+    const iat = Math.floor(new Date().getTime() / 1000);
+    const additionalClaims = this.additionalClaims || {};
+    const payload = Object.assign(
+      {
+        iss: this.iss,
+        scope: this.scope,
+        aud: GOOGLE_TOKEN_URL,
+        exp: iat + 3600,
+        iat,
+        sub: this.sub,
+      },
+      additionalClaims,
+    );
+    const privateKey = await importPKCS8(this.key, "RS256");
+    const signedJWT = await new SignJWT(payload)
+      .setProtectedHeader({ alg: "RS256" })
+      .sign(privateKey);
+    const body = new URLSearchParams();
+    body.append("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer");
+    body.append("assertion", signedJWT);
+    try {
+      // const res = await axios.post<TokenData>(GOOGLE_TOKEN_URL, body, {
+      //   headers: { "Content-Type": "application/x-www-form-urlencoded" },
+      //   timeout: 15000,
+      //   validateStatus: (status) => {
+      //     return status >= 200 && status < 300;
+      //   },
+      // });
+      // this.rawToken = res.data;
+
+      // uncomment below if prefer using fetch, but fetch will not follow HTTPS_PROXY
+      const res = await fetch(GOOGLE_TOKEN_URL, {
+        method: "POST",
+        body,
+        headers: { "Content-Type": "application/x-www-form-urlencoded" },
+      });
+      this.rawToken = (await res.json()) as TokenData;
+
+      this.expiresAt =
+        this.rawToken.expires_in === null ||
+        this.rawToken.expires_in === undefined
+          ? undefined
+          : (iat + this.rawToken.expires_in!) * 1000;
+      return this.rawToken;
+    } catch (e) {
+      this.rawToken = undefined;
+      this.tokenExpires = undefined;
+      if (e instanceof Error) {
+        throw Error("failed to get token: " + e.message);
+      } else {
+        throw Error("failed to get token: " + String(e));
+      }
+    }
+  }
+}

app/utils/model.ts

@@ -200,5 +200,6 @@ export function isModelAvailableInServer(
 ) {
   const fullName = `${modelName}@${providerName}`;
   const modelTable = collectModelTable(DEFAULT_MODELS, customModels);
+  //TODO: this always return false, because providerName's first letter is capitalized, but the providerName in modelTable is lowercase
   return modelTable[fullName]?.available === false;
 }