code:
import { Sha256 } from "@aws-crypto/sha256-js";
import { SignatureV4 } from "@smithy/signature-v4";
import { HttpRequest } from "@smithy/protocol-http";
import { AxiosError, AxiosHeaders, AxiosRequestConfig } from "axios";
import { URL } from "url";
import { config } from "../../../config";
import { getAwsBedrockModelFamily } from "../../models";
import { getAxiosInstance } from "../../network";
import { KeyCheckerBase } from "../key-checker-base";
import type { AwsBedrockKey, AwsBedrockKeyProvider } from "./provider";

const axios = getAxiosInstance();

type ParentModelId = string;
type AliasModelId = string;
type ModuleAliasTuple = [ParentModelId, ...AliasModelId[]];

const KNOWN_MODEL_IDS: ModuleAliasTuple[] = [
  ["anthropic.claude-instant-v1"],
  ["anthropic.claude-v2", "anthropic.claude-v2:1"],
  ["anthropic.claude-3-sonnet-20240229-v1:0"],
  ["anthropic.claude-3-haiku-20240307-v1:0"],
  ["anthropic.claude-3-5-haiku-20241022-v1:0"],
  ["anthropic.claude-3-opus-20240229-v1:0"],
  ["anthropic.claude-3-5-sonnet-20240620-v1:0"],
  ["anthropic.claude-3-5-sonnet-20241022-v2:0"],
  ["mistral.mistral-7b-instruct-v0:2"],
  ["mistral.mixtral-8x7b-instruct-v0:1"],
  ["mistral.mistral-large-2402-v1:0"],
  ["mistral.mistral-large-2407-v1:0"],
  ["mistral.mistral-small-2402-v1:0"], // Seems to return 400
];

const KEY_CHECK_BATCH_SIZE = 2; // AWS checker needs to do lots of concurrent requests so should lower the batch size
const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds
const KEY_CHECK_PERIOD = 90 * 60 * 1000; // 90 minutes
const AMZ_HOST =
  process.env.AMZ_HOST || "bedrock-runtime.%REGION%.amazonaws.com";
const GET_CALLER_IDENTITY_URL = `https://sts.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15`;
const GET_INVOCATION_LOGGING_CONFIG_URL = (region: string) =>
  `https://bedrock.${region}.amazonaws.com/logging/modelinvocations`;
const GET_LIST_INFERENCE_PROFILES_URL = (region: string) =>
  `https://bedrock.${region}.amazonaws.com/inference-profiles?maxResults=1000`;
const POST_INVOKE_MODEL_URL = (region: string, model: string) =>
  `https://${AMZ_HOST.replace("%REGION%", region)}/model/${model}/invoke`;
const TEST_MESSAGES = [
  { role: "user", content: "Hi!" },
  { role: "assistant", content: "Hello!" },
];

type AwsError = { error: {} };

type GetInferenceProfilesResponse = {
  inferenceProfileSummaries: {
    inferenceProfileId: string;
    inferenceProfileName: string;
    inferenceProfileArn: string;
    description?: string;
    createdAt?: string;
    updatedAt?: string;
    status: "ACTIVE" | unknown;
    type: "SYSTEM_DEFINED" | unknown;
    models: {
      modelArn?: string;
    }[];
  }[];
};

type GetLoggingConfigResponse = {
  loggingConfig: null | {
    cloudWatchConfig: null | unknown;
    s3Config: null | unknown;
    embeddingDataDeliveryEnabled: boolean;
    imageDataDeliveryEnabled: boolean;
    textDataDeliveryEnabled: boolean;
  };
};

type UpdateFn = typeof AwsBedrockKeyProvider.prototype.update;

export class AwsKeyChecker extends KeyCheckerBase<AwsBedrockKey> {
  constructor(keys: AwsBedrockKey[], updateKey: UpdateFn) {
    super(keys, {
      service: "aws",
      keyCheckPeriod: KEY_CHECK_PERIOD,
      minCheckInterval: MIN_CHECK_INTERVAL,
      keyCheckBatchSize: KEY_CHECK_BATCH_SIZE,
      updateKey,
    });
  }

  protected async testKeyOrFail(key: AwsBedrockKey) {
    const isInitialCheck = !key.lastChecked;

    if (isInitialCheck) {
      try {
        await this.checkInferenceProfiles(key);
      } catch (e) {
        const asError = e as AxiosError<AwsError>;
        const data = asError.response?.data;
        this.log.warn(
          { key: key.hash, error: e.message, data },
          "Cannot list inference profiles.\n\
Principal may be missing `AmazonBedrockFullAccess`, or has no policy allowing action `bedrock:ListInferenceProfiles` against resource `arn:aws:bedrock:*:*:inference-profile/*`.\n\
Requests will be made without inference profiles using on-demand quotas, which may be subject to more restrictive rate limits.\n\
See https://docs.aws.amazon.com/bedrock/latest/userguide/cross-region-inference-prereq.html."
        );
      }
    }

    // Perform checks for all parent model IDs
    // TODO: use allsettled
    const results = await Promise.all(
      KNOWN_MODEL_IDS.filter(([model]) =>
        // Skip checks for models that are disabled anyway
        config.allowedModelFamilies.includes(getAwsBedrockModelFamily(model))
      ).map(async ([model, ...aliases]) => ({
        models: [model, ...aliases],
        success: await this.invokeModel(model, key),
      }))
    );

    // Filter out models that are disabled
    const modelIds = results
      .filter(({ success }) => success)
      .flatMap(({ models }) => models);

    if (modelIds.length === 0) {
      this.log.warn(
        { key: key.hash },
        "Key does not have access to any models; disabling."
      );
      return this.updateKey(key.hash, { isDisabled: true });
    }

    this.updateKey(key.hash, {
      modelIds,
      modelFamilies: Array.from(
        new Set(modelIds.map(getAwsBedrockModelFamily))
      ),
    });

    this.log.info(
      {
        key: key.hash,
        logged: key.awsLoggingStatus,
        families: key.modelFamilies,
        models: key.modelIds,
      },
      "Checked key."
    );
  }

  protected handleAxiosError(key: AwsBedrockKey, error: AxiosError) {
    if (error.response && AwsKeyChecker.errorIsAwsError(error)) {
      const errorHeader = error.response.headers["x-amzn-errortype"] as string;
      const errorType = errorHeader.split(":")[0];
      switch (errorType) {
        case "AccessDeniedException":
          // Indicates that the principal's attached policy does not allow them
          // to perform the requested action.
          // How we handle this depends on whether the action was one that we
          // must be able to perform in order to use the key.
          const path = new URL(error.config?.url!).pathname;
          const data = error.response.data;
          this.log.warn(
            { key: key.hash, type: errorType, path, data },
            "Key can't perform a required action; disabling."
          );
          return this.updateKey(key.hash, { isDisabled: true });
        case "UnrecognizedClientException":
          // This is a 403 error that indicates the key is revoked.
          this.log.warn(
            { key: key.hash, errorType, error: error.response.data },
            "Key is revoked; disabling."
          );
          return this.updateKey(key.hash, {
            isDisabled: true,
            isRevoked: true,
          });
        case "ThrottlingException":
          // This is a 429 error that indicates the key is rate-limited, but
          // not necessarily disabled. Retry in 10 seconds.
          this.log.warn(
            { key: key.hash, errorType, error: error.response.data },
            "Key is rate limited. Rechecking in 30 seconds."
          );
          const next = Date.now() - (KEY_CHECK_PERIOD - 30 * 1000);
          return this.updateKey(key.hash, { lastChecked: next });
        case "ValidationException":
        default:
          // This indicates some issue that we did not account for, possibly
          // a new ValidationException type. This likely means our key checker
          // needs to be updated so we'll just let the key through and let it
          // fail when someone tries to use it if the error is fatal.
          this.log.error(
            { key: key.hash, errorType, error: error.response.data },
            "Encountered unexpected error while checking key. This may indicate a change in the API; please report this."
          );
          return this.updateKey(key.hash, { lastChecked: Date.now() });
      }
    }
    const { response } = error;
    const { headers, status, data } = response ?? {};
    this.log.error(
      { key: key.hash, status, headers, data, 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 });
  }

