code:
import { keyPool } from "../../../../shared/key-management";
import { ProxyReqMutator } from "../index";

export const addGoogleAIKey: ProxyReqMutator = (manager) => {
  const req = manager.request;
  const inboundValid =
    req.inboundApi === "openai" || req.inboundApi === "google-ai";
  const outboundValid = req.outboundApi === "google-ai";

  const serviceValid = req.service === "google-ai";
  if (!inboundValid || !outboundValid || !serviceValid) {
    throw new Error("addGoogleAIKey called on invalid request");
  }

  const model = req.body.model;
  const key = keyPool.get(model, "google-ai");
  manager.setKey(key);

  req.log.info(
    { key: key.hash, model, stream: req.isStreaming },
    "Assigned Google AI API key to request"
  );

  // https://generativelanguage.googleapis.com/v1beta/models/$MODEL_ID:generateContent?key=$API_KEY
  // https://generativelanguage.googleapis.com/v1beta/models/$MODEL_ID:streamGenerateContent?key=${API_KEY}
  const payload = { ...req.body, stream: undefined, model: undefined };

  // TODO: this isn't actually signed, so the manager api is a little unclear
  // with the ProxyReqManager refactor, it's probably no longer necesasry to
  // do this because we can modify the path using Manager.setPath.
  manager.setSignedRequest({
    method: "POST",
    protocol: "https:",
    hostname: "generativelanguage.googleapis.com",
    path: `/v1beta/models/${model}:${
      req.isStreaming ? "streamGenerateContent?alt=sse&" : "generateContent?"
    }key=${key.key}`,
    headers: {
      ["host"]: `generativelanguage.googleapis.com`,
      ["content-type"]: "application/json",
    },
    body: JSON.stringify(payload),
  });
};


code:
import { AnthropicChatMessage } from "../../../../shared/api-schemas";
import { containsImageContent } from "../../../../shared/api-schemas/anthropic";
import { Key, OpenAIKey, keyPool } from "../../../../shared/key-management";
import { isEmbeddingsRequest } from "../../common";
import { assertNever } from "../../../../shared/utils";
import { ProxyReqMutator } from "../index";

export const addKey: ProxyReqMutator = (manager) => {
  const req = manager.request;

  let assignedKey: Key;
  const { service, inboundApi, outboundApi, body } = req;

  if (!inboundApi || !outboundApi) {
    const err = new Error(
      "Request API format missing. Did you forget to add the request preprocessor to your router?"
    );
    req.log.error({ inboundApi, outboundApi, path: req.path }, err.message);
    throw err;
  }

  if (!body?.model) {
    throw new Error("You must specify a model with your request.");
  }

  let needsMultimodal = false;
  if (outboundApi === "anthropic-chat") {
    needsMultimodal = containsImageContent(
      body.messages as AnthropicChatMessage[]
    );
  }

  if (inboundApi === outboundApi) {
    assignedKey = keyPool.get(body.model, service, needsMultimodal);
  } else {
    switch (outboundApi) {
      // If we are translating between API formats we may need to select a model
      // for the user, because the provided model is for the inbound API.
      // TODO: This whole else condition is probably no longer needed since API
      // translation now reassigns the model earlier in the request pipeline.
      case "anthropic-text":
      case "anthropic-chat":
      case "mistral-ai":
      case "mistral-text":
      case "google-ai":
        assignedKey = keyPool.get(body.model, service);
        break;
      case "openai-text":
        assignedKey = keyPool.get("gpt-3.5-turbo-instruct", service);
        break;
      case "openai-image":
        assignedKey = keyPool.get("dall-e-3", service);
        break;
      case "openai":
        throw new Error(
          `Outbound API ${outboundApi} is not supported for ${inboundApi}`
        );
      default:
        assertNever(outboundApi);
    }
  }

  manager.setKey(assignedKey);
  req.log.info(
    { key: assignedKey.hash, model: body.model, inboundApi, outboundApi },
    "Assigned key to request"
  );

  // TODO: KeyProvider should assemble all necessary headers
  switch (assignedKey.service) {
    case "anthropic":
      manager.setHeader("X-API-Key", assignedKey.key);
      if (!manager.request.headers["anthropic-version"]) {
        manager.setHeader("anthropic-version", "2023-06-01");
      }
      break;
    case "openai":
      const key: OpenAIKey = assignedKey as OpenAIKey;
      if (key.organizationId && !key.key.includes("svcacct")) {
        manager.setHeader("OpenAI-Organization", key.organizationId);
      }
      manager.setHeader("Authorization", `Bearer ${assignedKey.key}`);
      break;
    case "mistral-ai":
      manager.setHeader("Authorization", `Bearer ${assignedKey.key}`);
      break;
    case "azure":
      const azureKey = assignedKey.key;
      manager.setHeader("api-key", azureKey);
      break;
    case "aws":
    case "gcp":
    case "google-ai":
      throw new Error("add-key should not be used for this service.");
    default:
      assertNever(assignedKey.service);
  }
};

/**
 * Special case for embeddings requests which don't go through the normal
 * request pipeline.
 */
export const addKeyForEmbeddingsRequest: ProxyReqMutator = (manager) => {
  const req = manager.request;
  if (!isEmbeddingsRequest(req)) {
    throw new Error(
      "addKeyForEmbeddingsRequest called on non-embeddings request"
    );
  }

  if (req.inboundApi !== "openai") {
    throw new Error("Embeddings requests must be from OpenAI");
  }

  manager.setBody({ input: req.body.input, model: "text-embedding-ada-002" });

  const key = keyPool.get("text-embedding-ada-002", "openai") as OpenAIKey;

  manager.setKey(key);
  req.log.info(
    { key: key.hash, toApi: req.outboundApi },
    "Assigned Turbo key to embeddings request"
  );

  manager.setHeader("Authorization", `Bearer ${key.key}`);
  if (key.organizationId) {
    manager.setHeader("OpenAI-Organization", key.organizationId);
  }
};

license: unknown


code:
import type { ProxyReqMutator } from "../index";

/** Finalize the rewritten request body. Must be the last mutator. */
export const finalizeBody: ProxyReqMutator = (manager) => {
  const req = manager.request;

  if (["POST", "PUT", "PATCH"].includes(req.method ?? "") && req.body) {
    // For image generation requests, remove stream flag.
    if (req.outboundApi === "openai-image") {
      delete req.body.stream;
    }
    // For anthropic text to chat requests, remove undefined prompt.
    if (req.outboundApi === "anthropic-chat") {
      delete req.body.prompt;
    }

    const serialized =
      typeof req.body === "string" ? req.body : JSON.stringify(req.body);
    manager.setHeader("Content-Length", String(Buffer.byteLength(serialized)));
    manager.setBody(serialized);
  }
};


code:
import { ProxyReqMutator } from "../index";

/**
 * For AWS/GCP/Azure/Google requests, the body is signed earlier in the request
 * pipeline, before the proxy middleware. This function just assigns the path
 * and headers to the proxy request.
 */
export const finalizeSignedRequest: ProxyReqMutator = (manager) => {
  const req = manager.request;
  if (!req.signedRequest) {
    throw new Error("Expected req.signedRequest to be set");
  }

  // The path depends on the selected model and the assigned key's region.
  manager.setPath(req.signedRequest.path);

  // Amazon doesn't want extra headers, so we need to remove all of them and
  // reassign only the ones specified in the signed request.
  const headers = req.signedRequest.headers;
  Object.keys(headers).forEach((key) => {
    manager.removeHeader(key);
  });
  Object.entries(req.signedRequest.headers).forEach(([key, value]) => {
    manager.setHeader(key, value);
  });
  const serialized =
    typeof req.signedRequest.body === "string"
      ? req.signedRequest.body
      : JSON.stringify(req.signedRequest.body);
  manager.setHeader("Content-Length", String(Buffer.byteLength(serialized)));
  manager.setBody(serialized);
};


code:
import express, { Request } from "express";
import { Sha256 } from "@aws-crypto/sha256-js";
import { SignatureV4 } from "@smithy/signature-v4";
import { HttpRequest } from "@smithy/protocol-http";
import {
  AnthropicV1TextSchema,
  AnthropicV1MessagesSchema,
} from "../../../../shared/api-schemas";
import { AwsBedrockKey, keyPool } from "../../../../shared/key-management";
import {
  AWSMistralV1ChatCompletionsSchema,
  AWSMistralV1TextCompletionsSchema,
} from "../../../../shared/api-schemas/mistral-ai";
import { ProxyReqMutator } from "../index";

const AMZ_HOST =
  process.env.AMZ_HOST || "bedrock-runtime.%REGION%.amazonaws.com";

/**
 * Signs an outgoing AWS request with the appropriate headers modifies the
 * request object in place to fix the path.
 * This happens AFTER request transformation.
 */
export const signAwsRequest: ProxyReqMutator = async (manager) => {
  const req = manager.request;
  const { model, stream } = req.body;
  const key = keyPool.get(model, "aws") as AwsBedrockKey;
  manager.setKey(key);

  const credential = getCredentialParts(req);
  const host = AMZ_HOST.replace("%REGION%", credential.region);

  // AWS only uses 2023-06-01 and does not actually check this header, but we
  // set it so that the stream adapter always selects the correct transformer.
  manager.setHeader("anthropic-version", "2023-06-01");

  // If our key has an inference profile compatible with the requested model,
  // we want to use the inference profile instead of the model ID when calling
  // InvokeModel as that will give us higher rate limits.
  const profile =
    key.inferenceProfileIds.find((p) => p.includes(model)) || model;

  // Uses the AWS SDK to sign a request, then modifies our HPM proxy request
  // with the headers generated by the SDK.
  const newRequest = new HttpRequest({
    method: "POST",
    protocol: "https:",
    hostname: host,
    path: `/model/${profile}/invoke${stream ? "-with-response-stream" : ""}`,
    headers: {
      ["Host"]: host,
      ["content-type"]: "application/json",
    },
    body: JSON.stringify(getStrictlyValidatedBodyForAws(req)),
  });

  if (stream) {
    newRequest.headers["x-amzn-bedrock-accept"] = "application/json";
  } else {
    newRequest.headers["accept"] = "*/*";
  }

  const { body, inboundApi, outboundApi } = req;
  req.log.info(
    { key: key.hash, model: body.model, profile, inboundApi, outboundApi },
    "Assigned AWS credentials to request"
  );

  manager.setSignedRequest(await sign(newRequest, getCredentialParts(req)));
};

type Credential = {
  accessKeyId: string;
  secretAccessKey: string;
  region: string;
};

function getCredentialParts(req: express.Request): Credential {
  const [accessKeyId, secretAccessKey, region] = req.key!.key.split(":");

  if (!accessKeyId || !secretAccessKey || !region) {
    req.log.error(
      { key: req.key!.hash },
      "AWS_CREDENTIALS isn't correctly formatted; refer to the docs"
    );
    throw new Error("The key assigned to this request is invalid.");
  }

  return { accessKeyId, secretAccessKey, region };
}

async function sign(request: HttpRequest, credential: Credential) {
  const { accessKeyId, secretAccessKey, region } = credential;

  const signer = new SignatureV4({
    sha256: Sha256,
    credentials: { accessKeyId, secretAccessKey },
    region,
    service: "bedrock",
  });

  return signer.sign(request);
}

