From ffddf93a8f9e370b29f265984cd0e7b45e6aa440 Mon Sep 17 00:00:00 2001
From: lloydzhou <lloydzhou@qq.com>
Date: Fri, 12 Jul 2024 12:00:25 +0800
Subject: [PATCH] support cloudflare ai gateway

---
 app/api/anthropic/[...path]/route.ts | 20 +++++++++++---------
 app/api/common.ts                    |  5 ++++-
 app/client/platforms/anthropic.ts    |  4 +++-
 app/client/platforms/openai.ts       |  4 +++-
 app/utils/cloudflare.ts              | 26 ++++++++++++++++++++++++++
 5 files changed, 47 insertions(+), 12 deletions(-)
 create mode 100644 app/utils/cloudflare.ts

diff --git a/app/api/anthropic/[...path]/route.ts b/app/api/anthropic/[...path]/route.ts
index 4621f5783..daaeb45b7 100644
--- a/app/api/anthropic/[...path]/route.ts
+++ b/app/api/anthropic/[...path]/route.ts
@@ -11,6 +11,7 @@ import { prettyObject } from "@/app/utils/format";
 import { NextRequest, NextResponse } from "next/server";
 import { auth } from "../../auth";
 import { isModelAvailableInServer } from "@/app/utils/model";
+import { cloudflareAIGatewayUrl } from "@/app/utils/cloudflare";
 
 const ALLOWD_PATH = new Set([Anthropic.ChatPath, Anthropic.ChatPath1]);
 
@@ -114,7 +115,8 @@ async function request(req: NextRequest) {
     10 * 60 * 1000,
   );
 
-  const fetchUrl = `${baseUrl}${path}`;
+  // try rebuild url, when using cloudflare ai gateway in server
+  const fetchUrl = cloudflareAIGatewayUrl(`${baseUrl}${path}`);
 
   const fetchOptions: RequestInit = {
     headers: {
@@ -164,17 +166,17 @@ async function request(req: NextRequest) {
       console.error(`[Anthropic] filter`, e);
     }
   }
-  console.log("[Anthropic request]", fetchOptions.headers, req.method);
+  // console.log("[Anthropic request]", fetchOptions.headers, req.method);
   try {
     const res = await fetch(fetchUrl, fetchOptions);
 
-    console.log(
-      "[Anthropic response]",
-      res.status,
-      "   ",
-      res.headers,
-      res.url,
-    );
+    // console.log(
+    //   "[Anthropic response]",
+    //   res.status,
+    //   "   ",
+    //   res.headers,
+    //   res.url,
+    // );
     // to prevent browser prompt for credentials
     const newHeaders = new Headers(res.headers);
     newHeaders.delete("www-authenticate");
diff --git a/app/api/common.ts b/app/api/common.ts
index ddd0722fc..125a46fdc 100644
--- a/app/api/common.ts
+++ b/app/api/common.ts
@@ -14,6 +14,7 @@ import prisma from "@/lib/prisma";
 import { getTokenLength } from "@/lib/utils";
 
 import { isModelAvailableInServer } from "../utils/model";
+import { cloudflareAIGatewayUrl } from "../utils/cloudflare";
 
 const serverConfig = getServerSideConfig();
 
@@ -109,8 +110,10 @@ export async function requestOpenai(
     }
   }
 
-  const fetchUrl = `${baseUrl}/${path}`;
+  // const fetchUrl = `${baseUrl}/${path}`;
   const jsonBody = await req.json();
+  const fetchUrl = cloudflareAIGatewayUrl(`${baseUrl}/${path}`);
+  console.log("fetchUrl", fetchUrl);
   const fetchOptions: RequestInit = {
     headers: {
       "Content-Type": "application/json",
diff --git a/app/client/platforms/anthropic.ts b/app/client/platforms/anthropic.ts
index 460e8f0bc..bf8faf837 100644
--- a/app/client/platforms/anthropic.ts
+++ b/app/client/platforms/anthropic.ts
@@ -12,6 +12,7 @@ import {
 import Locale from "../../locales";
 import { prettyObject } from "@/app/utils/format";
 import { getMessageTextContent, isVisionModel } from "@/app/utils";
+import { cloudflareAIGatewayUrl } from "@/app/utils/cloudflare";
 
 export type MultiBlockContent = {
   type: "image" | "text";
@@ -375,7 +376,8 @@ export class ClaudeApi implements LLMApi {
 
     baseUrl = trimEnd(baseUrl, "/");
 
-    return `${baseUrl}/${path}`;
+    // try rebuild url, when using cloudflare ai gateway in client
+    return cloudflareAIGatewayUrl(`${baseUrl}/${path}`);
   }
 }
 
diff --git a/app/client/platforms/openai.ts b/app/client/platforms/openai.ts
index c35ec7cc2..83ce5e235 100644
--- a/app/client/platforms/openai.ts
+++ b/app/client/platforms/openai.ts
@@ -12,6 +12,7 @@ import {
 } from "@/app/constant";
 import { useAccessStore, useAppConfig, useChatStore } from "@/app/store";
 import { collectModelsWithDefaultModel } from "@/app/utils/model";
+import { cloudflareAIGatewayUrl } from "@/app/utils/cloudflare";
 
 import {
   ChatOptions,
@@ -95,7 +96,8 @@ export class ChatGPTApi implements LLMApi {
 
     console.log("[Proxy Endpoint] ", baseUrl, path);
 
-    return [baseUrl, path].join("/");
+    // try rebuild url, when using cloudflare ai gateway in client
+    return cloudflareAIGatewayUrl([baseUrl, path].join("/"));
   }
 
   extractMessage(res: any) {
diff --git a/app/utils/cloudflare.ts b/app/utils/cloudflare.ts
new file mode 100644
index 000000000..5094640fc
--- /dev/null
+++ b/app/utils/cloudflare.ts
@@ -0,0 +1,26 @@
+export function cloudflareAIGatewayUrl(fetchUrl: string) {
+  // rebuild fetchUrl, if using cloudflare ai gateway
+  // document: https://developers.cloudflare.com/ai-gateway/providers/openai/
+
+  const paths = fetchUrl.split("/");
+  if ("gateway.ai.cloudflare.com" == paths[2]) {
+    // is cloudflare.com ai gateway
+    // https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/azure-openai/{resource_name}/{deployment_name}/chat/completions?api-version=2023-05-15'
+    if ("azure-openai" == paths[6]) {
+      // is azure gateway
+      return paths.slice(0, 8).concat(paths.slice(-3)).join("/"); // rebuild ai gateway azure_url
+    }
+    // https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/openai/chat/completions
+    if ("openai" == paths[6]) {
+      // is openai gateway
+      return paths.slice(0, 7).concat(paths.slice(-2)).join("/"); // rebuild ai gateway openai_url
+    }
+    // https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/anthropic/v1/messages \
+    if ("anthropic" == paths[6]) {
+      // is anthropic gateway
+      return paths.slice(0, 7).concat(paths.slice(-2)).join("/"); // rebuild ai gateway anthropic_url
+    }
+    // TODO: Amazon Bedrock, Groq, HuggingFace...
+  }
+  return fetchUrl;
+}