  /**
   * Attempt to invoke the given model with the given key.  Returns true if the
   * key has access to the model, false if it does not. Throws an error if the
   * key is disabled.
   */
  private async invokeModel(
    model: string,
    key: AwsBedrockKey
  ): Promise<boolean> {
    if (model.includes("claude")) {
      // If inference profiles are available, try testing model with them.
      // If they are not available or the invocation fails with the inference
      // profile, fall back to regular model ID.
      const { region } = AwsKeyChecker.getCredentialsFromKey(key);
      const continent = region.split("-")[0];
      const profile = key.inferenceProfileIds.find(
        (id) => `${continent}.${model}` === id
      );

      if (profile) {
        this.log.debug(
          { key: key.hash, model, profile },
          "Testing model via inference profile."
        );
        let result: boolean;
        try {
          result = await this.testClaudeModel(key, profile);
        } catch (e) {
          this.log.error(
            { key: key.hash, model, profile, error: e.message },
            "InvokeModel via inference profile returned an error; trying model ID directly."
          );
          result = false;
        }

        // If the profile worked, we'll return success. Caller will add the
        // model (not the profile) to the list of enabled models, but the
        // profile will be used when the key is used for inference.
        if (result) return true;
      }
      this.log.debug({ key: key.hash, model }, "Testing model via model ID.");
      return this.testClaudeModel(key, model);
    } else if (model.includes("mistral")) {
      return this.testMistralModel(key, model);
    }
    throw new Error("AwsKeyChecker#invokeModel: no implementation for model");
  }

  private async testClaudeModel(
    key: AwsBedrockKey,
    model: string
  ): Promise<boolean> {
    const creds = AwsKeyChecker.getCredentialsFromKey(key);
    // This is not a valid invocation payload, but a 400 response indicates that
    // the principal at least has permission to invoke the model.
    // A 403 response indicates that the model is not accessible -- if none of
    // the models are accessible, the key is effectively disabled.
    const payload = {
      max_tokens: -1,
      messages: TEST_MESSAGES,
      anthropic_version: "bedrock-2023-05-31",
    };
    const config: AxiosRequestConfig = {
      method: "POST",
      url: POST_INVOKE_MODEL_URL(creds.region, model),
      data: payload,
      validateStatus: (status) => [400, 403, 404, 429, 503].includes(status),
    };
    config.headers = new AxiosHeaders({
      "content-type": "application/json",
      accept: "*/*",
    });
    await AwsKeyChecker.signRequestForAws(config, key);
    const response = await axios.request(config);
    const { data, status, headers } = response;
    const errorType = (headers["x-amzn-errortype"] as string).split(":")[0];
    const errorMessage = data?.message;

    // 503 ServiceUnavailableException errors are usually due to temporary
    // outages in the AWS infrastructure. However, because a 503 response also
    // indicates that the key can invoke the model, we can treat this as a
    // successful response.
    if (status === 503 && errorType.match(/ServiceUnavailableException/i)) {
      this.log.warn(
        { key: key.hash, model, errorType, data, status, headers },
        "Model is accessible, but may be temporarily unavailable."
      );
      return true;
    }

    // 429 ThrottlingException can suggest the model is available but the key
    // is being rate limited. I think if a key does not have access to the
    // model, it cannot receive a 429 response, so this should be a success.
    if (status === 429) {
      if (errorType.match(/ThrottlingException/i)) {
        this.log.debug(
          { key: key.hash, model, errorType, data, status, headers },
          "Model is available but key is rate limited."
        );
        return true;
      } else {
        throw new AxiosError(
          `InvokeModel returned 429 of type ${errorType}`,
          `AWS_INVOKE_MODEL_RATE_LIMITED`,
          response.config,
          response.request,
          response
        );
      }
    }

    // This message indicates the key is valid but this particular model is not
    // accessible. Other 403s may indicate the key is not usable.
    if (
      status === 403 &&
      errorMessage?.match(/access to the model with the specified model ID/)
    ) {
      this.log.debug(
        { key: key.hash, model, errorType, data, status, headers },
        "Model is not available (principal does not have access)."
      );
      return false;
    }

    // ResourceNotFound typically indicates that the tested model cannot be used
    // on the configured region for this set of credentials.
    if (status === 404) {
      this.log.debug(
        { region: creds.region, model, key: key.hash },
        "Model is not available (not supported in this AWS region)."
      );
      return false;
    }

    // We're looking for a specific error type and message here
    // "ValidationException"
    const correctErrorType = errorType === "ValidationException";
    const correctErrorMessage = errorMessage?.match(/max_tokens/);
    if (!correctErrorType || !correctErrorMessage) {
      this.log.debug(
        { key: key.hash, model, errorType, data, status },
        "Model is not available (request rejected)."
      );
      return false;
    }

    this.log.debug(
      { key: key.hash, model, errorType, data, status },
      "Model is available."
    );
    return true;
  }

  private async testMistralModel(
    key: AwsBedrockKey,
    model: string
  ): Promise<boolean> {
    const creds = AwsKeyChecker.getCredentialsFromKey(key);

    const payload = {
      max_tokens: -1,
      prompt: "<s>[INST] What is your favourite condiment? [/INST]</s>",
    };
    const config: AxiosRequestConfig = {
      method: "POST",
      url: POST_INVOKE_MODEL_URL(creds.region, model),
      data: payload,
      validateStatus: (status) => [400, 403, 404].includes(status),
      headers: {
        "content-type": "application/json",
        accept: "*/*",
      },
    };
    await AwsKeyChecker.signRequestForAws(config, key);
    const response = await axios.request(config);
    const { data, status, headers } = response;
    const errorType = (headers["x-amzn-errortype"] as string).split(":")[0];
    const errorMessage = data?.message;

    if (status === 403 || status === 404) {
      this.log.debug(
        { key: key.hash, model, errorType, data, status },
        "Model is not available (no access or unsupported region)."
      );
      return false;
    }

    const isBadRequest = status === 400;
    const isValidationError = errorMessage?.match(/validation error/i);
    if (isBadRequest && !isValidationError) {
      this.log.debug(
        { key: key.hash, model, errorType, data, status, headers },
        "Model is not available (request rejected)."
      );
      return false;
    }

    this.log.debug(
      { key: key.hash, model, errorType, data, status },
      "Model is available."
    );
    return true;
  }

  private async checkInferenceProfiles(key: AwsBedrockKey) {
    const creds = AwsKeyChecker.getCredentialsFromKey(key);
    const req: AxiosRequestConfig = {
      method: "GET",
      url: GET_LIST_INFERENCE_PROFILES_URL(creds.region),
      headers: { accept: "application/json" },
    };
    await AwsKeyChecker.signRequestForAws(req, key);
    const { data } = await axios.request<GetInferenceProfilesResponse>(req);
    const { inferenceProfileSummaries } = data;
    const profileIds = inferenceProfileSummaries.map(
      (p) => p.inferenceProfileId
    );
    this.log.debug(
      { key: key.hash, profileIds, region: creds.region },
      "Inference profiles found."
    );
    this.updateKey(key.hash, { inferenceProfileIds: profileIds });
  }

  private async checkLoggingConfiguration(key: AwsBedrockKey) {
    if (config.allowAwsLogging) {
      // Don't check logging status if we're allowing it to reduce API calls.
      this.updateKey(key.hash, { awsLoggingStatus: "unknown" });
      return true;
    }

    const creds = AwsKeyChecker.getCredentialsFromKey(key);
    const req: AxiosRequestConfig = {
      method: "GET",
      url: GET_INVOCATION_LOGGING_CONFIG_URL(creds.region),
      headers: { accept: "application/json" },
      validateStatus: () => true,
    };
    await AwsKeyChecker.signRequestForAws(req, key);
    const { data, status, headers } =
      await axios.request<GetLoggingConfigResponse>(req);

    let result: AwsBedrockKey["awsLoggingStatus"] = "unknown";

    if (status === 200) {
      const { loggingConfig } = data;
      const loggingEnabled = !!loggingConfig?.textDataDeliveryEnabled;
      this.log.debug(
        { key: key.hash, loggingConfig, loggingEnabled },
        "AWS model invocation logging test complete."
      );
      result = loggingEnabled ? "enabled" : "disabled";
    } else {
      const errorType = (headers["x-amzn-errortype"] as string).split(":")[0];
      this.log.debug(
        { key: key.hash, errorType, data, status },
        "Can't determine AWS model invocation logging status."
      );
    }

    this.updateKey(key.hash, { awsLoggingStatus: result });
    return !!result;
  }

  static errorIsAwsError(error: AxiosError): error is AxiosError<AwsError> {
    const headers = error.response?.headers;
    if (!headers) return false;
    return !!headers["x-amzn-errortype"];
  }

  /** Given an Axios request, sign it with the given key. */
  static async signRequestForAws(
    axiosRequest: AxiosRequestConfig,
    key: AwsBedrockKey,
    awsService = "bedrock"
  ) {
    const creds = AwsKeyChecker.getCredentialsFromKey(key);
    const { accessKeyId, secretAccessKey, region } = creds;
    const { method, url: axUrl, headers: axHeaders, data } = axiosRequest;
    const url = new URL(axUrl!);

    let plainHeaders = {};
    if (axHeaders instanceof AxiosHeaders) {
      plainHeaders = axHeaders.toJSON();
    } else if (typeof axHeaders === "object") {
      plainHeaders = axHeaders;
    }

    const request = new HttpRequest({
      method,
      protocol: "https:",
      hostname: url.hostname,
      path: url.pathname,
      query: Object.fromEntries(url.searchParams),
      headers: { Host: url.hostname, ...plainHeaders },
    });

    if (data) {
      request.body = JSON.stringify(data);
    }

    const signer = new SignatureV4({
      sha256: Sha256,
      credentials: { accessKeyId, secretAccessKey },
      region,
      service: awsService,
    });
    const signedRequest = await signer.sign(request);
    axiosRequest.headers = signedRequest.headers;
  }