function getStrictlyValidatedBodyForAws(req: Readonly<Request>): unknown {
  // AWS uses vendor API formats but imposes additional (more strict) validation
  // rules, namely that extraneous parameters are not allowed. We will validate
  // using the vendor's zod schema but apply `.strip` to ensure that any
  // extraneous parameters are removed.
  let strippedParams: Record<string, unknown> = {};
  switch (req.outboundApi) {
    case "anthropic-text":
      strippedParams = AnthropicV1TextSchema.pick({
        prompt: true,
        max_tokens_to_sample: true,
        stop_sequences: true,
        temperature: true,
        top_k: true,
        top_p: true,
      })
        .strip()
        .parse(req.body);
      break;
    case "anthropic-chat":
      strippedParams = AnthropicV1MessagesSchema.pick({
        messages: true,
        system: true,
        max_tokens: true,
        stop_sequences: true,
        temperature: true,
        top_k: true,
        top_p: true,
      })
        .strip()
        .parse(req.body);
      strippedParams.anthropic_version = "bedrock-2023-05-31";
      break;
    case "mistral-ai":
      strippedParams = AWSMistralV1ChatCompletionsSchema.parse(req.body);
      break;
    case "mistral-text":
      strippedParams = AWSMistralV1TextCompletionsSchema.parse(req.body);
      break;
    default:
      throw new Error("Unexpected outbound API for AWS.");
  }
  return strippedParams;
}


code:
import { AnthropicV1MessagesSchema } from "../../../../shared/api-schemas";
import { GcpKey, keyPool } from "../../../../shared/key-management";
import { ProxyReqMutator } from "../index";
import {
  getCredentialsFromGcpKey,
  refreshGcpAccessToken,
} from "../../../../shared/key-management/gcp/oauth";

const GCP_HOST = process.env.GCP_HOST || "%REGION%-aiplatform.googleapis.com";

export const signGcpRequest: ProxyReqMutator = async (manager) => {
  const req = manager.request;
  const serviceValid = req.service === "gcp";
  if (!serviceValid) {
    throw new Error("addVertexAIKey called on invalid request");
  }

  if (!req.body?.model) {
    throw new Error("You must specify a model with your request.");
  }

  const { model } = req.body;
  const key: GcpKey = keyPool.get(model, "gcp") as GcpKey;

  if (!key.accessToken || Date.now() > key.accessTokenExpiresAt) {
    const [token, durationSec] = await refreshGcpAccessToken(key);
    keyPool.update(key, {
      accessToken: token,
      accessTokenExpiresAt: Date.now() + durationSec * 1000 * 0.95,
    } as GcpKey);
    // nb: key received by `get` is a clone and will not have the new access
    // token we just set, so it must be manually updated.
    key.accessToken = token;
  }

  manager.setKey(key);
  req.log.info({ key: key.hash, model }, "Assigned GCP key to request");

  // TODO: This should happen in transform-outbound-payload.ts
  // TODO: Support tools
  let strippedParams: Record<string, unknown>;
  strippedParams = AnthropicV1MessagesSchema.pick({
    messages: true,
    system: true,
    max_tokens: true,
    stop_sequences: true,
    temperature: true,
    top_k: true,
    top_p: true,
    stream: true,
  })
    .strip()
    .parse(req.body);
  strippedParams.anthropic_version = "vertex-2023-10-16";

  const credential = await getCredentialsFromGcpKey(key);

  const host = GCP_HOST.replace("%REGION%", credential.region);
  // GCP doesn't use the anthropic-version header, but we set it to ensure the
  // stream adapter selects the correct transformer.
  manager.setHeader("anthropic-version", "2023-06-01");

  manager.setSignedRequest({
    method: "POST",
    protocol: "https:",
    hostname: host,
    path: `/v1/projects/${credential.projectId}/locations/${credential.region}/publishers/anthropic/models/${model}:streamRawPredict`,
    headers: {
      ["host"]: host,
      ["content-type"]: "application/json",
      ["authorization"]: `Bearer ${key.accessToken}`,
    },
    body: JSON.stringify(strippedParams),
  });
};


code:
import { ProxyReqMutator } from "../index";

/**
 * Removes origin and referer headers before sending the request to the API for
 * privacy reasons.
 */
export const stripHeaders: ProxyReqMutator = (manager) => {
  manager.removeHeader("origin");
  manager.removeHeader("referer");

  // Some APIs refuse requests coming from browsers to discourage embedding
  // API keys in client-side code, so we must remove all CORS/fetch headers.
  Object.keys(manager.request.headers).forEach((key) => {
    if (key.startsWith("sec-")) {
      manager.removeHeader(key);
    }
  });

  manager.removeHeader("tailscale-user-login");
  manager.removeHeader("tailscale-user-name");
  manager.removeHeader("tailscale-headers-info");
  manager.removeHeader("tailscale-user-profile-pic");
  manager.removeHeader("cf-connecting-ip");
  manager.removeHeader("cf-ray");
  manager.removeHeader("cf-visitor");
  manager.removeHeader("cf-warp-tag-id");
  manager.removeHeader("forwarded");
  manager.removeHeader("true-client-ip");
  manager.removeHeader("x-forwarded-for");
  manager.removeHeader("x-forwarded-host");
  manager.removeHeader("x-forwarded-proto");
  manager.removeHeader("x-real-ip");
};


code:
import { RequestHandler } from "express";
import { ZodIssue } from "zod";
import { initializeSseStream } from "../../../shared/streaming";
import { classifyErrorAndSend } from "../common";
import {
  RequestPreprocessor,
  blockZoomerOrigins,
  countPromptTokens,
  languageFilter,
  setApiFormat,
  transformOutboundPayload,
  validateContextSize,
  validateModelFamily,
  validateVision,
  applyQuotaLimits,
} from ".";

type RequestPreprocessorOptions = {
  /**
   * Functions to run before the request body is transformed between API
   * formats. Use this to change the behavior of the transformation, such as for
   * endpoints which can accept multiple API formats.
   */
  beforeTransform?: RequestPreprocessor[];
  /**
   * Functions to run after the request body is transformed and token counts are
   * assigned. Use this to perform validation or other actions that depend on
   * the request body being in the final API format.
   */
  afterTransform?: RequestPreprocessor[];
};

/**
 * Returns a middleware function that processes the request body into the given
 * API format, and then sequentially runs the given additional preprocessors.
 * These should be used for validation and transformations that only need to
 * happen once per request.
 *
 * These run first in the request lifecycle, a single time per request before it
 * is added to the request queue. They aren't run again if the request is
 * re-attempted after a rate limit.
 *
 * To run functions against requests every time they are re-attempted, write a
 * ProxyReqMutator and pass it to createQueuedProxyMiddleware instead.
 */
export const createPreprocessorMiddleware = (
  apiFormat: Parameters<typeof setApiFormat>[0],
  { beforeTransform, afterTransform }: RequestPreprocessorOptions = {}
): RequestHandler => {
  const preprocessors: RequestPreprocessor[] = [
    setApiFormat(apiFormat),
    blockZoomerOrigins,
    ...(beforeTransform ?? []),
    transformOutboundPayload,
    countPromptTokens,
    languageFilter,
    ...(afterTransform ?? []),
    validateContextSize,
    validateVision,
    validateModelFamily,
    applyQuotaLimits,
  ];
  return async (...args) => executePreprocessors(preprocessors, args);
};

/**
 * Returns a middleware function that specifically prepares requests for
 * OpenAI's embeddings API. Tokens are not counted because embeddings requests
 * are basically free.
 */
export const createEmbeddingsPreprocessorMiddleware = (): RequestHandler => {
  const preprocessors: RequestPreprocessor[] = [
    setApiFormat({ inApi: "openai", outApi: "openai", service: "openai" }),
    (req) => void (req.promptTokens = req.outputTokens = 0),
  ];
  return async (...args) => executePreprocessors(preprocessors, args);
};

async function executePreprocessors(
  preprocessors: RequestPreprocessor[],
  [req, res, next]: Parameters<RequestHandler>
) {
  handleTestMessage(req, res, next);
  if (res.headersSent) return;

  try {
    for (const preprocessor of preprocessors) {
      await preprocessor(req);
    }
    next();
  } catch (error) {
    if (error.constructor.name === "ZodError") {
      const issues = error?.issues
        ?.map((issue: ZodIssue) => `${issue.path.join(".")}: ${issue.message}`)
        .join("; ");
      req.log.warn({ issues }, "Prompt failed preprocessor validation.");
    } else {
      req.log.error(error, "Error while executing request preprocessor");
    }

    // If the requested has opted into streaming, the client probably won't
    // handle a non-eventstream response, but we haven't initialized the SSE
    // stream yet as that is typically done later by the request queue. We'll
    // do that here and then call classifyErrorAndSend to use the streaming
    // error handler.
    const { stream } = req.body;
    const isStreaming = stream === "true" || stream === true;
    if (isStreaming && !res.headersSent) {
      initializeSseStream(res);
    }
    classifyErrorAndSend(error as Error, req, res);
  }
}

/**
 * Bypasses the API call and returns a test message response if the request body
 * is a known test message from SillyTavern. Otherwise these messages just waste
 * API request quota and confuse users when the proxy is busy, because ST always
 * makes them with `stream: false` (which is not allowed when the proxy is busy)
 */
const handleTestMessage: RequestHandler = (req, res) => {
  const { method, body } = req;
  if (method !== "POST") {
    return;
  }

  if (isTestMessage(body)) {
    req.log.info({ body }, "Received test message. Skipping API call.");
    res.json({
      id: "test-message",
      object: "chat.completion",
      created: Date.now(),
      model: body.model,
      // openai chat
      choices: [
        {
          message: { role: "assistant", content: "Hello!" },
          finish_reason: "stop",
          index: 0,
        },
      ],
      // anthropic text
      completion: "Hello!",
      // anthropic chat
      content: [{ type: "text", text: "Hello!" }],
      // gemini
      candidates: [
        {
          content: { parts: [{ text: "Hello!" }] },
          finishReason: "stop",
        },
      ],
      proxy_note:
        "SillyTavern connection test detected. Your prompt was not sent to the actual model and this response was generated by the proxy.",
    });
  }
};

function isTestMessage(body: any) {
  const { messages, prompt, contents } = body;

  if (messages) {
    return (
      messages.length === 1 &&
      messages[0].role === "user" &&
      messages[0].content === "Hi"
    );
  } else if (contents) {
    return contents.length === 1 && contents[0].parts[0]?.text === "Hi";
  } else {
    return (
      prompt?.trim() === "Human: Hi\n\nAssistant:" ||
      prompt?.startsWith("Hi\n\n")
    );
  }
}


code:
import { hasAvailableQuota } from "../../../../shared/users/user-store";
import { isImageGenerationRequest, isTextGenerationRequest } from "../../common";
import { RequestPreprocessor } from "../index";

export class QuotaExceededError extends Error {
  public quotaInfo: any;
  constructor(message: string, quotaInfo: any) {
    super(message);
    this.name = "QuotaExceededError";
    this.quotaInfo = quotaInfo;
  }
}

export const applyQuotaLimits: RequestPreprocessor = (req) => {
  const subjectToQuota =
    isTextGenerationRequest(req) || isImageGenerationRequest(req);
  if (!subjectToQuota || !req.user) return;

  const requestedTokens = (req.promptTokens ?? 0) + (req.outputTokens ?? 0);
  if (
    !hasAvailableQuota({
      userToken: req.user.token,
      model: req.body.model,
      api: req.outboundApi,
      requested: requestedTokens,
    })
  ) {
    throw new QuotaExceededError(
      "You have exceeded your proxy token quota for this model.",
      {
        quota: req.user.tokenLimits,
        used: req.user.tokenCounts,
        requested: requestedTokens,
      }
    );
  }
};


code:
import { RequestPreprocessor } from "../index";

const DISALLOWED_ORIGIN_SUBSTRINGS = "janitorai.com,janitor.ai".split(",");

class ZoomerForbiddenError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "ZoomerForbiddenError";
  }
}

/**
 * Blocks requests from Janitor AI users with a fake, scary error message so I
 * stop getting emails asking for tech support.
 */
