Merge c30fd63173 into 0c3d4462ca

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 =
-    serverConfig.anthropicUrl || serverConfig.baseUrl || ANTHROPIC_BASE_URL;
+  let path = serverConfig.vertexAIUrl
+    ? ""
+    : `${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":
-        req.headers.get("x-goog-api-key") ||
-        (req.headers.get("Authorization") ?? "").replace("Bearer ", ""),
+      [authHeaderName]: authValue,
     },
     method: req.method,
     body: req.body,