  static getCredentialsFromKey(key: AwsBedrockKey) {
    const [accessKeyId, secretAccessKey, region] = key.key.split(":");
    if (!accessKeyId || !secretAccessKey || !region) {
      throw new Error("Invalid AWS Bedrock key");
    }
    return { accessKeyId, secretAccessKey, region };
  }
}


code:
import crypto from "crypto";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { PaymentRequiredError } from "../../errors";
import { AwsBedrockModelFamily, getAwsBedrockModelFamily } from "../../models";
import { createGenericGetLockoutPeriod, Key, KeyProvider } from "..";
import { prioritizeKeys } from "../prioritize-keys";
import { AwsKeyChecker } from "./checker";

type AwsBedrockKeyUsage = {
  [K in AwsBedrockModelFamily as `${K}Tokens`]: number;
};

export interface AwsBedrockKey extends Key, AwsBedrockKeyUsage {
  readonly service: "aws";
  readonly modelFamilies: AwsBedrockModelFamily[];
  /**
   * The confirmed logging status of this key. This is "unknown" until we
   * receive a response from the AWS API. Keys which are logged, or not
   * confirmed as not being logged, won't be used unless ALLOW_AWS_LOGGING is
   * set.
   */
  awsLoggingStatus: "unknown" | "disabled" | "enabled";
  modelIds: string[];
  inferenceProfileIds: string[];
}

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

