Use Cloudflare AI gateway

.gitignore

@@ -1,3 +1,6 @@
+# Configuration
+wrangler.toml
+
 # Logs
 logs
 *.log

README.md

@@ -9,6 +9,16 @@ By accessing the OpenAI API through the Cloudflare worker, clients from all coun
 
 ## Usage
 
+### Setup AI Gateway
+
+Create an AI Gateway on Cloudflare. See <https://developers.cloudflare.com/ai-gateway/get-started/creating-gateway/>.
+
+After creating the AI Gateway, copy the OpenAI API endpoint URL.
+It should be displayed in the AI Gateway dashboard once you have created the gateway.
+
+In this project directory, make a copy of file `wrangler.toml.sample` and rename it to `wrangler.toml`.
+Edit the file and replace the existing value of `API_BASE` with the copied OpenAI API endpoint URL.
+
 ### [Setting up the secrets to the Cloudflare worker](https://developers.cloudflare.com/workers/platform/environment-variables/#add-secrets-to-your-project)
 
 To set up secrets for the Cloudflare worker, follow these steps:

src/key.mjs

@@ -9,15 +9,13 @@ program
   .version("0.0.1")
   .argument("<user-name>", "the user name of who is going to use this key")
   .action((userName) => {
-    userName = userName.replace(" ", "_");
-    console.log(`${prefix}${userName}-${tenRandomChars()}`);
+    console.log(generateOpenAiApiKey(userName));
   });
 
 program.parse();
 
 
-function tenRandomChars() {
-  const length = 10;
+function randomChars(length) {
   const characters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
   const charactersLength = characters.length;
   let result = '';
@@ -28,3 +26,7 @@ function tenRandomChars() {
   return result;
 }
 
+function generateOpenAiApiKey(userName) {
+  const first20chars = userName.replace(" ", "").padEnd(20, "0");
+  return `sk-${first20chars}T3BlbkFJ${randomChars(20)}`;
+}

src/worker.js

@@ -8,13 +8,17 @@
  * Learn more at https://developers.cloudflare.com/workers/
  */
 
-const api_base = "https://api.openai.com"
 export default {
 	async fetch(request, env, ctx) {
 		if (request.method === 'OPTIONS') {
 			return handleOPTIONS(request)
 		}
 
+		// Block request from China
+		// if (request.headers.get("CF-IPCountry") === "CN") {
+		//   return new Response("Forbidden country", { status: 503 });
+		// }
+
 		const accessKeys = env.ACCESS_KEYS.split(",");
 		const Authorization = request.headers.get("Authorization");
 		if (!Authorization) {
@@ -25,25 +29,26 @@ export default {
 			return new Response("Invalid Access Key", { status: 401 });
 		}
 
-		if (!accessKeys.includes(accessKey)) {
-			return new Response("Invalid Access Key", { status: 401 });
-		}
-
-		const newHeaders = new Headers(request.headers)
+		const newHeaders = new Headers()
 		newHeaders.set("Authorization", `Bearer ${env.OPENAI_API_KEY}`);
+		newHeaders.set("Content-Type", "application/json");
 
 		let pathname = new URL(request.url).pathname;
-		if (!pathname.startsWith("/v1")) {
+		if (pathname.startsWith("/v1")) {
 			// some sdk will put /v1 at the begin, some will not
-			pathname = "/v1" + pathname;
+			// remove the /v1 if it exists, because it isn't needed for AI gateway
+			pathname = pathname.slice(3);
 		}
-		const newRequest = new Request(api_base + pathname, {
+		const newRequest = new Request(env.API_BASE + pathname, {
 			method: request.method,
 			headers: newHeaders,
 			body: request.body
 		})
 
-		let response = await fetch(newRequest);
+		const response = await fetch(newRequest);
+		// if (response.status !== 200) {
+		// 	console.log(await response.text());
+		// }
 		return response;
 	},
 };

wrangler.toml.sample

@@ -0,0 +1,13 @@
+name = "openai-proxy"
+main = "src/worker.js"
+compatibility_date = "2023-05-25"
+
+[vars]
+# Copy the OpenAI API Endpoint from the AI Gateway page
+API_BASE = "https://gateway.ai.cloudflare.com/v1/xxxxxxxxxxxxxxx/name/openai"
+
+# The necessary secrets are:
+# - ACCESS_KEYS
+# - OPENAI_API_KEY
+# Run `wrangler secret put <NAME>` for each of these
+