export const blockZoomerOrigins: RequestPreprocessor = (req) => {
  const origin = req.headers.origin || req.headers.referer;
  if (origin && DISALLOWED_ORIGIN_SUBSTRINGS.some((s) => origin.includes(s))) {
    // Venus-derivatives send a test prompt to check if the proxy is working.
    // We don't want to block that just yet.
    if (req.body.messages[0]?.content === "Just say TEST") {
      return;
    }

    throw new ZoomerForbiddenError(
      `Your access was terminated due to violation of our policies, please check your email for more information. If you believe this is in error and would like to appeal, please contact us through our help center at help.openai.com.`
    );
  }
};


code:
import { RequestPreprocessor } from "../index";
import { countTokens } from "../../../../shared/tokenization";
import { assertNever } from "../../../../shared/utils";
import {
  GoogleAIChatMessage,
  MistralAIChatMessage,
  OpenAIChatMessage,
} from "../../../../shared/api-schemas";

/**
 * Given a request with an already-transformed body, counts the number of
 * tokens and assigns the count to the request.
 */
export const countPromptTokens: RequestPreprocessor = async (req) => {
  const service = req.outboundApi;
  let result;

  switch (service) {
    case "openai": {
      req.outputTokens = req.body.max_completion_tokens || req.body.max_tokens;
      const prompt: OpenAIChatMessage[] = req.body.messages;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "openai-text": {
      req.outputTokens = req.body.max_tokens;
      const prompt: string = req.body.prompt;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "anthropic-chat": {
      req.outputTokens = req.body.max_tokens;
      let system = req.body.system ?? "";
      if (Array.isArray(system)) {
        system = system
          .map((m: { type: string; text: string }) => m.text)
          .join("\n");
      }
      const prompt = { system, messages: req.body.messages };
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "anthropic-text": {
      req.outputTokens = req.body.max_tokens_to_sample;
      const prompt: string = req.body.prompt;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "google-ai": {
      req.outputTokens = req.body.generationConfig.maxOutputTokens;
      const prompt: GoogleAIChatMessage[] = req.body.contents;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "mistral-ai":
    case "mistral-text": {
      req.outputTokens = req.body.max_tokens;
      const prompt: string | MistralAIChatMessage[] =
        req.body.messages ?? req.body.prompt;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "openai-image": {
      req.outputTokens = 1;
      result = await countTokens({ req, service });
      break;
    }
    default:
      assertNever(service);
  }

  req.promptTokens = result.token_count;

  req.log.debug({ result: result }, "Counted prompt tokens.");
  req.tokenizerInfo = req.tokenizerInfo ?? {};
  req.tokenizerInfo = { ...req.tokenizerInfo, ...result };
};


code:
import { Request } from "express";
import { z } from "zod";
import { config } from "../../../../config";
import { assertNever } from "../../../../shared/utils";
import { RequestPreprocessor } from "../index";
import { BadRequestError } from "../../../../shared/errors";
import {
  MistralAIChatMessage,
  OpenAIChatMessage,
  flattenAnthropicMessages,
} from "../../../../shared/api-schemas";
import { GoogleAIV1GenerateContentSchema } from "../../../../shared/api-schemas/google-ai";

const rejectedClients = new Map<string, number>();

setInterval(() => {
  rejectedClients.forEach((count, ip) => {
    if (count > 0) {
      rejectedClients.set(ip, Math.floor(count / 2));
    } else {
      rejectedClients.delete(ip);
    }
  });
}, 30000);

/**
 * Block requests containing blacklisted phrases. Repeated rejections from the
 * same IP address will be throttled.
 */
export const languageFilter: RequestPreprocessor = async (req) => {
  if (!config.rejectPhrases.length) return;

  const prompt = getPromptFromRequest(req);
  const match = config.rejectPhrases.find((phrase) =>
    prompt.match(new RegExp(phrase, "i"))
  );

  if (match) {
    const ip = req.ip;
    const rejections = (rejectedClients.get(req.ip) || 0) + 1;
    const delay = Math.min(60000, Math.pow(2, rejections - 1) * 1000);
    rejectedClients.set(ip, rejections);
    req.log.warn(
      { match, ip, rejections, delay },
      "Prompt contains rejected phrase"
    );
    await new Promise((resolve) => {
      req.res!.once("close", resolve);
      setTimeout(resolve, delay);
    });
    throw new BadRequestError(config.rejectMessage);
  }
};

/*
TODO: this is not type safe and does not raise errors if request body zod schema
is changed.
*/
function getPromptFromRequest(req: Request) {
  const service = req.outboundApi;
  const body = req.body;
  switch (service) {
    case "anthropic-chat":
      return flattenAnthropicMessages(body.messages);
    case "openai":
    case "mistral-ai":
      return body.messages
        .map((msg: OpenAIChatMessage | MistralAIChatMessage) => {
          const text = Array.isArray(msg.content)
            ? msg.content
                .map((c) => {
                  if ("text" in c) return c.text;
                })
                .join()
            : msg.content;
          return `${msg.role}: ${text}`;
        })
        .join("\n\n");
    case "anthropic-text":
    case "openai-text":
    case "openai-image":
    case "mistral-text":
      return body.prompt;
    case "google-ai": {
      const b = body as z.infer<typeof GoogleAIV1GenerateContentSchema>;
      return [
        b.systemInstruction?.parts.map((p) => p.text),
        ...b.contents.flatMap((c) => c.parts.map((p) => p.text)),
      ].join("\n");
    }
    default:
      assertNever(service);
  }
}


code:
import { Request } from "express";
import { APIFormat } from "../../../../shared/key-management";
import { LLMService } from "../../../../shared/models";
import { RequestPreprocessor } from "../index";

export const setApiFormat = (api: {
  /**
   * The API format the user made the request in and expects the response to be
   * in.
   */
  inApi: Request["inboundApi"];
  /**
   * The API format the proxy will make the request in and expects the response
   * to be in. If different from `inApi`, the proxy will transform the user's
   * request body to this format, and will transform the response body or stream
   * events from this format.
   */
  outApi: APIFormat;
  /**
   * The service the request will be sent to, which determines authentication
   * and possibly the streaming transport.
   */
  service: LLMService;
}): RequestPreprocessor => {
  return function configureRequestApiFormat(req) {
    req.inboundApi = api.inApi;
    req.outboundApi = api.outApi;
    req.service = api.service;
  };
};


code:
import { Request } from "express";
import {
  API_REQUEST_VALIDATORS,
  API_REQUEST_TRANSFORMERS,
} from "../../../../shared/api-schemas";
import { BadRequestError } from "../../../../shared/errors";
import { fixMistralPrompt } from "../../../../shared/api-schemas/mistral-ai";
import {
  isImageGenerationRequest,
  isTextGenerationRequest,
} from "../../common";
import { RequestPreprocessor } from "../index";

/** Transforms an incoming request body to one that matches the target API. */
export const transformOutboundPayload: RequestPreprocessor = async (req) => {
  const alreadyTransformed = req.retryCount > 0;
  const notTransformable =
    !isTextGenerationRequest(req) && !isImageGenerationRequest(req);

  if (alreadyTransformed) {
    return;
  } else if (notTransformable) {
    // This is probably an indication of a bug in the proxy.
    const { inboundApi, outboundApi, method, path } = req;
    req.log.warn(
      { inboundApi, outboundApi, method, path },
      "`transformOutboundPayload` called on a non-transformable request."
    );
    return;
  }

  applyMistralPromptFixes(req);

  // Native prompts are those which were already provided by the client in the
  // target API format. We don't need to transform them.
  const isNativePrompt = req.inboundApi === req.outboundApi;
  if (isNativePrompt) {
    const result = API_REQUEST_VALIDATORS[req.inboundApi].parse(req.body);
    req.body = result;
    return;
  }

  // Prompt requires translation from one API format to another.
  const transformation = `${req.inboundApi}->${req.outboundApi}` as const;
  const transFn = API_REQUEST_TRANSFORMERS[transformation];

  if (transFn) {
    req.log.info({ transformation }, "Transforming request...");
    req.body = await transFn(req);
    return;
  }

  throw new BadRequestError(
    `${transformation} proxying is not supported. Make sure your client is configured to send requests in the correct format and to the correct endpoint.`
  );
};

// handles weird cases that don't fit into our abstractions
function applyMistralPromptFixes(req: Request): void {
  if (req.inboundApi === "mistral-ai") {
    // Mistral Chat is very similar to OpenAI but not identical and many clients
    // don't properly handle the differences. We will try to validate the
    // mistral prompt and try to fix it if it fails. It will be re-validated
    // after this function returns.
    const result = API_REQUEST_VALIDATORS["mistral-ai"].parse(req.body);
    req.body.messages = fixMistralPrompt(result.messages);
    req.log.info(
      { n: req.body.messages.length, prev: result.messages.length },
      "Applied Mistral chat prompt fixes."
    );

    // If the prompt relies on `prefix: true` for the last message, we need to
    // convert it to a text completions request because AWS Mistral support for
    // this feature is broken.
    // On Mistral La Plateforme, we can't do this because they don't expose
    // a text completions endpoint.
    const { messages } = req.body;
    const lastMessage = messages && messages[messages.length - 1];
    if (lastMessage?.role === "assistant" && req.service === "aws") {
      // enable prefix if client forgot, otherwise the template will insert an
      // eos token which is very unlikely to be what the client wants.
      lastMessage.prefix = true;
      req.outboundApi = "mistral-text";
      req.log.info(
        "Native Mistral chat prompt relies on assistant message prefix. Converting to text completions request."
      );
    }
  }
}


code:
import { Request } from "express";
import { z } from "zod";
import { config } from "../../../../config";
import { assertNever } from "../../../../shared/utils";
import { RequestPreprocessor } from "../index";

const CLAUDE_MAX_CONTEXT = config.maxContextTokensAnthropic;
const OPENAI_MAX_CONTEXT = config.maxContextTokensOpenAI;
// todo: make configurable
const GOOGLE_AI_MAX_CONTEXT = 2048000;
const MISTRAL_AI_MAX_CONTENT = 131072;

/**
 * Assigns `req.promptTokens` and `req.outputTokens` based on the request body
 * and outbound API format, which combined determine the size of the context.
 * If the context is too large, an error is thrown.
 * This preprocessor should run after any preprocessor that transforms the
 * request body.
 */
export const validateContextSize: RequestPreprocessor = async (req) => {
  assertRequestHasTokenCounts(req);
  const promptTokens = req.promptTokens;
  const outputTokens = req.outputTokens;
  const contextTokens = promptTokens + outputTokens;
  const model = req.body.model;

  let proxyMax: number;
  switch (req.outboundApi) {
    case "openai":
    case "openai-text":
      proxyMax = OPENAI_MAX_CONTEXT;
      break;
    case "anthropic-chat":
    case "anthropic-text":
      proxyMax = CLAUDE_MAX_CONTEXT;
      break;
    case "google-ai":
      proxyMax = GOOGLE_AI_MAX_CONTEXT;
      break;
    case "mistral-ai":
    case "mistral-text":
      proxyMax = MISTRAL_AI_MAX_CONTENT;
      break;
    case "openai-image":
      return;
    default:
      assertNever(req.outboundApi);
  }
  proxyMax ||= Number.MAX_SAFE_INTEGER;

  if (req.user?.type === "special") {
    req.log.debug("Special user, not enforcing proxy context limit.");
    proxyMax = Number.MAX_SAFE_INTEGER;
  }

  let modelMax: number;
  if (model.match(/gpt-3.5-turbo-16k/)) {
    modelMax = 16384;
  } else if (model.match(/^gpt-4o/)) {
    modelMax = 128000;
  } else if (model.match(/^chatgpt-4o/)) {
    modelMax = 128000;
  } else if (model.match(/gpt-4-turbo(-\d{4}-\d{2}-\d{2})?$/)) {
    modelMax = 131072;
  } else if (model.match(/gpt-4-turbo(-preview)?$/)) {
    modelMax = 131072;
  } else if (model.match(/gpt-4-(0125|1106)(-preview)?$/)) {
    modelMax = 131072;
  } else if (model.match(/^gpt-4(-\d{4})?-vision(-preview)?$/)) {
    modelMax = 131072;
  } else if (model.match(/^o1-mini(-\d{4}-\d{2}-\d{2})?$/)) {
    modelMax = 128000;
  } else if (model.match(/^o1(-preview)?(-\d{4}-\d{2}-\d{2})?$/)) {
    modelMax = 128000;
  } else if (model.match(/gpt-3.5-turbo/)) {
    modelMax = 16384;
  } else if (model.match(/gpt-4-32k/)) {
    modelMax = 32768;
  } else if (model.match(/gpt-4/)) {
    modelMax = 8192;
  } else if (model.match(/^claude-(?:instant-)?v1(?:\.\d)?-100k/)) {
    modelMax = 100000;
  } else if (model.match(/^claude-(?:instant-)?v1(?:\.\d)?$/)) {
    modelMax = 9000;
  } else if (model.match(/^claude-2\.0/)) {
    modelMax = 100000;
  } else if (model.match(/^claude-2/)) {
    modelMax = 200000;
  } else if (model.match(/^claude-3/)) {
    modelMax = 200000;
  } else if (model.match(/^gemini-/)) {
    modelMax = 1024000;
  } else if (model.match(/^anthropic\.claude-3/)) {
    modelMax = 200000;
  } else if (model.match(/^anthropic\.claude-v2:\d/)) {
    modelMax = 200000;
  } else if (model.match(/^anthropic\.claude/)) {
    modelMax = 100000;
  } else if (model.match(/tral/)) {
    // catches mistral, mixtral, codestral, mathstral, etc. mistral models have
    // no name convention and wildly different context windows so this is a
    // catch-all
    modelMax = MISTRAL_AI_MAX_CONTENT;
  } else {
    req.log.warn({ model }, "Unknown model, using 200k token limit.");
    modelMax = 200000;
  }

  const finalMax = Math.min(proxyMax, modelMax);
  z.object({
    tokens: z
      .number()
      .int()
      .max(finalMax, {
        message: `Your request exceeds the context size limit. (max: ${finalMax} tokens, requested: ${promptTokens} prompt + ${outputTokens} output = ${contextTokens} context tokens)`,
      }),
  }).parse({ tokens: contextTokens });

  req.log.debug(
    { promptTokens, outputTokens, contextTokens, modelMax, proxyMax },
    "Prompt size validated"
  );

  req.tokenizerInfo.prompt_tokens = promptTokens;
  req.tokenizerInfo.completion_tokens = outputTokens;
  req.tokenizerInfo.max_model_tokens = modelMax;
  req.tokenizerInfo.max_proxy_tokens = proxyMax;
};

function assertRequestHasTokenCounts(
  req: Request
): asserts req is Request & { promptTokens: number; outputTokens: number } {
  z.object({
    promptTokens: z.number().int().min(1),
    outputTokens: z.number().int().min(1),
  })
    .nonstrict()
    .parse({ promptTokens: req.promptTokens, outputTokens: req.outputTokens });
}


code:
import { config } from "../../../../config";
import { ForbiddenError } from "../../../../shared/errors";
import { getModelFamilyForRequest } from "../../../../shared/models";
import { RequestPreprocessor } from "../index";

/**
 * Ensures the selected model family is enabled by the proxy configuration.
 */
export const validateModelFamily: RequestPreprocessor = (req) => {
  const family = getModelFamilyForRequest(req);
  if (!config.allowedModelFamilies.includes(family)) {
    throw new ForbiddenError(
      `Model family '${family}' is not enabled on this proxy`
    );
  }
};


code:
import { config } from "../../../../config";
import { assertNever } from "../../../../shared/utils";
import { RequestPreprocessor } from "../index";
import { containsImageContent as containsImageContentOpenAI } from "../../../../shared/api-schemas/openai";
import { containsImageContent as containsImageContentAnthropic } from "../../../../shared/api-schemas/anthropic";
import { ForbiddenError } from "../../../../shared/errors";

/**
 * Rejects prompts containing images if multimodal prompts are disabled.
 */
export const validateVision: RequestPreprocessor = async (req) => {
  if (req.service === undefined) {
    throw new Error("Request service must be set before validateVision");
  }

  if (req.user?.type === "special") return;
  if (config.allowedVisionServices.includes(req.service)) return;

  // vision not allowed for req's service, block prompts with images
  let hasImage = false;
  switch (req.outboundApi) {
    case "openai":
      hasImage = containsImageContentOpenAI(req.body.messages);
      break;
    case "anthropic-chat":
      hasImage = containsImageContentAnthropic(req.body.messages);
      break;
    case "anthropic-text":
    case "google-ai":
    case "mistral-ai":
    case "mistral-text":
    case "openai-image":
    case "openai-text":
      return;
    default:
      assertNever(req.outboundApi);
  }

  if (hasImage) {
    throw new ForbiddenError(
      "Prompts containing images are not permitted. Disable 'Send Inline Images' in your client and try again."
    );
  }
};


code:
import { Request, Response } from "express";
import http from "http";
import ProxyServer from "http-proxy";
import { Readable } from "stream";
import {
  createProxyMiddleware,
  Options,
  debugProxyErrorsPlugin,
  proxyEventsPlugin,
} from "http-proxy-middleware";
import { ProxyReqMutator, stripHeaders } from "./index";
import { createOnProxyResHandler, ProxyResHandlerWithBody } from "../response";
import { createQueueMiddleware } from "../../queue";
import { getHttpAgents } from "../../../shared/network";
import { classifyErrorAndSend } from "../common";

/**
 * Options for the `createQueuedProxyMiddleware` factory function.
 */
type ProxyMiddlewareFactoryOptions = {
  /**
   * Functions which receive a ProxyReqManager and can modify the request before
   * it is proxied. The modifications will be automatically reverted if the
   * request needs to be returned to the queue.
   */
  mutations?: ProxyReqMutator[];
  /**
   * The target URL to proxy requests to. This can be a string or a function
   * which accepts the request and returns a string.
   */
  target: string | Options<Request>["router"];
  /**
   * A function which receives the proxy response and the JSON-decoded request
   * body. Only fired for non-streaming responses; streaming responses are
   * handled in `handle-streaming-response.ts`.
   */
  blockingResponseHandler?: ProxyResHandlerWithBody;
};

/**
 * Returns a middleware function that accepts incoming requests and places them
 * into the request queue. When the request is dequeued, it is proxied to the
 * target URL using the given options and middleware. Non-streaming responses
 * are handled by the given `blockingResponseHandler`.
 */
export function createQueuedProxyMiddleware({
  target,
  mutations,
  blockingResponseHandler,
}: ProxyMiddlewareFactoryOptions) {
  const hpmTarget = typeof target === "string" ? target : "https://setbyrouter";
  const hpmRouter = typeof target === "function" ? target : undefined;

  const [httpAgent, httpsAgent] = getHttpAgents();
  const agent = hpmTarget.startsWith("http:") ? httpAgent : httpsAgent;

  const proxyMiddleware = createProxyMiddleware<Request, Response>({
    target: hpmTarget,
    router: hpmRouter,
    agent,
    changeOrigin: true,
    toProxy: true,
    selfHandleResponse: typeof blockingResponseHandler === "function",
    // Disable HPM logger plugin (requires re-adding the other default plugins).
    // Contrary to name, debugProxyErrorsPlugin is not just for debugging and
    // fixes several error handling/connection close issues in http-proxy core.
    ejectPlugins: true,
    // Inferred (via Options<express.Request>) as Plugin<express.Request>, but
    // the default plugins only allow http.IncomingMessage for TReq. They are
    // compatible with express.Request, so we can use them. `Plugin` type is not
    // exported for some reason.
    plugins: [
      debugProxyErrorsPlugin,
      pinoLoggerPlugin,
      proxyEventsPlugin,
    ] as any,
    on: {
      proxyRes: createOnProxyResHandler(
        blockingResponseHandler ? [blockingResponseHandler] : []
      ),
      error: classifyErrorAndSend,
    },
    buffer: ((req: Request) => {
      // This is a hack/monkey patch and is not part of the official
      // http-proxy-middleware package. See patches/http-proxy+1.18.1.patch.
      let payload = req.body;
      if (typeof payload === "string") {
        payload = Buffer.from(payload);
      }
      const stream = new Readable();
      stream.push(payload);
      stream.push(null);
      return stream;
    }) as any,
  });

  return createQueueMiddleware({
    mutations: [stripHeaders, ...(mutations ?? [])],
    proxyMiddleware,
  });
}

type ProxiedResponse = http.IncomingMessage & Response & any;
function pinoLoggerPlugin(proxyServer: ProxyServer<Request>) {
  proxyServer.on("error", (err, req, res, target) => {
    req.log.error(
      { originalUrl: req.originalUrl, targetUrl: String(target), err },
      "Error occurred while proxying request to target"
    );
  });
  proxyServer.on("proxyReq", (proxyReq, req) => {
    const { protocol, host, path } = proxyReq;
    req.log.info(
      {
        from: req.originalUrl,
        to: `${protocol}//${host}${path}`,
      },
      "Sending request to upstream API..."
    );
  });
  proxyServer.on("proxyRes", (proxyRes: ProxiedResponse, req, _res) => {
    const { protocol, host, path } = proxyRes.req;
    req.log.info(
      {
        target: `${protocol}//${host}${path}`,
        status: proxyRes.statusCode,
        contentType: proxyRes.headers["content-type"],
        contentEncoding: proxyRes.headers["content-encoding"],
        contentLength: proxyRes.headers["content-length"],
        transferEncoding: proxyRes.headers["transfer-encoding"],
      },
      "Got response from upstream API."
    );
  });
}


code:
import { Request } from "express";
import { Key } from "../../../shared/key-management";
import { assertNever } from "../../../shared/utils";

/**
 * Represents a change to the request that will be reverted if the request
 * fails.
 */
interface ProxyReqMutation {
  target: "header" | "path" | "body" | "api-key" | "signed-request";
  key?: string;
  originalValue: any | undefined;
}

/**
 * Manages a request's headers, body, and path, allowing them to be modified
 * before the request is proxied and automatically reverted if the request
 * needs to be retried.
 */
export class ProxyReqManager {
  private req: Request;
  private mutations: ProxyReqMutation[] = [];

  /**
   * A read-only proxy of the request object. Avoid changing any properties
   * here as they will persist across retries.
   */
  public readonly request: Readonly<Request>;

  constructor(req: Request) {
    this.req = req;

    this.request = new Proxy(req, {
      get: (target, prop) => {
        if (typeof prop === "string") return target[prop as keyof Request];
        return undefined;
      },
    });
  }

  setHeader(name: string, newValue: string): void {
    const originalValue = this.req.get(name);
    this.mutations.push({ target: "header", key: name, originalValue });
    this.req.headers[name.toLowerCase()] = newValue;
  }

  removeHeader(name: string): void {
    const originalValue = this.req.get(name);
    this.mutations.push({ target: "header", key: name, originalValue });
    delete this.req.headers[name.toLowerCase()];
  }

  setBody(newBody: any): void {
    const originalValue = this.req.body;
    this.mutations.push({ target: "body", key: "body", originalValue });
    this.req.body = newBody;
  }

  setKey(newKey: Key): void {
    const originalValue = this.req.key;
    this.mutations.push({ target: "api-key", key: "key", originalValue });
    this.req.key = newKey;
  }

  setPath(newPath: string): void {
    const originalValue = this.req.path;
    this.mutations.push({ target: "path", key: "path", originalValue });
    this.req.url = newPath;
  }

  setSignedRequest(newSignedRequest: typeof this.req.signedRequest): void {
    const originalValue = this.req.signedRequest;
    this.mutations.push({ target: "signed-request", key: "signedRequest", originalValue });
    this.req.signedRequest = newSignedRequest;
  }

  hasChanged(): boolean {
    return this.mutations.length > 0;
  }

  revert(): void {
    for (const mutation of this.mutations.reverse()) {
      switch (mutation.target) {
        case "header":
          if (mutation.originalValue === undefined) {
            delete this.req.headers[mutation.key!.toLowerCase()];
            continue;
          } else {
            this.req.headers[mutation.key!.toLowerCase()] =
              mutation.originalValue;
          }
          break;
        case "path":
          this.req.url = mutation.originalValue;
          break;
        case "body":
          this.req.body = mutation.originalValue;
          break;
        case "api-key":
          // We don't reset the key here because it's not a property of the
          // inbound request, so we'd only ever be reverting it to null.
          break;
        case "signed-request":
          this.req.signedRequest = mutation.originalValue;
          break;
        default:
          assertNever(mutation.target);
      }
    }
    this.mutations = [];
  }
}


code:
import util from "util";
import zlib from "zlib";
import { PassThrough } from "stream";

const BUFFER_DECODER_MAP = {
  gzip: util.promisify(zlib.gunzip),
  deflate: util.promisify(zlib.inflate),
  br: util.promisify(zlib.brotliDecompress),
  text: (data: Buffer) => data,
};

const STREAM_DECODER_MAP = {
  gzip: zlib.createGunzip,
  deflate: zlib.createInflate,
  br: zlib.createBrotliDecompress,
  text: () => new PassThrough(),
};

type SupportedContentEncoding = keyof typeof BUFFER_DECODER_MAP;
const isSupportedContentEncoding = (
  encoding: string
): encoding is SupportedContentEncoding => encoding in BUFFER_DECODER_MAP;

export async function decompressBuffer(buf: Buffer, encoding: string = "text") {
  if (isSupportedContentEncoding(encoding)) {
    return (await BUFFER_DECODER_MAP[encoding](buf)).toString();
  }
  throw new Error(`Unsupported content-encoding: ${encoding}`);
}

export function getStreamDecompressor(encoding: string = "text") {
  if (isSupportedContentEncoding(encoding)) {
    return STREAM_DECODER_MAP[encoding]();
  }
  throw new Error(`Unsupported content-encoding: ${encoding}`);
}


code:
import express from "express";
import { APIFormat } from "../../../shared/key-management";
import { assertNever } from "../../../shared/utils";
import { initializeSseStream } from "../../../shared/streaming";
import http from "http";

/**
 * Returns a Markdown-formatted message that renders semi-nicely in most chat
 * frontends. For example:
 *
 * **Proxy error (HTTP 404 Not Found)**
 * The proxy encountered an error while trying to send your prompt to the upstream service. Further technical details are provided below.
 * ***
 * *The requested Claude model might not exist, or the key might not be provisioned for it.*
 * ```
 * {
 *  "type": "error",
 *  "error": {
 *    "type": "not_found_error",
 *    "message": "model: some-invalid-model-id",
 *   },
 *  "proxy_note": "The requested Claude model might not exist, or the key might not be provisioned for it."
 * }
 * ```
 */
function getMessageContent(params: {
  title: string;
  message: string;
  obj?: Record<string, any>;
}) {
  const { title, message, obj } = params;
  const note = obj?.proxy_note || obj?.error?.message || "";
  const header = `### **${title}**`;
  const friendlyMessage = note ? `${message}\n\n----\n\n*${note}*` : message;

  const serializedObj = obj
    ? ["```", JSON.stringify(obj, null, 2), "```"].join("\n")
    : "";

  const { stack } = JSON.parse(JSON.stringify(obj ?? {}));
  let prettyTrace = "";
  if (stack && obj) {
    prettyTrace = [
      "Include this trace when reporting an issue.",
      "```",
      stack,
      "```",
    ].join("\n");
    delete obj.stack;
  }

  return [
    header,
    friendlyMessage,
    serializedObj,
    prettyTrace,
    "<!-- oai-proxy-error -->",
  ].join("\n\n");
}

type ErrorGeneratorOptions = {
  format: APIFormat | "unknown";
  title: string;
  message: string;
  obj?: Record<string, any>;
  reqId: string | number | object;
  model?: string;
  statusCode?: number;
};

/**
 * Very crude inference of the request format based on the request body. Don't
 * rely on this to be very accurate.
 */
function tryInferFormat(body: any): APIFormat | "unknown" {
  if (typeof body !== "object" || !body.model) {
    return "unknown";
  }

  if (body.model.includes("gpt")) {
    return "openai";
  }

  if (body.model.includes("mistral")) {
    return "mistral-ai";
  }

  if (body.model.includes("claude")) {
    return body.messages?.length ? "anthropic-chat" : "anthropic-text";
  }

  if (body.model.includes("gemini")) {
    return "google-ai";
  }

  return "unknown";
}

/**
 * Redacts the hostname from the error message if it contains a DNS resolution
 * error. This is to avoid leaking upstream hostnames on DNS resolution errors,
 * as those may contain sensitive information about the proxy's configuration.
 */
function redactHostname(options: ErrorGeneratorOptions): ErrorGeneratorOptions {
  if (!options.message.includes("getaddrinfo")) return options;

  const redacted = { ...options };
  redacted.message = "Could not resolve hostname";

  if (typeof redacted.obj?.error === "object") {
    redacted.obj = {
      ...redacted.obj,
      error: { message: "Could not resolve hostname" },
    };
  }

  return redacted;
}

/**
 * Generates an appropriately-formatted error response and sends it to the
 * client over their requested transport (blocking or SSE stream).
 */
export function sendErrorToClient(params: {
  options: ErrorGeneratorOptions;
  req: express.Request;
  res: express.Response;
}) {
  const { req, res } = params;
  const options = redactHostname(params.options);
  const { statusCode, message, title, obj: details } = options;

  // Since we want to send the error in a format the client understands, we
  // need to know the request format. `setApiFormat` might not have been called
  // yet, so we'll try to infer it from the request body.
  const format =
    options.format === "unknown" ? tryInferFormat(req.body) : options.format;
  if (format === "unknown") {
    // Early middleware error (auth, rate limit) so we can only send something
    // generic.
    const code = statusCode || 400;
    const hasDetails = details && Object.keys(details).length > 0;
    return res.status(code).json({
      error: {
        message,
        type: http.STATUS_CODES[code]!.replace(/\s+/g, "_").toLowerCase(),
      },
      ...(hasDetails ? { details } : {}),
    });
  }

  // Cannot modify headers if client opted into streaming and made it into the
  // proxy request queue, because that immediately starts an SSE stream.
  if (!res.headersSent) {
    res.setHeader("x-oai-proxy-error", title);
    res.setHeader("x-oai-proxy-error-status", statusCode || 500);
  }

  // By this point, we know the request format. To get the error to display in
  // chat clients' UIs, we'll send it as a 200 response as a spoofed completion
  // from the language model. Depending on whether the client is streaming, we
  // will either send an SSE event or a JSON response.
  const isStreaming = req.isStreaming || String(req.body.stream) === "true";
  if (isStreaming) {
    // User can have opted into streaming but not made it into the queue yet,
    // in which case the stream must be started first.
    if (!res.headersSent) {
      initializeSseStream(res);
    }
    res.write(buildSpoofedSSE({ ...options, format }));
    res.write(`data: [DONE]\n\n`);
    res.end();
  } else {
    res.status(200).json(buildSpoofedCompletion({ ...options, format }));
  }
}

/**
 * Returns a non-streaming completion object that looks like it came from the
 * service that the request is being proxied to. Used to send error messages to
 * the client and have them look like normal responses, for clients with poor
 * error handling.
 */
export function buildSpoofedCompletion({
  format,
  title,
  message,
  obj,
  reqId,
  model = "unknown",
}: ErrorGeneratorOptions & { format: Exclude<APIFormat, "unknown"> }) {
  const id = String(reqId);
  const content = getMessageContent({ title, message, obj });

  switch (format) {
    case "openai":
    case "mistral-ai":
      return {
        id: "error-" + id,
        object: "chat.completion",
        created: Date.now(),
        model,
        usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
        choices: [
          {
            message: { role: "assistant", content },
            finish_reason: title,
            index: 0,
          },
        ],
      };
    case "mistral-text":
      return {
        outputs: [{ text: content, stop_reason: title }],
        model,
      };
    case "openai-text":
      return {
        id: "error-" + id,
        object: "text_completion",
        created: Date.now(),
        model,
        usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
        choices: [
          { text: content, index: 0, logprobs: null, finish_reason: title },
        ],
      };
    case "anthropic-text":
      return {
        id: "error-" + id,
        type: "completion",
        completion: content,
        stop_reason: title,
        stop: null,
        model,
      };
    case "anthropic-chat":
      return {
        id: "error-" + id,
        type: "message",
        role: "assistant",
        content: [{ type: "text", text: content }],
        model,
        stop_reason: title,
        stop_sequence: null,
      };
    case "google-ai":
      return {
        candidates: [
          {
            content: { parts: [{ text: content }], role: "model" },
            finishReason: title,
            index: 0,
            tokenCount: null,
            safetyRatings: [],
          },
        ],
      };
    case "openai-image":
      return obj;
    default:
      assertNever(format);
  }
}

/**
 * Returns an SSE message that looks like a completion event for the service
 * that the request is being proxied to. Used to send error messages to the
 * client in the middle of a streaming request.
 */
export function buildSpoofedSSE({
  format,
  title,
  message,
  obj,
  reqId,
  model = "unknown",
}: ErrorGeneratorOptions & { format: Exclude<APIFormat, "unknown"> }) {
  const id = String(reqId);
  const content = getMessageContent({ title, message, obj });

  let event;

  switch (format) {
    case "openai":
    case "mistral-ai":
      event = {
        id: "chatcmpl-" + id,
        object: "chat.completion.chunk",
        created: Date.now(),
        model,
        choices: [{ delta: { content }, index: 0, finish_reason: title }],
      };
      break;
    case "mistral-text":
      event = {
        outputs: [{ text: content, stop_reason: title }],
      };
      break;
    case "openai-text":
      event = {
        id: "cmpl-" + id,
        object: "text_completion",
        created: Date.now(),
        choices: [
          { text: content, index: 0, logprobs: null, finish_reason: title },
        ],
        model,
      };
      break;
    case "anthropic-text":
      event = {
        completion: content,
        stop_reason: title,
        truncated: false,
        stop: null,
        model,
        log_id: "proxy-req-" + id,
      };
      break;
    case "anthropic-chat":
      event = {
        type: "content_block_delta",
        index: 0,
        delta: { type: "text_delta", text: content },
      };
      break;
    case "google-ai":
      // TODO: google ai supports two streaming transports, SSE and JSON.
      // we currently only support SSE.
      // return JSON.stringify({
      event = {
        candidates: [
          {
            content: { parts: [{ text: content }], role: "model" },
            finishReason: title,
            index: 0,
            tokenCount: null,
            safetyRatings: [],
          },
        ],
      };
      break;
    case "openai-image":
      return JSON.stringify(obj);
    default:
      assertNever(format);
  }

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

  // ugh.
  if (format === "anthropic-chat") {
    return (
      [
        [
          "event: message_start",
          `data: ${JSON.stringify({
            type: "message_start",
            message: {
              id: "error-" + id,
              type: "message",
              role: "assistant",
              content: [],
              model,
            },
          })}`,
        ].join("\n"),
        [
          "event: content_block_start",
          `data: ${JSON.stringify({
            type: "content_block_start",
            index: 0,
            content_block: { type: "text", text: "" },
          })}`,
        ].join("\n"),
        ["event: content_block_delta", `data: ${JSON.stringify(event)}`].join(
          "\n"
        ),
        [
          "event: content_block_stop",
          `data: ${JSON.stringify({ type: "content_block_stop", index: 0 })}`,
        ].join("\n"),
        [
          "event: message_delta",
          `data: ${JSON.stringify({
            type: "message_delta",
            delta: { stop_reason: title, stop_sequence: null, usage: null },
          })}`,
        ],
        [
          "event: message_stop",
          `data: ${JSON.stringify({ type: "message_stop" })}`,
        ].join("\n"),
      ].join("\n\n") + "\n\n"
    );
  }

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


code:
import { sendProxyError } from "../common";
import type { RawResponseBodyHandler } from "./index";
import { decompressBuffer } from "./compression";

/**
 * Handles the response from the upstream service and decodes the body if
 * necessary. If the response is JSON, it will be parsed and returned as an
 * object. Otherwise, it will be returned as a string. Does not handle streaming
 * responses.
 * @throws {Error} Unsupported content-encoding or invalid application/json body
 */
export const handleBlockingResponse: RawResponseBodyHandler = async (
  proxyRes,
  req,
  res
) => {
  if (req.isStreaming) {
    const err = new Error(
      "handleBlockingResponse called for a streaming request."
    );
    req.log.error({ stack: err.stack, api: req.inboundApi }, err.message);
    throw err;
  }

  return new Promise((resolve, reject) => {
    let chunks: Buffer[] = [];
    proxyRes.on("data", (chunk) => chunks.push(chunk));
    proxyRes.on("end", async () => {
      const contentEncoding = proxyRes.headers["content-encoding"];
      const contentType = proxyRes.headers["content-type"];
      let body: string | Buffer = Buffer.concat(chunks);
      const rejectWithMessage = function (msg: string, err: Error) {
        const error = `${msg} (${err.message})`;
        req.log.warn(
          { msg: error, stack: err.stack },
          "Error in blocking response handler"
        );
        sendProxyError(req, res, 500, "Internal Server Error", { error });
        return reject(error);
      };

      try {
        body = await decompressBuffer(body, contentEncoding);
      } catch (e) {
        return rejectWithMessage(`Could not decode response body`, e);
      }

      try {
        return resolve(tryParseAsJson(body, contentType));
      } catch (e) {
        return rejectWithMessage("API responded with invalid JSON", e);
      }
    });
  });
};

function tryParseAsJson(body: string, contentType?: string) {
  // If the response is declared as JSON, it must parse or we will throw
  if (contentType?.includes("application/json")) {
    return JSON.parse(body);
  }
  // If it's not declared as JSON, some APIs we'll try to parse it as JSON
  // anyway since some APIs return the wrong content-type header in some cases.
  // If it fails to parse, we'll just return the raw body without throwing.
  try {
    return JSON.parse(body);
  } catch (e) {
    return body;
  }
}


code:
import express from "express";
import { pipeline, Readable, Transform } from "stream";
import { StringDecoder } from "string_decoder";
import { promisify } from "util";
import type { logger } from "../../../logger";
import { BadRequestError, RetryableError } from "../../../shared/errors";
import { APIFormat, keyPool } from "../../../shared/key-management";
import {
  copySseResponseHeaders,
  initializeSseStream,
} from "../../../shared/streaming";
import { reenqueueRequest } from "../../queue";
import type { RawResponseBodyHandler } from ".";
import { handleBlockingResponse } from "./handle-blocking-response";
import { buildSpoofedSSE, sendErrorToClient } from "./error-generator";
import { getAwsEventStreamDecoder } from "./streaming/aws-event-stream-decoder";
import { EventAggregator } from "./streaming/event-aggregator";
import { SSEMessageTransformer } from "./streaming/sse-message-transformer";
import { SSEStreamAdapter } from "./streaming/sse-stream-adapter";
import { getStreamDecompressor } from "./compression";

const pipelineAsync = promisify(pipeline);

/**
 * `handleStreamedResponse` consumes a streamed response from the upstream API,
 * decodes chunk-by-chunk into a stream of events, transforms those events into
 * the client's requested format, and forwards the result to the client.
 *
 * After the entire stream has been consumed, it resolves with the full response
 * body so that subsequent middleware in the chain can process it as if it were
 * a non-streaming response (to count output tokens, track usage, etc).
 *
 * In the event of an error, the request's streaming flag is unset and the
 * request is bounced back to the non-streaming response handler. If the error
 * is retryable, that handler will re-enqueue the request and also reset the
 * streaming flag. Unfortunately the streaming flag is set and unset in multiple
 * places, so it's hard to keep track of.
 */
export const handleStreamedResponse: RawResponseBodyHandler = async (
  proxyRes,
  req,
  res
) => {
  const { headers, statusCode } = proxyRes;
  if (!req.isStreaming) {
    throw new Error("handleStreamedResponse called for non-streaming request.");
  }

  if (statusCode! > 201) {
    req.isStreaming = false;
    req.log.warn(
      { statusCode },
      `Streaming request returned error status code. Falling back to non-streaming response handler.`
    );
    return handleBlockingResponse(proxyRes, req, res);
  }

  req.log.debug({ headers }, `Starting to proxy SSE stream.`);

  // Typically, streaming will have already been initialized by the request
  // queue to send heartbeat pings.
  if (!res.headersSent) {
    copySseResponseHeaders(proxyRes, res);
    initializeSseStream(res);
  }

  const prefersNativeEvents = req.inboundApi === req.outboundApi;
  const streamOptions = {
    contentType: headers["content-type"],
    api: req.outboundApi,
    logger: req.log,
  };

  // While the request is streaming, aggregator collects all events so that we
  // can compile them into a single response object and publish that to the
  // remaining middleware. Because we have an OpenAI transformer for every
  // supported format, EventAggregator always consumes OpenAI events so that we
  // only have to write one aggregator (OpenAI input) for each output format.
  const aggregator = new EventAggregator(req);

  const decompressor = getStreamDecompressor(headers["content-encoding"]);
  // Decoder reads from the response bytes to produce a stream of plaintext.
  const decoder = getDecoder({ ...streamOptions, input: proxyRes });
  // Adapter consumes the decoded text and produces server-sent events so we
  // have a standard event format for the client and to translate between API
  // message formats.
  const adapter = new SSEStreamAdapter(streamOptions);
  // Transformer converts server-sent events from one vendor's API message
  // format to another.
  const transformer = new SSEMessageTransformer({
    inputFormat: req.outboundApi, // The format of the upstream service's events
    outputFormat: req.inboundApi, // The format the client requested
    inputApiVersion: String(req.headers["anthropic-version"]),
    logger: req.log,
    requestId: String(req.id),
    requestedModel: req.body.model,
  })
    .on("originalMessage", (msg: string) => {
      if (prefersNativeEvents) res.write(msg);
    })
    .on("data", (msg) => {
      if (!prefersNativeEvents) res.write(`data: ${JSON.stringify(msg)}\n\n`);
      aggregator.addEvent(msg);
    });

  try {
    await Promise.race([
      handleAbortedStream(req, res),
      pipelineAsync(proxyRes, decompressor, decoder, adapter, transformer),
    ]);
    req.log.debug(`Finished proxying SSE stream.`);
    res.end();
    return aggregator.getFinalResponse();
  } catch (err) {
    if (err instanceof RetryableError) {
      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
    } else if (err instanceof BadRequestError) {
      sendErrorToClient({
        req,
        res,
        options: {
          format: req.inboundApi,
          title: "Proxy streaming error (Bad Request)",
          message: `The API returned an error while streaming your request. Your prompt might not be formatted correctly.\n\n*${err.message}*`,
          reqId: req.id,
          model: req.body?.model,
        },
      });
    } else {
      const { message, stack, lastEvent } = err;
      const eventText = JSON.stringify(lastEvent, null, 2) ?? "undefined";
      const errorEvent = buildSpoofedSSE({
        format: req.inboundApi,
        title: "Proxy stream error",
        message: "An unexpected error occurred while streaming the response.",
        obj: { message, stack, lastEvent: eventText },
        reqId: req.id,
        model: req.body?.model,
      });
      res.write(errorEvent);
      res.write(`data: [DONE]\n\n`);
      res.end();
    }

    // At this point the response is closed. If the request resulted in any
    // tokens being consumed (suggesting a mid-stream error), we will resolve
    // and continue the middleware chain so tokens can be counted.
    if (aggregator.hasEvents()) {
      return aggregator.getFinalResponse();
    } else {
      // If there is nothing, then this was a completely failed prompt that
      // will not have billed any tokens. Throw to stop the middleware chain.
      throw err;
    }
  }
};

function handleAbortedStream(req: express.Request, res: express.Response) {
  return new Promise<void>((resolve) =>
    res.on("close", () => {
      if (!res.writableEnded) {
        req.log.info("Client prematurely closed connection during stream.");
      }
      resolve();
    })
  );
}

function getDecoder(options: {
  input: Readable;
  api: APIFormat;
  logger: typeof logger;
  contentType?: string;
}) {
  const { contentType, input, logger } = options;
  if (contentType?.includes("application/vnd.amazon.eventstream")) {
    return getAwsEventStreamDecoder({ input, logger });
  } else if (contentType?.includes("application/json")) {
    throw new Error("JSON streaming not supported, request SSE instead");
  } else {
    // Ensures split chunks across multi-byte characters are handled correctly.
    const stringDecoder = new StringDecoder("utf8");
    return new Transform({
      readableObjectMode: true,
      writableObjectMode: false,
      transform(chunk, _encoding, callback) {
        const text = stringDecoder.write(chunk);
        if (text) this.push(text);
        callback();
      },
    });
  }
}


code:
/* This file is fucking horrendous, sorry */
// TODO: extract all per-service error response handling into its own modules
import { Request, Response } from "express";
import * as http from "http";
import { config } from "../../../config";
import { HttpError, RetryableError } from "../../../shared/errors";
import { keyPool } from "../../../shared/key-management";
import { getOpenAIModelFamily } from "../../../shared/models";
import { countTokens } from "../../../shared/tokenization";
import {
  incrementPromptCount,
  incrementTokenCount,
} from "../../../shared/users/user-store";
import { assertNever } from "../../../shared/utils";
import { reenqueueRequest, trackWaitTime } from "../../queue";
import { refundLastAttempt } from "../../rate-limit";
import {
  getCompletionFromBody,
  isImageGenerationRequest,
  isTextGenerationRequest,
  sendProxyError,
} from "../common";
import { handleBlockingResponse } from "./handle-blocking-response";
import { handleStreamedResponse } from "./handle-streamed-response";
import { logPrompt } from "./log-prompt";
import { logEvent } from "./log-event";
import { saveImage } from "./save-image";

/**
 * Either decodes or streams the entire response body and then resolves with it.
 * @returns The response body as a string or parsed JSON object depending on the
 * response's content-type.
 */
export type RawResponseBodyHandler = (
  proxyRes: http.IncomingMessage,
  req: Request,
  res: Response
) => Promise<string | Record<string, any>>;

export type ProxyResHandlerWithBody = (
  proxyRes: http.IncomingMessage,
  req: Request,
  res: Response,
  /**
   * This will be an object if the response content-type is application/json,
   * or if the response is a streaming response. Otherwise it will be a string.
   */
  body: string | Record<string, any>
) => Promise<void>;
export type ProxyResMiddleware = ProxyResHandlerWithBody[] | undefined;

/**
 * Returns a on.proxyRes handler that executes the given middleware stack after
 * the common proxy response handlers have processed the response and decoded
 * the body.  Custom middleware won't execute if the response is determined to
 * be an error from the upstream service as the response will be taken over by
 * the common error handler.
 *
 * For streaming responses, the handleStream middleware will block remaining
 * middleware from executing as it consumes the stream and forwards events to
 * the client. Once the stream is closed, the finalized body will be attached
 * to res.body and the remaining middleware will execute.
 *
 * @param apiMiddleware - Custom middleware to execute after the common response
 * handlers. These *only* execute for non-streaming responses, so should be used
 * to transform non-streaming responses into the desired format.
 */
export const createOnProxyResHandler = (apiMiddleware: ProxyResMiddleware) => {
  return async (
    proxyRes: http.IncomingMessage,
    req: Request,
    res: Response
  ) => {
    // Proxied request has by now been sent to the upstream API, so we revert
    // tracked mutations that were only needed to send the request.
    // This generally means path adjustment, headers, and body serialization.
    if (req.changeManager) {
      req.changeManager.revert();
    }

    const initialHandler = req.isStreaming
      ? handleStreamedResponse
      : handleBlockingResponse;
    let lastMiddleware = initialHandler.name;

    if (Buffer.isBuffer(req.body)) {
      req.body = JSON.parse(req.body.toString());
    }

    try {
      const body = await initialHandler(proxyRes, req, res);
      const middlewareStack: ProxyResMiddleware = [];

      if (req.isStreaming) {
        // Handlers for streaming requests must never write to the response.
        middlewareStack.push(
          trackKeyRateLimit,
          countResponseTokens,
          incrementUsage,
          logPrompt,
          logEvent
        );
      } else {
        middlewareStack.push(
          trackKeyRateLimit,
          injectProxyInfo,
          handleUpstreamErrors,
          countResponseTokens,
          incrementUsage,
          copyHttpHeaders,
          saveImage,
          logPrompt,
          logEvent,
          ...(apiMiddleware ?? [])
        );
      }

      for (const middleware of middlewareStack) {
        lastMiddleware = middleware.name;
        await middleware(proxyRes, req, res, body);
      }

      trackWaitTime(req);
    } catch (error) {
      // Hack: if the error is a retryable rate-limit error, the request has
      // been re-enqueued and we can just return without doing anything else.
      if (error instanceof RetryableError) {
        return;
      }

      // Already logged and responded to the client by handleUpstreamErrors
      if (error instanceof HttpError) {
        if (!res.writableEnded) res.end();
        return;
      }

      const { stack, message } = error;
      const details = { stack, message, lastMiddleware, key: req.key?.hash };
      const description = `Error while executing proxy response middleware: ${lastMiddleware} (${message})`;

      if (res.headersSent) {
        req.log.error(details, description);
        if (!res.writableEnded) res.end();
        return;
      } else {
        req.log.error(details, description);
        res
          .status(500)
          .json({ error: "Internal server error", proxy_note: description });
      }
    }
  };
};

type ProxiedErrorPayload = {
  error?: Record<string, any>;
  message?: string;
  proxy_note?: string;
};

/**
 * Handles non-2xx responses from the upstream service.  If the proxied response
 * is an error, this will respond to the client with an error payload and throw
 * an error to stop the middleware stack.
 * On 429 errors, if request queueing is enabled, the request will be silently
 * re-enqueued.  Otherwise, the request will be rejected with an error payload.
 * @throws {HttpError} On HTTP error status code from upstream service
 */
const handleUpstreamErrors: ProxyResHandlerWithBody = async (
  proxyRes,
  req,
  res,
  body
) => {
  const statusCode = proxyRes.statusCode || 500;
  const statusMessage = proxyRes.statusMessage || "Internal Server Error";
  const service = req.key!.service;
  // Not an error, continue to next response handler
  if (statusCode < 400) return;

  // Parse the error response body
  let errorPayload: ProxiedErrorPayload;
  try {
    assertJsonResponse(body);
    errorPayload = body;
  } catch (parseError) {
    const strBody = String(body).slice(0, 128);
    req.log.error({ statusCode, strBody }, "Error body is not JSON");

    const details = {
      error: parseError.message,
      status: statusCode,
      statusMessage,
      proxy_note: `Proxy got back an error, but it was not in JSON format. This is likely a temporary problem with the upstream service. Response body: ${strBody}`,
    };

    sendProxyError(req, res, statusCode, statusMessage, details);
    throw new HttpError(statusCode, parseError.message);
  }

  // Extract the error type from the response body depending on the service
  if (service === "gcp") {
    if (Array.isArray(errorPayload)) {
      errorPayload = errorPayload[0];
    }
  }
  const errorType =
    errorPayload.error?.code ||
    errorPayload.error?.type ||
    getAwsErrorType(proxyRes.headers["x-amzn-errortype"]);

  req.log.warn(
    { statusCode, statusMessage, errorType, errorPayload, key: req.key?.hash },
    `API returned an error.`
  );

  // Try to convert response body to a ProxiedErrorPayload with message/type
  if (service === "aws") {
    errorPayload.error = { message: errorPayload.message, type: errorType };
    delete errorPayload.message;
  } else if (service === "gcp") {
    if (errorPayload.error?.code) {
      errorPayload.error = {
        message: errorPayload.error.message,
        type: errorPayload.error.status || errorPayload.error.code,
      };
    }
  }

  // Figure out what to do with the error
  // TODO: separate error handling for each service
  if (statusCode === 400) {
    switch (service) {
      case "openai":
      case "mistral-ai":
      case "azure":
        const filteredCodes = ["content_policy_violation", "content_filter"];
        if (filteredCodes.includes(errorPayload.error?.code)) {
          errorPayload.proxy_note = `Request was filtered by the upstream API's content moderation system. Modify your prompt and try again.`;
          refundLastAttempt(req);
        } else if (errorPayload.error?.code === "billing_hard_limit_reached") {
          // For some reason, some models return this 400 error instead of the
          // same 429 billing error that other models return.
          await handleOpenAIRateLimitError(req, errorPayload);
        } else {
          errorPayload.proxy_note = `The upstream API rejected the request. Check the error message for details.`;
        }
        break;
      case "anthropic":
      case "aws":
      case "gcp":
        await handleAnthropicAwsBadRequestError(req, errorPayload);
        break;
      case "google-ai":
        await handleGoogleAIBadRequestError(req, errorPayload);
        break;
      default:
        assertNever(service);
    }
  } else if (statusCode === 401) {
    // Key is invalid or was revoked
    keyPool.disable(req.key!, "revoked");
    errorPayload.proxy_note = `Assigned API key is invalid or revoked, please try again.`;
  } else if (statusCode === 403) {
    switch (service) {
      case "anthropic":
        if (
          errorType === "permission_error" &&
          errorPayload.error?.message?.toLowerCase().includes("multimodal")
        ) {
          keyPool.update(req.key!, { allowsMultimodality: false });
          await reenqueueRequest(req);
          throw new RetryableError(
            "Claude request re-enqueued because key does not support multimodality."
          );
        } else {
          keyPool.disable(req.key!, "revoked");
          errorPayload.proxy_note = `Assigned API key is invalid or revoked, please try again.`;
        }
        return;
      case "aws":
        switch (errorType) {
          case "UnrecognizedClientException":
            // Key is invalid.
            keyPool.disable(req.key!, "revoked");
            errorPayload.proxy_note = `Assigned API key is invalid or revoked, please try again.`;
            break;
          case "AccessDeniedException":
            const isModelAccessError =
              errorPayload.error?.message?.includes(`specified model ID`);
            if (!isModelAccessError) {
              req.log.error(
                { key: req.key?.hash, model: req.body?.model },
                "Disabling key due to AccessDeniedException when invoking model. If credentials are valid, check IAM permissions."
              );
              keyPool.disable(req.key!, "revoked");
            }
            errorPayload.proxy_note = `API key doesn't have access to the requested resource. Model ID: ${req.body?.model}`;
            break;
          default:
            errorPayload.proxy_note = `Received 403 error. Key may be invalid.`;
        }
        return;
      case "mistral-ai":
      case "gcp":
        keyPool.disable(req.key!, "revoked");
        errorPayload.proxy_note = `Assigned API key is invalid or revoked, please try again.`;
        return;
    }
  } else if (statusCode === 429) {
    switch (service) {
      case "openai":
        await handleOpenAIRateLimitError(req, errorPayload);
        break;
      case "anthropic":
        await handleAnthropicRateLimitError(req, errorPayload);
        break;
      case "aws":
        await handleAwsRateLimitError(req, errorPayload);
        break;
      case "gcp":
        await handleGcpRateLimitError(req, errorPayload);
        break;
      case "azure":
      case "mistral-ai":
        await handleAzureRateLimitError(req, errorPayload);
        break;
      case "google-ai":
        await handleGoogleAIRateLimitError(req, errorPayload);
        break;
      default:
        assertNever(service);
    }
  } else if (statusCode === 404) {
    // Most likely model not found
    switch (service) {
      case "openai":
        if (errorType === "model_not_found") {
          const requestedModel = req.body.model;
          const modelFamily = getOpenAIModelFamily(requestedModel);
          errorPayload.proxy_note = `The key assigned to your prompt does not support the requested model (${requestedModel}, family: ${modelFamily}).`;
          req.log.error(
            { key: req.key?.hash, model: requestedModel, modelFamily },
            "Prompt was routed to a key that does not support the requested model."
          );
        }
        break;
      case "anthropic":
      case "google-ai":
      case "mistral-ai":
      case "aws":
      case "gcp":
      case "azure":
        errorPayload.proxy_note = `The key assigned to your prompt does not support the requested model.`;
        break;
      default:
        assertNever(service);
    }
  } else if (statusCode === 503) {
    switch (service) {
      case "aws":
        if (
          errorType === "ServiceUnavailableException" &&
          errorPayload.error?.message?.match(/too many connections/i)
        ) {
          errorPayload.proxy_note = `The requested AWS Bedrock model is overloaded. Try again in a few minutes, or try another model.`;
        }
        break;
      default:
        errorPayload.proxy_note = `Upstream service unavailable. Try again later.`;
        break;
    }
  } else {
    errorPayload.proxy_note = `Unrecognized error from upstream service.`;
  }

  // Redact the OpenAI org id from the error message
  if (errorPayload.error?.message) {
    errorPayload.error.message = errorPayload.error.message.replace(
      /org-.{24}/gm,
      "org-xxxxxxxxxxxxxxxxxxx"
    );
  }

  // Send the error to the client
  sendProxyError(req, res, statusCode, statusMessage, errorPayload);

  // Re-throw the error to bubble up to onProxyRes's handler for logging
  throw new HttpError(statusCode, errorPayload.error?.message);
};

async function handleAnthropicAwsBadRequestError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const { error } = errorPayload;
  const isMissingPreamble = error?.message.startsWith(
    `prompt must start with "\n\nHuman:" turn`
  );

  // Some keys mandate a \n\nHuman: preamble, which we can add and retry
  if (isMissingPreamble) {
    req.log.warn(
      { key: req.key?.hash },
      "Request failed due to missing preamble. Key will be marked as such for subsequent requests."
    );
    keyPool.update(req.key!, { requiresPreamble: true });
    await reenqueueRequest(req);
    throw new RetryableError("Claude request re-enqueued to add preamble.");
  }

  // {"type":"error","error":{"type":"invalid_request_error","message":"Usage blocked until 2024-03-01T00:00:00+00:00 due to user specified spend limits."}}
  // {"type":"error","error":{"type":"invalid_request_error","message":"Your credit balance is too low to access the Claude API. Please go to Plans & Billing to upgrade or purchase credits."}}
  const isOverQuota =
    error?.message?.match(/usage blocked until/i) ||
    error?.message?.match(/credit balance is too low/i);
  if (isOverQuota) {
    req.log.warn(
      { key: req.key?.hash, message: error?.message },
      "Anthropic key has hit spending limit and will be disabled."
    );
    keyPool.disable(req.key!, "quota");
    errorPayload.proxy_note = `Assigned key has hit its spending limit. ${error?.message}`;
    return;
  }

  const isDisabled =
    error?.message?.match(/organization has been disabled/i) ||
    error?.message?.match(/^operation not allowed/i);
  if (isDisabled) {
    req.log.warn(
      { key: req.key?.hash, message: error?.message },
      "Anthropic/AWS key has been disabled."
    );
    keyPool.disable(req.key!, "revoked");
    errorPayload.proxy_note = `Assigned key has been disabled. (${error?.message})`;
    return;
  }

  errorPayload.proxy_note = `Unrecognized error from the API. (${error?.message})`;
}

async function handleAnthropicRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  if (errorPayload.error?.type === "rate_limit_error") {
    keyPool.markRateLimited(req.key!);
    await reenqueueRequest(req);
    throw new RetryableError("Claude rate-limited request re-enqueued.");
  } else {
    errorPayload.proxy_note = `Unrecognized 429 Too Many Requests error from the API.`;
  }
}

async function handleAwsRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const errorType = errorPayload.error?.type;
  switch (errorType) {
    case "ThrottlingException":
      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
      throw new RetryableError("AWS rate-limited request re-enqueued.");
    case "ModelNotReadyException":
      errorPayload.proxy_note = `The requested model is overloaded. Try again in a few seconds.`;
      break;
    default:
      errorPayload.proxy_note = `Unrecognized rate limit error from AWS. (${errorType})`;
  }
}

async function handleGcpRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  if (errorPayload.error?.type === "RESOURCE_EXHAUSTED") {
    keyPool.markRateLimited(req.key!);
    await reenqueueRequest(req);
    throw new RetryableError("GCP rate-limited request re-enqueued.");
  } else {
    errorPayload.proxy_note = `Unrecognized 429 Too Many Requests error from GCP.`;
  }
}

async function handleOpenAIRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
): Promise<Record<string, any>> {
  const type = errorPayload.error?.type;
  switch (type) {
    case "insufficient_quota":
    case "invalid_request_error": // this is the billing_hard_limit_reached error seen in some cases
      // Billing quota exceeded (key is dead, disable it)
      keyPool.disable(req.key!, "quota");
      errorPayload.proxy_note = `Assigned key's quota has been exceeded. Please try again.`;
      break;
    case "access_terminated":
      // Account banned (key is dead, disable it)
      keyPool.disable(req.key!, "revoked");
      errorPayload.proxy_note = `Assigned key has been banned by OpenAI for policy violations. Please try again.`;
      break;
    case "billing_not_active":
      // Key valid but account billing is delinquent
      keyPool.disable(req.key!, "quota");
      errorPayload.proxy_note = `Assigned key has been disabled due to delinquent billing. Please try again.`;
      break;
    case "requests":
    case "tokens":
      keyPool.markRateLimited(req.key!);
      if (errorPayload.error?.message?.match(/on requests per day/)) {
        // This key has a very low rate limit, so we can't re-enqueue it.
        errorPayload.proxy_note = `Assigned key has reached its per-day request limit for this model. Try another model.`;
        break;
      }

      // Per-minute request or token rate limit is exceeded, which we can retry
      await reenqueueRequest(req);
      throw new RetryableError("Rate-limited request re-enqueued.");
    default:
      errorPayload.proxy_note = `This is likely a temporary error with the API. Try again in a few seconds.`;
      break;
  }
  return errorPayload;
}

async function handleAzureRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const code = errorPayload.error?.code;
  switch (code) {
    case "429":
      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
      throw new RetryableError("Rate-limited request re-enqueued.");
    default:
      errorPayload.proxy_note = `Unrecognized rate limit error from Azure (${code}). Please report this.`;
      break;
  }
}

//{"error":{"code":400,"message":"API Key not found. Please pass a valid API key.","status":"INVALID_ARGUMENT","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"API_KEY_INVALID","domain":"googleapis.com","metadata":{"service":"generativelanguage.googleapis.com"}}]}}
//{"error":{"code":400,"message":"Gemini API free tier is not available in your country. Please enable billing on your project in Google AI Studio.","status":"FAILED_PRECONDITION"}}
async function handleGoogleAIBadRequestError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const error = errorPayload.error || {};
  // google changes this shit every few months
  // i don't want to deal with it
  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);
  if (keyDeadMsgs.some((msg) => text.match(msg))) {
    req.log.warn(
      { key: req.key?.hash, error: text },
      "Google API key appears to be inoperative."
    );
    keyPool.disable(req.key!, "revoked");
    errorPayload.proxy_note = `Assigned API key cannot be used.`;
  } else {
    req.log.warn(
      { key: req.key?.hash, error: text },
      "Unknown Google API error."
    );
    errorPayload.proxy_note = `Unrecognized error from Google AI.`;
  }

  // const { message, status, details } = error;
  //
  // if (status === "INVALID_ARGUMENT") {
  //   const reason = details?.[0]?.reason;
  //   if (reason === "API_KEY_INVALID") {
  //     req.log.warn(
  //       { key: req.key?.hash, status, reason, msg: error.message },
  //       "Received `API_KEY_INVALID` error from Google AI. Check the configured API key."
  //     );
  //     keyPool.disable(req.key!, "revoked");
  //     errorPayload.proxy_note = `Assigned API key is invalid.`;
  //   }
  // } else if (status === "FAILED_PRECONDITION") {
  //   if (message.match(/please enable billing/i)) {
  //     req.log.warn(
  //       { key: req.key?.hash, status, msg: error.message },
  //       "Cannot use key due to billing restrictions."
  //     );
  //     keyPool.disable(req.key!, "revoked");
  //     errorPayload.proxy_note = `Assigned API key cannot be used.`;
  //   }
  // } else {
  //   req.log.warn(
  //     { key: req.key?.hash, status, msg: error.message },
  //     "Received unexpected 400 error from Google AI."
  //   );
  // }
}