export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
  readonly service = "aws";

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

  constructor() {
    const keyConfig = config.awsCredentials?.trim();
    if (!keyConfig) {
      this.log.warn(
        "AWS_CREDENTIALS is not set. AWS Bedrock API will not be available."
      );
      return;
    }
    let bareKeys: string[];
    bareKeys = [...new Set(keyConfig.split(",").map((k) => k.trim()))];
    for (const key of bareKeys) {
      const newKey: AwsBedrockKey = {
        key,
        service: this.service,
        modelFamilies: ["aws-claude"],
        isDisabled: false,
        isRevoked: false,
        promptCount: 0,
        lastUsed: 0,
        rateLimitedAt: 0,
        rateLimitedUntil: 0,
        awsLoggingStatus: "unknown",
        hash: `aws-${crypto
          .createHash("sha256")
          .update(key)
          .digest("hex")
          .slice(0, 8)}`,
        lastChecked: 0,
        modelIds: ["anthropic.claude-3-sonnet-20240229-v1:0"],
        inferenceProfileIds: [],
        ["aws-claudeTokens"]: 0,
        ["aws-claude-opusTokens"]: 0,
        ["aws-mistral-tinyTokens"]: 0,
        ["aws-mistral-smallTokens"]: 0,
        ["aws-mistral-mediumTokens"]: 0,
        ["aws-mistral-largeTokens"]: 0,
      };
      this.keys.push(newKey);
    }
    this.log.info({ keyCount: this.keys.length }, "Loaded AWS Bedrock keys.");
  }

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

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

  public get(model: string) {
    let neededVariantId = model;
    // This function accepts both Anthropic/Mistral IDs and AWS IDs.
    // Generally all AWS model IDs are supersets of the original vendor IDs.
    // Claude 2 is the only model that breaks this convention; Anthropic calls
    // it claude-2 but AWS calls it claude-v2.
    if (model.includes("claude-2")) neededVariantId = "claude-v2";
    const neededFamily = getAwsBedrockModelFamily(model);

    const availableKeys = this.keys.filter((k) => {
      // Select keys which
      return (
        // are enabled
        !k.isDisabled &&
        // are not logged, unless policy allows it
        (config.allowAwsLogging || k.awsLoggingStatus !== "enabled") &&
        // have access to the model family we need
        k.modelFamilies.includes(neededFamily) &&
        // have access to the specific variant we need
        k.modelIds.some((m) => m.includes(neededVariantId))
      );
    });

    this.log.debug(
      {
        requestedModel: model,
        selectedVariant: neededVariantId,
        selectedFamily: neededFamily,
        totalKeys: this.keys.length,
        availableKeys: availableKeys.length,
      },
      "Selecting AWS key"
    );

    if (availableKeys.length === 0) {
      throw new PaymentRequiredError(
        `No AWS Bedrock keys available for model ${model}`
      );
    }

    /**
     * Comparator for prioritizing keys on inference profile compatibility.
     * Requests made via inference profiles have higher rate limits so we want
     * to use keys with compatible inference profiles first.
     */
    const hasInferenceProfile = (
      a: AwsBedrockKey,
      b: AwsBedrockKey
    ) => {
      const aMatch = +a.inferenceProfileIds.some((p) => p.includes(model));
      const bMatch = +b.inferenceProfileIds.some((p) => p.includes(model));
      return aMatch - bMatch;
    };

    const selectedKey = prioritizeKeys(availableKeys, hasInferenceProfile)[0];
    selectedKey.lastUsed = Date.now();
    this.throttle(selectedKey.hash);
    return { ...selectedKey };
  }

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

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

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

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

  getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys);

  /**
   * 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() {
    this.keys.forEach(({ hash }) =>
      this.update(hash, { lastChecked: 0, isDisabled: false, isRevoked: false })
    );
    this.checker?.scheduleNextCheck();
  }

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

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

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

license: unknown


code:
import { AxiosError } from "axios";
import { getAzureOpenAIModelFamily } from "../../models";
import { getAxiosInstance } from "../../network";
import { KeyCheckerBase } from "../key-checker-base";
import type { AzureOpenAIKey, AzureOpenAIKeyProvider } from "./provider";

const axios = getAxiosInstance();

const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds
const KEY_CHECK_PERIOD = 60 * 60 * 1000; // 1 hour
const AZURE_HOST = process.env.AZURE_HOST || "%RESOURCE_NAME%.openai.azure.com";
const POST_CHAT_COMPLETIONS = (resourceName: string, deploymentId: string) =>
  `https://${AZURE_HOST.replace(
    "%RESOURCE_NAME%",
    resourceName
  )}/openai/deployments/${deploymentId}/chat/completions?api-version=2023-09-01-preview`;

type AzureError = {
  error: {
    message: string;
    type: string | null;
    param: string;
    code: string;
    status: number;
  };
};
type UpdateFn = typeof AzureOpenAIKeyProvider.prototype.update;

export class AzureOpenAIKeyChecker extends KeyCheckerBase<AzureOpenAIKey> {
  constructor(keys: AzureOpenAIKey[], updateKey: UpdateFn) {
    super(keys, {
      service: "azure",
      keyCheckPeriod: KEY_CHECK_PERIOD,
      minCheckInterval: MIN_CHECK_INTERVAL,
      recurringChecksEnabled: true,
      updateKey,
    });
  }

  protected async testKeyOrFail(key: AzureOpenAIKey) {
    const model = await this.testModel(key);
    this.log.info({ key: key.hash, deploymentModel: model }, "Checked key.");
    this.updateKey(key.hash, { modelFamilies: [model] });
  }

  protected handleAxiosError(key: AzureOpenAIKey, error: AxiosError) {
    if (error.response && AzureOpenAIKeyChecker.errorIsAzureError(error)) {
      const data = error.response.data;
      const errorType = data.error.code || data.error.type;
      switch (errorType) {
        case "DeploymentNotFound":
          this.log.warn(
            { key: key.hash, errorType, error: error.response.data },
            "Key is revoked or deployment ID is incorrect. Disabling key."
          );
          return this.updateKey(key.hash, {
            isDisabled: true,
            isRevoked: true,
          });
        case "401":
          this.log.warn(
            { key: key.hash, errorType, error: error.response.data },
            "Key is disabled or incorrect. Disabling key."
          );
          return this.updateKey(key.hash, {
            isDisabled: true,
            isRevoked: true,
          });
        case "429":
          const headers = error.response.headers;
          const retryAfter = Number(headers["retry-after"] || 0);
          if (retryAfter > 3600) {
            this.log.warn(
              { key: key.hash, errorType, error: error.response.data, headers },
              "Key has an excessive rate limit and will be disabled."
            );
            return this.updateKey(key.hash, { isDisabled: true });
          }
          this.log.warn(
            { key: key.hash, errorType, error: error.response.data, headers },
            "Key is rate limited. Rechecking key in 1 minute."
          );
          this.updateKey(key.hash, { lastChecked: Date.now() });
          setTimeout(async () => {
            this.log.info(
              { key: key.hash },
              "Rechecking Azure key after rate limit."
            );
            await this.checkKey(key);
          }, 1000 * 60);
          return;
        default:
          const { data: errorData, status: errorStatus } = error.response;
          this.log.error(
            { key: key.hash, errorType, errorData, errorStatus },
            "Unknown Azure API error while checking key. Please report this."
          );
          return this.updateKey(key.hash, { lastChecked: Date.now() });
      }
    }

    const { response, code } = error;
    if (code === "ENOTFOUND") {
      this.log.warn(
        { key: key.hash, error: error.message },
        "Resource name is probably incorrect. Disabling key."
      );
      return this.updateKey(key.hash, { isDisabled: true, isRevoked: true });
    }

    const { headers, status, data } = response ?? {};
    this.log.error(
      { key: key.hash, status, headers, data, error: error.stack },
      "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 });
  }

  private async testModel(key: AzureOpenAIKey) {
    const { apiKey, deploymentId, resourceName } =
      AzureOpenAIKeyChecker.getCredentialsFromKey(key);
    const url = POST_CHAT_COMPLETIONS(resourceName, deploymentId);
    const testRequest = {
      max_tokens: 1,
      stream: false,
      messages: [{ role: "user", content: "" }],
    };
    const response = await axios.post(url, testRequest, {
      headers: { "Content-Type": "application/json", "api-key": apiKey },
      validateStatus: (status) => status === 200 || status === 400,
    });
    const { data } = response;

    // We allow one 400 condition, OperationNotSupported, which is returned when
    // we try to invoke /chat/completions on dall-e-3. This is expected and
    // indicates a DALL-E deployment.
    if (response.status === 400) {
      if (data.error.code === "OperationNotSupported") return "azure-dall-e";
      throw new AxiosError(
        `Unexpected error when testing deployment ${deploymentId}`,
        "AZURE_TEST_ERROR",
        response.config,
        response.request,
        response
      );
    }

    const family = getAzureOpenAIModelFamily(data.model);
    this.updateKey(key.hash, { modelIds: [data.model] });

    // Azure returns "gpt-4" even for GPT-4 Turbo, so we need further checks.
    // Otherwise we can use the model family Azure returned.
    if (family !== "azure-gpt4") {
      return family;
    }

    // Try to send an oversized prompt. GPT-4 Turbo can handle this but regular
    // GPT-4 will return a Bad Request error.
    const contextText = {
      max_tokens: 9000,
      stream: false,
      temperature: 0,
      seed: 0,
      messages: [{ role: "user", content: "" }],
    };
    const { data: contextTest, status } = await axios.post(url, contextText, {
      headers: { "Content-Type": "application/json", "api-key": apiKey },
      validateStatus: (status) => status === 400 || status === 200,
    });
    const code = contextTest.error?.code;
    this.log.debug({ code, status }, "Performed Azure GPT4 context size test.");

    if (code === "context_length_exceeded") return "azure-gpt4";
    return "azure-gpt4-turbo";
  }

  static errorIsAzureError(error: AxiosError): error is AxiosError<AzureError> {
    const data = error.response?.data as any;
    return data?.error?.code || data?.error?.type;
  }

  static getCredentialsFromKey(key: AzureOpenAIKey) {
    const [resourceName, deploymentId, apiKey] = key.key.split(":");
    if (!resourceName || !deploymentId || !apiKey) {
      throw new Error(
        "Invalid Azure credential format. Refer to .env.example and ensure your credentials are in the format RESOURCE_NAME:DEPLOYMENT_ID:API_KEY with commas between each credential set."
      );
    }
    return { resourceName, deploymentId, apiKey };
  }
}


code:
import crypto from "crypto";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { PaymentRequiredError } from "../../errors";
import {
  AzureOpenAIModelFamily,
  getAzureOpenAIModelFamily,
} from "../../models";
import { createGenericGetLockoutPeriod, Key, KeyProvider } from "..";
import { prioritizeKeys } from "../prioritize-keys";
import { AzureOpenAIKeyChecker } from "./checker";

type AzureOpenAIKeyUsage = {
  [K in AzureOpenAIModelFamily as `${K}Tokens`]: number;
};

export interface AzureOpenAIKey extends Key, AzureOpenAIKeyUsage {
  readonly service: "azure";
  readonly modelFamilies: AzureOpenAIModelFamily[];
  contentFiltering: boolean;
  modelIds: string[];
}

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

export class AzureOpenAIKeyProvider implements KeyProvider<AzureOpenAIKey> {
  readonly service = "azure";

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

  constructor() {
    const keyConfig = config.azureCredentials;
    if (!keyConfig) {
      this.log.warn(
        "AZURE_CREDENTIALS is not set. Azure OpenAI API will not be available."
      );
      return;
    }
    let bareKeys: string[];
    bareKeys = [...new Set(keyConfig.split(",").map((k) => k.trim()))];
    for (const key of bareKeys) {
      const newKey: AzureOpenAIKey = {
        key,
        service: this.service,
        modelFamilies: ["azure-gpt4"],
        isDisabled: false,
        isRevoked: false,
        promptCount: 0,
        lastUsed: 0,
        rateLimitedAt: 0,
        rateLimitedUntil: 0,
        contentFiltering: false,
        hash: `azu-${crypto
          .createHash("sha256")
          .update(key)
          .digest("hex")
          .slice(0, 8)}`,
        lastChecked: 0,
        "azure-turboTokens": 0,
        "azure-gpt4Tokens": 0,
        "azure-gpt4-32kTokens": 0,
        "azure-gpt4-turboTokens": 0,
        "azure-gpt4oTokens": 0,
        "azure-o1Tokens": 0,
        "azure-o1-miniTokens": 0,
        "azure-dall-eTokens": 0,
        modelIds: [],
      };
      this.keys.push(newKey);
    }
    this.log.info({ keyCount: this.keys.length }, "Loaded Azure OpenAI keys.");
  }

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

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

  public get(model: string) {
    const neededFamily = getAzureOpenAIModelFamily(model);
    const availableKeys = this.keys.filter(
      (k) => !k.isDisabled && k.modelFamilies.includes(neededFamily)
    );
    if (availableKeys.length === 0) {
      throw new PaymentRequiredError(
        `No keys available for model family '${neededFamily}'.`
      );
    }

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

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

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

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

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

  getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys);

  /**
   * 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() {
    this.keys.forEach(({ hash }) =>
      this.update(hash, { lastChecked: 0, isDisabled: false, isRevoked: false })
    );
    this.checker?.scheduleNextCheck();
  }

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

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

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


code:
import { AxiosError } from "axios";
import { GcpModelFamily } from "../../models";
import { getAxiosInstance } from "../../network";
import { KeyCheckerBase } from "../key-checker-base";
import type { GcpKey, GcpKeyProvider } from "./provider";
import { getCredentialsFromGcpKey, refreshGcpAccessToken } from "./oauth";

const axios = getAxiosInstance();

const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds
const KEY_CHECK_PERIOD = 90 * 60 * 1000; // 90 minutes
const GCP_HOST = process.env.GCP_HOST || "%REGION%-aiplatform.googleapis.com";
const POST_STREAM_RAW_URL = (project: string, region: string, model: string) =>
  `https://${GCP_HOST.replace(
    "%REGION%",
    region
  )}/v1/projects/${project}/locations/${region}/publishers/anthropic/models/${model}:streamRawPredict`;
const TEST_MESSAGES = [
  { role: "user", content: "Hi!" },
  { role: "assistant", content: "Hello!" },
];

type UpdateFn = typeof GcpKeyProvider.prototype.update;

export class GcpKeyChecker extends KeyCheckerBase<GcpKey> {
  constructor(keys: GcpKey[], updateKey: UpdateFn) {
    super(keys, {
      service: "gcp",
      keyCheckPeriod: KEY_CHECK_PERIOD,
      minCheckInterval: MIN_CHECK_INTERVAL,
      recurringChecksEnabled: false,
      updateKey,
    });
  }

  protected async testKeyOrFail(key: GcpKey) {
    let checks: Promise<boolean>[] = [];
    const isInitialCheck = !key.lastChecked;
    if (isInitialCheck) {
      await this.maybeRefreshAccessToken(key);
      checks = [
        this.invokeModel("claude-3-haiku@20240307", key, true),
        this.invokeModel("claude-3-sonnet@20240229", key, true),
        this.invokeModel("claude-3-opus@20240229", key, true),
        this.invokeModel("claude-3-5-sonnet@20240620", key, true),
      ];

      const [sonnet, haiku, opus, sonnet35] = await Promise.all(checks);

      this.log.debug(
        { key: key.hash, sonnet, haiku, opus, sonnet35 },
        "GCP model initial tests complete."
      );

      const families: GcpModelFamily[] = [];
      if (sonnet || sonnet35 || haiku) families.push("gcp-claude");
      if (opus) families.push("gcp-claude-opus");

      if (families.length === 0) {
        this.log.warn(
          { key: key.hash },
          "Key does not have access to any models; disabling."
        );
        return this.updateKey(key.hash, { isDisabled: true });
      }

      this.updateKey(key.hash, {
        sonnetEnabled: sonnet,
        haikuEnabled: haiku,
        sonnet35Enabled: sonnet35,
        modelFamilies: families,
      });
    } else {
      await this.maybeRefreshAccessToken(key);
      if (key.haikuEnabled) {
        await this.invokeModel("claude-3-haiku@20240307", key, false);
      } else if (key.sonnetEnabled) {
        await this.invokeModel("claude-3-sonnet@20240229", key, false);
      } else if (key.sonnet35Enabled) {
        await this.invokeModel("claude-3-5-sonnet@20240620", key, false);
      } else {
        await this.invokeModel("claude-3-opus@20240229", key, false);
      }

      this.updateKey(key.hash, { lastChecked: Date.now() });
      this.log.debug({ key: key.hash }, "GCP key check complete.");
    }

    this.log.info(
      { key: key.hash, families: key.modelFamilies },
      "Checked key."
    );
  }

  protected handleAxiosError(key: GcpKey, error: AxiosError) {
    if (error.response && GcpKeyChecker.errorIsGcpError(error)) {
      const { status, data } = error.response;
      if (status === 400 || status === 401 || status === 403) {
        this.log.warn(
          { key: key.hash, error: data },
          "Key is invalid or revoked. Disabling key."
        );
        this.updateKey(key.hash, { isDisabled: true, isRevoked: true });
      } else if (status === 429) {
        this.log.warn(
          { key: key.hash, error: data },
          "Key is rate limited. Rechecking in a minute."
        );
        const next = Date.now() - (KEY_CHECK_PERIOD - 60 * 1000);
        this.updateKey(key.hash, { lastChecked: next });
      } 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;
    }
    const { response, cause } = error;
    const { headers, status, data } = response ?? {};
    this.log.error(
      { key: key.hash, status, headers, data, cause, 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 });
  }

  private async maybeRefreshAccessToken(key: GcpKey) {
    if (key.accessToken && key.accessTokenExpiresAt >= Date.now()) {
      return;
    }

    this.log.info({ key: key.hash }, "Refreshing GCP access token...");
    const [token, durationSec] = await refreshGcpAccessToken(key);
    this.updateKey(key.hash, {
      accessToken: token,
      accessTokenExpiresAt: Date.now() + durationSec * 1000 * 0.95,
    });
  }

  /**
   * Attempt to invoke the given model with the given key.  Returns true if the
   * key has access to the model, false if it does not. Throws an error if the
   * key is disabled.
   */
  private async invokeModel(model: string, key: GcpKey, initial: boolean) {
    const creds = await getCredentialsFromGcpKey(key);
    try {
      await this.maybeRefreshAccessToken(key);
    } catch (e) {
      this.log.error(
        { key: key.hash, error: e.message },
        "Could not test key due to error while getting access token."
      );
      return false;
    }

    const payload = {
      max_tokens: 1,
      messages: TEST_MESSAGES,
      anthropic_version: "vertex-2023-10-16",
    };
    const { data, status } = await axios.post(
      POST_STREAM_RAW_URL(creds.projectId, creds.region, model),
      payload,
      {
        headers: GcpKeyChecker.getRequestHeaders(key.accessToken),
        validateStatus: initial
          ? () => true
          : (status: number) => status >= 200 && status < 300,
      }
    );
    this.log.debug({ key: key.hash, data }, "Response from GCP");

    if (initial) {
      return (
        (status >= 200 && status < 300) || status === 429 || status === 529
      );
    }

    return true;
  }

  static errorIsGcpError(error: AxiosError): error is AxiosError {
    const data = error.response?.data as any;
    if (Array.isArray(data)) {
      return data.length > 0 && data[0]?.error?.message;
    } else {
      return data?.error?.message;
    }
  }

  static getRequestHeaders(accessToken: string) {
    return {
      Authorization: `Bearer ${accessToken}`,
      "Content-Type": "application/json",
    };
  }
}


