code:
import crypto from "crypto";
import { Key, KeyProvider, Model } from "..";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { MistralAIModelFamily, getMistralAIModelFamily } from "../../models";
import { MistralAIKeyChecker } from "./checker";

export type MistralAIModel =
  | "mistral-tiny"
  | "mistral-small"
  | "mistral-medium";

export type MistralAIKeyUpdate = Omit<
  Partial<MistralAIKey>,
  | "key"
  | "hash"
  | "lastUsed"
  | "promptCount"
  | "rateLimitedAt"
  | "rateLimitedUntil"
>;

type MistralAIKeyUsage = {
  [K in MistralAIModelFamily as `${K}Tokens`]: number;
};

export interface MistralAIKey extends Key, MistralAIKeyUsage {
  readonly service: "mistral-ai";
  readonly modelFamilies: MistralAIModelFamily[];
  /** The time at which this key was last rate limited. */
  rateLimitedAt: number;
  /** The time until which this key is rate limited. */
  rateLimitedUntil: number;
}

/**
 * Upon being rate limited, a key will be locked out for this many milliseconds
 * while we wait for other concurrent requests to finish.
 */
const RATE_LIMIT_LOCKOUT = 2000;
/**
 * Upon assigning a key, we will wait this many milliseconds before allowing it
 * to be used again. This is to prevent the queue from flooding a key with too
 * many requests while we wait to learn whether previous ones succeeded.
 */
const KEY_REUSE_DELAY = 500;

export class MistralAIKeyProvider implements KeyProvider<MistralAIKey> {
  readonly service = "mistral-ai";

  private keys: MistralAIKey[] = [];
  private checker?: MistralAIKeyChecker;
  private log = logger.child({ module: "key-provider", service: this.service });

  constructor() {
    const keyConfig = config.mistralAIKey?.trim();
    if (!keyConfig) {
      this.log.warn(
        "MISTRAL_AI_KEY is not set. Mistral AI API will not be available."
      );
      return;
    }
    let bareKeys: string[];
    bareKeys = [...new Set(keyConfig.split(",").map((k) => k.trim()))];
    for (const key of bareKeys) {
      const newKey: MistralAIKey = {
        key,
        service: this.service,
        modelFamilies: ["mistral-tiny", "mistral-small", "mistral-medium"],
        isDisabled: false,
        isRevoked: false,
        promptCount: 0,
        lastUsed: 0,
        rateLimitedAt: 0,
        rateLimitedUntil: 0,
        hash: `mst-${crypto
          .createHash("sha256")
          .update(key)
          .digest("hex")
          .slice(0, 8)}`,
        lastChecked: 0,
        "mistral-tinyTokens": 0,
        "mistral-smallTokens": 0,
        "mistral-mediumTokens": 0,
      };
      this.keys.push(newKey);
    }
    this.log.info({ keyCount: this.keys.length }, "Loaded Mistral AI keys.");
  }

  public init() {
    if (config.checkKeys) {
      const updateFn = this.update.bind(this);
      this.checker = new MistralAIKeyChecker(this.keys, updateFn);
      this.checker.start();
    }
  }

  public list() {
    return this.keys.map((k) => Object.freeze({ ...k, key: undefined }));
  }

  public get(_model: Model) {
    const availableKeys = this.keys.filter((k) => !k.isDisabled);
    if (availableKeys.length === 0) {
      throw new Error("No Mistral AI keys available");
    }

    // (largely copied from the OpenAI provider, without trial key support)
    // Select a key, from highest priority to lowest priority:
    // 1. Keys which are not rate limited
    //    a. If all keys were rate limited recently, select the least-recently
    //       rate limited key.
    // 3. Keys which have not been used in the longest time

    const now = Date.now();

    const keysByPriority = availableKeys.sort((a, b) => {
      const aRateLimited = now - a.rateLimitedAt < RATE_LIMIT_LOCKOUT;
      const bRateLimited = now - b.rateLimitedAt < RATE_LIMIT_LOCKOUT;

      if (aRateLimited && !bRateLimited) return 1;
      if (!aRateLimited && bRateLimited) return -1;
      if (aRateLimited && bRateLimited) {
        return a.rateLimitedAt - b.rateLimitedAt;
      }

      return a.lastUsed - b.lastUsed;
    });

    const selectedKey = keysByPriority[0];
    selectedKey.lastUsed = now;
    this.throttle(selectedKey.hash);
    return { ...selectedKey };
  }

  public disable(key: MistralAIKey) {
    const keyFromPool = this.keys.find((k) => k.hash === key.hash);
    if (!keyFromPool || keyFromPool.isDisabled) return;
    keyFromPool.isDisabled = true;
    this.log.warn({ key: key.hash }, "Key disabled");
  }

  public update(hash: string, update: Partial<MistralAIKey>) {
    const keyFromPool = this.keys.find((k) => k.hash === hash)!;
    Object.assign(keyFromPool, { lastChecked: Date.now(), ...update });
  }

  public available() {
    return this.keys.filter((k) => !k.isDisabled).length;
  }

  public incrementUsage(hash: string, model: string, tokens: number) {
    const key = this.keys.find((k) => k.hash === hash);
    if (!key) return;
    key.promptCount++;
    const family = getMistralAIModelFamily(model);
    key[`${family}Tokens`] += tokens;
  }

  public getLockoutPeriod() {
    const activeKeys = this.keys.filter((k) => !k.isDisabled);
    // Don't lock out if there are no keys available or the queue will stall.
    // Just let it through so the add-key middleware can throw an error.
    if (activeKeys.length === 0) return 0;

    const now = Date.now();
    const rateLimitedKeys = activeKeys.filter((k) => now < k.rateLimitedUntil);
    const anyNotRateLimited = rateLimitedKeys.length < activeKeys.length;

    if (anyNotRateLimited) return 0;

    // If all keys are rate-limited, return the time until the first key is
    // ready.
    return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now));
  }

  /**
   * This is called when we receive a 429, which means there are already five
   * concurrent requests running on this key. We don't have any information on
   * when these requests will resolve, so all we can do is wait a bit and try
   * again. We will lock the key for 2 seconds after getting a 429 before
   * retrying in order to give the other requests a chance to finish.
   */
  public markRateLimited(keyHash: string) {
    this.log.debug({ key: keyHash }, "Key rate limited");
    const key = this.keys.find((k) => k.hash === keyHash)!;
    const now = Date.now();
    key.rateLimitedAt = now;
    key.rateLimitedUntil = now + RATE_LIMIT_LOCKOUT;
  }

  public recheck() {}

  /**
   * Applies a short artificial delay to the key upon dequeueing, in order to
   * prevent it from being immediately assigned to another request before the
   * current one can be dispatched.
   **/
  private throttle(hash: string) {
    const now = Date.now();
    const key = this.keys.find((k) => k.hash === hash)!;

    const currentRateLimit = key.rateLimitedUntil;
    const nextRateLimit = now + KEY_REUSE_DELAY;

    key.rateLimitedAt = now;
    key.rateLimitedUntil = Math.max(currentRateLimit, nextRateLimit);
  }
}


code:
import axios, { AxiosError } from "axios";
import type { OpenAIModelFamily } from "../../models";
import { KeyCheckerBase } from "../key-checker-base";
import type { OpenAIKey, OpenAIKeyProvider } from "./provider";
import { getOpenAIModelFamily } from "../../models";

const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds
const KEY_CHECK_PERIOD = 60 * 60 * 1000; // 1 hour
const POST_CHAT_COMPLETIONS_URL = "https://api.openai.com/v1/chat/completions";
const GET_MODELS_URL = "https://api.openai.com/v1/models";
const GET_ORGANIZATIONS_URL = "https://api.openai.com/v1/organizations";

type GetModelsResponse = {
  data: [{ id: string }];
};

type GetOrganizationsResponse = {
  data: [{ id: string; is_default: boolean }];
};

type OpenAIError = {
  error: { type: string; code: string; param: unknown; message: string };
};

type CloneFn = typeof OpenAIKeyProvider.prototype.clone;
type UpdateFn = typeof OpenAIKeyProvider.prototype.update;

export class OpenAIKeyChecker extends KeyCheckerBase<OpenAIKey> {
  private readonly cloneKey: CloneFn;

  constructor(keys: OpenAIKey[], cloneFn: CloneFn, updateKey: UpdateFn) {
    super(keys, {
      service: "openai",
      keyCheckPeriod: KEY_CHECK_PERIOD,
      minCheckInterval: MIN_CHECK_INTERVAL,
      recurringChecksEnabled: false,
      updateKey,
    });
    this.cloneKey = cloneFn;
  }

  protected async testKeyOrFail(key: OpenAIKey) {
    // We only need to check for provisioned models on the initial check.
    const isInitialCheck = !key.lastChecked;
    if (isInitialCheck) {
      const [provisionedModels, livenessTest] = await Promise.all([
        this.getProvisionedModels(key),
        this.testLiveness(key),
        this.maybeCreateOrganizationClones(key),
      ]);
      const updates = {
        modelFamilies: provisionedModels,
        isTrial: livenessTest.rateLimit <= 250,
      };
      this.updateKey(key.hash, updates);
    } else {
      // No updates needed as models and trial status generally don't change.
      const [_livenessTest] = await Promise.all([this.testLiveness(key)]);
      this.updateKey(key.hash, {});
    }
    this.log.info(
      { key: key.hash, models: key.modelFamilies, trial: key.isTrial },
      "Checked key."
    );
  }

  private async getProvisionedModels(
    key: OpenAIKey
  ): Promise<OpenAIModelFamily[]> {
    const opts = { headers: OpenAIKeyChecker.getHeaders(key) };
    const { data } = await axios.get<GetModelsResponse>(GET_MODELS_URL, opts);
    const models = data.data;

    const families = new Set<OpenAIModelFamily>();
    models.forEach(({ id }) => families.add(getOpenAIModelFamily(id, "turbo")));

    // as of 2023-11-18, many keys no longer return the dalle3 model but still
    // have access to it via the api for whatever reason.
    // if (families.has("dall-e") && !models.find(({ id }) => id === "dall-e-3")) {
    //   families.delete("dall-e");
    // }

    // We want to update the key's model families here, but we don't want to
    // update its `lastChecked` timestamp because we need to let the liveness
    // check run before we can consider the key checked.

    const familiesArray = [...families];
    const keyFromPool = this.keys.find((k) => k.hash === key.hash)!;
    this.updateKey(key.hash, {
      modelFamilies: familiesArray,
      lastChecked: keyFromPool.lastChecked,
    });
    return familiesArray;
  }

  private async maybeCreateOrganizationClones(key: OpenAIKey) {
    if (key.organizationId) return; // already cloned
    const opts = { headers: { Authorization: `Bearer ${key.key}` } };
    const { data } = await axios.get<GetOrganizationsResponse>(
      GET_ORGANIZATIONS_URL,
      opts
    );
    const organizations = data.data;
    const defaultOrg = organizations.find(({ is_default }) => is_default);
    this.updateKey(key.hash, { organizationId: defaultOrg?.id });
    if (organizations.length <= 1) return undefined;

    this.log.info(
      { parent: key.hash, organizations: organizations.map((org) => org.id) },
      "Key is associated with multiple organizations; cloning key for each organization."
    );

    const ids = organizations
      .filter(({ is_default }) => !is_default)
      .map(({ id }) => id);
    this.cloneKey(key.hash, ids);

    // It's possible that the keychecker may be stopped if all non-cloned keys
    // happened to be unusable, in which case this clnoe will never be checked
    // unless we restart the keychecker.
    if (!this.timeout) {
      this.log.warn(
        { parent: key.hash },
        "Restarting key checker to check cloned keys."
      );
      this.scheduleNextCheck();
    }
  }

