src/keys.ts

/* 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 };