code:
import crypto from "crypto";
import type { GcpKey } from "./provider";
import { getAxiosInstance } from "../../network";
import { logger } from "../../../logger";

const axios = getAxiosInstance();
const log = logger.child({ module: "gcp-oauth" });

const authUrl = "https://www.googleapis.com/oauth2/v4/token";
const scope = "https://www.googleapis.com/auth/cloud-platform";

type GoogleAuthResponse = {
  access_token: string;
  scope: string;
  token_type: "Bearer";
  expires_in: number;
};

type GoogleAuthError = {
  error:
    | "unauthorized_client"
    | "access_denied"
    | "admin_policy_enforced"
    | "invalid_client"
    | "invalid_grant"
    | "invalid_scope"
    | "disabled_client"
    | "org_internal";
  error_description: string;
};

export async function refreshGcpAccessToken(
  key: GcpKey
): Promise<[string, number]> {
  log.info({ key: key.hash }, "Entering GCP OAuth flow...");
  const { clientEmail, privateKey } = await getCredentialsFromGcpKey(key);

  // https://developers.google.com/identity/protocols/oauth2/service-account#authorizingrequests
  const jwt = await createSignedJWT(clientEmail, privateKey);
  log.info({ key: key.hash }, "Signed JWT, exchanging for access token...");
  const res = await axios.post<GoogleAuthResponse | GoogleAuthError>(
    authUrl,
    {
      grant_type: "urn:ietf:params:oauth:grant-type:jwt-bearer",
      assertion: jwt,
    },
    {
      headers: { "Content-Type": "application/x-www-form-urlencoded" },
      validateStatus: () => true,
    }
  );
  const status = res.status;
  const headers = res.headers;
  const data = res.data;

  if ("error" in data || status >= 400) {
    log.error(
      { key: key.hash, status, headers, data },
      "Error from Google Identity API while getting access token."
    );
    throw new Error(
      `Google Identity API returned error: ${(data as GoogleAuthError).error}`
    );
  }

  log.info({ key: key.hash, exp: data.expires_in }, "Got access token.");
  return [data.access_token, data.expires_in];
}

export async function getCredentialsFromGcpKey(key: GcpKey) {
  const [projectId, clientEmail, region, rawPrivateKey] = key.key.split(":");
  if (!projectId || !clientEmail || !region || !rawPrivateKey) {
    log.error(
      { key: key.hash },
      "Cannot parse GCP credentials. Ensure they are in the format PROJECT_ID:CLIENT_EMAIL:REGION:PRIVATE_KEY, and ensure no whitespace or newlines are in the private key."
    );
    throw new Error("Cannot parse GCP credentials.");
  }

  if (!key.privateKey) {
    await importPrivateKey(key, rawPrivateKey);
  }

  return { projectId, clientEmail, region, privateKey: key.privateKey! };
}

async function createSignedJWT(
  email: string,
  pkey: crypto.webcrypto.CryptoKey
) {
  const issued = Math.floor(Date.now() / 1000);
  const expires = issued + 600;

  const header = { alg: "RS256", typ: "JWT" };

  const payload = {
    iss: email,
    aud: authUrl,
    iat: issued,
    exp: expires,
    scope,
  };

  const encodedHeader = urlSafeBase64Encode(JSON.stringify(header));
  const encodedPayload = urlSafeBase64Encode(JSON.stringify(payload));

  const unsignedToken = `${encodedHeader}.${encodedPayload}`;

  const signature = await crypto.subtle.sign(
    "RSASSA-PKCS1-v1_5",
    pkey,
    new TextEncoder().encode(unsignedToken)
  );

  const encodedSignature = urlSafeBase64Encode(signature);
  return `${unsignedToken}.${encodedSignature}`;
}