  protected handleAxiosError(key: OpenAIKey, error: AxiosError) {
    if (error.response && OpenAIKeyChecker.errorIsOpenAIError(error)) {
      const { status, data } = error.response;
      if (status === 401) {
        this.log.warn(
          { key: key.hash, error: data },
          "Key is invalid or revoked. Disabling key."
        );
        this.updateKey(key.hash, {
          isDisabled: true,
          isRevoked: true,
          modelFamilies: ["turbo"],
        });
      } else if (status === 429) {
        switch (data.error.type) {
          case "insufficient_quota":
          case "billing_not_active":
          case "access_terminated":
            const isRevoked = data.error.type === "access_terminated";
            const isOverQuota = !isRevoked;
            const modelFamilies: OpenAIModelFamily[] = isRevoked
              ? ["turbo"]
              : key.modelFamilies;
            this.log.warn(
              { key: key.hash, rateLimitType: data.error.type, error: data },
              "Key returned a non-transient 429 error. Disabling key."
            );
            this.updateKey(key.hash, {
              isDisabled: true,
              isRevoked,
              isOverQuota,
              modelFamilies,
            });
            break;
          case "requests":
            // If we hit the text completion rate limit on a trial key, it is
            // likely being used by many proxies. We will disable the key since
            // it's just going to be constantly rate limited.
            const isTrial =
              Number(error.response.headers["x-ratelimit-limit-requests"]) <=
              250;

            if (isTrial) {
              this.log.warn(
                { key: key.hash, error: data },
                "Trial key is rate limited on text completion endpoint. This indicates the key is being used by several proxies at once and is not likely to be usable. Disabling key."
              );
              this.updateKey(key.hash, {
                isTrial,
                isDisabled: true,
                isOverQuota: true,
                modelFamilies: ["turbo"],
                lastChecked: Date.now(),
              });
            } else {
              this.log.warn(
                { key: key.hash, error: data },
                "Non-trial key is rate limited on text completion endpoint. This is unusual and may indicate a bug. Assuming key is operational."
              );
              this.updateKey(key.hash, { lastChecked: Date.now() });
            }
            break;
          case "tokens":
            // Hitting a token rate limit, even on a trial key, actually implies
            // that the key is valid and can generate completions, so we will
            // treat this as effectively a successful `testLiveness` call.
            this.log.info(
              { key: key.hash },
              "Key is currently `tokens` rate limited; assuming it is operational."
            );
            this.updateKey(key.hash, { lastChecked: Date.now() });
            break;
          default:
            this.log.error(
              { key: key.hash, rateLimitType: data.error.type, error: data },
              "Encountered unexpected rate limit error class while checking key. This may indicate a change in the API; please report this."
            );
            // We don't know what this error means, so we just let the key
            // through and maybe it will fail when someone tries to use it.
            this.updateKey(key.hash, { lastChecked: Date.now() });
        }
      } else {
        this.log.error(
          { key: key.hash, status, error: data },
          "Encountered unexpected error status while checking key. This may indicate a change in the API; please report this."
        );
        this.updateKey(key.hash, { lastChecked: Date.now() });
      }
      return;
    }
    this.log.error(
      { key: key.hash, error: error.message },
      "Network error while checking key; trying this key again in a minute."
    );
    const oneMinute = 60 * 1000;
    const next = Date.now() - (KEY_CHECK_PERIOD - oneMinute);
    this.updateKey(key.hash, { lastChecked: next });
  }

  /**
   * Tests whether the key is valid and has quota remaining. The request we send
   * is actually not valid, but keys which are revoked or out of quota will fail
   * with a 401 or 429 error instead of the expected 400 Bad Request error.
   * This lets us avoid test keys without spending any quota.
   *
   * We use the rate limit header to determine whether it's a trial key.
   */
  private async testLiveness(key: OpenAIKey): Promise<{ rateLimit: number }> {
    // What the hell this is doing:

    // OpenAI enforces separate rate limits for chat and text completions. Trial
    // keys have extremely low rate limits of 200 per day per API type. In order
    // to avoid wasting more valuable chat quota, we send an (invalid) chat
    // request to Babbage (a text completion model). Even though our request is
    // to the chat endpoint, we get text rate limit headers back because the
    // requested model determines the rate limit used, not the endpoint.

    // Once we have headers, we can determine:
    // 1. Is the key revoked? (401, OAI doesn't even validate the request)
    // 2. Is the key out of quota? (400, OAI will still validate the request)
    // 3. Is the key a trial key? (400, x-ratelimit-limit-requests: 200)

    // This might still cause issues if too many proxies are running a train on
    // the same trial key and even the text completion quota is exhausted, but
    // it should work better than the alternative.

    const payload = {
      model: "babbage-002",
      max_tokens: -1,
      messages: [{ role: "user", content: "" }],
    };
    const { headers, data } = await axios.post<OpenAIError>(
      POST_CHAT_COMPLETIONS_URL,
      payload,
      {
        headers: OpenAIKeyChecker.getHeaders(key),
        validateStatus: (status) => status === 400,
      }
    );
    const rateLimitHeader = headers["x-ratelimit-limit-requests"];
    const rateLimit = parseInt(rateLimitHeader) || 3500; // trials have 200

    // invalid_request_error is the expected error
    if (data.error.type !== "invalid_request_error") {
      this.log.warn(
        { key: key.hash, error: data },
        "Unexpected 400 error class while checking key; assuming key is valid, but this may indicate a change in the API."
      );
    }
    return { rateLimit };
  }

  static errorIsOpenAIError(
    error: AxiosError
  ): error is AxiosError<OpenAIError> {
    const data = error.response?.data as any;
    return data?.error?.type;
  }

  static getHeaders(key: OpenAIKey) {
    return {
      Authorization: `Bearer ${key.key}`,
      ...(key.organizationId && { "OpenAI-Organization": key.organizationId }),
    };
  }
}

license: unknown


code:
/* 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 as a comma-separated list of keys. */
import crypto from "crypto";
import http from "http";
import { Key, KeyProvider, Model } from "../index";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { OpenAIKeyChecker } from "./checker";
import { getOpenAIModelFamily, OpenAIModelFamily } from "../../models";

export type OpenAIModel =
  | "gpt-3.5-turbo"
  | "gpt-3.5-turbo-instruct"
  | "gpt-4"
  | "gpt-4-32k"
  | "gpt-4-1106"
  | "text-embedding-ada-002"
  | "dall-e-2"
  | "dall-e-3"

// Flattening model families instead of using a nested object for easier
// cloning.
type OpenAIKeyUsage = {
  [K in OpenAIModelFamily as `${K}Tokens`]: number;
};

export interface OpenAIKey extends Key, OpenAIKeyUsage {
  readonly service: "openai";
  modelFamilies: OpenAIModelFamily[];
  /**
   * Some keys are assigned to multiple organizations, each with their own quota
   * limits. We clone the key for each organization and track usage/disabled
   * status separately.
   */
  organizationId?: string;
  /** Whether this is a free trial key. These are prioritized over paid keys if they can fulfill the request. */
  isTrial: boolean;
  /** Set when key check returns a non-transient 429. */
  isOverQuota: boolean;
  /** The time at which this key was last rate limited. */
  rateLimitedAt: number;
  /**
   * Last known X-RateLimit-Requests-Reset header from OpenAI, converted to a
   * number.
   * Formatted as a `\d+(m|s)` string denoting the time until the limit resets.
   * Specifically, it seems to indicate the time until the key's quota will be
   * fully restored; the key may be usable before this time as the limit is a
   * rolling window.
   *
   * Requests which return a 429 do not count against the quota.
   *
   * Requests which fail for other reasons (e.g. 401) count against the quota.
   */
  rateLimitRequestsReset: number;
  /**
   * Last known X-RateLimit-Tokens-Reset header from OpenAI, converted to a
   * number.
   * Appears to follow the same format as `rateLimitRequestsReset`.
   *
   * Requests which fail do not count against the quota as they do not consume
   * tokens.
   */
  rateLimitTokensReset: number;
  /**
   * This key's maximum request rate for GPT-4, per minute.
   */
  gpt4Rpm: number;
}

export type OpenAIKeyUpdate = Omit<
  Partial<OpenAIKey>,
  "key" | "hash" | "promptCount"
>;

/**
 * Upon assigning a key, we will wait this many milliseconds before allowing it
 * to be used again. This is to prevent the queue from flooding a key with too
 * many requests while we wait to learn whether previous ones succeeded.
 */
const KEY_REUSE_DELAY = 1000;

export class OpenAIKeyProvider implements KeyProvider<OpenAIKey> {
  readonly service = "openai" as const;

  private keys: OpenAIKey[] = [];
  private checker?: OpenAIKeyChecker;
  private log = logger.child({ module: "key-provider", service: this.service });

  constructor() {
    const keyString = config.openaiKey?.trim();
    if (!keyString) {
      this.log.warn("OPENAI_KEY is not set. OpenAI API will not be available.");
      return;
    }
    let bareKeys: string[];
    bareKeys = keyString.split(",").map((k) => k.trim());
    bareKeys = [...new Set(bareKeys)];
    for (const k of bareKeys) {
      const newKey: OpenAIKey = {
        key: k,
        service: "openai" as const,
        modelFamilies: [
          "turbo" as const,
          "gpt4" as const,
          "gpt4-turbo" as const,
        ],
        isTrial: false,
        isDisabled: false,
        isRevoked: false,
        isOverQuota: false,
        lastUsed: 0,
        lastChecked: 0,
        promptCount: 0,
        hash: `oai-${crypto
          .createHash("sha256")
          .update(k)
          .digest("hex")
          .slice(0, 8)}`,
        rateLimitedAt: 0,
        rateLimitRequestsReset: 0,
        rateLimitTokensReset: 0,
        turboTokens: 0,
        gpt4Tokens: 0,
        "gpt4-32kTokens": 0,
        "gpt4-turboTokens": 0,
        "dall-eTokens": 0,
        gpt4Rpm: 0,
      };
      this.keys.push(newKey);
    }
    this.log.info({ keyCount: this.keys.length }, "Loaded OpenAI keys.");
  }

  public init() {
    if (config.checkKeys) {
      const cloneFn = this.clone.bind(this);
      const updateFn = this.update.bind(this);
      this.checker = new OpenAIKeyChecker(this.keys, cloneFn, updateFn);
      this.checker.start();
    }
  }

  /**
   * Returns a list of all keys, with the key field removed.
   * Don't mutate returned keys, use a KeyPool method instead.
   **/
  public list() {
    return this.keys.map((key) => {
      return Object.freeze({
        ...key,
        key: undefined,
      });
    });
  }

  public get(model: Model) {
    const neededFamily = getOpenAIModelFamily(model);
    const excludeTrials = model === "text-embedding-ada-002";

    const availableKeys = this.keys.filter(
      // Allow keys which
      (key) =>
        !key.isDisabled && // are not disabled
        key.modelFamilies.includes(neededFamily) && // have access to the model
        (!excludeTrials || !key.isTrial) // and are not trials (if applicable)
    );

    if (availableKeys.length === 0) {
      throw new Error(`No keys available for model family '${neededFamily}'.`);
    }

    // Select a key, from highest priority to lowest priority:
    // 1. Keys which are not rate limited
    //    a. We ignore rate limits from >30 seconds ago
    //    b. If all keys were rate limited in the last minute, select the
    //       least recently rate limited key
    // 2. Keys which are trials
    // 3. Keys which do *not* have access to GPT-4-32k
    // 4. Keys which have not been used in the longest time

    const now = Date.now();
    const rateLimitThreshold = 30 * 1000;

    const keysByPriority = availableKeys.sort((a, b) => {
      // TODO: this isn't quite right; keys are briefly artificially rate-
      // limited when they are selected, so this will deprioritize keys that
      // may not actually be limited, simply because they were used recently.
      // This should be adjusted to use a new `rateLimitedUntil` field instead
      // of `rateLimitedAt`.
      const aRateLimited = now - a.rateLimitedAt < rateLimitThreshold;
      const bRateLimited = now - b.rateLimitedAt < rateLimitThreshold;

      if (aRateLimited && !bRateLimited) return 1;
      if (!aRateLimited && bRateLimited) return -1;
      if (aRateLimited && bRateLimited) {
        return a.rateLimitedAt - b.rateLimitedAt;
      }
      // Neither key is rate limited, continue

      if (a.isTrial && !b.isTrial) return -1;
      if (!a.isTrial && b.isTrial) return 1;
      // Neither or both keys are trials, continue

      const aHas32k = a.modelFamilies.includes("gpt4-32k");
      const bHas32k = b.modelFamilies.includes("gpt4-32k");
      if (aHas32k && !bHas32k) return 1;
      if (!aHas32k && bHas32k) return -1;
      // Neither or both keys have 32k, continue

      return a.lastUsed - b.lastUsed;
    });

    const selectedKey = keysByPriority[0];
    selectedKey.lastUsed = now;
    this.throttle(selectedKey.hash);
    return { ...selectedKey };
  }

