nai-degen
131 lines
4.0 KiB
TypeScript
131 lines
4.0 KiB
TypeScript
/* 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 };
|