async function importPrivateKey(key: GcpKey, rawPrivateKey: string) {
  log.info({ key: key.hash }, "Importing GCP private key...");
  const privateKey = rawPrivateKey
    .replace(
      /-----BEGIN PRIVATE KEY-----|-----END PRIVATE KEY-----|\r|\n|\\n/g,
      ""
    )
    .trim();
  const binaryKey = Buffer.from(privateKey, "base64");
  key.privateKey = await crypto.subtle.importKey(
    "pkcs8",
    binaryKey,
    { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" },
    true,
    ["sign"]
  );
  log.info({ key: key.hash }, "GCP private key imported.");
}

function urlSafeBase64Encode(data: string | ArrayBuffer): string {
  let base64: string;
  if (typeof data === "string") {
    base64 = btoa(
      encodeURIComponent(data).replace(/%([0-9A-F]{2})/g, (match, p1) =>
        String.fromCharCode(parseInt("0x" + p1, 16))
      )
    );
  } else {
    base64 = btoa(String.fromCharCode(...new Uint8Array(data)));
  }
  return base64.replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
}


code:
import crypto from "crypto";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { PaymentRequiredError } from "../../errors";
import { GcpModelFamily, getGcpModelFamily } from "../../models";
import { createGenericGetLockoutPeriod, Key, KeyProvider } from "..";
import { prioritizeKeys } from "../prioritize-keys";
import { GcpKeyChecker } from "./checker";

type GcpKeyUsage = {
  [K in GcpModelFamily as `${K}Tokens`]: number;
};

export interface GcpKey extends Key, GcpKeyUsage {
  readonly service: "gcp";
  readonly modelFamilies: GcpModelFamily[];
  sonnetEnabled: boolean;
  haikuEnabled: boolean;
  sonnet35Enabled: boolean;

  privateKey?: crypto.webcrypto.CryptoKey;
  /** Cached access token for GCP APIs. */
  accessToken: string;
  accessTokenExpiresAt: number;
}

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

export class GcpKeyProvider implements KeyProvider<GcpKey> {
  readonly service = "gcp";

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

  constructor() {
    const keyConfig = config.gcpCredentials?.trim();
    if (!keyConfig) {
      this.log.warn(
        "GCP_CREDENTIALS is not set. GCP API will not be available."
      );
      return;
    }
    let bareKeys: string[];
    bareKeys = [...new Set(keyConfig.split(",").map((k) => k.trim()))];
    for (const key of bareKeys) {
      const newKey: GcpKey = {
        key,
        service: this.service,
        modelFamilies: ["gcp-claude"],
        isDisabled: false,
        isRevoked: false,
        promptCount: 0,
        lastUsed: 0,
        rateLimitedAt: 0,
        rateLimitedUntil: 0,
        hash: `gcp-${crypto
          .createHash("sha256")
          .update(key)
          .digest("hex")
          .slice(0, 8)}`,
        lastChecked: 0,
        sonnetEnabled: true,
        haikuEnabled: false,
        sonnet35Enabled: false,
        accessToken: "",
        accessTokenExpiresAt: 0,
        ["gcp-claudeTokens"]: 0,
        ["gcp-claude-opusTokens"]: 0,
      };
      this.keys.push(newKey);
    }
    this.log.info({ keyCount: this.keys.length }, "Loaded GCP keys.");
  }

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

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

  public get(model: string) {
    const neededFamily = getGcpModelFamily(model);

    // this is a horrible mess
    // each of these should be separate model families, but adding model
    // families is not low enough friction for the rate at which gcp claude
    // model variants are added.
    const needsSonnet35 =
      model.includes("claude-3-5-sonnet") && neededFamily === "gcp-claude";
    const needsSonnet =
      !needsSonnet35 &&
      model.includes("sonnet") &&
      neededFamily === "gcp-claude";
    const needsHaiku = model.includes("haiku") && neededFamily === "gcp-claude";

    const availableKeys = this.keys.filter((k) => {
      return (
        !k.isDisabled &&
        (k.sonnetEnabled || !needsSonnet) && // sonnet and haiku are both under gcp-claude, while opus is not
        (k.haikuEnabled || !needsHaiku) &&
        (k.sonnet35Enabled || !needsSonnet35) &&
        k.modelFamilies.includes(neededFamily)
      );
    });

    this.log.debug(
      {
        model,
        neededFamily,
        needsSonnet,
        needsHaiku,
        needsSonnet35,
        availableKeys: availableKeys.length,
        totalKeys: this.keys.length,
      },
      "Selecting GCP key"
    );

    if (availableKeys.length === 0) {
      throw new PaymentRequiredError(
        `No GCP keys available for model ${model}`
      );
    }

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

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

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

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

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

  getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys);

  /**
   * 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() {
    this.keys.forEach(({ hash }) =>
      this.update(hash, { lastChecked: 0, isDisabled: false, isRevoked: false })
    );
    this.checker?.scheduleNextCheck();
  }

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

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

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


code:
import { AxiosError } from "axios";
import { GoogleAIModelFamily, getGoogleAIModelFamily } from "../../models";
import { getAxiosInstance } from "../../network";
import { KeyCheckerBase } from "../key-checker-base";
import type { GoogleAIKey, GoogleAIKeyProvider } from "./provider";

const axios = getAxiosInstance();

const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds
const KEY_CHECK_PERIOD = 6 * 60 * 60 * 1000; // 3 hours
const LIST_MODELS_URL =
  "https://generativelanguage.googleapis.com/v1beta/models";
const GENERATE_CONTENT_URL =
  "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=%KEY%";

type ListModelsResponse = {
  models: {
    name: string;
    baseModelId: string;
    version: string;
    displayName: string;
    description: string;
    inputTokenLimit: number;
    outputTokenLimit: number;
    supportedGenerationMethods: string[];
    temperature: number;
    maxTemperature: number;
    topP: number;
    topK: number;
  }[];
  nextPageToken: string;
};

type UpdateFn = typeof GoogleAIKeyProvider.prototype.update;

export class GoogleAIKeyChecker extends KeyCheckerBase<GoogleAIKey> {
  constructor(keys: GoogleAIKey[], updateKey: UpdateFn) {
    super(keys, {
      service: "google-ai",
      keyCheckPeriod: KEY_CHECK_PERIOD,
      minCheckInterval: MIN_CHECK_INTERVAL,
      recurringChecksEnabled: true,
      updateKey,
    });
  }

  protected async testKeyOrFail(key: GoogleAIKey) {
    const provisionedModels = await this.getProvisionedModels(key);
    await this.testGenerateContent(key);

    const updates = { modelFamilies: provisionedModels };
    this.updateKey(key.hash, updates);
    this.log.info(
      { key: key.hash, models: key.modelFamilies, ids: key.modelIds.length },
      "Checked key."
    );
  }

  private async getProvisionedModels(
    key: GoogleAIKey
  ): Promise<GoogleAIModelFamily[]> {
    const { data } = await axios.get<ListModelsResponse>(
      `${LIST_MODELS_URL}?pageSize=1000&key=${key.key}`
    );
    const models = data.models;

    const ids = new Set<string>();
    const families = new Set<GoogleAIModelFamily>();
    models.forEach(({ name }) => {
      families.add(getGoogleAIModelFamily(name));
      ids.add(name);
    });

    const familiesArray = Array.from(families);
    this.updateKey(key.hash, {
      modelFamilies: familiesArray,
      modelIds: Array.from(ids),
    });

    return familiesArray;
  }

  private async testGenerateContent(key: GoogleAIKey) {
    const payload = {
      contents: [{ parts: { text: "hello" }, role: "user" }],
      tools: [],
      safetySettings: [],
      generationConfig: { maxOutputTokens: 1 },
    };
    await axios.post(
      GENERATE_CONTENT_URL.replace("%KEY%", key.key),
      payload,
      { validateStatus: (status) => status === 200 }
    );
  }

  protected handleAxiosError(key: GoogleAIKey, error: AxiosError): void {
    if (error.response && GoogleAIKeyChecker.errorIsGoogleAIError(error)) {
      const httpStatus = error.response.status;
      const { code, message, status, details } = error.response.data.error;

      switch (httpStatus) {
        case 400: {
          const keyDeadMsgs = [
            /please enable billing/i,
            /API key not valid/i,
            /API key expired/i,
            /pass a valid API/i,
          ];
          const text = JSON.stringify(error.response.data.error);
          if (text.match(keyDeadMsgs.join("|"))) {
            this.log.warn(
              { key: key.hash, error: text },
              "Key check returned a non-transient 400 error. Disabling key."
            );
            this.updateKey(key.hash, { isDisabled: true, isRevoked: true });
            return;
          }
          break;
        }
        case 401:
        case 403:
          this.log.warn(
            { key: key.hash, status, code, message, details },
            "Key check returned Forbidden/Unauthorized error. Disabling key."
          );
          this.updateKey(key.hash, { isDisabled: true, isRevoked: true });
          return;
        case 429: {
          const text = JSON.stringify(error.response.data.error);

          const keyDeadMsgs = [
            /GenerateContentRequestsPerMinutePerProjectPerRegion/i,
            /"quota_limit_value":"0"/i,
          ];
          if (text.match(keyDeadMsgs.join("|"))) {
            this.log.warn(
              { key: key.hash, error: text },
              "Key check returned a non-transient 429 error. Disabling key."
            );
            this.updateKey(key.hash, { isDisabled: true, isRevoked: true });
            return;
          }

          this.log.warn(
            { key: key.hash, status, code, message, details },
            "Key is rate limited. Rechecking key in 1 minute."
          );
          const next = Date.now() - (KEY_CHECK_PERIOD - 60 * 1000);
          this.updateKey(key.hash, { lastChecked: next });
          return;
        }
      }

      this.log.error(
        { key: key.hash, status, code, message, details },
        "Encountered unexpected error status while checking key. This may indicate a change in the API; please report this."
      );
      return this.updateKey(key.hash, { lastChecked: Date.now() });
    }

    this.log.error(
      { key: key.hash, error: error.message },
      "Network error while checking key; trying this key again in a minute."
    );
    const oneMinute = 10 * 1000;
    const next = Date.now() - (KEY_CHECK_PERIOD - oneMinute);
    return this.updateKey(key.hash, { lastChecked: next });
  }

  static errorIsGoogleAIError(
    error: AxiosError
  ): error is AxiosError<GoogleAIError> {
    const data = error.response?.data as any;
    return data?.error?.code || data?.error?.status;
  }
}

type GoogleAIError = {
  error: {
    code: string;
    message: string;
    status: string;
    details: any[];
  };
};


code:
import crypto from "crypto";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { PaymentRequiredError } from "../../errors";
import { getGoogleAIModelFamily, type GoogleAIModelFamily } from "../../models";
import { createGenericGetLockoutPeriod, Key, KeyProvider } from "..";
import { prioritizeKeys } from "../prioritize-keys";
import { GoogleAIKeyChecker } from "./checker";

// Note that Google AI is not the same as Vertex AI, both are provided by
// Google but Vertex is the GCP product for enterprise, while Google API is a
// development/hobbyist product. They use completely different APIs and keys.
// https://ai.google.dev/docs/migrate_to_cloud

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

type GoogleAIKeyUsage = {
  [K in GoogleAIModelFamily as `${K}Tokens`]: number;
};

export interface GoogleAIKey extends Key, GoogleAIKeyUsage {
  readonly service: "google-ai";
  readonly modelFamilies: GoogleAIModelFamily[];
  /** All detected model IDs on this key. */
  modelIds: string[];
}

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

export class GoogleAIKeyProvider implements KeyProvider<GoogleAIKey> {
  readonly service = "google-ai";

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

  constructor() {
    const keyConfig = config.googleAIKey?.trim();
    if (!keyConfig) {
      this.log.warn(
        "GOOGLE_AI_KEY is not set. Google AI API will not be available."
      );
      return;
    }
    let bareKeys: string[];
    bareKeys = [...new Set(keyConfig.split(",").map((k) => k.trim()))];
    for (const key of bareKeys) {
      const newKey: GoogleAIKey = {
        key,
        service: this.service,
        modelFamilies: ["gemini-pro"],
        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,
        "gemini-flashTokens": 0,
        "gemini-proTokens": 0,
        "gemini-ultraTokens": 0,
        modelIds: [],
      };
      this.keys.push(newKey);
    }
    this.log.info({ keyCount: this.keys.length }, "Loaded Google AI keys.");
  }

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

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

  public get(model: string) {
    const neededFamily = getGoogleAIModelFamily(model);
    const availableKeys = this.keys.filter(
      (k) => !k.isDisabled && k.modelFamilies.includes(neededFamily)
    );
    if (availableKeys.length === 0) {
      throw new PaymentRequiredError("No Google AI keys available");
    }

    const keysByPriority = prioritizeKeys(availableKeys);

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

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

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

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

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

  getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys);

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

  public recheck() {}

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

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

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


code:
import type { LLMService, ModelFamily } from "../models";
import { KeyPool } from "./key-pool";

/** The request and response format used by a model's API. */
export type APIFormat =
  | "openai"
  | "openai-text"
  | "openai-image"
  | "anthropic-chat" // Anthropic's newer messages array format
  | "anthropic-text" // Legacy flat string prompt format
  | "google-ai"
  | "mistral-ai"
  | "mistral-text"

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, meaning its quota has been exceeded or it has been revoked. */
  isDisabled: boolean;
  /** Whether this key specifically has been revoked. */
  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;
  /** The time at which this key was last rate limited. */
  rateLimitedAt: number;
  /** The time until which this key is rate limited. */
  rateLimitedUntil: number;
}

/*
KeyPool and KeyProvider's similarities are a relic of the old design where
there was only a single KeyPool for OpenAI keys. Now that there are multiple
supported services, the service-specific functionality has been moved to
KeyProvider and KeyPool is just a wrapper around multiple KeyProviders,
delegating to the appropriate one based on the model requested.

Existing code will continue to call methods on KeyPool, which routes them to
the appropriate KeyProvider or returns data aggregated across all KeyProviders
for service-agnostic functionality.
*/

export interface KeyProvider<T extends Key = Key> {
  readonly service: LLMService;
  init(): void;
  get(model: string): T;
  list(): Omit<T, "key">[];
  disable(key: T): void;
  update(hash: string, update: Partial<T>): void;
  available(): number;
  incrementUsage(hash: string, model: string, tokens: number): void;
  getLockoutPeriod(model: ModelFamily): number;
  markRateLimited(hash: string): void;
  recheck(): void;
}

export function createGenericGetLockoutPeriod<T extends Key>(
  getKeys: () => T[]
) {
  return function (this: unknown, family?: ModelFamily): number {
    const keys = getKeys();
    const activeKeys = keys.filter(
      (k) => !k.isDisabled && (!family || k.modelFamilies.includes(family))
    );

    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;

    return Math.min(...activeKeys.map((k) => k.rateLimitedUntil - now));
  };
}

export const keyPool = new KeyPool();
export { AnthropicKey } from "./anthropic/provider";
export { AwsBedrockKey } from "./aws/provider";
export { GcpKey } from "./gcp/provider";
export { AzureOpenAIKey } from "./azure/provider";
export { GoogleAIKey } from "././google-ai/provider";
export { MistralAIKey } from "./mistral-ai/provider";
export { OpenAIKey } from "./openai/provider";


code:
import pino from "pino";
import { logger } from "../../logger";
import { Key } from "./index";
import { AxiosError } from "axios";

type KeyCheckerOptions<TKey extends Key = Key> = {
  service: string;
  keyCheckPeriod: number;
  minCheckInterval: number;
  keyCheckBatchSize?: number;
  recurringChecksEnabled?: boolean;
  updateKey: (hash: string, props: Partial<TKey>) => void;
};

export abstract class KeyCheckerBase<TKey extends Key> {
  protected readonly service: string;
  protected readonly recurringChecksEnabled: boolean;
  /** Minimum time in between any two key checks. */
  protected readonly minCheckInterval: number;
  /**
   * Minimum time in between checks for a given key. Because we can no longer
   * read quota usage, there is little reason to check a single key more often
   * than this.
   */
  protected readonly keyCheckPeriod: number;
  /** Maximum number of keys to check simultaneously. */
  protected readonly keyCheckBatchSize: number;
  protected readonly updateKey: (hash: string, props: Partial<TKey>) => void;
  protected readonly keys: TKey[] = [];
  protected log: pino.Logger;
  protected timeout?: NodeJS.Timeout;
  protected lastCheck = 0;

  protected constructor(keys: TKey[], opts: KeyCheckerOptions<TKey>) {
    this.keys = keys;
    this.keyCheckPeriod = opts.keyCheckPeriod;
    this.minCheckInterval = opts.minCheckInterval;
    this.recurringChecksEnabled = opts.recurringChecksEnabled ?? true;
    this.keyCheckBatchSize = opts.keyCheckBatchSize ?? 12;
    this.updateKey = opts.updateKey;
    this.service = opts.service;
    this.log = logger.child({ module: "key-checker", service: opts.service });
  }

  public start() {
    this.log.info("Starting key checker...");
    this.timeout = setTimeout(() => this.scheduleNextCheck(), 0);
  }

  public stop() {
    if (this.timeout) {
      this.log.debug("Stopping key checker...");
      clearTimeout(this.timeout);
    }
  }

  /**
   * Schedules the next check. If there are still keys yet to be checked, it
   * will schedule a check immediately for the next unchecked key. Otherwise,
   * it will schedule a check for the least recently checked key, respecting
   * the minimum check interval.
   */
  public scheduleNextCheck() {
    // Gives each concurrent check a correlation ID to make logs less confusing.
    const callId = Math.random().toString(36).slice(2, 8);
    const timeoutId = this.timeout?.[Symbol.toPrimitive]?.();
    const checkLog = this.log.child({ callId, timeoutId });

    const enabledKeys = this.keys.filter((key) => !key.isDisabled);
    const uncheckedKeys = enabledKeys.filter((key) => !key.lastChecked);
    const numEnabled = enabledKeys.length;
    const numUnchecked = uncheckedKeys.length;

    clearTimeout(this.timeout);
    this.timeout = undefined;

    if (!numEnabled) {
      checkLog.warn("All keys are disabled. Stopping.");
      return;
    }

    checkLog.debug({ numEnabled, numUnchecked }, "Scheduling next check...");

    if (numUnchecked > 0) {
      const keycheckBatch = uncheckedKeys.slice(0, this.keyCheckBatchSize);

      this.timeout = setTimeout(async () => {
        try {
          await Promise.all(keycheckBatch.map((key) => this.checkKey(key)));
        } catch (error) {
          checkLog.error({ error }, "Error checking one or more keys.");
        }
        checkLog.info("Batch complete.");
        this.scheduleNextCheck();
      }, 250);

      checkLog.info(
        {
          batch: keycheckBatch.map((k) => k.hash),
          remaining: uncheckedKeys.length - keycheckBatch.length,
          newTimeoutId: this.timeout?.[Symbol.toPrimitive]?.(),
        },
        "Scheduled batch of initial checks."
      );
      return;
    }

    if (!this.recurringChecksEnabled) {
      checkLog.info(
        "Initial checks complete and recurring checks are disabled for this service. Stopping."
      );
      return;
    }

    // Schedule the next check for the oldest key.
    const oldestKey = enabledKeys.reduce((oldest, key) =>
      key.lastChecked < oldest.lastChecked ? key : oldest
    );

    // Don't check any individual key too often.
    // Don't check anything at all more frequently than some minimum interval
    // even if keys still need to be checked.
    const nextCheck = Math.max(
      oldestKey.lastChecked + this.keyCheckPeriod,
      this.lastCheck + this.minCheckInterval
    );

    const baseDelay = nextCheck - Date.now();
    const jitter = (Math.random() - 0.5) * baseDelay * 0.5;
    const jitteredDelay = Math.max(1000, baseDelay + jitter);

    this.timeout = setTimeout(
      () => this.checkKey(oldestKey).then(() => this.scheduleNextCheck()),
      jitteredDelay
    );
    checkLog.debug(
      { key: oldestKey.hash, nextCheck: new Date(nextCheck), jitteredDelay },
      "Scheduled next recurring check."
    );
  }

  public async checkKey(key: TKey): Promise<void> {
    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...");

    try {
      await this.testKeyOrFail(key);
    } catch (error) {
      this.updateKey(key.hash, {});
      this.handleAxiosError(key, error as AxiosError);
    }

    this.lastCheck = Date.now();
  }

  protected abstract testKeyOrFail(key: TKey): Promise<void>;

  protected abstract handleAxiosError(key: TKey, error: AxiosError): void;
}


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 { LLMService, MODEL_FAMILY_SERVICE, ModelFamily } from "../models";
import { Key, KeyProvider } from "./index";
import { AnthropicKeyProvider, AnthropicKeyUpdate } from "./anthropic/provider";
import { OpenAIKeyProvider, OpenAIKeyUpdate } from "./openai/provider";
import { GoogleAIKeyProvider  } from "./google-ai/provider";
import { AwsBedrockKeyProvider } from "./aws/provider";
import { GcpKeyProvider, GcpKey } from "./gcp/provider";
import { AzureOpenAIKeyProvider } from "./azure/provider";
import { MistralAIKeyProvider } from "./mistral-ai/provider";

type AllowedPartial = OpenAIKeyUpdate | AnthropicKeyUpdate | Partial<GcpKey>;

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

  constructor() {
    this.keyProviders.push(new OpenAIKeyProvider());
    this.keyProviders.push(new AnthropicKeyProvider());
    this.keyProviders.push(new GoogleAIKeyProvider());
    this.keyProviders.push(new MistralAIKeyProvider());
    this.keyProviders.push(new AwsBedrockKeyProvider());
    this.keyProviders.push(new GcpKeyProvider());
    this.keyProviders.push(new AzureOpenAIKeyProvider());
  }

  public init() {
    this.keyProviders.forEach((provider) => provider.init());
    const availableKeys = this.available("all");
    if (availableKeys === 0) {
      throw new Error(
        "No keys loaded. Ensure that at least one key is configured."
      );
    }
    this.scheduleRecheck();
  }

  public get(model: string, service?: LLMService, multimodal?: boolean): Key {
    // hack for some claude requests needing keys with particular permissions
    // even though they use the same models as the non-multimodal requests
    if (multimodal) {
      model += "-multimodal";
    }

    const queryService = service || this.getServiceForModel(model);
    return this.getKeyProvider(queryService).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 instanceof AnthropicKeyProvider
    ) {
      service.update(key.hash, { isOverQuota: reason === "quota" });
    }
  }

  /**
   * Updates a key in the keypool with the given properties.
   *
   * Be aware that the `key` argument may not be the same object instance as the
   * one in the keypool (such as if it is a clone received via `KeyPool.get` in
   * which case you are responsible for updating your clone with the new
   * properties.
   */
  public update(key: Key, props: AllowedPartial): void {
    const service = this.getKeyProvider(key.service);
    service.update(key.hash, props);
  }

  public available(model: string | "all" = "all"): number {
    return this.keyProviders.reduce((sum, provider) => {
      const includeProvider =
        model === "all" || this.getServiceForModel(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(family: ModelFamily): number {
    const service = MODEL_FAMILY_SERVICE[family];
    return this.getKeyProvider(service).getLockoutPeriod(family);
  }

  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 getServiceForModel(model: string): LLMService {
    if (
      model.startsWith("gpt") ||
      model.startsWith("text-embedding-ada") ||
      model.startsWith("dall-e")
    ) {
      // https://platform.openai.com/docs/models/model-endpoint-compatibility
      return "openai";
    } else if (model.startsWith("claude-")) {
      // https://console.anthropic.com/docs/api/reference#parameters
      if (!model.includes('@')) {
        return "anthropic";
      } else {
        return "gcp";
      }
    } else if (model.includes("gemini")) {
      // https://developers.generativeai.google.com/models/language
      return "google-ai";
    } else if (model.includes("mistral")) {
      // https://docs.mistral.ai/platform/endpoints
      return "mistral-ai";
    } 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";
    } else if (model.startsWith("azure")) {
      return "azure";
    }
    throw new Error(`Unknown service for model '${model}'`);
  }

  private getKeyProvider(service: LLMService): KeyProvider {
    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.");
      this.recheck("openai");
      this.recheck("google-ai");
    });
    logger.info(
      { rule: crontab, next: job.nextInvocation() },
      "Scheduled periodic key recheck job"
    );
    this.recheckJobs.openai = job;
  }
}