  /** Called by the key checker to update key information. */
  public update(keyHash: string, update: OpenAIKeyUpdate) {
    const keyFromPool = this.keys.find((k) => k.hash === keyHash)!;
    Object.assign(keyFromPool, { lastChecked: Date.now(), ...update });
  }

  /** Called by the key checker to create clones of keys for the given orgs. */
  public clone(keyHash: string, newOrgIds: string[]) {
    const keyFromPool = this.keys.find((k) => k.hash === keyHash)!;
    const clones = newOrgIds.map((orgId) => {
      const clone: OpenAIKey = {
        ...keyFromPool,
        organizationId: orgId,
        isDisabled: false,
        isRevoked: false,
        isOverQuota: false,
        hash: `oai-${crypto
          .createHash("sha256")
          .update(keyFromPool.key + orgId)
          .digest("hex")
          .slice(0, 8)}`,
        lastChecked: 0, // Force re-check in case the org has different models
      };
      this.log.info(
        { cloneHash: clone.hash, parentHash: keyFromPool.hash, orgId },
        "Cloned organization key"
      );
      return clone;
    });
    this.keys.push(...clones);
  }

  /** Disables a key, or does nothing if the key isn't in this pool. */
  public disable(key: Key) {
    const keyFromPool = this.keys.find((k) => k.hash === key.hash);
    if (!keyFromPool || keyFromPool.isDisabled) return;
    this.update(key.hash, { isDisabled: true });
    this.log.warn({ key: key.hash }, "Key disabled");
  }

  public available() {
    return this.keys.filter((k) => !k.isDisabled).length;
  }

  /**
   * Given a model, returns the period until a key will be available to service
   * the request, or returns 0 if a key is ready immediately.
   */
  public getLockoutPeriod(family: OpenAIModelFamily): number {
    const activeKeys = this.keys.filter(
      (key) => !key.isDisabled && key.modelFamilies.includes(family)
    );

    // Don't lock out if there are no keys available or the queue will stall.
    // Just let it through so the add-key middleware can throw an error.
    if (activeKeys.length === 0) return 0;

    // A key is rate-limited if its `rateLimitedAt` plus the greater of its
    // `rateLimitRequestsReset` and `rateLimitTokensReset` is after the
    // current time.

    // If there are any keys that are not rate-limited, we can fulfill requests.
    const now = Date.now();
    const rateLimitedKeys = activeKeys.filter((key) => {
      const resetTime = Math.max(
        key.rateLimitRequestsReset,
        key.rateLimitTokensReset
      );
      return now < key.rateLimitedAt + Math.min(20000, resetTime);
    }).length;
    const anyNotRateLimited = rateLimitedKeys < activeKeys.length;

    if (anyNotRateLimited) {
      return 0;
    }

    // If all keys are rate-limited, return the time until the first key is
    // ready. We don't want to wait longer than 10 seconds because rate limits
    // are a rolling window and keys may become available sooner than the stated
    // reset time.
    return Math.min(
      ...activeKeys.map((key) => {
        const resetTime = Math.max(
          key.rateLimitRequestsReset,
          key.rateLimitTokensReset
        );
        return key.rateLimitedAt + Math.min(20000, resetTime) - now;
      })
    );
  }

  public markRateLimited(keyHash: string) {
    this.log.debug({ key: keyHash }, "Key rate limited");
    const key = this.keys.find((k) => k.hash === keyHash)!;
    key.rateLimitedAt = Date.now();
    // DALL-E requests do not send headers telling us when the rate limit will
    // be reset so we need to set a fallback value here.  Other models will have
    // this overwritten by the `updateRateLimits` method.
    key.rateLimitRequestsReset = 20000;
  }

  public incrementUsage(keyHash: string, model: string, tokens: number) {
    const key = this.keys.find((k) => k.hash === keyHash);
    if (!key) return;
    key.promptCount++;
    key[`${getOpenAIModelFamily(model)}Tokens`] += tokens;
  }

  public updateRateLimits(keyHash: string, headers: http.IncomingHttpHeaders) {
    const key = this.keys.find((k) => k.hash === keyHash)!;
    const requestsReset = headers["x-ratelimit-reset-requests"];
    const tokensReset = headers["x-ratelimit-reset-tokens"];

    if (typeof requestsReset === "string") {
      key.rateLimitRequestsReset = getResetDurationMillis(requestsReset);
    }

    if (typeof tokensReset === "string") {
      key.rateLimitTokensReset = getResetDurationMillis(tokensReset);
    }

    if (!requestsReset && !tokensReset) {
      this.log.warn({ key: key.hash }, `No ratelimit headers; skipping update`);
      return;
    }
  }

  public recheck() {
    this.keys.forEach((key) => {
      this.update(key.hash, {
        isRevoked: false,
        isOverQuota: false,
        isDisabled: false,
        lastChecked: 0,
      });
    });
    this.checker?.scheduleNextCheck();
  }

  /**
   * Called when a key is selected for a request, briefly disabling it to
   * avoid spamming the API with requests while we wait to learn whether this
   * key is already rate limited.
   */
  private throttle(hash: string) {
    const now = Date.now();
    const key = this.keys.find((k) => k.hash === hash)!;

    const currentRateLimit =
      Math.max(key.rateLimitRequestsReset, key.rateLimitTokensReset) +
      key.rateLimitedAt;
    const nextRateLimit = now + KEY_REUSE_DELAY;

    // Don't throttle if the key is already naturally rate limited.
    if (currentRateLimit > nextRateLimit) return;

    key.rateLimitedAt = Date.now();
    key.rateLimitRequestsReset = KEY_REUSE_DELAY;
  }
}

// wip
function calculateRequestsPerMinute(headers: http.IncomingHttpHeaders) {
  const requestsLimit = headers["x-ratelimit-limit-requests"];
  const requestsReset = headers["x-ratelimit-reset-requests"];

  if (typeof requestsLimit !== "string" || typeof requestsReset !== "string") {
    return 0;
  }

  const limit = parseInt(requestsLimit, 10);
  const reset = getResetDurationMillis(requestsReset);

  // If `reset` is less than one minute, OpenAI specifies the `limit` as an
  // integer representing requests per minute.  Otherwise it actually means the
  // requests per day.
  const isPerMinute = reset < 60000;
  if (isPerMinute) return limit;
  return limit / 1440;
}

/**
 * Converts reset string ("14m25s", "21.0032s", "14ms" or "21ms") to a number of
 * milliseconds.
 **/
function getResetDurationMillis(resetDuration?: string): number {
  const match = resetDuration?.match(
    /(?:(\d+)m(?!s))?(?:(\d+(?:\.\d+)?)s)?(?:(\d+)ms)?/
  );

  if (match) {
    const [, minutes, seconds, milliseconds] = match.map(Number);

    const minutesToMillis = (minutes || 0) * 60 * 1000;
    const secondsToMillis = (seconds || 0) * 1000;
    const millisecondsValue = milliseconds || 0;

    return minutesToMillis + secondsToMillis + millisecondsValue;
  }

  return 0;
}


code:
// Don't import any other project files here as this is one of the first modules
// loaded and it will cause circular imports.

import pino from "pino";
import type { Request } from "express";

/**
 * The service that a model is hosted on. Distinct from `APIFormat` because some
 * services have interoperable APIs (eg Anthropic/AWS, OpenAI/Azure).
 */
export type LLMService =
  | "openai"
  | "anthropic"
  | "google-ai"
  | "mistral-ai"
  | "aws"
  | "azure";

export type OpenAIModelFamily =
  | "turbo"
  | "gpt4"
  | "gpt4-32k"
  | "gpt4-turbo"
  | "dall-e";
export type AnthropicModelFamily = "claude";
export type GoogleAIModelFamily = "gemini-pro";
export type MistralAIModelFamily =
  | "mistral-tiny"
  | "mistral-small"
  | "mistral-medium";
export type AwsBedrockModelFamily = "aws-claude";
export type AzureOpenAIModelFamily = `azure-${Exclude<
  OpenAIModelFamily,
  "dall-e"
>}`;
export type ModelFamily =
  | OpenAIModelFamily
  | AnthropicModelFamily
  | GoogleAIModelFamily
  | MistralAIModelFamily
  | AwsBedrockModelFamily
  | AzureOpenAIModelFamily;

export const MODEL_FAMILIES = (<A extends readonly ModelFamily[]>(
  arr: A & ([ModelFamily] extends [A[number]] ? unknown : never)
) => arr)([
  "turbo",
  "gpt4",
  "gpt4-32k",
  "gpt4-turbo",
  "dall-e",
  "claude",
  "gemini-pro",
  "mistral-tiny",
  "mistral-small",
  "mistral-medium",
  "aws-claude",
  "azure-turbo",
  "azure-gpt4",
  "azure-gpt4-32k",
  "azure-gpt4-turbo",
] as const);

export const LLM_SERVICES = (<A extends readonly LLMService[]>(
  arr: A & ([LLMService] extends [A[number]] ? unknown : never)
) => arr)([
  "openai",
  "anthropic",
  "google-ai",
  "mistral-ai",
  "aws",
  "azure",
] as const);

export const OPENAI_MODEL_FAMILY_MAP: { [regex: string]: OpenAIModelFamily } = {
  "^gpt-4-1106(-preview)?$": "gpt4-turbo",
  "^gpt-4(-\\d{4})?-vision(-preview)?$": "gpt4-turbo",
  "^gpt-4-32k-\\d{4}$": "gpt4-32k",
  "^gpt-4-32k$": "gpt4-32k",
  "^gpt-4-\\d{4}$": "gpt4",
  "^gpt-4$": "gpt4",
  "^gpt-3.5-turbo": "turbo",
  "^text-embedding-ada-002$": "turbo",
  "^dall-e-\\d{1}$": "dall-e",
};

export const MODEL_FAMILY_SERVICE: {
  [f in ModelFamily]: LLMService;
} = {
  turbo: "openai",
  gpt4: "openai",
  "gpt4-turbo": "openai",
  "gpt4-32k": "openai",
  "dall-e": "openai",
  claude: "anthropic",
  "aws-claude": "aws",
  "azure-turbo": "azure",
  "azure-gpt4": "azure",
  "azure-gpt4-32k": "azure",
  "azure-gpt4-turbo": "azure",
  "gemini-pro": "google-ai",
  "mistral-tiny": "mistral-ai",
  "mistral-small": "mistral-ai",
  "mistral-medium": "mistral-ai",
};

pino({ level: "debug" }).child({ module: "startup" });

export function getOpenAIModelFamily(
  model: string,
  defaultFamily: OpenAIModelFamily = "gpt4"
): OpenAIModelFamily {
  for (const [regex, family] of Object.entries(OPENAI_MODEL_FAMILY_MAP)) {
    if (model.match(regex)) return family;
  }
  return defaultFamily;
}

export function getClaudeModelFamily(model: string): ModelFamily {
  if (model.startsWith("anthropic.")) return getAwsBedrockModelFamily(model);
  return "claude";
}

export function getGoogleAIModelFamily(_model: string): ModelFamily {
  return "gemini-pro";
}

export function getMistralAIModelFamily(model: string): MistralAIModelFamily {
  switch (model) {
    case "mistral-tiny":
    case "mistral-small":
    case "mistral-medium":
      return model;
    default:
      return "mistral-tiny";
  }
}

export function getAwsBedrockModelFamily(_model: string): ModelFamily {
  return "aws-claude";
}

export function getAzureOpenAIModelFamily(
  model: string,
  defaultFamily: AzureOpenAIModelFamily = "azure-gpt4"
): AzureOpenAIModelFamily {
  // Azure model names omit periods.  addAzureKey also prepends "azure-" to the
  // model name to route the request the correct keyprovider, so we need to
  // remove that as well.
  const modified = model
    .replace("gpt-35-turbo", "gpt-3.5-turbo")
    .replace("azure-", "");
  for (const [regex, family] of Object.entries(OPENAI_MODEL_FAMILY_MAP)) {
    if (modified.match(regex)) {
      return `azure-${family}` as AzureOpenAIModelFamily;
    }
  }
  return defaultFamily;
}

