From e1606e3f66b35d5ffa4e58d186ad205ff59f791b Mon Sep 17 00:00:00 2001
From: nai-degen <44111-khanon@users.noreply.gitgud.io>
Date: Sat, 8 Apr 2023 04:12:15 -0500
Subject: [PATCH] adds basic key management

---
 src/keys.ts | 102 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 102 insertions(+)
 create mode 100644 src/keys.ts

diff --git a/src/keys.ts b/src/keys.ts
new file mode 100644
index 0000000..f2b8259
--- /dev/null
+++ b/src/keys.ts
@@ -0,0 +1,102 @@
+/* Manages OpenAI API keys. Tracks usage, disables expired keys, and provides
+round-robin access to keys. Keys are stored in the OPENAI_KEY environment
+variable, either as a single key, or a base64-encoded JSON array of keys.*/
+import { logger } from "./logger";
+import crypto from "crypto";
+
+/** Represents a key stored in the OPENAI_KEY environment variable. */
+type KeySchema = {
+  /** The OpenAI API key itself. */
+  key: string;
+  /** Whether this is a free trial key. These are prioritized over paid keys if they can fulfill the request. */
+  isTrial?: boolean;
+  /** Whether this key has been provisioned for GPT-4. */
+  isGpt4?: boolean;
+};
+
+/** Runtime information about a key. */
+type Key = KeySchema & {
+  /** Whether this key is currently disabled. We set this if we get a 429 or 401 response from OpenAI. */
+  isDisabled?: boolean;
+  /** Threshold at which a warning email will be sent by OpenAI. */
+  softLimit?: number;
+  /** Threshold at which the key will be disabled because it has reached the user-defined limit. */
+  hardLimit?: number;
+  /** The maximum quota allocated to this key by OpenAI. */
+  systemHardLimit?: number;
+  /** The current usage of this key. */
+  usage?: number;
+  /** The time at which this key was last used. */
+  lastUsed: number;
+  /** Key hash for displaying usage in the dashboard. */
+  hash: string;
+};
+
+const keys: Key[] = [];
+
+function init() {
+  const keyString = process.env.OPENAI_KEY;
+  if (!keyString) {
+    throw new Error("OPENAI_KEY environment variable is not set");
+  }
+  let keyList: KeySchema[];
+  try {
+    keyList = JSON.parse(Buffer.from(keyString, "base64").toString());
+  } catch (err) {
+    // We don't actually know if bare keys are paid/GPT-4 so we assume they are
+    keyList = [{ key: keyString, isTrial: false, isGpt4: true }];
+  }
+  for (const key of keyList) {
+    keys.push({
+      ...key,
+      isDisabled: false,
+      softLimit: 0,
+      hardLimit: 0,
+      systemHardLimit: 0,
+      usage: 0,
+      lastUsed: 0,
+      hash: crypto
+        .createHash("sha256")
+        .update(key.key)
+        .digest("hex")
+        .slice(0, 6),
+    });
+  }
+}
+
+function list() {
+  return keys.map((key) => ({
+    ...key,
+    key: undefined,
+  }));
+}
+
+function getKey(model: string) {
+  const needsGpt4Key = model.startsWith("gpt-4");
+  const availableKeys = keys.filter(
+    (key) => !key.isDisabled && (!needsGpt4Key || key.isGpt4)
+  );
+  if (availableKeys.length === 0) {
+    let message = "No keys available. Please add more keys.";
+    if (needsGpt4Key) {
+      message =
+        "No GPT-4 keys available. Please add more keys or use a non-GPT-4 model.";
+    }
+    logger.error(message);
+    throw new Error(message);
+  }
+
+  // Prioritize trial keys
+  const trialKeys = availableKeys.filter((key) => key.isTrial);
+  if (trialKeys.length > 0) {
+    logger.info("Using trial key", { key: trialKeys[0].hash });
+    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 });
+  return oldestKey;
+}
+
+export { init, list, getKey };