code:
import { AxiosError } from "axios";
import { MistralAIModelFamily, getMistralAIModelFamily } from "../../models";
import { getAxiosInstance } from "../../network";
import { KeyCheckerBase } from "../key-checker-base";
import type { MistralAIKey, MistralAIKeyProvider } from "./provider";

const axios = getAxiosInstance();

const MIN_CHECK_INTERVAL = 3 * 1000; // 3 seconds
const KEY_CHECK_PERIOD = 60 * 60 * 1000; // 1 hour
const GET_MODELS_URL = "https://api.mistral.ai/v1/models";

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

type MistralAIError = {
  message: string;
  request_id: string;
};

type UpdateFn = typeof MistralAIKeyProvider.prototype.update;

export class MistralAIKeyChecker extends KeyCheckerBase<MistralAIKey> {
  constructor(keys: MistralAIKey[], updateKey: UpdateFn) {
    super(keys, {
      service: "mistral-ai",
      keyCheckPeriod: KEY_CHECK_PERIOD,
      minCheckInterval: MIN_CHECK_INTERVAL,
      recurringChecksEnabled: false,
      updateKey,
    });
  }

  protected async testKeyOrFail(key: MistralAIKey) {
    // We only need to check for provisioned models on the initial check.
    const isInitialCheck = !key.lastChecked;
    if (isInitialCheck) {
      const provisionedModels = await this.getProvisionedModels(key);
      const updates = {
        modelFamilies: provisionedModels,
      };
      this.updateKey(key.hash, updates);
    }
    this.log.info({ key: key.hash, models: key.modelFamilies }, "Checked key.");
  }

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

