loads keys on startup

.env.example

@@ -0,0 +1,20 @@
+# Copy this file to .env and fill in the values.
+
+# Uncomment the following line and replace the value with your own secret key
+# to control access to the proxy server
+# PROXY_KEY=your-secret-key
+
+# Set your OpenAI API key below.
+OPENAI_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+
+# You can also set a base-64 encoded JSON array of keys matching this schema:
+# Trial keys will be prioritized. GPT4-enabled keys are necessary to use GPT-4.
+# For example:
+# [
+#  { "key": "your-openai-key-1", "isTrial": true, "isGpt4": false },
+#  { "key": "your-openai-key-2", "isTrial": false, "isGpt4": false },
+#  { "key": "your-openai-key-3", "isTrial": false, "isGpt4": true }
+# ]
+# Encoded in base-64, this would look like:
+# OPENAI_KEYS=WwogeyAia2V5IjogInlvdXItb3BlbmFpLWtleS0xIiwgImlzVHJpYWwiOiB0cnVlLCAiaXNHcHQ0IjogZmFsc2UgfSwKIHsgImtleSI6ICJ5b3VyLW9wZW5haS1rZXktMiIsICJpc1RyaWFsIjogZmFsc2UsICJpc0dwdDQiOiBmYWxzZSB9LAogeyAia2V5IjogInlvdXItb3BlbmFpLWtleS0zIiwgImlzVHJpYWwiOiBmYWxzZSwgImlzR3B0NCI6IHRydWUgfQpd
+  

.gitignore

@@ -1,3 +1,4 @@
 .vscode
+.env
 build
 node_modules

package-lock.json

@@ -10,6 +10,7 @@
       "license": "MIT",
       "dependencies": {
         "cors": "^2.8.5",
+        "dotenv": "^16.0.3",
         "express": "^4.18.2",
         "http-proxy-middleware": "^2.0.6",
         "pino-http": "^8.3.3"
@@ -495,6 +496,14 @@
         "node": ">=0.3.1"
       }
     },
+    "node_modules/dotenv": {
+      "version": "16.0.3",
+      "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-16.0.3.tgz",
+      "integrity": "sha512-7GO6HghkA5fYG9TYnNxi14/7K9f5occMlp3zXAuSxn7CKCxt9xbNWG7yF8hTCSUchlfWSe3uLmlPfigevRItzQ==",
+      "engines": {
+        "node": ">=12"
+      }
+    },
     "node_modules/ee-first": {
       "version": "1.1.1",
       "resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz",

package.json

@@ -11,6 +11,7 @@
   "license": "MIT",
   "dependencies": {
     "cors": "^2.8.5",
+    "dotenv": "^16.0.3",
     "express": "^4.18.2",
     "http-proxy-middleware": "^2.0.6",
     "pino-http": "^8.3.3"

src/keys.ts

@@ -36,13 +36,15 @@ const keyPool: Key[] = [];
 
 function init() {
   const keyString = process.env.OPENAI_KEY;
-  if (!keyString) {
+  if (!keyString?.trim()) {
     throw new Error("OPENAI_KEY environment variable is not set");
   }
   let keyList: KeySchema[];
   try {
-    keyList = JSON.parse(Buffer.from(keyString, "base64").toString());
+    const decoded = Buffer.from(keyString, "base64").toString();
+    keyList = JSON.parse(decoded) as KeySchema[];
   } catch (err) {
+    console.log("Key is not base64-encoded JSON, assuming it's a bare key");
     // We don't actually know if bare keys are paid/GPT-4 so we assume they are
     keyList = [{ key: keyString, isTrial: false, isGpt4: true }];
   }
@@ -63,7 +65,7 @@ function init() {
     };
     keyPool.push(newKey);
 
-    logger.info("Key added", { key: newKey.hash });
+    logger.info({ key: newKey.hash }, "Key added");
   }
   // TODO: check each key's usage upon startup.
 }
@@ -79,7 +81,7 @@ function disable(key: Key) {
   const keyFromPool = keyPool.find((k) => k.key === key.key)!;
   if (keyFromPool.isDisabled) return;
   keyFromPool.isDisabled = true;
-  logger.warn("Key disabled", { key: key.hash });
+  logger.warn({ key: key.hash }, "Key disabled");
 }
 
 function anyAvailable() {
@@ -104,13 +106,13 @@ function get(model: string) {
   // Prioritize trial keys
   const trialKeys = availableKeys.filter((key) => key.isTrial);
   if (trialKeys.length > 0) {
-    logger.info("Using trial key", { key: trialKeys[0].hash });
+    logger.info({ key: trialKeys[0].hash }, "Using trial key");
     return trialKeys[0];
   }
 
   // Otherwise, return the oldest key
   const oldestKey = availableKeys.sort((a, b) => a.lastUsed - b.lastUsed)[0];
-  logger.info("Using key", { key: oldestKey.hash });
+  logger.info({ key: oldestKey.hash }, "Assigning key to request.");
   oldestKey.lastUsed = Date.now();
   return { ...oldestKey };
 }

src/openai.ts

@@ -45,7 +45,7 @@ const handleResponse = (
       try {
         errorPayload = JSON.parse(body);
       } catch (err) {
-        logger.error(errorPayload.error, { error: err });
+        logger.error({ error: err }, errorPayload.error);
         res.status(statusCode).json(errorPayload);
         return;
       }
@@ -71,8 +71,8 @@ const handleResponse = (
           errorPayload.proxy_note = message;
         } else {
           logger.warn(
-            `OpenAI rate limit exceeded or model overloaded. Keyhash ${req.key?.hash}`,
-            { errorCode: errorPayload.error?.type }
+            { errorCode: errorPayload.error?.type },
+            `OpenAI rate limit exceeded or model overloaded. Keyhash ${req.key?.hash}`
           );
         }
       }

src/server.ts

@@ -1,8 +1,11 @@
+import dotenv from "dotenv";
+dotenv.config();
 import express from "express";
 import cors from "cors";
 import pinoHttp from "pino-http";
 
 import { logger } from "./logger";
+import { keys } from "./keys";
 import { proxy } from "./proxy";
 
 const PORT = process.env.PORT || 7860;
@@ -18,6 +21,10 @@ app.use(
 
 app.use("/", proxy);
 
+app.use((_req: unknown, res: express.Response) => {
+  res.status(404).json({ error: "Not found" });
+});
+
 app.use((err: any, _req: unknown, res: express.Response, _next: unknown) => {
   if (err.status) {
     res.status(err.status).json({ error: err.message });
@@ -29,4 +36,5 @@ app.use((err: any, _req: unknown, res: express.Response, _next: unknown) => {
 
 app.listen(PORT, () => {
   logger.info(`Server listening on port ${PORT}`);
+  keys.init();
 });