code:
import crypto from "crypto";
import type * as http from "http";
import os from "os";
import schedule from "node-schedule";
import { config } from "../../config";
import { logger } from "../../logger";
import { KeyProviderBase } from "./key-provider-base";
import { getSerializer } from "./serializers";
import { FirebaseKeyStore, MemoryKeyStore } from "./stores";
import { AnthropicKeyProvider } from "./anthropic/provider";
import { OpenAIKeyProvider } from "./openai/provider";
import { GooglePalmKeyProvider } from "./palm/provider";
import { AwsBedrockKeyProvider } from "./aws/provider";
import { Key, KeyStore, LLMService, Model, ServiceToKey } from "./types";

export class KeyPool {
  private keyProviders: KeyProviderBase[] = [];
  private recheckJobs: Partial<Record<LLMService, schedule.Job | null>> = {
    openai: null,
  };

  constructor() {
    this.keyProviders.push(
      new OpenAIKeyProvider(createKeyStore("openai")),
      new AnthropicKeyProvider(createKeyStore("anthropic")),
      new GooglePalmKeyProvider(createKeyStore("google-palm")),
      new AwsBedrockKeyProvider(createKeyStore("aws"))
    );
  }

  public async init() {
    await Promise.all(this.keyProviders.map((p) => p.init()));

    const availableKeys = this.available("all");
    if (availableKeys === 0) {
      throw new Error("No keys loaded, the application cannot start.");
    }
    this.scheduleRecheck();
  }

  public get(model: Model): Key {
    const service = this.getService(model);
    return this.getKeyProvider(service).get(model);
  }

  public list(): Omit<Key, "key">[] {
    return this.keyProviders.flatMap((provider) => provider.list());
  }

  /**
   * Marks a key as disabled for a specific reason. `revoked` should be used
   * to indicate a key that can never be used again, while `quota` should be
   * used to indicate a key that is still valid but has exceeded its quota.
   */
  public disable(key: Key, reason: "quota" | "revoked"): void {
    const service = this.getKeyProvider(key.service);
    service.disable(key);
    service.update(key.hash, { isRevoked: reason === "revoked" });
    if (service instanceof OpenAIKeyProvider) {
      service.update(key.hash, { isOverQuota: reason === "quota" });
    }
  }

  public update<T extends Key>(key: T, props: Partial<T>): void {
    const service = this.getKeyProvider(key.service);
    service.update(key.hash, props);
  }

  public available(model: Model | "all" = "all"): number {
    return this.keyProviders.reduce((sum, provider) => {
      const includeProvider =
        model === "all" || this.getService(model) === provider.service;
      return sum + (includeProvider ? provider.available() : 0);
    }, 0);
  }

  public incrementUsage(key: Key, model: string, tokens: number): void {
    const provider = this.getKeyProvider(key.service);
    provider.incrementUsage(key.hash, model, tokens);
  }

  public getLockoutPeriod(model: Model): number {
    const service = this.getService(model);
    return this.getKeyProvider(service).getLockoutPeriod(model);
  }

  public markRateLimited(key: Key): void {
    const provider = this.getKeyProvider(key.service);
    provider.markRateLimited(key.hash);
  }

  public updateRateLimits(key: Key, headers: http.IncomingHttpHeaders): void {
    const provider = this.getKeyProvider(key.service);
    if (provider instanceof OpenAIKeyProvider) {
      provider.updateRateLimits(key.hash, headers);
    }
  }

  public recheck(service: LLMService): void {
    if (!config.checkKeys) {
      logger.info("Skipping key recheck because key checking is disabled");
      return;
    }

    const provider = this.getKeyProvider(service);
    provider.recheck();
  }

  private getService(model: Model): LLMService {
    if (model.startsWith("gpt") || model.startsWith("text-embedding-ada")) {
      // https://platform.openai.com/docs/models/model-endpoint-compatibility
      return "openai";
    } else if (model.startsWith("claude-")) {
      // https://console.anthropic.com/docs/api/reference#parameters
      return "anthropic";
    } else if (model.includes("bison")) {
      // https://developers.generativeai.google.com/models/language
      return "google-palm";
    } else if (model.startsWith("anthropic.claude")) {
      // AWS offers models from a few providers
      // https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html
      return "aws";
    }
    throw new Error(`Unknown service for model '${model}'`);
  }

  private getKeyProvider(service: LLMService): KeyProviderBase {
    return this.keyProviders.find((provider) => provider.service === service)!;
  }