export function assertIsKnownModelFamily(
  modelFamily: string
): asserts modelFamily is ModelFamily {
  if (!MODEL_FAMILIES.includes(modelFamily as ModelFamily)) {
    throw new Error(`Unknown model family: ${modelFamily}`);
  }
}

export function getModelFamilyForRequest(req: Request): ModelFamily {
  if (req.modelFamily) return req.modelFamily;
  // There is a single request queue, but it is partitioned by model family.
  // Model families are typically separated on cost/rate limit boundaries so
  // they should be treated as separate queues.
  const model = req.body.model ?? "gpt-3.5-turbo";
  let modelFamily: ModelFamily;

  // Weird special case for AWS/Azure because they serve multiple models from
  // different vendors, even if currently only one is supported.
  if (req.service === "aws") {
    modelFamily = getAwsBedrockModelFamily(model);
  } else if (req.service === "azure") {
    modelFamily = getAzureOpenAIModelFamily(model);
  } else {
    switch (req.outboundApi) {
      case "anthropic":
        modelFamily = getClaudeModelFamily(model);
        break;
      case "openai":
      case "openai-text":
      case "openai-image":
        modelFamily = getOpenAIModelFamily(model);
        break;
      case "google-ai":
        modelFamily = getGoogleAIModelFamily(model);
        break;
      case "mistral-ai":
        modelFamily = getMistralAIModelFamily(model);
        break;
      default:
        assertNever(req.outboundApi);
    }
  }

  return (req.modelFamily = modelFamily);
}

function assertNever(x: never): never {
  throw new Error(`Called assertNever with argument ${x}.`);
}


code: export * as sheets from "./sheets";


code:
/* Google Sheets backend for prompt logger.  Upon every flush, this backend
writes the batch to a Sheets spreadsheet. If the sheet becomes too large, it
will create a new sheet and continue writing there.

This is essentially a really shitty ORM for Sheets. Absolutely no concurrency
support because it relies on local state to match up with the remote state. */

import { google, sheets_v4 } from "googleapis";
import type { CredentialBody } from "google-auth-library";
import type { GaxiosResponse } from "googleapis-common";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { PromptLogEntry } from "..";

// There is always a sheet called __index__ which contains a list of all the
// other sheets. We use this rather than iterating over all the sheets in case
// the user needs to manually work with the spreadsheet.
// If no __index__ sheet exists, we will assume that the spreadsheet is empty
// and create one.

type IndexSheetModel = {
  /**
   * Stored in cell B2. Set on startup; if it changes, we assume that another
   * instance of the proxy is writing to the spreadsheet and stop.
   */
  lockId: string;
  /**
   * Data starts at row 4. Row 1-3 are headers
   */
  rows: { logSheetName: string; createdAt: string; rowCount: number }[];
};

type LogSheetModel = {
  sheetName: string;
  rows: {
    model: string;
    endpoint: string;
    promptRaw: string;
    promptFlattened: string;
    response: string;
  }[];
};

const MAX_ROWS_PER_SHEET = 2000;
const log = logger.child({ module: "sheets" });

let sheetsClient: sheets_v4.Sheets | null = null;
/** Called when log backend aborts to tell the log queue to stop. */
let stopCallback: (() => void) | null = null;
/** Lock/synchronization ID for this session. */
let lockId = Math.random().toString(36).substring(2, 15);
/** In-memory cache of the index sheet. */
let indexSheet: IndexSheetModel | null = null;
/** In-memory cache of the active log sheet. */
let activeLogSheet: LogSheetModel | null = null;

/**
 * Loads the __index__ sheet into memory. By default, asserts that the lock ID
 * has not changed since the start of the session.
 */
const loadIndexSheet = async (assertLockId = true) => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  log.info({ assertLockId }, "Loading __index__ sheet.");
  const res = await client.spreadsheets.values.get({
    spreadsheetId: spreadsheetId,
    range: "__index__!A1:D",
    majorDimension: "ROWS",
  });
  const data = assertData(res);
  if (!data.values || data.values[2][0] !== "logSheetName") {
    log.error({ values: data.values }, "Unexpected format for __index__ sheet");
    throw new Error("Unexpected format for __index__ sheet");
  }

  if (assertLockId) {
    const lockIdCell = data.values[1][1];
    if (lockIdCell !== lockId) {
      log.error(
        { receivedLock: lockIdCell, expectedLock: lockId },
        "Another instance of the proxy is writing to the spreadsheet; stopping."
      );
      stop();
      throw new Error(`Lock ID assertion failed`);
    }
  }

  const rows = data.values.slice(3).map((row) => {
    return {
      logSheetName: row[0],
      createdAt: row[1],
      rowCount: row[2],
    };
  });
  indexSheet = { lockId, rows };
};

/** Creates empty __index__ sheet for a new spreadsheet. */
const createIndexSheet = async () => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  log.info("Creating empty __index__ sheet.");
  const res = await client.spreadsheets.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      requests: [
        {
          addSheet: {
            properties: {
              title: "__index__",
              gridProperties: { rowCount: 1, columnCount: 3 },
            },
          },
        },
      ],
    },
  });
  assertData(res);
  indexSheet = { lockId, rows: [] };
  await writeIndexSheet();
};

/** Writes contents of in-memory indexSheet to the remote __index__ sheet. */
const writeIndexSheet = async () => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  const headerRows = [
    ["Don't edit this sheet while the server is running.", "", ""],
    ["Lock ID", lockId, ""],
    ["logSheetName", "createdAt", "rowCount"],
  ];
  const contentRows = indexSheet!.rows.map((row) => {
    return [row.logSheetName, row.createdAt, row.rowCount];
  });
  log.info("Persisting __index__ sheet.");
  await client.spreadsheets.values.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      valueInputOption: "RAW",
      data: [
        { range: "__index__!A1:D", values: [...headerRows, ...contentRows] },
      ],
    },
  });
};

/** Creates a new log sheet, adds it to the index, and sets it as active. */
const createLogSheet = async () => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  // Sheet name format is Log_YYYYMMDD_HHMMSS
  const sheetName = `Log_${new Date()
    .toISOString()
    // YYYY-MM-DDTHH:MM:SS.sssZ -> YYYYMMDD_HHMMSS
    .replace(/[-:.]/g, "")
    .replace(/T/, "_")
    .substring(0, 15)}`;

  log.info({ sheetName }, "Creating new log sheet.");
  const res = await client.spreadsheets.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      requests: [
        {
          addSheet: {
            properties: {
              title: sheetName,
              gridProperties: { rowCount: MAX_ROWS_PER_SHEET, columnCount: 5 },
            },
          },
        },
      ],
    },
  });
  assertData(res);
  // Increase row/column size and wrap text for readability.
  const sheetId = res.data.replies![0].addSheet!.properties!.sheetId;
  await client.spreadsheets.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      requests: [
        {
          repeatCell: {
            range: { sheetId },
            cell: {
              userEnteredFormat: {
                wrapStrategy: "WRAP",
                verticalAlignment: "TOP",
              },
            },
            fields: "*",
          },
        },
        {
          updateDimensionProperties: {
            range: {
              sheetId,
              dimension: "COLUMNS",
              startIndex: 3,
              endIndex: 5,
            },
            properties: { pixelSize: 500 },
            fields: "pixelSize",
          },
        },
        {
          updateDimensionProperties: {
            range: {
              sheetId,
              dimension: "ROWS",
              startIndex: 1,
            },
            properties: { pixelSize: 200 },
            fields: "pixelSize",
          },
        },
      ],
    },
  });
  await client.spreadsheets.values.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      valueInputOption: "RAW",
      data: [
        {
          range: `${sheetName}!A1:E`,
          values: [
            ["model", "endpoint", "prompt json", "prompt string", "response"],
          ],
        },
      ],
    },
  });
  indexSheet!.rows.push({
    logSheetName: sheetName,
    createdAt: new Date().toISOString(),
    rowCount: 0,
  });
  await writeIndexSheet();
  activeLogSheet = { sheetName, rows: [] };
};

export const appendBatch = async (batch: PromptLogEntry[]) => {
  if (!activeLogSheet) {
    // Create a new log sheet if we don't have one yet.
    await createLogSheet();
  } else {
    // Check lock to ensure we're the only instance writing to the spreadsheet.
    await loadIndexSheet(true);
  }

  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  const sheetName = activeLogSheet!.sheetName;
  const newRows = batch.map((entry) => {
    return [
      entry.model,
      entry.endpoint,
      entry.promptRaw.slice(-50000),
      entry.promptFlattened.slice(-50000),
      entry.response.slice(0, 50000),
    ];
  });
  log.info({ sheetName, rowCount: newRows.length }, "Appending log batch.");
  const data = await client.spreadsheets.values.append({
    spreadsheetId: spreadsheetId,
    range: `${sheetName}!A1:D`,
    valueInputOption: "RAW",
    requestBody: { values: newRows, majorDimension: "ROWS" },
  });
  assertData(data);
  if (data.data.updates && data.data.updates.updatedRows) {
    const newRowCount = data.data.updates.updatedRows;
    log.info({ sheetName, rowCount: newRowCount }, "Successfully appended.");
    activeLogSheet!.rows = activeLogSheet!.rows.concat(
      newRows.map((row) => ({
        model: row[0],
        endpoint: row[1],
        promptRaw: row[2],
        promptFlattened: row[3],
        response: row[4],
      }))
    );
  } else {
    // We didn't receive an error but we didn't get any updates either.
    // We may need to create a new sheet and throw to make the queue retry the
    // batch.
    log.warn(
      { sheetName, rowCount: newRows.length },
      "No updates received from append. Creating new sheet and retrying."
    );
    await createLogSheet();
    throw new Error("No updates received from append.");
  }
  await finalizeBatch();
};

const finalizeBatch = async () => {
  const sheetName = activeLogSheet!.sheetName;
  const rowCount = activeLogSheet!.rows.length;
  const indexRow = indexSheet!.rows.find(
    ({ logSheetName }) => logSheetName === sheetName
  )!;
  indexRow.rowCount = rowCount;
  if (rowCount >= MAX_ROWS_PER_SHEET) {
    await createLogSheet(); // Also updates index sheet
  } else {
    await writeIndexSheet();
  }
  log.info({ sheetName, rowCount }, "Batch finalized.");
};

type LoadLogSheetArgs = {
  sheetName: string;
  /** The starting row to load. If omitted, loads all rows (expensive). */
  fromRow?: number;
};

/** Not currently used. */
export const loadLogSheet = async ({
  sheetName,
  fromRow = 2, // omit header row
}: LoadLogSheetArgs) => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;

  const range = `${sheetName}!A${fromRow}:E`;
  const res = await client.spreadsheets.values.get({
    spreadsheetId: spreadsheetId,
    range,
  });
  const data = assertData(res);
  const values = data.values || [];
  const rows = values.slice(1).map((row) => {
    return {
      model: row[0],
      endpoint: row[1],
      promptRaw: row[2],
      promptFlattened: row[3],
      response: row[4],
    };
  });
  activeLogSheet = { sheetName, rows };
};

export const init = async (onStop: () => void) => {
  if (sheetsClient) {
    return;
  }
  if (!config.googleSheetsKey || !config.googleSheetsSpreadsheetId) {
    throw new Error(
      "Missing required Google Sheets config. Refer to documentation for setup instructions."
    );
  }

  log.info("Initializing Google Sheets backend.");
  const encodedCreds = config.googleSheetsKey;
  // encodedCreds is a base64-encoded JSON key from the GCP console.
  const creds: CredentialBody = JSON.parse(
    Buffer.from(encodedCreds, "base64").toString("utf8").trim()
  );
  const auth = new google.auth.GoogleAuth({
    scopes: ["https://www.googleapis.com/auth/spreadsheets"],
    credentials: creds,
  });
  sheetsClient = google.sheets({ version: "v4", auth });
  stopCallback = onStop;

  const sheetId = config.googleSheetsSpreadsheetId;
  const res = await sheetsClient.spreadsheets.get({
    spreadsheetId: sheetId,
  });
  if (!res.data) {
    const { status, statusText, headers } = res;
    log.error(
      {
        res: { status, statusText, headers },
        creds: {
          client_email: creds.client_email?.slice(0, 5) + "********",
          private_key: creds.private_key?.slice(0, 5) + "********",
        },
        sheetId: config.googleSheetsSpreadsheetId,
      },
      "Could not connect to Google Sheets."
    );
    stop();
    throw new Error("Could not connect to Google Sheets.");
  } else {
    const sheetTitle = res.data.properties?.title;
    log.info({ sheetId, sheetTitle }, "Connected to Google Sheets.");
  }

  // Load or create the index sheet and write the lockId to it.
  try {
    log.info("Loading index sheet.");
    await loadIndexSheet(false);
    await writeIndexSheet();
  } catch (e) {
    log.warn({ error: e.message }, "Could not load index sheet. Creating a new one.");
    await createIndexSheet();
  }
};