//{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"}
//
async function handleGoogleAIRateLimitError(
  req: Request,
  errorPayload: ProxiedErrorPayload
) {
  const status = errorPayload.error?.status;
  const text = JSON.stringify(errorPayload.error);

  // sometimes they block keys by rate limiting them to 0 requests per minute
  // for some indefinite period of time
  const keyDeadMsgs = [
    /GenerateContentRequestsPerMinutePerProjectPerRegion/i,
    /"quota_limit_value":"0"/i,
  ];

  switch (status) {
    case "RESOURCE_EXHAUSTED": {
      if (keyDeadMsgs.every((msg) => text.match(msg))) {
        req.log.warn(
          { key: req.key?.hash, error: text },
          "Google API key appears to be temporarily inoperative and will be disabled."
        );
        keyPool.disable(req.key!, "revoked");
        errorPayload.proxy_note = `Assigned API key cannot be used.`;
        return;
      }

      keyPool.markRateLimited(req.key!);
      await reenqueueRequest(req);
      throw new RetryableError("Rate-limited request re-enqueued.");
    }
    default:
      errorPayload.proxy_note = `Unrecognized rate limit error from Google AI (${status}). Please report this.`;
      break;
  }
}

const incrementUsage: ProxyResHandlerWithBody = async (_proxyRes, req) => {
  if (isTextGenerationRequest(req) || isImageGenerationRequest(req)) {
    const model = req.body.model;
    const tokensUsed = req.promptTokens! + req.outputTokens!;
    req.log.debug(
      {
        model,
        tokensUsed,
        promptTokens: req.promptTokens,
        outputTokens: req.outputTokens,
      },
      `Incrementing usage for model`
    );
    keyPool.incrementUsage(req.key!, model, tokensUsed);
    if (req.user) {
      incrementPromptCount(req.user.token);
      incrementTokenCount(req.user.token, model, req.outboundApi, tokensUsed);
    }
  }
};