  /**
   * Schedules a periodic recheck of OpenAI keys, which runs every 8 hours on
   * a schedule offset by the server's hostname.
   */
  private scheduleRecheck(): void {
    const machineHash = crypto
      .createHash("sha256")
      .update(os.hostname())
      .digest("hex");
    const offset = parseInt(machineHash, 16) % 7;
    const hour = [0, 8, 16].map((h) => h + offset).join(",");
    const crontab = `0 ${hour} * * *`;

    const job = schedule.scheduleJob(crontab, () => {
      const next = job.nextInvocation();
      logger.info({ next }, "Performing periodic recheck of OpenAI keys");
      this.recheck("openai");
    });
    logger.info(
      { rule: crontab, next: job.nextInvocation() },
      "Scheduled periodic key recheck job"
    );
    this.recheckJobs.openai = job;
  }
}

function createKeyStore<S extends LLMService>(
  service: S
): KeyStore<ServiceToKey[S]> {
  const serializer = getSerializer(service);

  switch (config.persistenceProvider) {
    case "memory":
      return new MemoryKeyStore(service, serializer);
    case "firebase_rtdb":
      return new FirebaseKeyStore(service, serializer);
    default:
      throw new Error(`Unknown store type: ${config.persistenceProvider}`);
  }
}

export let keyPool: KeyPool;

export async function init() {
  keyPool = new KeyPool();
  await keyPool.init();
}


code:
import { logger } from "../../logger";
import { Key, KeyStore, LLMService, Model } from "./types";

export abstract class KeyProviderBase<K extends Key = Key> {
  public abstract readonly service: LLMService;

  protected abstract readonly keys: K[];
  protected abstract log: typeof logger;
  protected readonly store: KeyStore<K>;

  public constructor(keyStore: KeyStore<K>) {
    this.store = keyStore;
  }

  public abstract init(): Promise<void>;

  public addKey(key: K): void {
    this.keys.push(key);
    this.store.add(key);
  }

  public abstract get(model: Model): K;

  /**
   * Returns a list of all keys, with the actual key value removed. Don't
   * mutate the returned objects; use `update` instead to ensure the changes
   * are synced to the key store.
   */
  public list(): Omit<K, "key">[] {
    return this.keys.map((k) => Object.freeze({ ...k, key: undefined }));
  }

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

  public update(hash: string, update: Partial<K>, force = false): void {
    const key = this.keys.find((k) => k.hash === hash);
    if (!key) {
      throw new Error(`No key with hash ${hash}`);
    }

    Object.assign(key, { lastChecked: Date.now(), ...update });
    this.store.update(hash, update, force);
  }

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

  public abstract incrementUsage(
    hash: string,
    model: string,
    tokens: number
  ): void;

  public abstract getLockoutPeriod(model: Model): number;

  public abstract markRateLimited(hash: string): void;

  public abstract recheck(): void;
}

license: unknown


code:
import { Key, KeySerializer, SerializedKey } from "./types";

export abstract class KeySerializerBase<K extends Key>
  implements KeySerializer<K>
{
  protected constructor(protected serializableFields: (keyof K)[]) {}

  serialize(keyObj: K): SerializedKey {
    return {
      ...Object.fromEntries(
        this.serializableFields
          .map((f) => [f, keyObj[f]])
          .filter(([, v]) => v !== undefined)
      ),
      key: keyObj.key,
    };
  }

  partialSerialize(key: string, update: Partial<K>): Partial<SerializedKey> {
    return {
      ...Object.fromEntries(
        this.serializableFields
          .map((f) => [f, update[f]])
          .filter(([, v]) => v !== undefined)
      ),
      key,
    };
  }

  abstract deserialize(serializedKey: SerializedKey): K;
}


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

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;
  private readonly updateKey: UpdateFn;

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

  protected async checkKey(key: OpenAIKey) {
    if (key.isDisabled) {
      this.log.warn({ key: key.hash }, "Skipping check for disabled key.");
      this.scheduleNextCheck();
      return;
    }

    this.log.debug({ key: key.hash }, "Checking key...");
    let isInitialCheck = !key.lastChecked;
    try {
      // We only need to check for provisioned models on the initial check.
      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 },
        "Key check complete."
      );
    } catch (error) {
      // touch the key so we don't check it again for a while
      this.updateKey(key.hash, {});
      this.handleAxiosError(key, error as AxiosError);
    }

    this.lastCheck = Date.now();
    // Only enqueue the next check if this wasn't a startup check, since those
    // are batched together elsewhere.
    if (!isInitialCheck) {
      this.log.info(
        { key: key.hash },
        "Recurring keychecks are disabled, no-op."
      );
      // this.scheduleNextCheck();
    }
  }

  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: OpenAIModelFamily[] = [];
    if (models.some(({ id }) => id.startsWith("gpt-3.5-turbo"))) {
      families.push("turbo");
    }

    if (models.some(({ id }) => id.startsWith("gpt-4"))) {
      families.push("gpt4");
    }

    if (models.some(({ id }) => id.startsWith("gpt-4-32k"))) {
      families.push("gpt4-32k");
    }

    // 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 keyFromPool = this.keys.find((k) => k.hash === key.hash)!;
    this.updateKey(key.hash, {
      modelFamilies: families,
      lastChecked: keyFromPool.lastChecked,
    });
    return families;
  }

  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);
  }

  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 }),
    };
  }
}