/** Called during some unrecoverable error to tell the log queue to stop. */
function stop() {
  log.warn("Stopping Google Sheets backend.");
  if (stopCallback) {
    stopCallback();
  }
  sheetsClient = null;
}

function assertData<T = sheets_v4.Schema$ValueRange>(res: GaxiosResponse<T>) {
  if (!res.data) {
    const { status, statusText, headers } = res;
    log.error(
      { res: { status, statusText, headers } },
      "Unexpected response from Google Sheets API."
    );
  }
  return res.data!;
}


code:
/* Logs prompts and model responses to a persistent storage backend, if enabled.
Since the proxy is generally deployed to free-tier services, our options for
persistent storage are pretty limited. We'll use Google Sheets as a makeshift
database for now.

Due to the limitations of Google Sheets, we'll queue up log entries and flush
them to the API periodically. */

export interface PromptLogEntry {
  model: string;
  endpoint: string;
  /** JSON prompt passed to the model */
  promptRaw: string;
  /** Prompt with user and assistant messages flattened into a single string */
  promptFlattened: string;
  response: string;
  // TODO: temperature, top_p, top_k, etc.
}

export * as logQueue from "./log-queue";


code:
/* Queues incoming prompts/responses and periodically flushes them to configured
 * logging backend. */

import { logger } from "../../logger";
import { PromptLogEntry } from ".";
import { sheets } from "./backends";

const FLUSH_INTERVAL = 1000 * 10;
const MAX_BATCH_SIZE = 25;

const queue: PromptLogEntry[] = [];
const log = logger.child({ module: "log-queue" });

let started = false;
let timeoutId: NodeJS.Timeout | null = null;
let retrying = false;
let consecutiveFailedBatches = 0;

export const enqueue = (payload: PromptLogEntry) => {
  if (!started) {
    log.warn("Log queue not started, discarding incoming log entry.");
    return;
  }
  queue.push(payload);
};

export const flush = async () => {
  if (!started) {
    return;
  }

  if (queue.length > 0) {
    const batchSize = Math.min(MAX_BATCH_SIZE, queue.length);
    const nextBatch = queue.splice(0, batchSize);
    log.info({ size: nextBatch.length }, "Submitting new batch.");
    try {
      await sheets.appendBatch(nextBatch);
      retrying = false;
      consecutiveFailedBatches = 0;
    } catch (e: any) {
      if (retrying) {
        log.error(
          { message: e.message, stack: e.stack },
          "Failed twice to flush batch, discarding."
        );
        retrying = false;
        consecutiveFailedBatches++;
      } else {
        // Put the batch back at the front of the queue and try again
        log.warn(
          { message: e.message, stack: e.stack },
          "Failed to flush batch. Retrying."
        );
        queue.unshift(...nextBatch);
        retrying = true;
        setImmediate(() => flush());
        return;
      }
    }
  }

  const useHalfInterval = queue.length > MAX_BATCH_SIZE / 2;
  scheduleFlush(useHalfInterval);
};

export const start = async () => {
  try {
    await sheets.init(() => stop());
    log.info("Logging backend initialized.");
    started = true;
  } catch (e) {
    log.error({ error: e.message }, "Could not initialize logging backend.");
    return;
  }
  scheduleFlush();
};

export const stop = () => {
  if (timeoutId) {
    clearTimeout(timeoutId);
  }
  log.info("Stopping log queue.");
  started = false;
};

const scheduleFlush = (halfInterval = false) => {
  if (consecutiveFailedBatches > 3) {
    // TODO: may cause memory issues on busy servers, though if we crash that
    // may actually fix the problem with logs randomly not being flushed.
    const oneMinute = 60 * 1000;
    const maxBackoff = 10 * oneMinute;
    const backoff = Math.min(consecutiveFailedBatches * oneMinute, maxBackoff);
    timeoutId = setTimeout(() => {
      flush();
    }, backoff);
    log.warn(
      { consecutiveFailedBatches, backoffMs: backoff },
      "Failed to flush 3 batches in a row, pausing for a few minutes."
    );
    return;
  }

  if (halfInterval) {
    log.warn(
      { queueSize: queue.length },
      "Queue is falling behind, switching to faster flush interval."
    );
  }

  timeoutId = setTimeout(
    () => {
      flush();
    },
    halfInterval ? FLUSH_INTERVAL / 2 : FLUSH_INTERVAL
  );
};


code:
import { config } from "../config";
import { ModelFamily } from "./models";

// technically slightly underestimates, because completion tokens cost more
// than prompt tokens but we don't track those separately right now
export function getTokenCostUsd(model: ModelFamily, tokens: number) {
  let cost = 0;
  switch (model) {
    case "gpt4-turbo":
      cost = 0.00001;
      break;
    case "gpt4-32k":
      cost = 0.00006;
      break;
    case "gpt4":
      cost = 0.00003;
      break;
    case "turbo":
      cost = 0.000001;
      break;
    case "dall-e":
      cost = 0.00001;
      break;
    case "aws-claude":
    case "claude":
      cost = 0.00001102;
      break;
    case "mistral-tiny":
      cost = 0.00000031;
      break;
    case "mistral-small":
      cost = 0.00000132;
      break;
    case "mistral-medium":
      cost = 0.0000055;
      break;
  }
  return cost * Math.max(0, tokens);
}

export function prettyTokens(tokens: number): string {
  const absTokens = Math.abs(tokens);
  if (absTokens < 1000) {
    return tokens.toString();
  } else if (absTokens < 1000000) {
    return (tokens / 1000).toFixed(1) + "k";
  } else if (absTokens < 1000000000) {
    return (tokens / 1000000).toFixed(2) + "m";
  } else {
    return (tokens / 1000000000).toFixed(3) + "b";
  }
}

export function getCostSuffix(cost: number) {
  if (!config.showTokenCosts) return "";
  return ` ($${cost.toFixed(2)})`;
}


code:
import { Response } from "express";
import { IncomingMessage } from "http";
import { assertNever } from "./utils";
import { APIFormat } from "./key-management";

export function initializeSseStream(res: Response) {
  res.statusCode = 200;
  res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
  res.setHeader("Cache-Control", "no-cache");
  res.setHeader("Connection", "keep-alive");
  res.setHeader("X-Accel-Buffering", "no"); // nginx-specific fix
  res.flushHeaders();
}

/**
 * Copies headers received from upstream API to the SSE response, excluding
 * ones we need to set ourselves for SSE to work.
 */
export function copySseResponseHeaders(
  proxyRes: IncomingMessage,
  res: Response
) {
  const toOmit = [
    "content-length",
    "content-encoding",
    "transfer-encoding",
    "content-type",
    "connection",
    "cache-control",
  ];
  for (const [key, value] of Object.entries(proxyRes.headers)) {
    if (!toOmit.includes(key) && value) {
      res.setHeader(key, value);
    }
  }
}

/**
 * Returns an SSE message that looks like a completion event for the service
 * that the request is being proxied to. Used to send error messages to the
 * client in the middle of a streaming request.
 */
