File size: 4,104 Bytes
5bddbd8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
/* 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. */
export 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 number of prompts that have been sent with this key. */
promptCount: number;
/** The time at which this key was last used. */
lastUsed: number;
/** Key hash for displaying usage in the dashboard. */
hash: string;
};
const keyPool: Key[] = [];
function init() {
const keyString = process.env.OPENAI_KEY;
if (!keyString?.trim()) {
throw new Error("OPENAI_KEY environment variable is not set");
}
let keyList: KeySchema[];
try {
const decoded = Buffer.from(keyString, "base64").toString();
keyList = JSON.parse(decoded) as KeySchema[];
} catch (err) {
logger.info("OPENAI_KEY is not base64-encoded JSON, assuming 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 }];
}
for (const key of keyList) {
const newKey = {
...key,
isDisabled: false,
softLimit: 0,
hardLimit: 0,
systemHardLimit: 0,
usage: 0,
lastUsed: 0,
promptCount: 0,
hash: crypto
.createHash("sha256")
.update(key.key)
.digest("hex")
.slice(0, 6),
};
keyPool.push(newKey);
logger.info({ key: newKey.hash }, "Key added");
}
// TODO: check each key's usage upon startup.
}
function list() {
return keyPool.map((key) => ({
...key,
key: undefined,
}));
}
function disable(key: Key) {
const keyFromPool = keyPool.find((k) => k.key === key.key)!;
if (keyFromPool.isDisabled) return;
keyFromPool.isDisabled = true;
logger.warn({ key: key.hash }, "Key disabled");
}
function anyAvailable() {
return keyPool.some((key) => !key.isDisabled);
}
function get(model: string) {
const needsGpt4Key = model.startsWith("gpt-4");
const availableKeys = keyPool.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({ key: trialKeys[0].hash }, "Using trial key");
trialKeys[0].lastUsed = Date.now();
return trialKeys[0];
}
// Otherwise, return the oldest key
const oldestKey = availableKeys.sort((a, b) => a.lastUsed - b.lastUsed)[0];
logger.info({ key: oldestKey.hash }, "Assigning key to request.");
oldestKey.lastUsed = Date.now();
return { ...oldestKey };
}
function incrementPrompt(keyHash?: string) {
if (!keyHash) return;
const key = keyPool.find((k) => k.hash === keyHash)!;
key.promptCount++;
}
export const keys = { init, list, get, anyAvailable, disable, incrementPrompt };
|