code:
import crypto from "crypto";
import { IncomingHttpHeaders } from "http";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { getOpenAIModelFamily, OpenAIModelFamily } from "../../models";
import { Key, Model } from "../types";
import { OpenAIKeyChecker } from "./checker";
import { KeyProviderBase } from "../key-provider-base";

const KEY_REUSE_DELAY = 1000;

export const OPENAI_SUPPORTED_MODELS = [
  "gpt-3.5-turbo",
  "gpt-3.5-turbo-instruct",
  "gpt-4",
  "gpt-4-32k",
  "text-embedding-ada-002",
] as const;
export type OpenAIModel = (typeof OPENAI_SUPPORTED_MODELS)[number];

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

export class OpenAIKeyProvider extends KeyProviderBase<OpenAIKey> {
  readonly service = "openai" as const;

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

  public async init() {
    const storeName = this.store.constructor.name;
    const loadedKeys = await this.store.load();

    // TODO: after key management UI, keychecker should always be enabled
    // because keys may be added after initialization.

    if (loadedKeys.length === 0) {
      return this.log.warn({ via: storeName }, "No OpenAI keys found.");
    }

    this.keys.push(...loadedKeys);
    this.log.info(
      { count: this.keys.length, via: storeName },
      "Loaded OpenAI keys."
    );

    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();
    }
  }

  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 active keys available for ${neededFamily} models.`);
    }

    if (!config.allowedModelFamilies.includes(neededFamily)) {
      throw new Error(
        `Proxy operator has disabled access to ${neededFamily} models.`
      );
    }

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

    // logger.debug(
    //   {
    //     byPriority: keysByPriority.map((k) => ({
    //       hash: k.hash,
    //       isRateLimited: now - k.rateLimitedAt < rateLimitThreshold,
    //       modelFamilies: k.modelFamilies,
    //     })),
    //   },
    //   "Keys sorted by priority"
    // );

    const selectedKey = keysByPriority[0];
    selectedKey.lastUsed = now;

    // When a key is selected, we rate-limit it for a brief period of time to
    // prevent the queue processor from immediately flooding it with requests
    // while the initial request is still being processed (which is when we will
    // get new rate limit headers).
    // Instead, we will let a request through every second until the key
    // becomes fully saturated and locked out again.
    selectedKey.rateLimitedAt = now;
    selectedKey.rateLimitRequestsReset = KEY_REUSE_DELAY;
    return { ...selectedKey };
  }

  /** 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,
        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;
    });
    clones.forEach((clone) => this.addKey(clone));
  }

  /**
   * 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(model: Model = "gpt-4"): number {
    const neededFamily = getOpenAIModelFamily(model);
    const activeKeys = this.keys.filter(
      (key) => !key.isDisabled && key.modelFamilies.includes(neededFamily)
    );

    if (activeKeys.length === 0) {
      // If there are no active keys for this model we can't fulfill requests.
      // We'll return 0 to let the request through and return an error,
      // otherwise the request will be stuck in the queue forever.
      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 + 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.
    return Math.min(
      ...activeKeys.map((key) => {
        const resetTime = Math.max(
          key.rateLimitRequestsReset,
          key.rateLimitTokensReset
        );
        return key.rateLimitedAt + 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();
  }

  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: IncomingHttpHeaders) {
    const key = this.keys.find((k) => k.hash === keyHash)!;
    const requestsReset = headers["x-ratelimit-reset-requests"];
    const tokensReset = headers["x-ratelimit-reset-tokens"];

    // Sometimes OpenAI only sends one of the two rate limit headers, it's
    // unclear why.

    if (requestsReset && typeof requestsReset === "string") {
      this.log.debug(
        { key: key.hash, requestsReset },
        `Updating rate limit requests reset time`
      );
      key.rateLimitRequestsReset = getResetDurationMillis(requestsReset);
    }

    if (tokensReset && typeof tokensReset === "string") {
      this.log.debug(
        { key: key.hash, tokensReset },
        `Updating rate limit tokens reset time`
      );
      key.rateLimitTokensReset = getResetDurationMillis(tokensReset);
    }

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

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

/**
 * Converts reset string ("21.0032s" or "21ms") to a number of milliseconds.
 * Result is clamped to 10s even though the API returns up to 60s, because the
 * API returns the time until the entire quota is reset, even if a key may be
 * able to fulfill requests before then due to partial resets.
 **/
function getResetDurationMillis(resetDuration?: string): number {
  const match = resetDuration?.match(/(\d+(\.\d+)?)(s|ms)/);
  if (match) {
    const [, time, , unit] = match;
    const value = parseFloat(time);
    const result = unit === "s" ? value * 1000 : value;
    return Math.min(result, 10000);
  }
  return 0;
}


code:
import crypto from "crypto";
import type { OpenAIKey, SerializedKey } from "../index";
import { KeySerializerBase } from "../key-serializer-base";

const SERIALIZABLE_FIELDS: (keyof OpenAIKey)[] = [
  "key",
  "service",
  "hash",
  "organizationId",
  "promptCount",
  "gpt4Tokens",
  "gpt4-32kTokens",
  "turboTokens",
];
export type SerializedOpenAIKey = SerializedKey &
  Partial<Pick<OpenAIKey, (typeof SERIALIZABLE_FIELDS)[number]>>;

export class OpenAIKeySerializer extends KeySerializerBase<OpenAIKey> {
  constructor() {
    super(SERIALIZABLE_FIELDS);
  }

  deserialize({ key, ...rest }: SerializedOpenAIKey): OpenAIKey {
    return {
      key,
      service: "openai",
      modelFamilies: ["turbo" as const, "gpt4" as const],
      isTrial: false,
      isDisabled: false,
      isRevoked: false,
      isOverQuota: false,
      lastUsed: 0,
      lastChecked: 0,
      promptCount: 0,
      hash: `oai-${crypto
        .createHash("sha256")
        .update(key)
        .digest("hex")
        .slice(0, 8)}`,
      rateLimitedAt: 0,
      rateLimitRequestsReset: 0,
      rateLimitTokensReset: 0,
      turboTokens: 0,
      gpt4Tokens: 0,
      "gpt4-32kTokens": 0,
      ...rest,
    };
  }
}


code:
import { logger } from "../../../logger";
import type { GooglePalmModelFamily } from "../../models";
import { KeyProviderBase } from "../key-provider-base";
import { Key } from "../types";

const RATE_LIMIT_LOCKOUT = 2000;
const KEY_REUSE_DELAY = 500;

// https://developers.generativeai.google.com/models/language
export const GOOGLE_PALM_SUPPORTED_MODELS = ["text-bison-001"] as const;
export type GooglePalmModel = (typeof GOOGLE_PALM_SUPPORTED_MODELS)[number];

type GooglePalmKeyUsage = {
  [K in GooglePalmModelFamily as `${K}Tokens`]: number;
};

export interface GooglePalmKey extends Key, GooglePalmKeyUsage {
  readonly service: "google-palm";
  readonly modelFamilies: GooglePalmModelFamily[];
  /** The time at which this key was last rate limited. */
  rateLimitedAt: number;
  /** The time until which this key is rate limited. */
  rateLimitedUntil: number;
}

export class GooglePalmKeyProvider extends KeyProviderBase<GooglePalmKey> {
  readonly service = "google-palm";

  protected keys: GooglePalmKey[] = [];
  protected log = logger.child({ module: "key-provider", service: this.service });

  public async init() {
    const storeName = this.store.constructor.name;
    const loadedKeys = await this.store.load();

    if (loadedKeys.length === 0) {
      return this.log.warn({ via: storeName }, "No Google PaLM keys found.");
    }

    this.keys.push(...loadedKeys);
    this.log.info(
      { count: this.keys.length, via: storeName },
      "Loaded PaLM keys."
    );
  }

  public get(_model: GooglePalmModel) {
    const availableKeys = this.keys.filter((k) => !k.isDisabled);
    if (availableKeys.length === 0) {
      throw new Error("No Google PaLM 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;
    selectedKey.rateLimitedAt = now;
    // Intended to throttle the queue processor as otherwise it will just
    // flood the API with requests and we want to wait a sec to see if we're
    // going to get a rate limit error on this key.
    selectedKey.rateLimitedUntil = now + KEY_REUSE_DELAY;
    return { ...selectedKey };
  }

  public incrementUsage(hash: string, _model: string, tokens: number) {
    const key = this.keys.find((k) => k.hash === hash);
    if (!key) return;
    key.promptCount++;
    key.bisonTokens += tokens;
  }

  public getLockoutPeriod(_model: GooglePalmModel) {
    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() {}
}


code:
import crypto from "crypto";
import type { GooglePalmKey, SerializedKey } from "../index";
import { KeySerializerBase } from "../key-serializer-base";

const SERIALIZABLE_FIELDS: (keyof GooglePalmKey)[] = [
  "key",
  "service",
  "hash",
  "promptCount",
  "bisonTokens",
];
export type SerializedGooglePalmKey = SerializedKey &
  Partial<Pick<GooglePalmKey, (typeof SERIALIZABLE_FIELDS)[number]>>;

export class GooglePalmKeySerializer extends KeySerializerBase<GooglePalmKey> {
  constructor() {
    super(SERIALIZABLE_FIELDS);
  }

  deserialize(serializedKey: SerializedGooglePalmKey): GooglePalmKey {
    const { key, ...rest } = serializedKey;
    return {
      key,
      service: "google-palm" as const,
      modelFamilies: ["bison"],
      isDisabled: false,
      isRevoked: false,
      promptCount: 0,
      lastUsed: 0,
      rateLimitedAt: 0,
      rateLimitedUntil: 0,
      hash: `plm-${crypto
        .createHash("sha256")
        .update(key)
        .digest("hex")
        .slice(0, 8)}`,
      lastChecked: 0,
      bisonTokens: 0,
      ...rest,
    };
  }
}


code:
import { assertNever } from "../utils";
import {
  Key,
  KeySerializer,
  LLMService,
  SerializedKey,
  ServiceToKey,
} from "./types";
import { OpenAIKeySerializer } from "./openai/serializer";
import { AnthropicKeySerializer } from "./anthropic/serializer";
import { GooglePalmKeySerializer } from "./palm/serializer";
import { AwsBedrockKeySerializer } from "./aws/serializer";

export function assertSerializedKey(k: any): asserts k is SerializedKey {
  if (typeof k !== "object" || !k || typeof (k as any).key !== "string") {
    throw new Error("Invalid serialized key data");
  }
}

export function getSerializer<S extends LLMService>(
  service: S
): KeySerializer<ServiceToKey[S]>;
export function getSerializer(service: LLMService): KeySerializer<Key> {
  switch (service) {
    case "openai":
      return new OpenAIKeySerializer();
    case "anthropic":
      return new AnthropicKeySerializer();
    case "google-palm":
      return new GooglePalmKeySerializer();
    case "aws":
      return new AwsBedrockKeySerializer();
    default:
      assertNever(service);
  }
}


code:
import firebase from "firebase-admin";
import { config, getFirebaseApp } from "../../../config";
import { logger } from "../../../logger";
import { assertSerializedKey } from "../serializers";
import type {
  Key,
  KeySerializer,
  KeyStore,
  LLMService,
  SerializedKey,
} from "../types";
import { MemoryKeyStore } from "./index";

export class FirebaseKeyStore<K extends Key> implements KeyStore<K> {
  private readonly db: firebase.database.Database;
  private readonly log: typeof logger;
  private readonly pendingUpdates: Map<string, Partial<SerializedKey>>;
  private readonly root: string;
  private readonly serializer: KeySerializer<K>;
  private readonly service: LLMService;
  private flushInterval: NodeJS.Timeout | null = null;
  private keysRef: firebase.database.Reference | null = null;

  constructor(
    service: LLMService,
    serializer: KeySerializer<K>,
    app = getFirebaseApp()
  ) {
    this.db = firebase.database(app);
    this.log = logger.child({ module: "firebase-key-store", service });
    this.root = `keys/${config.firebaseRtdbRoot.toLowerCase()}/${service}`;
    this.serializer = serializer;
    this.service = service;
    this.pendingUpdates = new Map();
    this.scheduleFlush();
  }

  public async load(isMigrating = false): Promise<K[]> {
    const keysRef = this.db.ref(this.root);
    const snapshot = await keysRef.once("value");
    const keys = snapshot.val();
    this.keysRef = keysRef;

    if (!keys) {
      if (isMigrating) return [];
      this.log.warn("No keys found in Firebase. Migrating from environment.");
      await this.migrate();
      return this.load(true);
    }

    return Object.values(keys).map((k) => {
      assertSerializedKey(k);
      return this.serializer.deserialize(k);
    });
  }

  public add(key: K) {
    const serialized = this.serializer.serialize(key);
    this.pendingUpdates.set(key.hash, serialized);
    this.forceFlush();
  }

  public update(id: string, update: Partial<K>, force = false) {
    const existing = this.pendingUpdates.get(id) ?? {};
    Object.assign(existing, this.serializer.partialSerialize(id, update));
    this.pendingUpdates.set(id, existing);
    if (force) this.forceFlush();
  }

  private forceFlush() {
    if (this.flushInterval) clearInterval(this.flushInterval);
    this.flushInterval = setTimeout(() => this.flush(), 0);
  }

  private scheduleFlush() {
    if (this.flushInterval) clearInterval(this.flushInterval);
    this.flushInterval = setInterval(() => this.flush(), 1000 * 60 * 5);
  }

  private async flush() {
    if (!this.keysRef) {
      this.log.warn(
        { pendingUpdates: this.pendingUpdates.size },
        "Database not loaded yet. Skipping flush."
      );
      return this.scheduleFlush();
    }

    if (this.pendingUpdates.size === 0) {
      this.log.debug("No pending key updates to flush.");
      return this.scheduleFlush();
    }

    const updates: Record<string, Partial<SerializedKey>> = {};
    this.pendingUpdates.forEach((v, k) => (updates[k] = v));
    this.pendingUpdates.clear();
    console.log(updates);

    await this.keysRef.update(updates);

    this.log.debug(
      { count: Object.keys(updates).length },
      "Flushed pending key updates."
    );
    this.scheduleFlush();
  }

  private async migrate(): Promise<SerializedKey[]> {
    const keysRef = this.db.ref(this.root);
    const envStore = new MemoryKeyStore<K>(this.service, this.serializer);
    const keys = await envStore.load();

    if (keys.length === 0) {
      this.log.warn("No keys found in environment or Firebase.");
      return [];
    }

    const updates: Record<string, SerializedKey> = {};
    keys.forEach((k) => (updates[k.hash] = this.serializer.serialize(k)));
    await keysRef.update(updates);

    this.log.info({ count: keys.length }, "Migrated keys from environment.");
    return Object.values(updates);
  }
}


code:
export { FirebaseKeyStore } from "./firebase";
export { MemoryKeyStore } from "./memory";


code:
import { assertNever } from "../../utils";
import { Key, KeySerializer, KeyStore, LLMService } from "../types";

export class MemoryKeyStore<K extends Key> implements KeyStore<K> {
  private readonly env: string;
  private readonly serializer: KeySerializer<K>;

  constructor(service: LLMService, serializer: KeySerializer<K>) {
    switch (service) {
      case "anthropic":
        this.env = "ANTHROPIC_KEY";
        break;
      case "openai":
        this.env = "OPENAI_KEY";
        break;
      case "google-palm":
        this.env = "GOOGLE_PALM_KEY";
        break;
      case "aws":
        this.env = "AWS_CREDENTIALS";
        break;
      default:
        assertNever(service);
    }
    this.serializer = serializer;
  }

  public async load() {
    let envKeys: string[];
    envKeys = [
      ...new Set(process.env[this.env]?.split(",").map((k) => k.trim())),
    ];
    return envKeys
      .filter((k) => k)
      .map((k) => this.serializer.deserialize({ key: k }));
  }

  public add() {}

  public update() {}
}


code:
import type { OpenAIKey, OpenAIModel } from "./openai/provider";
import type { AnthropicKey, AnthropicModel } from "./anthropic/provider";
import type { GooglePalmKey, GooglePalmModel } from "./palm/provider";
import type { AwsBedrockKey, AwsBedrockModel } from "./aws/provider";
import type { ModelFamily } from "../models";

/** The request and response format used by a model's API. */
export type APIFormat = "openai" | "anthropic" | "google-palm" | "openai-text";
/**
 * The service that a model is hosted on; distinct because services like AWS
 * provide APIs from other service providers, but have their own authentication
 * and key management.
 */
export type LLMService = "openai" | "anthropic" | "google-palm" | "aws";

export type Model =
  | OpenAIModel
  | AnthropicModel
  | GooglePalmModel
  | AwsBedrockModel;

type AllKeys = OpenAIKey | AnthropicKey | GooglePalmKey | AwsBedrockKey;
export type ServiceToKey = {
  [K in AllKeys["service"]]: Extract<AllKeys, { service: K }>;
};
export type SerializedKey = { key: string };

export interface Key {
  /** The API key itself. Never log this, use `hash` instead. */
  readonly key: string;
  /** The service that this key is for. */
  service: LLMService;
  /** The model families that this key has access to. */
  modelFamilies: ModelFamily[];
  /** Whether this key is currently disabled for some reason. */
  isDisabled: boolean;
  /**
   * Whether this key specifically has been revoked. This is different from
   * `isDisabled` because a key can be disabled for other reasons, such as
   * exceeding its quota. A revoked key is assumed to be permanently disabled,
   * and KeyStore implementations should not return it when loading keys.
   */
  isRevoked: boolean;
  /** The number of prompts that have been sent with this key. */
  promptCount: number;
  /** The time at which this key was last used. */
  lastUsed: number;
  /** The time at which this key was last checked. */
  lastChecked: number;
  /** Hash of the key, for logging and to find the key in the pool. */
  hash: string;
}

export interface KeySerializer<K> {
  serialize(keyObj: K): SerializedKey;
  deserialize(serializedKey: SerializedKey): K;
  partialSerialize(key: string, update: Partial<K>): Partial<SerializedKey>;
}

export interface KeyStore<K extends Key> {
  load(): Promise<K[]>;
  add(key: K): void;
  update(id: string, update: Partial<K>, force?: boolean): void;
}


code:
import { logger } from "../logger";

export type OpenAIModelFamily = "turbo" | "gpt4" | "gpt4-32k";
export type AnthropicModelFamily = "claude";
export type GooglePalmModelFamily = "bison";
export type AwsBedrockModelFamily = "aws-claude";
export type ModelFamily =
  | OpenAIModelFamily
  | AnthropicModelFamily
  | GooglePalmModelFamily
  | AwsBedrockModelFamily;

export const MODEL_FAMILIES = (<A extends readonly ModelFamily[]>(
  arr: A & ([ModelFamily] extends [A[number]] ? unknown : never)
) => arr)([
  "turbo",
  "gpt4",
  "gpt4-32k",
  "claude",
  "bison",
  "aws-claude",
] as const);

export const OPENAI_MODEL_FAMILY_MAP: { [regex: string]: OpenAIModelFamily } = {
  "^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",
};

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

export function getClaudeModelFamily(_model: string): ModelFamily {
  return "claude";
}

export function getGooglePalmModelFamily(model: string): ModelFamily {
  if (model.match(/^\w+-bison-\d{3}$/)) return "bison";
  const stack = new Error().stack;
  logger.warn({ model, stack }, "Unmapped PaLM model family");
  return "bison";
}

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

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


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(0, 50000),
      entry.promptFlattened.slice(0, 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.info("Creating new index sheet.");
    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(e, "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 { 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-32k":
      cost = 0.00006;
      break;
    case "gpt4":
      cost = 0.00003;
      break;
    case "turbo":
      cost = 0.0000015;
      break;
    case "aws-claude":
    case "claude":
      cost = 0.00001102;
      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(2) + "b";
  }
}


code:
import { Request, Response } from "express";
import { IncomingMessage } from "http";
import { assertNever } from "./utils";

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 buildFakeSse(
  type: string,
  string: string,
  req: Request
) {
  let fakeEvent;
  const content = `\`\`\`\n[${type}: ${string}]\n\`\`\`\n`;

  switch (req.inboundApi) {
    case "openai":
      fakeEvent = {
        id: "chatcmpl-" + req.id,
        object: "chat.completion.chunk",
        created: Date.now(),
        model: req.body?.model,
        choices: [{ delta: { content }, index: 0, finish_reason: type }]
      };
      break;
    case "openai-text":
      fakeEvent = {
        id: "cmpl-" + req.id,
        object: "text_completion",
        created: Date.now(),
        choices: [
          { text: content, index: 0, logprobs: null, finish_reason: type }
        ],
        model: req.body?.model
      };
      break;
    case "anthropic":
      fakeEvent = {
        completion: content,
        stop_reason: type,
        truncated: false, // I've never seen this be true
        stop: null,
        model: req.body?.model,
        log_id: "proxy-req-" + req.id
      };
      break;
    case "google-palm":
      throw new Error("PaLM not supported as an inbound API format");
    default:
      assertNever(req.inboundApi);
  }

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

  return `data: ${JSON.stringify(fakeEvent)}\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 { OpenAIPromptMessage } from "./openai";
export { init, countTokens } from "./tokenizer";


code:
import { Tiktoken } from "tiktoken/lite";
import cl100k_base from "tiktoken/encoders/cl100k_base.json";

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 function getTokenCount(
  prompt: string | OpenAIPromptMessage[],
  model: string
) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  const gpt4 = model.startsWith("gpt-4");

  const tokensPerMessage = gpt4 ? 3 : 4;
  const tokensPerName = gpt4 ? 1 : -1; // turbo omits role if name is present

  let numTokens = 0;

  for (const message of prompt) {
    numTokens += tokensPerMessage;
    for (const key of Object.keys(message)) {
      {
        const value = message[key as keyof OpenAIPromptMessage];
        if (!value || typeof value !== "string") continue;
        // Break if we get a huge message or exceed the token limit to prevent
        // DoS.
        // 100k tokens allows for future 100k GPT-4 models and 500k characters
        // is just a sanity check
        if (value.length > 500000 || numTokens > 100000) {
          numTokens = 100000;
          return {
            tokenizer: "tiktoken (prompt length limit exceeded)",
            token_count: numTokens,
          };
        }

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

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

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

export type OpenAIPromptMessage = {
  name?: string;
  content: string;
  role: string;
};


code:
import { Request } from "express";
import { assertNever } from "../utils";
import {
  init as initClaude,
  getTokenCount as getClaudeTokenCount,
} from "./claude";
import {
  init as initOpenAi,
  getTokenCount as getOpenAITokenCount,
  OpenAIPromptMessage,
} from "./openai";
import { APIFormat } from "../key-management";

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

/** 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: OpenAIPromptMessage[]; completion?: never; service: "openai" }
  | {
      prompt: string;
      completion?: never;
      service: "openai-text" | "anthropic" | "google-palm";
    }
  | { prompt?: never; completion: string; service: APIFormat }
);

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 {
        ...getOpenAITokenCount(prompt ?? completion, req.body.model),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "google-palm":
      // TODO: Can't find a tokenization library for PaLM. There is an API
      // endpoint for it but it adds significant latency to the request.
      return {
        ...getOpenAITokenCount(prompt ?? completion, req.body.model),
        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),
  claude: z.number().optional().default(0),
  bison: 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 { MODEL_FAMILIES, ModelFamily } from "../models";
import { logger } from "../../logger";
import { User, UserTokenCounts, UserUpdate } from "./schema";

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

const INITIAL_TOKENS: Required<UserTokenCounts> = {
  turbo: 0,
  gpt4: 0,
  "gpt4-32k": 0,
  claude: 0,
  bison: 0,
  "aws-claude": 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.persistenceProvider }, "Initializing user store...");
  if (config.persistenceProvider === "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.persistenceProvider === "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,
  consumption: number
) {
  const user = users.get(token);
  if (!user) return;
  const modelFamily = getModelFamilyForQuotaUsage(model);
  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) {
  const user = users.get(token);
  if (!user || user.disabledAt) return;
  if (!user.ip.includes(ip)) user.ip.push(ip);

  const configIpLimit = user.maxIps ?? config.maxIpsPerUser;
  const ipLimit =
    user.type === "special" || !configIpLimit ? Infinity : configIpLimit;
  if (user.ip.length > ipLimit) {
    disableUser(token, "IP address limit exceeded.");
    return;
  }

  user.lastUsedAt = Date.now();
  usersToFlush.add(token);
  return user;
}

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

  const modelFamily = getModelFamilyForQuotaUsage(model);
  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"
  );
}

// TODO: use key-management/models.ts for family mapping
function getModelFamilyForQuotaUsage(model: string): ModelFamily {
  if (model.includes("32k")) {
    return "gpt4-32k";
  }
  if (model.startsWith("gpt-4")) {
    return "gpt4";
  }
  if (model.startsWith("gpt-3.5")) {
    return "turbo";
  }
  if (model.includes("bison")) {
    return "bison";
  }
  if (model.startsWith("claude")) {
    return "claude";
  }
  if(model.startsWith("anthropic.claude")) {
    return "aws-claude";
  }
  throw new Error(`Unknown quota model family for model ${model}`);
}

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 { 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 },
});

const withSession = [cookieParserMiddleware, sessionMiddleware];

export { withSession };


code:
import type { HttpRequest } from "@smithy/types";
import { Express } from "express-serve-static-core";
import { APIFormat, Key, LLMService } from "../shared/key-management";
import { User } from "../shared/users/schema";

declare global {
  namespace Express {
    interface Request {
      key?: Key;
      service?: LLMService;
      /** Denotes the format of the user's submitted request. */
      inboundApi: APIFormat;
      /** Denotes the format of the request being proxied to the API. */
      outboundApi: APIFormat;
      /** If the request comes from a RisuAI.xyz user, this is their token. */
      risuToken?: string;
      user?: User;
      isStreaming?: boolean;
      startTime: number;
      retryCount: number;
      queueOutTime?: number;
      onAborted?: () => void;
      proceed: () => void;
      heartbeatInterval?: NodeJS.Timeout;
      promptTokens?: number;
      outputTokens?: number;
      // TODO: remove later
      debug: Record<string, any>;
      signedRequest: HttpRequest;
    }
  }
}

declare module "express-session" {
  interface SessionData {
    adminToken?: string;
    userToken?: string;
    csrf?: string;
    flash?: { type: string; message: string };
  }
}


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);
  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.");
  }

  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");
});

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) %>