    const families = new Set<MistralAIModelFamily>();
    models.forEach(({ id }) => families.add(getMistralAIModelFamily(id)));

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

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

  protected handleAxiosError(key: MistralAIKey, error: AxiosError) {
    if (error.response && MistralAIKeyChecker.errorIsMistralAIError(error)) {
      const { status, data } = error.response;
      if ([401, 403].includes(status)) {
        this.log.warn(
          { key: key.hash, error: data, status },
          "Key is invalid or revoked. Disabling key."
        );
        this.updateKey(key.hash, {
          isDisabled: true,
          isRevoked: true,
          modelFamilies: ["mistral-tiny"],
        });
      } 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 });
  }

  static errorIsMistralAIError(
    error: AxiosError
  ): error is AxiosError<MistralAIError> {
    const data = error.response?.data as any;
    return data?.message && data?.request_id;
  }

  static getHeaders(key: MistralAIKey) {
    return {
      Authorization: `Bearer ${key.key}`,
    };
  }
}


code:
import crypto from "crypto";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { HttpError } from "../../errors";
import { MistralAIModelFamily, getMistralAIModelFamily } from "../../models";
import { createGenericGetLockoutPeriod, Key, KeyProvider } from "..";
import { prioritizeKeys } from "../prioritize-keys";
import { MistralAIKeyChecker } from "./checker";

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

export interface MistralAIKey extends Key, MistralAIKeyUsage {
  readonly service: "mistral-ai";
  readonly modelFamilies: MistralAIModelFamily[];
}

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

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

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

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

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

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

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

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

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

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

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

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

  getLockoutPeriod = createGenericGetLockoutPeriod(() => this.keys);

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

  public recheck() {}

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

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

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