const countResponseTokens: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  _res,
  body
) => {
  if (req.outboundApi === "openai-image") {
    req.outputTokens = req.promptTokens;
    req.promptTokens = 0;
    return;
  }

  // This function is prone to breaking if the upstream API makes even minor
  // changes to the response format, especially for SSE responses. If you're
  // seeing errors in this function, check the reassembled response body from
  // handleStreamedResponse to see if the upstream API has changed.
  try {
    assertJsonResponse(body);
    const service = req.outboundApi;
    const completion = getCompletionFromBody(req, body);
    const tokens = await countTokens({ req, completion, service });

    if (req.service === "openai" || req.service === "azure") {
      // O1 consumes (a significant amount of) invisible tokens for the chain-
      // of-thought reasoning. We have no way to count these other than to check
      // the response body.
      tokens.reasoning_tokens =
        body.usage?.completion_tokens_details?.reasoning_tokens;
    }

    req.log.debug(
      { service, prevOutputTokens: req.outputTokens, tokens },
      `Counted tokens for completion`
    );
    if (req.tokenizerInfo) {
      req.tokenizerInfo.completion_tokens = tokens;
    }

    req.outputTokens = tokens.token_count + (tokens.reasoning_tokens ?? 0);
  } catch (error) {
    req.log.warn(
      error,
      "Error while counting completion tokens; assuming `max_output_tokens`"
    );
    // req.outputTokens will already be set to `max_output_tokens` from the
    // prompt counting middleware, so we don't need to do anything here.
  }
};