export function makeCompletionSSE({
  format,
  title,
  message,
  obj,
  reqId,
  model = "unknown",
}: {
  format: APIFormat;
  title: string;
  message: string;
  obj?: object;
  reqId: string | number | object;
  model?: string;
}) {
  const id = String(reqId);
  const content = `\n\n**${title}**\n${message}${
    obj ? `\n\`\`\`\n${JSON.stringify(obj, null, 2)}\n\`\`\`\n` : ""
  }`;

  let event;

  switch (format) {
    case "openai":
    case "mistral-ai":
      event = {
        id: "chatcmpl-" + id,
        object: "chat.completion.chunk",
        created: Date.now(),
        model,
        choices: [{ delta: { content }, index: 0, finish_reason: title }],
      };
      break;
    case "openai-text":
      event = {
        id: "cmpl-" + id,
        object: "text_completion",
        created: Date.now(),
        choices: [
          { text: content, index: 0, logprobs: null, finish_reason: title },
        ],
        model,
      };
      break;
    case "anthropic":
      event = {
        completion: content,
        stop_reason: title,
        truncated: false,
        stop: null,
        model,
        log_id: "proxy-req-" + id,
      };
      break;
    case "google-ai":
      return JSON.stringify({
        candidates: [
          {
            content: { parts: [{ text: content }], role: "model" },
            finishReason: title,
            index: 0,
            tokenCount: null,
            safetyRatings: [],
          },
        ],
      });
    case "openai-image":
      throw new Error(`SSE not supported for ${format} requests`);
    default:
      assertNever(format);
  }

  if (format === "anthropic") {
    return (
      ["event: completion", `data: ${JSON.stringify(event)}`].join("\n") +
      "\n\n"
    );
  }

  return `data: ${JSON.stringify(event)}\n\n`;
}


code:
import { getTokenizer } from "@anthropic-ai/tokenizer";
import { Tiktoken } from "tiktoken/lite";

let encoder: Tiktoken;

export function init() {
  // they export a `countTokens` function too but it instantiates a new
  // tokenizer every single time and it is not fast...
  encoder = getTokenizer();
  return true;
}

export function getTokenCount(prompt: string, _model: string) {
  // Don't try tokenizing if the prompt is massive to prevent DoS.
  // 500k characters should be sufficient for all supported models.
  if (prompt.length > 500000) {
    return {
      tokenizer: "length fallback",
      token_count: 100000,
    };
  }

  return {
    tokenizer: "@anthropic-ai/tokenizer",
    token_count: encoder.encode(prompt.normalize("NFKC"), "all").length,
  };
}


code: export { init, countTokens } from "./tokenizer";


code:
import { MistralAIChatMessage } from "../../proxy/middleware/request/preprocessors/transform-outbound-payload.js";
import * as tokenizer from "./mistral-tokenizer-js";

export function init() {
  tokenizer.initializemistralTokenizer();
  return true;
}

export function getTokenCount(prompt: MistralAIChatMessage[] | string) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  let chunks = [];
  for (const message of prompt) {
    switch (message.role) {
      case "system":
        chunks.push(message.content);
        break;
      case "assistant":
        chunks.push(message.content + "</s>");
        break;
      case "user":
        chunks.push("[INST] " + message.content + " [/INST]");
        break;
    }
  }
  return getTextTokenCount(chunks.join(" "));
}

function getTextTokenCount(prompt: string) {
  // Don't try tokenizing if the prompt is massive to prevent DoS.
  // 500k characters should be sufficient for all supported models.
  if (prompt.length > 500000) {
    return {
      tokenizer: "length fallback",
      token_count: 100000,
    };
  }

  return {
    tokenizer: "mistral-tokenizer-js",
    token_count: tokenizer.encode(prompt.normalize("NFKC"))!.length,
  };
}


code:
import { Tiktoken } from "tiktoken/lite";
import cl100k_base from "tiktoken/encoders/cl100k_base.json";
import { logger } from "../../logger";
import { libSharp } from "../file-storage";
import type {
  GoogleAIChatMessage,
  OpenAIChatMessage,
} from "../../proxy/middleware/request/preprocessors/transform-outbound-payload";
import { z } from "zod";

const log = logger.child({ module: "tokenizer", service: "openai" });
const GPT4_VISION_SYSTEM_PROMPT_SIZE = 170;

let encoder: Tiktoken;

export function init() {
  encoder = new Tiktoken(
    cl100k_base.bpe_ranks,
    cl100k_base.special_tokens,
    cl100k_base.pat_str
  );
  return true;
}

// Tested against:
// https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb

export async function getTokenCount(
  prompt: string | OpenAIChatMessage[],
  model: string
) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  const oldFormatting = model.startsWith("turbo-0301");
  const vision = model.includes("vision");

  const tokensPerMessage = oldFormatting ? 4 : 3;
  const tokensPerName = oldFormatting ? -1 : 1; // older formatting replaces role with name if name is present

  let numTokens = vision ? GPT4_VISION_SYSTEM_PROMPT_SIZE : 0;

  for (const message of prompt) {
    numTokens += tokensPerMessage;
    for (const key of Object.keys(message)) {
      {
        let textContent: string = "";
        const value = message[key as keyof OpenAIChatMessage];

        if (!value) continue;

        if (Array.isArray(value)) {
          for (const item of value) {
            if (item.type === "text") {
              textContent += item.text;
            } else if (item.type === "image_url") {
              const { url, detail } = item.image_url;
              const cost = await getGpt4VisionTokenCost(url, detail);
              numTokens += cost ?? 0;
            }
          }
        } else {
          textContent = value;
        }

        // Break if we get a huge message or exceed the token limit to prevent
        // DoS.
        // 200k tokens allows for future 200k GPT-4 models and 500k characters
        // is just a sanity check
        if (textContent.length > 500000 || numTokens > 200000) {
          numTokens = 200000;
          return {
            tokenizer: "tiktoken (prompt length limit exceeded)",
            token_count: numTokens,
          };
        }

        numTokens += encoder.encode(textContent).length;
        if (key === "name") {
          numTokens += tokensPerName;
        }
      }
    }
  }
  numTokens += 3; // every reply is primed with <|start|>assistant<|message|>
  return { tokenizer: "tiktoken", token_count: numTokens };
}

async function getGpt4VisionTokenCost(
  url: string,
  detail: "auto" | "low" | "high" = "auto"
) {
  // For now we do not allow remote images as the proxy would have to download
  // them, which is a potential DoS vector.
  if (!url.startsWith("data:image/")) {
    throw new Error(
      "Remote images are not supported. Add the image to your prompt as a base64 data URL."
    );
  }

  const base64Data = url.split(",")[1];
  const buffer = Buffer.from(base64Data, "base64");
  const image = libSharp(buffer);
  const metadata = await image.metadata();

  if (!metadata || !metadata.width || !metadata.height) {
    throw new Error("Prompt includes an image that could not be parsed");
  }

  const { width, height } = metadata;

  let selectedDetail: "low" | "high";
  if (detail === "auto") {
    const threshold = 512 * 512;
    const imageSize = width * height;
    selectedDetail = imageSize > threshold ? "high" : "low";
  } else {
    selectedDetail = detail;
  }

  // https://platform.openai.com/docs/guides/vision/calculating-costs
  if (selectedDetail === "low") {
    log.info(
      { width, height, tokens: 85 },
      "Using fixed GPT-4-Vision token cost for low detail image"
    );
    return 85;
  }

  let newWidth = width;
  let newHeight = height;
  if (width > 2048 || height > 2048) {
    const aspectRatio = width / height;
    if (width > height) {
      newWidth = 2048;
      newHeight = Math.round(2048 / aspectRatio);
    } else {
      newHeight = 2048;
      newWidth = Math.round(2048 * aspectRatio);
    }
  }

  if (newWidth < newHeight) {
    newHeight = Math.round((newHeight / newWidth) * 768);
    newWidth = 768;
  } else {
    newWidth = Math.round((newWidth / newHeight) * 768);
    newHeight = 768;
  }

  const tiles = Math.ceil(newWidth / 512) * Math.ceil(newHeight / 512);
  const tokens = 170 * tiles + 85;

  log.info(
    { width, height, newWidth, newHeight, tiles, tokens },
    "Calculated GPT-4-Vision token cost for high detail image"
  );
  return tokens;
}

function getTextTokenCount(prompt: string) {
  if (prompt.length > 500000) {
    return {
      tokenizer: "length fallback",
      token_count: 100000,
    };
  }

  return {
    tokenizer: "tiktoken",
    token_count: encoder.encode(prompt).length,
  };
}

// Model	Resolution	Price
// DALL·E 3	1024×1024	$0.040 / image
// 1024×1792, 1792×1024	$0.080 / image
// DALL·E 3 HD	1024×1024	$0.080 / image
// 1024×1792, 1792×1024	$0.120 / image
// DALL·E 2	1024×1024	$0.020 / image
// 512×512	$0.018 / image
// 256×256	$0.016 / image

export const DALLE_TOKENS_PER_DOLLAR = 100000;

/**
 * OpenAI image generation with DALL-E doesn't use tokens but everything else
 * in the application does. There is a fixed cost for each image generation
 * request depending on the model and selected quality/resolution parameters,
 * which we convert to tokens at a rate of 100000 tokens per dollar.
 */
export function getOpenAIImageCost(params: {
  model: "dall-e-2" | "dall-e-3";
  quality: "standard" | "hd";
  resolution: "512x512" | "256x256" | "1024x1024" | "1024x1792" | "1792x1024";
  n: number | null;
}) {
  const { model, quality, resolution, n } = params;
  const usd = (() => {
    switch (model) {
      case "dall-e-2":
        switch (resolution) {
          case "512x512":
            return 0.018;
          case "256x256":
            return 0.016;
          case "1024x1024":
            return 0.02;
          default:
            throw new Error("Invalid resolution");
        }
      case "dall-e-3":
        switch (resolution) {
          case "1024x1024":
            return quality === "standard" ? 0.04 : 0.08;
          case "1024x1792":
          case "1792x1024":
            return quality === "standard" ? 0.08 : 0.12;
          default:
            throw new Error("Invalid resolution");
        }
      default:
        throw new Error("Invalid image generation model");
    }
  })();

  const tokens = (n ?? 1) * (usd * DALLE_TOKENS_PER_DOLLAR);

  return {
    tokenizer: `openai-image cost`,
    token_count: Math.ceil(tokens),
  };
}

export function estimateGoogleAITokenCount(prompt: string | GoogleAIChatMessage[]) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  const tokensPerMessage = 3;

  let numTokens = 0;
  for (const message of prompt) {
    numTokens += tokensPerMessage;
    numTokens += encoder.encode(message.parts[0].text).length;
  }

  numTokens += 3;

  return {
    tokenizer: "tiktoken (google-ai estimate)",
    token_count: numTokens,
  };
}


code:
import { Request } from "express";
import type {
  GoogleAIChatMessage,
  MistralAIChatMessage,
  OpenAIChatMessage,
} from "../../proxy/middleware/request/preprocessors/transform-outbound-payload";
import { assertNever } from "../utils";
import {
  init as initClaude,
  getTokenCount as getClaudeTokenCount,
} from "./claude";
import {
  init as initOpenAi,
  getTokenCount as getOpenAITokenCount,
  getOpenAIImageCost,
  estimateGoogleAITokenCount,
} from "./openai";
import {
  init as initMistralAI,
  getTokenCount as getMistralAITokenCount,
} from "./mistral";
import { APIFormat } from "../key-management";

export async function init() {
  initClaude();
  initOpenAi();
  initMistralAI();
}

/** Tagged union via `service` field of the different types of requests that can
 * be made to the tokenization service, for both prompts and completions */
type TokenCountRequest = { req: Request } & (
  | { prompt: OpenAIChatMessage[]; completion?: never; service: "openai" }
  | {
      prompt: string;
      completion?: never;
      service: "openai-text" | "anthropic" | "google-ai";
    }
  | { prompt?: GoogleAIChatMessage[]; completion?: never; service: "google-ai" }
  | { prompt: MistralAIChatMessage[]; completion?: never; service: "mistral-ai" }
  | { prompt?: never; completion: string; service: APIFormat }
  | { prompt?: never; completion?: never; service: "openai-image" }
);

type TokenCountResult = {
  token_count: number;
  tokenizer: string;
  tokenization_duration_ms: number;
};

export async function countTokens({
  req,
  service,
  prompt,
  completion,
}: TokenCountRequest): Promise<TokenCountResult> {
  const time = process.hrtime();
  switch (service) {
    case "anthropic":
      return {
        ...getClaudeTokenCount(prompt ?? completion, req.body.model),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "openai":
    case "openai-text":
      return {
        ...(await getOpenAITokenCount(prompt ?? completion, req.body.model)),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "openai-image":
      return {
        ...getOpenAIImageCost({
          model: req.body.model,
          quality: req.body.quality,
          resolution: req.body.size,
          n: parseInt(req.body.n, 10) || null,
        }),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "google-ai":
      // TODO: Can't find a tokenization library for Gemini. There is an API
      // endpoint for it but it adds significant latency to the request.
      return {
        ...estimateGoogleAITokenCount(prompt ?? (completion || [])),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "mistral-ai":
      return {
        ...getMistralAITokenCount(prompt ?? completion),
        tokenization_duration_ms: getElapsedMs(time),
      };
    default:
      assertNever(service);
  }
}

function getElapsedMs(time: [number, number]) {
  const diff = process.hrtime(time);
  return diff[0] * 1000 + diff[1] / 1e6;
}


code:
import { ZodType, z } from "zod";
import type { ModelFamily } from "../models";
import { makeOptionalPropsNullable } from "../utils";

export const tokenCountsSchema: ZodType<UserTokenCounts> = z.object({
  turbo: z.number().optional().default(0),
  gpt4: z.number().optional().default(0),
  "gpt4-32k": z.number().optional().default(0),
  "gpt4-turbo": z.number().optional().default(0),
  "dall-e": z.number().optional().default(0),
  claude: z.number().optional().default(0),
  "gemini-pro": z.number().optional().default(0),
  "aws-claude": z.number().optional().default(0),
});

export const UserSchema = z
  .object({
    /** User's personal access token. */
    token: z.string(),
    /** IP addresses the user has connected from. */
    ip: z.array(z.string()),
    /** User's nickname. */
    nickname: z.string().max(80).optional(),
    /**
     * The user's privilege level.
     * - `normal`: Default role. Subject to usual rate limits and quotas.
     * - `special`: Special role. Higher quotas and exempt from
     *   auto-ban/lockout.
     **/
    type: z.enum(["normal", "special", "temporary"]),
    /** Number of prompts the user has made. */
    promptCount: z.number(),
    /**
     * @deprecated Use `tokenCounts` instead.
     * Never used; retained for backwards compatibility.
     */
    tokenCount: z.any().optional(),
    /** Number of tokens the user has consumed, by model family. */
    tokenCounts: tokenCountsSchema,
    /** Maximum number of tokens the user can consume, by model family. */
    tokenLimits: tokenCountsSchema,
    /** Time at which the user was created. */
    createdAt: z.number(),
    /** Time at which the user last connected. */
    lastUsedAt: z.number().optional(),
    /** Time at which the user was disabled, if applicable. */
    disabledAt: z.number().optional(),
    /** Reason for which the user was disabled, if applicable. */
    disabledReason: z.string().optional(),
    /** Time at which the user will expire and be disabled (for temp users). */
    expiresAt: z.number().optional(),
    /** The user's maximum number of IP addresses; supercedes global max. */
    maxIps: z.coerce.number().int().min(0).optional(),
    /** Private note about the user. */
    adminNote: z.string().optional(),
  })
  .strict();

/**
 * Variant of `UserSchema` which allows for partial updates, and makes any
 * optional properties on the base schema nullable. Null values are used to
 * indicate that the property should be deleted from the user object.
 */
export const UserPartialSchema = makeOptionalPropsNullable(UserSchema)
  .partial()
  .extend({ token: z.string() });

export type UserTokenCounts = {
  [K in ModelFamily]?: number;
};
export type User = z.infer<typeof UserSchema>;
export type UserUpdate = z.infer<typeof UserPartialSchema>;


code:
/**
 * Basic user management. Handles creation and tracking of proxy users, personal
 * access tokens, and quota management. Supports in-memory and Firebase Realtime
 * Database persistence stores.
 *
 * Users are identified solely by their personal access token. The token is
 * used to authenticate the user for all proxied requests.
 */

import admin from "firebase-admin";
import schedule from "node-schedule";
import { v4 as uuid } from "uuid";
import { config, getFirebaseApp } from "../../config";
import {
  getAzureOpenAIModelFamily,
  getClaudeModelFamily,
  getGoogleAIModelFamily,
  getMistralAIModelFamily,
  getOpenAIModelFamily,
  MODEL_FAMILIES,
  ModelFamily,
} from "../models";
import { logger } from "../../logger";
import { User, UserTokenCounts, UserUpdate } from "./schema";
import { APIFormat } from "../key-management";
import { assertNever } from "../utils";

const log = logger.child({ module: "users" });

const INITIAL_TOKENS: Required<UserTokenCounts> = {
  turbo: 0,
  gpt4: 0,
  "gpt4-32k": 0,
  "gpt4-turbo": 0,
  "dall-e": 0,
  claude: 0,
  "gemini-pro": 0,
  "mistral-tiny": 0,
  "mistral-small": 0,
  "mistral-medium": 0,
  "aws-claude": 0,
  "azure-turbo": 0,
  "azure-gpt4": 0,
  "azure-gpt4-turbo": 0,
  "azure-gpt4-32k": 0,
};

const users: Map<string, User> = new Map();
const usersToFlush = new Set<string>();
let quotaRefreshJob: schedule.Job | null = null;
let userCleanupJob: schedule.Job | null = null;

export async function init() {
  log.info({ store: config.gatekeeperStore }, "Initializing user store...");
  if (config.gatekeeperStore === "firebase_rtdb") {
    await initFirebase();
  }
  if (config.quotaRefreshPeriod) {
    const crontab = getRefreshCrontab();
    quotaRefreshJob = schedule.scheduleJob(crontab, refreshAllQuotas);
    if (!quotaRefreshJob) {
      throw new Error(
        "Unable to schedule quota refresh. Is QUOTA_REFRESH_PERIOD set correctly?"
      );
    }
    log.debug(
      { nextRefresh: quotaRefreshJob.nextInvocation() },
      "Scheduled token quota refresh."
    );
  }

  userCleanupJob = schedule.scheduleJob("* * * * *", cleanupExpiredTokens);

  log.info("User store initialized.");
}

/**
 * Creates a new user and returns their token. Optionally accepts parameters
 * for setting an expiry date and/or token limits for temporary users.
 **/
export function createUser(createOptions?: {
  type?: User["type"];
  expiresAt?: number;
  tokenLimits?: User["tokenLimits"];
}) {
  const token = uuid();
  const newUser: User = {
    token,
    ip: [],
    type: "normal",
    promptCount: 0,
    tokenCounts: { ...INITIAL_TOKENS },
    tokenLimits: createOptions?.tokenLimits ?? { ...config.tokenQuota },
    createdAt: Date.now(),
  };

  if (createOptions?.type === "temporary") {
    Object.assign(newUser, {
      type: "temporary",
      expiresAt: createOptions.expiresAt,
    });
  } else {
    Object.assign(newUser, { type: createOptions?.type ?? "normal" });
  }

  users.set(token, newUser);
  usersToFlush.add(token);
  return token;
}

/** Returns the user with the given token if they exist. */
export function getUser(token: string) {
  return users.get(token);
}

/** Returns a list of all users. */
export function getUsers() {
  return Array.from(users.values()).map((user) => ({ ...user }));
}

/**
 * Upserts the given user. Intended for use with the /admin API for updating
 * arbitrary fields on a user; use the other functions in this module for
 * specific use cases. `undefined` values are left unchanged. `null` will delete
 * the property from the user.
 *
 * Returns the upserted user.
 */
export function upsertUser(user: UserUpdate) {
  const existing: User = users.get(user.token) ?? {
    token: user.token,
    ip: [],
    type: "normal",
    promptCount: 0,
    tokenCounts: { ...INITIAL_TOKENS },
    tokenLimits: { ...config.tokenQuota },
    createdAt: Date.now(),
  };

  const updates: Partial<User> = {};

  for (const field of Object.entries(user)) {
    const [key, value] = field as [keyof User, any]; // already validated by zod
    if (value === undefined || key === "token") continue;
    if (value === null) {
      delete existing[key];
    } else {
      updates[key] = value;
    }
  }

  // TODO: Write firebase migration to backfill new fields
  if (updates.tokenCounts) {
    for (const family of MODEL_FAMILIES) {
      updates.tokenCounts[family] ??= 0;
    }
  }
  if (updates.tokenLimits) {
    for (const family of MODEL_FAMILIES) {
      updates.tokenLimits[family] ??= 0;
    }
  }

  users.set(user.token, Object.assign(existing, updates));
  usersToFlush.add(user.token);

  // Immediately schedule a flush to the database if we're using Firebase.
  if (config.gatekeeperStore === "firebase_rtdb") {
    setImmediate(flushUsers);
  }

  return users.get(user.token);
}

/** Increments the prompt count for the given user. */
export function incrementPromptCount(token: string) {
  const user = users.get(token);
  if (!user) return;
  user.promptCount++;
  usersToFlush.add(token);
}

/** Increments token consumption for the given user and model. */
export function incrementTokenCount(
  token: string,
  model: string,
  api: APIFormat,
  consumption: number
) {
  const user = users.get(token);
  if (!user) return;
  const modelFamily = getModelFamilyForQuotaUsage(model, api);
  const existing = user.tokenCounts[modelFamily] ?? 0;
  user.tokenCounts[modelFamily] = existing + consumption;
  usersToFlush.add(token);
}

/**
 * Given a user's token and IP address, authenticates the user and adds the IP
 * to the user's list of IPs. Returns the user if they exist and are not
 * disabled, otherwise returns undefined.
 */
export function authenticate(
  token: string,
  ip: string
): { user?: User; result: "success" | "disabled" | "not_found" | "limited" } {
  const user = users.get(token);
  if (!user) return { result: "not_found" };
  if (user.disabledAt) return { result: "disabled" };

  const newIp = !user.ip.includes(ip);

  const userLimit = user.maxIps ?? config.maxIpsPerUser;
  const enforcedLimit =
    user.type === "special" || !userLimit ? Infinity : userLimit;

  if (newIp && user.ip.length >= enforcedLimit) {
    if (config.maxIpsAutoBan) {
      user.ip.push(ip);
      disableUser(token, "IP address limit exceeded.");
      return { result: "disabled" };
    }
    return { result: "limited" };
  } else if (newIp) {
    user.ip.push(ip);
  }

  user.lastUsedAt = Date.now();
  usersToFlush.add(token);
  return { user, result: "success" };
}

export function hasAvailableQuota({
  userToken,
  model,
  api,
  requested,
}: {
  userToken: string;
  model: string;
  api: APIFormat;
  requested: number;
}) {
  const user = users.get(userToken);
  if (!user) return false;
  if (user.type === "special") return true;

  const modelFamily = getModelFamilyForQuotaUsage(model, api);
  const { tokenCounts, tokenLimits } = user;
  const tokenLimit = tokenLimits[modelFamily];

  if (!tokenLimit) return true;

  const tokensConsumed = (tokenCounts[modelFamily] ?? 0) + requested;
  return tokensConsumed < tokenLimit;
}

export function refreshQuota(token: string) {
  const user = users.get(token);
  if (!user) return;
  const { tokenCounts, tokenLimits } = user;
  const quotas = Object.entries(config.tokenQuota) as [ModelFamily, number][];
  quotas
    // If a quota is not configured, don't touch any existing limits a user may
    // already have been assigned manually.
    .filter(([, quota]) => quota > 0)
    .forEach(
      ([model, quota]) =>
        (tokenLimits[model] = (tokenCounts[model] ?? 0) + quota)
    );
  usersToFlush.add(token);
}

export function resetUsage(token: string) {
  const user = users.get(token);
  if (!user) return;
  const { tokenCounts } = user;
  const counts = Object.entries(tokenCounts) as [ModelFamily, number][];
  counts.forEach(([model]) => (tokenCounts[model] = 0));
  usersToFlush.add(token);
}

/** Disables the given user, optionally providing a reason. */
export function disableUser(token: string, reason?: string) {
  const user = users.get(token);
  if (!user) return;
  user.disabledAt = Date.now();
  user.disabledReason = reason;
  usersToFlush.add(token);
}

export function getNextQuotaRefresh() {
  if (!quotaRefreshJob) return "never (manual refresh only)";
  return quotaRefreshJob.nextInvocation().getTime();
}

/**
 * Cleans up expired temporary tokens by disabling tokens past their access
 * expiry date and permanently deleting tokens three days after their access
 * expiry date.
 */
function cleanupExpiredTokens() {
  const now = Date.now();
  let disabled = 0;
  let deleted = 0;
  for (const user of users.values()) {
    if (user.type !== "temporary") continue;
    if (user.expiresAt && user.expiresAt < now && !user.disabledAt) {
      disableUser(user.token, "Temporary token expired.");
      disabled++;
    }
    if (user.disabledAt && user.disabledAt + 72 * 60 * 60 * 1000 < now) {
      users.delete(user.token);
      usersToFlush.add(user.token);
      deleted++;
    }
  }
  log.trace({ disabled, deleted }, "Expired tokens cleaned up.");
}

function refreshAllQuotas() {
  let count = 0;
  for (const user of users.values()) {
    if (user.type === "temporary") continue;
    refreshQuota(user.token);
    count++;
  }
  log.info(
    { refreshed: count, nextRefresh: quotaRefreshJob!.nextInvocation() },
    "Token quotas refreshed."
  );
}

// TODO: Firebase persistence is pretend right now and just polls the in-memory
// store to sync it with Firebase when it changes. Will refactor to abstract
// persistence layer later so we can support multiple stores.
let firebaseTimeout: NodeJS.Timeout | undefined;

async function initFirebase() {
  log.info("Connecting to Firebase...");
  const app = getFirebaseApp();
  const db = admin.database(app);
  const usersRef = db.ref("users");
  const snapshot = await usersRef.once("value");
  const users: Record<string, User> | null = snapshot.val();
  firebaseTimeout = setInterval(flushUsers, 20 * 1000);
  if (!users) {
    log.info("No users found in Firebase.");
    return;
  }
  for (const token in users) {
    upsertUser(users[token]);
  }
  usersToFlush.clear();
  const numUsers = Object.keys(users).length;
  log.info({ users: numUsers }, "Loaded users from Firebase");
}

async function flushUsers() {
  const app = getFirebaseApp();
  const db = admin.database(app);
  const usersRef = db.ref("users");
  const updates: Record<string, User> = {};
  const deletions = [];

  for (const token of usersToFlush) {
    const user = users.get(token);
    if (!user) {
      deletions.push(token);
      continue;
    }
    updates[token] = user;
  }

  usersToFlush.clear();

  const numUpdates = Object.keys(updates).length + deletions.length;
  if (numUpdates === 0) {
    return;
  }

  await usersRef.update(updates);
  await Promise.all(deletions.map((token) => usersRef.child(token).remove()));
  log.info(
    { users: Object.keys(updates).length, deletions: deletions.length },
    "Flushed changes to Firebase"
  );
}

function getModelFamilyForQuotaUsage(
  model: string,
  api: APIFormat
): ModelFamily {
  // TODO: this seems incorrect
  if (model.includes("azure")) return getAzureOpenAIModelFamily(model);

  switch (api) {
    case "openai":
    case "openai-text":
    case "openai-image":
      return getOpenAIModelFamily(model);
    case "anthropic":
      return getClaudeModelFamily(model);
    case "google-ai":
      return getGoogleAIModelFamily(model);
    case "mistral-ai":
      return getMistralAIModelFamily(model);
    default:
      assertNever(api);
  }
}

function getRefreshCrontab() {
  switch (config.quotaRefreshPeriod!) {
    case "hourly":
      return "0 * * * *";
    case "daily":
      return "0 0 * * *";
    default:
      return config.quotaRefreshPeriod ?? "0 0 * * *";
  }
}


code:
import { Query } from "express-serve-static-core";
import sanitize from "sanitize-html";
import { z } from "zod";

export function parseSort(sort: Query["sort"]) {
  if (!sort) return null;
  if (typeof sort === "string") return sort.split(",");
  if (Array.isArray(sort)) return sort.splice(3) as string[];
  return null;
}

export function sortBy(fields: string[], asc = true) {
  return (a: any, b: any) => {
    for (const field of fields) {
      if (a[field] !== b[field]) {
        // always sort nulls to the end
        if (a[field] == null) return 1;
        if (b[field] == null) return -1;

        const valA = Array.isArray(a[field]) ? a[field].length : a[field];
        const valB = Array.isArray(b[field]) ? b[field].length : b[field];

        const result = valA < valB ? -1 : 1;
        return asc ? result : -result;
      }
    }
    return 0;
  };
}

export function paginate(set: unknown[], page: number, pageSize: number = 20) {
  const p = Math.max(1, Math.min(page, Math.ceil(set.length / pageSize)));
  return {
    page: p,
    items: set.slice((p - 1) * pageSize, p * pageSize),
    pageSize,
    pageCount: Math.ceil(set.length / pageSize),
    totalCount: set.length,
    nextPage: p * pageSize < set.length ? p + 1 : null,
    prevPage: p > 1 ? p - 1 : null,
  };
}

export function sanitizeAndTrim(
  input?: string | null,
  options: sanitize.IOptions = {
    allowedTags: [],
    allowedAttributes: {},
  }
) {
  return sanitize((input ?? "").trim(), options);
}

// https://github.com/colinhacks/zod/discussions/2050#discussioncomment-5018870
export function makeOptionalPropsNullable<Schema extends z.AnyZodObject>(
  schema: Schema
) {
  const entries = Object.entries(schema.shape) as [
    keyof Schema["shape"],
    z.ZodTypeAny
  ][];
  const newProps = entries.reduce(
    (acc, [key, value]) => {
      acc[key] =
        value instanceof z.ZodOptional ? value.unwrap().nullable() : value;
      return acc;
    },
    {} as {
      [key in keyof Schema["shape"]]: Schema["shape"][key] extends z.ZodOptional<
        infer T
      >
        ? z.ZodNullable<T>
        : Schema["shape"][key];
    }
  );
  return z.object(newProps);
}

export function redactIp(ip: string) {
  const ipv6 = ip.includes(":");
  return ipv6 ? "redacted:ipv6" : ip.replace(/\.\d+\.\d+$/, ".xxx.xxx");
}

export function assertNever(x: never): never {
  throw new Error(`Called assertNever with argument ${x}.`);
}


code:
<% if (flashData) {
  let flashStyle = { title: "", style: "" };
  switch (flashData.type) {
    case "success":
      flashStyle.title = "✅ Success:";
      flashStyle.style = "color: green; background-color: #ddffee; padding: 1em";
      break;
    case "error":
      flashStyle.title = "⚠ Error:";
      flashStyle.style = "color: red; background-color: #eedddd; padding: 1em";
      break;
    case "warning":
      flashStyle.title = "⚠ Alert:";
      flashStyle.style = "color: darkorange; background-color: #ffeecc; padding: 1em";
      break;
    case "info":
      flashStyle.title = "ℹ Notice:";
      flashStyle.style = "color: blue; background-color: #ddeeff; padding: 1em";
      break;
  }
%>
<p style="<%= flashStyle.style %>">
  <%= flashStyle.title %> <%= flashData.message %>

<% } %>


code:
<!DOCTYPE html>
<html>
  <head>
    <meta charset="utf-8">
    <meta name="csrf-token" content="<%= csrfToken %>">
    <title><%= title %></title>
    <style>
      a:hover {
        background-color: #e0e6f6;
      }

      a:visited:hover {
        background-color: #e7e0f6;
      }

      .pagination {
        list-style-type: none;
        padding: 0;
      }
      .pagination li {
        display: inline-block;
      }
      .pagination li a {
        display: block;
        padding: 0.5em 1em;
        text-decoration: none;
      }
      .pagination li.active a {
        background-color: #58739c;
        color: #fff;
      }

      table {
        border-collapse: collapse;
        border: 1px solid #ccc;
      }
      table.striped tr:nth-child(even) {
        background-color: #eaeaea
      }
      table td, table th {
        border: 1px solid #ccc;
        padding: 0.25em 0.5em;
      }
      th.active {
        background-color: #e0e6f6;
      }
      td.actions {
        padding: 0;
        width: 0;
        text-align: center;
      }
      td.actions a {
        text-decoration: none;
        background-color: transparent;
        padding: 0.5em;
        height: 100%;
        width: 100%;
      }
      td.actions:hover {
        background-color: #e0e6f6;
      }

      @media (max-width: 600px) {
        table {
          width: 100%;
        }
        table td, table th {
          display: block;
          width: 100%;
        }
      }
    </style>
  </head>
  <body style="font-family: sans-serif; background-color: #f0f0f0; padding: 1em;">
    <%- include("partials/shared_flash", { flashData: flash }) %>


code:
  Page Size
  <select id="pageSize" onchange="setPageSize(this.value)" style="margin-bottom: 1rem;">
    <option value="10" <% if (pageSize === 10) { %>selected<% } %>>10</option>
    <option value="20" <% if (pageSize === 20) { %>selected<% } %>>20</option>
    <option value="50" <% if (pageSize === 50) { %>selected<% } %>>50</option>
    <option value="100" <% if (pageSize === 100) { %>selected<% } %>>100</option>
    <option value="200" <% if (pageSize === 200) { %>selected<% } %>>200</option>
  </select>


code:
Next refresh: <time><%- nextQuotaRefresh %></time>

      Model Family
      Usage
      <% if (showTokenCosts) { %>
      Cost
      <% } %>
      Limit
      Remaining
      Refresh Amount

    <% Object.entries(quota).forEach(([key, limit]) => { %>

      <%- key %>
      <%- prettyTokens(user.tokenCounts[key]) %>
      <% if (showTokenCosts) { %>
      $<%- tokenCost(key, user.tokenCounts[key]).toFixed(2) %>
      <% } %>
      <% if (!user.tokenLimits[key]) { %>
      unlimited
      <% } else { %>
      <%- prettyTokens(user.tokenLimits[key]) %>
      <%- prettyTokens(user.tokenLimits[key] - user.tokenCounts[key]) %>
      <% } %>
      <% if (user.type === "temporary") { %>
      N/A
      <% } else { %>
      <%- prettyTokens(quota[key]) %>
      <% } %>

    <% }) %>


code:
Show all (<%- user.ip.length %>)

  <% user.ip.forEach((ip) => { %>

<%- shouldRedact ? redactIp(ip) : ip %>

  <% }) %>


code:
import cookieParser from "cookie-parser";
import expressSession from "express-session";
import MemoryStore from "memorystore";
import { config, COOKIE_SECRET } from "../config";

const ONE_WEEK = 1000 * 60 * 60 * 24 * 7;

const cookieParserMiddleware = cookieParser(COOKIE_SECRET);

const sessionMiddleware = expressSession({
  secret: COOKIE_SECRET,
  resave: false,
  saveUninitialized: false,
  store: new (MemoryStore(expressSession))({ checkPeriod: ONE_WEEK }),
  cookie: {
    sameSite: "strict",
    maxAge: ONE_WEEK,
    signed: true,
    secure: !config.useInsecureCookies,
  },
});

const withSession = [cookieParserMiddleware, sessionMiddleware];

export { withSession };


code:
import express, { Router } from "express";
import { injectCsrfToken, checkCsrfToken } from "../shared/inject-csrf";
import { selfServiceRouter } from "./web/self-service";
import { injectLocals } from "../shared/inject-locals";
import { withSession } from "../shared/with-session";

const userRouter = Router();

userRouter.use(
  express.json({ limit: "1mb" }),
  express.urlencoded({ extended: true, limit: "1mb" })
);
userRouter.use(withSession);
userRouter.use(injectCsrfToken, checkCsrfToken);
userRouter.use(injectLocals);

userRouter.use(selfServiceRouter);

userRouter.use(
  (
    err: Error,
    _req: express.Request,
    res: express.Response,
    _next: express.NextFunction
  ) => {
    const data: any = { message: err.message, stack: err.stack, status: 500 };
    res.status(500).render("user_error", { ...data, flash: null });
  }
);

export { userRouter };


code:
import { Router } from "express";
import { UserPartialSchema } from "../../shared/users/schema";
import * as userStore from "../../shared/users/user-store";
import { ForbiddenError, UserInputError } from "../../shared/errors";
import { sanitizeAndTrim } from "../../shared/utils";
import { config } from "../../config";

const router = Router();

router.use((req, res, next) => {
  if (req.session.userToken) {
    res.locals.currentSelfServiceUser =
      userStore.getUser(req.session.userToken) || null;
  }
  next();
});

router.get("/", (_req, res) => {
  res.redirect("/");
});

router.get("/lookup", (_req, res) => {
  const ipLimit =
    (res.locals.currentSelfServiceUser?.maxIps ?? config.maxIpsPerUser) || 0;
  res.render("user_lookup", {
    user: res.locals.currentSelfServiceUser,
    ipLimit,
  });
});

router.post("/lookup", (req, res) => {
  const token = req.body.token;
  const user = userStore.getUser(token);
  req.log.info(
    { token: truncateToken(token), success: !!user },
    "User self-service lookup"
  );
  if (!user) {
    req.session.flash = { type: "error", message: "Invalid user token." };
    return res.redirect("/user/lookup");
  }
  req.session.userToken = user.token;
  return res.redirect("/user/lookup");
});

router.post("/edit-nickname", (req, res) => {
  const existing = res.locals.currentSelfServiceUser;

  if (!existing) {
    throw new ForbiddenError("Not logged in.");
  } else if (!config.allowNicknameChanges || existing.disabledAt) {
    throw new ForbiddenError("Nickname changes are not allowed.");
  } else if (!config.maxIpsAutoBan && !existing.ip.includes(req.ip)) {
    throw new ForbiddenError(
      "Nickname changes are only allowed from registered IPs."
    );
  }

  const schema = UserPartialSchema.pick({ nickname: true })
    .strict()
    .transform((v) => ({ nickname: sanitizeAndTrim(v.nickname) }));

  const result = schema.safeParse(req.body);
  if (!result.success) {
    throw new UserInputError(result.error.message);
  }

  const newNickname = result.data.nickname || null;
  userStore.upsertUser({ token: existing.token, nickname: newNickname });
  req.session.flash = { type: "success", message: "Nickname updated." };
  return res.redirect("/user/lookup");
});

function truncateToken(token: string) {
  const sliceLength = Math.max(Math.floor(token.length / 8), 1);
  return `${token.slice(0, sliceLength)}...${token.slice(-sliceLength)}`;
}

export { router as selfServiceRouter };


code:
<%- include("partials/shared_header", { title: "Error" }) %>

⚠ Error <%= status %>: <%= message %>

<%= stack %>

      Go Back


code:
<!DOCTYPE html>

  <meta charset="UTF-8">
  <meta name="robots" content="noindex" />
  <title><%= title %></title>

    <%= pageHeader %>

Service Info

<%= JSON.stringify(serviceInfo, null, 2) %>


code:
<%- include("partials/shared_header", { title: "User Token Lookup" }) %>

User Token Lookup

Provide your user token to check your usage and quota information.

<form action="/user/lookup" method="post">
  <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
  User Token
  <input type="password" name="token" value="<%= user?.token %>" />
  <input type="submit" value="Lookup" />
</form>
<% if (user) { %>

<% if (user.type === "temporary" && Boolean(user.disabledAt)) { %>
<%- include("partials/shared_flash", { flashData: {
      type: "info",
      message: "This temporary user token has expired and is no longer usable. These records will be deleted soon.",
    } }) %>
<% } else if (user.disabledAt) { %>
<%- include("partials/shared_flash", { flashData: {
      type: "warning",
      message: "This user token has been disabled." + (user.disabledReason ? ` Reason: ${user.disabledReason}` : ""),
    } }) %>
<% } %>

      User Token
       <%- "..." + user.token.slice(-5) %>

      Nickname
      <%- user.nickname ?? "none" %>

        ✏

      Type
      <%- user.type %>

      Prompts
      <%- user.promptCount %>

      Created At
      <%- user.createdAt %>

      Last Used At
      <%- user.lastUsedAt || "never" %>

      IPs<%- ipLimit ? ` (max ${ipLimit})` : "" %>
      <%- include("partials/shared_user_ip_list", { user, shouldRedact: true }) %>

    <% if (user.type === "temporary") { %>

      Expires At
      <%- user.expiresAt %>

    <% } %>

Quota Information

<%- include("partials/shared_quota-info", { quota, user }) %>

<form id="edit-nickname-form" style="display: none" action="/user/edit-nickname" method="post">
  <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
  <input type="hidden" name="nickname" value="<%= user.nickname %>" />
</form>

<% } %> <%- include("partials/user_footer") %>