Merge c30fd63173 into 0c3d4462ca

.env.template

@@ -66,4 +66,14 @@ ANTHROPIC_API_VERSION=
 ANTHROPIC_URL=
 
 ### (optional)
-WHITE_WEBDAV_ENDPOINTS=
+WHITE_WEBDAV_ENDPOINTS=
+
+# (optional)
+# Default: Empty
+# Google Cloud Vertex AI full url, set if you want to use Vertex AI.
+VERTEX_AI_URL=
+
+# (optional)
+# Default: Empty
+# Text content of Google Cloud service account JSON key, set if you want to use Vertex AI.
+GOOGLE_CLOUD_JSON_KEY=''

app/api/anthropic.ts

@@ -11,6 +11,7 @@ import { NextRequest, NextResponse } from "next/server";
 import { auth } from "./auth";
 import { isModelAvailableInServer } from "@/app/utils/model";
 import { cloudflareAIGatewayUrl } from "@/app/utils/cloudflare";
+import { getGCloudToken } from "./common";
 
 const ALLOWD_PATH = new Set([Anthropic.ChatPath, Anthropic.ChatPath1]);
 
@@ -67,10 +68,20 @@ async function request(req: NextRequest) {
     serverConfig.anthropicApiKey ||
     "";
 
-  let path = `${req.nextUrl.pathname}`.replaceAll(ApiPath.Anthropic, "");
+  // adjust header and url when using vertex ai
+  if (serverConfig.isVertexAI) {
+    authHeaderName = "Authorization";
+    const gCloudToken = await getGCloudToken();
+    authValue = `Bearer ${gCloudToken}`;
+  }
 
-  let baseUrl =
+  let path = serverConfig.vertexAIUrl
-    serverConfig.anthropicUrl || serverConfig.baseUrl || ANTHROPIC_BASE_URL;
+    ? ""
+    : `${req.nextUrl.pathname}`.replaceAll(ApiPath.Anthropic, "");
+
+  let baseUrl = serverConfig.vertexAIUrl
+    ? serverConfig.vertexAIUrl
+    : serverConfig.anthropicUrl || serverConfig.baseUrl || ANTHROPIC_BASE_URL;
 
   if (!baseUrl.startsWith("http")) {
     baseUrl = `https://${baseUrl}`;
@@ -112,13 +123,16 @@ async function request(req: NextRequest) {
     signal: controller.signal,
   };
 
-  // #1815 try to refuse some request to some models
+  // #1815 try to refuse some request to some models or tick json body for vertex ai
   if (serverConfig.customModels && req.body) {
     try {
       const clonedBody = await req.text();
       fetchOptions.body = clonedBody;
 
-      const jsonBody = JSON.parse(clonedBody) as { model?: string };
+      const jsonBody = JSON.parse(clonedBody) as {
+        model?: string;
+        anthropic_version?: string;
+      };
 
       // not undefined and is false
       if (
@@ -138,6 +152,14 @@ async function request(req: NextRequest) {
           },
         );
       }
+
+      // tick json body for vertex ai and update fetch options
+      if (serverConfig.isVertexAI) {
+        delete jsonBody.model;
+        jsonBody.anthropic_version =
+          serverConfig.anthropicApiVersion || "vertex-2023-10-16";
+        fetchOptions.body = JSON.stringify(jsonBody);
+      }
     } catch (e) {
       console.error(`[Anthropic] filter`, e);
     }

app/api/common.ts

@@ -2,6 +2,7 @@ import { NextRequest, NextResponse } from "next/server";
 import { getServerSideConfig } from "../config/server";
 import { OPENAI_BASE_URL, ServiceProvider } from "../constant";
 import { cloudflareAIGatewayUrl } from "../utils/cloudflare";
+import { GoogleToken } from "../utils/gtoken";
 import { getModelProvider, isModelAvailableInServer } from "../utils/model";
 
 const serverConfig = getServerSideConfig();
@@ -185,3 +186,25 @@ export async function requestOpenai(req: NextRequest) {
     clearTimeout(timeoutId);
   }
 }
+
+let gTokenClient: GoogleToken | undefined;
+
+/**
+ * Get access token for google cloud,
+ * requires GOOGLE_CLOUD_JSON_KEY to be set
+ * @returns access token for google cloud
+ */
+export async function getGCloudToken() {
+  if (!gTokenClient) {
+    if (!serverConfig.googleCloudJsonKey)
+      throw new Error("GOOGLE_CLOUD_JSON_KEY is not set");
+    const keys = JSON.parse(serverConfig.googleCloudJsonKey);
+    gTokenClient = new GoogleToken({
+      email: keys.client_email,
+      key: keys.private_key,
+      scope: ["https://www.googleapis.com/auth/cloud-platform"],
+    });
+  }
+  const credentials = await gTokenClient?.getToken();
+  return credentials?.access_token;
+}

app/api/google.ts

@@ -3,6 +3,7 @@ import { auth } from "./auth";
 import { getServerSideConfig } from "@/app/config/server";
 import { ApiPath, GEMINI_BASE_URL, ModelProvider } from "@/app/constant";
 import { prettyObject } from "@/app/utils/format";
+import { getGCloudToken } from "./common";
 
 const serverConfig = getServerSideConfig();
 
@@ -29,7 +30,9 @@ export async function handle(
 
   const apiKey = token ? token : serverConfig.googleApiKey;
 
-  if (!apiKey) {
+  // When using Vertex AI, the API key is not required.
+  // Instead, a GCloud token will be used later in the request.
+  if (!apiKey && !serverConfig.isVertexAI) {
     return NextResponse.json(
       {
         error: true,
@@ -73,7 +76,9 @@ async function request(req: NextRequest, apiKey: string) {
 
   let baseUrl = serverConfig.googleUrl || GEMINI_BASE_URL;
 
-  let path = `${req.nextUrl.pathname}`.replaceAll(ApiPath.Google, "");
+  let path = serverConfig.vertexAIUrl
+    ? ""
+    : `${req.nextUrl.pathname}`.replaceAll(ApiPath.Google, "");
 
   if (!baseUrl.startsWith("http")) {
     baseUrl = `https://${baseUrl}`;
@@ -92,18 +97,30 @@ async function request(req: NextRequest, apiKey: string) {
     },
     10 * 60 * 1000,
   );
-  const fetchUrl = `${baseUrl}${path}${
+
-    req?.nextUrl?.searchParams?.get("alt") === "sse" ? "?alt=sse" : ""
+  let authHeaderName = "x-goog-api-key";
-  }`;
+  let authValue =
+    req.headers.get(authHeaderName) ||
+    (req.headers.get("Authorization") ?? "").replace("Bearer ", "");
+
+  // adjust header and url when use with vertex ai
+  if (serverConfig.vertexAIUrl) {
+    authHeaderName = "Authorization";
+    const gCloudToken = await getGCloudToken();
+    authValue = `Bearer ${gCloudToken}`;
+  }
+  const fetchUrl = serverConfig.vertexAIUrl
+    ? serverConfig.vertexAIUrl
+    : `${baseUrl}${path}${
+        req?.nextUrl?.searchParams?.get("alt") === "sse" ? "?alt=sse" : ""
+      }`;
 
   console.log("[Fetch Url] ", fetchUrl);
   const fetchOptions: RequestInit = {
     headers: {
       "Content-Type": "application/json",
       "Cache-Control": "no-store",
-      "x-goog-api-key":
+      [authHeaderName]: authValue,
-        req.headers.get("x-goog-api-key") ||
-        (req.headers.get("Authorization") ?? "").replace("Bearer ", ""),
     },
     method: req.method,
     body: req.body,

app/client/api.ts

@@ -317,7 +317,14 @@ export function getHeaders(ignoreHeaders: boolean = false) {
 
   if (bearerToken) {
     headers[authHeader] = bearerToken;
-  } else if (isEnabledAccessControl && validString(accessStore.accessCode)) {
+  }
+  // ensure access code is being sent when access control is enabled,
+  // this will fix an issue where the access code is not being sent when provider is google, azure or anthropic
+  if (
+    isEnabledAccessControl &&
+    validString(accessStore.accessCode) &&
+    authHeader !== "Authorization"
+  ) {
     headers["Authorization"] = getBearerToken(
       ACCESS_CODE_PREFIX + accessStore.accessCode,
     );

app/config/server.ts

@@ -44,6 +44,10 @@ declare global {
       ANTHROPIC_API_KEY?: string;
       ANTHROPIC_API_VERSION?: string;
 
+      // google cloud vertex ai only
+      VERTEX_AI_URL?: string; // https://{loc}-aiaiplatfor.googleapis.com/v1/{project}/locations/{loc}/models/{model}/versions/{version}:predict
+      GOOGLE_CLOUD_JSON_KEY?: string; // service account json key content
+
       // baidu only
       BAIDU_URL?: string;
       BAIDU_API_KEY?: string;
@@ -149,6 +153,7 @@ export const getServerSideConfig = () => {
   const isGoogle = !!process.env.GOOGLE_API_KEY;
   const isAnthropic = !!process.env.ANTHROPIC_API_KEY;
   const isTencent = !!process.env.TENCENT_API_KEY;
+  const isVertexAI = !!process.env.VERTEX_AI_URL;
 
   const isBaidu = !!process.env.BAIDU_API_KEY;
   const isBytedance = !!process.env.BYTEDANCE_API_KEY;
@@ -192,6 +197,10 @@ export const getServerSideConfig = () => {
     anthropicApiVersion: process.env.ANTHROPIC_API_VERSION,
     anthropicUrl: process.env.ANTHROPIC_URL,
 
+    isVertexAI,
+    vertexAIUrl: process.env.VERTEX_AI_URL,
+    googleCloudJsonKey: process.env.GOOGLE_CLOUD_JSON_KEY,
+
     isBaidu,
     baiduUrl: process.env.BAIDU_URL,
     baiduApiKey: getApiKey(process.env.BAIDU_API_KEY),

app/layout.tsx

@@ -7,6 +7,7 @@ import type { Metadata, Viewport } from "next";
 import { SpeedInsights } from "@vercel/speed-insights/next";
 import { getServerSideConfig } from "./config/server";
 import { GoogleTagManager, GoogleAnalytics } from "@next/third-parties/google";
+import { Analytics } from "@vercel/analytics/react";
 const serverConfig = getServerSideConfig();
 
 export const metadata: Metadata = {
@@ -53,6 +54,7 @@ export default function RootLayout({
         {serverConfig?.isVercel && (
           <>
             <SpeedInsights />
+            <Analytics />
           </>
         )}
         {serverConfig?.gtmId && (

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;
 }

package.json

@@ -24,7 +24,7 @@
     "@hello-pangea/dnd": "^16.5.0",
     "@next/third-parties": "^14.1.0",
     "@svgr/webpack": "^6.5.1",
-    "@vercel/analytics": "^0.1.11",
+    "@vercel/analytics": "^1.3.2",
     "@vercel/speed-insights": "^1.0.2",
     "axios": "^1.7.5",
     "clsx": "^2.1.1",
@@ -33,6 +33,7 @@
     "heic2any": "^0.0.4",
     "html-to-image": "^1.11.11",
     "idb-keyval": "^6.2.1",
+    "jose": "^5.9.6",
     "lodash-es": "^4.17.21",
     "markdown-to-txt": "^2.0.1",
     "mermaid": "^10.6.1",

test/model-provider.test.ts

@@ -1,4 +1,4 @@
-import { getModelProvider } from "../app/utils/model";
+import { getModelProvider, isModelAvailableInServer } from "../app/utils/model";
 
 describe("getModelProvider", () => {
   test("should return model and provider when input contains '@'", () => {
@@ -29,3 +29,15 @@ describe("getModelProvider", () => {
     expect(provider).toBeUndefined();
   });
 });
+
+describe("isModelAvailableInServer", () => {
+  test("works when model null", () => {
+    const jsonBody = JSON.parse("{}") as { model?: string };
+    const result = isModelAvailableInServer(
+      "gpt-3.5-turbo@OpenAI",
+      jsonBody.model as string,
+      "OpenAI",
+    );
+    expect(result).toBe(false);
+  });
+});

yarn.lock

@@ -2481,10 +2481,12 @@
     "@typescript-eslint/types" "6.4.0"
     eslint-visitor-keys "^3.4.1"
 
-"@vercel/analytics@^0.1.11":
+"@vercel/analytics@^1.3.2":
-  version "0.1.11"
+  version "1.3.2"
-  resolved "https://registry.yarnpkg.com/@vercel/analytics/-/analytics-0.1.11.tgz#727a0ac655a4a89104cdea3e6925476470299428"
+  resolved "https://mirrors.cloud.tencent.com/npm/@vercel/analytics/-/analytics-1.3.2.tgz#e7a8e22c83a7945e069960bab172308498b12b4e"
-  integrity sha512-mj5CPR02y0BRs1tN3oZcBNAX9a8NxsIUl9vElDPcqxnMfP0RbRc9fI9Ud7+QDg/1Izvt5uMumsr+6YsmVHcyuw==
+  integrity sha512-n/Ws7skBbW+fUBMeg+jrT30+GP00jTHvCcL4fuVrShuML0uveEV/4vVUdvqEVnDgXIGfLm0GXW5EID2mCcRXhg==
+  dependencies:
+    server-only "^0.0.1"
 
 "@vercel/speed-insights@^1.0.2":
   version "1.0.2"
@@ -5782,6 +5784,11 @@ jest@^29.7.0:
     import-local "^3.0.2"
     jest-cli "^29.7.0"
 
+jose@^5.9.6:
+  version "5.9.6"
+  resolved "https://mirrors.cloud.tencent.com/npm/jose/-/jose-5.9.6.tgz#77f1f901d88ebdc405e57cce08d2a91f47521883"
+  integrity sha512-AMlnetc9+CV9asI19zHmrgS/WYsWUwCn2R7RzlbJWD7F9eWYUTGyBmU9o6PxngtLGOiDGPRu+Uc4fhKzbpteZQ==
+
 "js-tokens@^3.0.0 || ^4.0.0", js-tokens@^4.0.0:
   version "4.0.0"
   resolved "https://registry.yarnpkg.com/js-tokens/-/js-tokens-4.0.0.tgz#19203fb59991df98e3a287050d4647cdeaf32499"
@@ -7569,6 +7576,11 @@ serialize-javascript@^6.0.1:
   dependencies:
     randombytes "^2.1.0"
 
+server-only@^0.0.1:
+  version "0.0.1"
+  resolved "https://mirrors.cloud.tencent.com/npm/server-only/-/server-only-0.0.1.tgz#0f366bb6afb618c37c9255a314535dc412cd1c9e"
+  integrity sha512-qepMx2JxAa5jjfzxG79yPPq+8BuFToHd1hm7kI+Z4zAq1ftQiP7HcxMhDDItrbtwVeLg/cY2JnKnrcFkmiswNA==
+
 shebang-command@^2.0.0:
   version "2.0.0"
   resolved "https://registry.yarnpkg.com/shebang-command/-/shebang-command-2.0.0.tgz#ccd0af4f8835fbdc265b82461aaf0c36663f34ea"