const trackKeyRateLimit: ProxyResHandlerWithBody = async (proxyRes, req) => {
  keyPool.updateRateLimits(req.key!, proxyRes.headers);
};

const omittedHeaders = new Set<string>([
  // Omit content-encoding because we will always decode the response body
  "content-encoding",
  // Omit transfer-encoding because we are using response.json which will
  // set a content-length header, which is not valid for chunked responses.
  "transfer-encoding",
  // Don't set cookies from upstream APIs because proxied requests are stateless
  "set-cookie",
  "openai-organization",
  "x-request-id",
  "cf-ray",
]);
const copyHttpHeaders: ProxyResHandlerWithBody = async (
  proxyRes,
  _req,
  res
) => {
  Object.keys(proxyRes.headers).forEach((key) => {
    if (omittedHeaders.has(key)) return;
    res.setHeader(key, proxyRes.headers[key] as string);
  });
};

/**
 * Injects metadata into the response, such as the tokenizer used, logging
 * status, upstream API endpoint used, and whether the input prompt was modified
 * or transformed.
 * Only used for non-streaming requests.
 */
const injectProxyInfo: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  const { service, inboundApi, outboundApi, tokenizerInfo } = req;
  const native = inboundApi === outboundApi;
  const info: any = {
    logged: config.promptLogging,
    tokens: tokenizerInfo,
    service,
    in_api: inboundApi,
    out_api: outboundApi,
    prompt_transformed: !native,
  };

  if (req.query?.debug?.length) {
    info.final_request_body = req.signedRequest?.body || req.body;
  }

  if (typeof body === "object") {
    body.proxy = info;
  }
};

function getAwsErrorType(header: string | string[] | undefined) {
  const val = String(header).match(/^(\w+):?/)?.[1];
  return val || String(header);
}

function assertJsonResponse(body: any): asserts body is Record<string, any> {
  if (typeof body !== "object") {
    throw new Error(`Expected response to be an object, got ${typeof body}`);
  }
}