code:
import { Request, Response, RequestHandler, Router } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addKey,
  addAnthropicPreamble,
  createPreprocessorMiddleware,
  finalizeBody,
  createOnProxyReqHandler,
} from "./middleware/request";
import {
  ProxyResHandlerWithBody,
  createOnProxyResHandler,
} from "./middleware/response";
import { sendErrorToClient } from "./middleware/response/error-generator";

let modelsCache: any = null;
let modelsCacheTime = 0;

const getModelsResponse = () => {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60) {
    return modelsCache;
  }

  if (!config.anthropicKey) return { object: "list", data: [] };

  const claudeVariants = [
    "claude-v1",
    "claude-v1-100k",
    "claude-instant-v1",
    "claude-instant-v1-100k",
    "claude-v1.3",
    "claude-v1.3-100k",
    "claude-v1.2",
    "claude-v1.0",
    "claude-instant-v1.1",
    "claude-instant-v1.1-100k",
    "claude-instant-v1.0",
    "claude-2",
    "claude-2.0",
    "claude-2.1",
    "claude-3-opus-20240229",
    "claude-3-sonnet-20240229",
  ];

  const models = claudeVariants.map((id) => ({
    id,
    object: "model",
    created: new Date().getTime(),
    owned_by: "anthropic",
    permission: [],
    root: "claude",
    parent: null,
  }));

  modelsCache = { object: "list", data: models };
  modelsCacheTime = new Date().getTime();

  return modelsCache;
};

const handleModelRequest: RequestHandler = (_req, res) => {
  res.status(200).json(getModelsResponse());
};

/** Only used for non-streaming requests. */
const anthropicResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.inboundApi === "openai") {
    req.log.info("Transforming Anthropic text to OpenAI format");
    body = transformAnthropicTextResponseToOpenAI(body, req);
  }

  if (
    req.inboundApi === "anthropic-text" &&
    req.outboundApi === "anthropic-chat"
  ) {
    req.log.info("Transforming Anthropic text to Anthropic chat format");
    body = transformAnthropicChatResponseToAnthropicText(body);
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  res.status(200).json(body);
};

export function transformAnthropicChatResponseToAnthropicText(
  anthropicBody: Record<string, any>
): Record<string, any> {
  return {
    type: "completion",
    id: "trans-" + anthropicBody.id,
    completion: anthropicBody.content
      .map((part: { type: string; text: string }) =>
        part.type === "text" ? part.text : ""
      )
      .join(""),
    stop_reason: anthropicBody.stop_reason,
    stop: anthropicBody.stop_sequence,
    model: anthropicBody.model,
    usage: anthropicBody.usage,
  };
}

/**
 * Transforms a model response from the Anthropic API to match those from the
 * OpenAI API, for users using Claude via the OpenAI-compatible endpoint. This
 * is only used for non-streaming requests as streaming requests are handled
 * on-the-fly.
 */
function transformAnthropicTextResponseToOpenAI(
  anthropicBody: Record<string, any>,
  req: Request
): Record<string, any> {
  const totalTokens = (req.promptTokens ?? 0) + (req.outputTokens ?? 0);
  return {
    id: "ant-" + anthropicBody.log_id,
    object: "chat.completion",
    created: Date.now(),
    model: anthropicBody.model,
    usage: {
      prompt_tokens: req.promptTokens,
      completion_tokens: req.outputTokens,
      total_tokens: totalTokens,
    },
    choices: [
      {
        message: {
          role: "assistant",
          content: anthropicBody.completion?.trim(),
        },
        finish_reason: anthropicBody.stop_reason,
        index: 0,
      },
    ],
  };
}

const anthropicProxy = createQueueMiddleware({
  proxyMiddleware: createProxyMiddleware({
    target: "https://api.anthropic.com",
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    on: {
      proxyReq: createOnProxyReqHandler({
        pipeline: [addKey, addAnthropicPreamble, finalizeBody],
      }),
      proxyRes: createOnProxyResHandler([anthropicResponseHandler]),
      error: handleProxyError,
    },
    // Abusing pathFilter to rewrite the paths dynamically.
    pathFilter: (pathname, req) => {
      const isText = req.outboundApi === "anthropic-text";
      const isChat = req.outboundApi === "anthropic-chat";
      if (isChat && pathname === "/v1/complete") {
        req.url = "/v1/messages";
      }
      if (isText && pathname === "/v1/chat/completions") {
        req.url = "/v1/complete";
      }
      if (isChat && ["sonnet", "opus"].includes(req.params.type)) {
        req.url = "/v1/messages";
      }
      return true;
    },
  }),
});

const nativeTextPreprocessor = createPreprocessorMiddleware({
  inApi: "anthropic-text",
  outApi: "anthropic-text",
  service: "anthropic",
});

const textToChatPreprocessor = createPreprocessorMiddleware({
  inApi: "anthropic-text",
  outApi: "anthropic-chat",
  service: "anthropic",
});

/**
 * Routes text completion prompts to anthropic-chat if they need translation
 * (claude-3 based models do not support the old text completion endpoint).
 */
const claudeTextCompletionRouter: RequestHandler = (req, res, next) => {
  if (req.body.model?.startsWith("claude-3")) {
    textToChatPreprocessor(req, res, next);
  } else {
    nativeTextPreprocessor(req, res, next);
  }
};

const anthropicRouter = Router();
anthropicRouter.get("/v1/models", handleModelRequest);
// Anthropic text completion endpoint. Dynamic routing based on model.
anthropicRouter.post(
  "/v1/complete",
  ipLimiter,
  claudeTextCompletionRouter,
  anthropicProxy
);
// Native Anthropic chat completion endpoint.
anthropicRouter.post(
  "/v1/messages",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "anthropic-chat",
    outApi: "anthropic-chat",
    service: "anthropic",
  }),
  anthropicProxy
);
// OpenAI-to-Anthropic Text compatibility endpoint.
anthropicRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "openai", outApi: "anthropic-text", service: "anthropic" },
    { afterTransform: [maybeReassignModel] }
  ),
  anthropicProxy
);
// Temporary force Anthropic Text to Anthropic Chat for frontends which do not
// yet support the new model. Forces claude-3. Will be removed once common
// frontends have been updated.
anthropicRouter.post(
  "/v1/:type(sonnet|opus)/:action(complete|messages)",
  ipLimiter,
  handleCompatibilityRequest,
  createPreprocessorMiddleware({
    inApi: "anthropic-text",
    outApi: "anthropic-chat",
    service: "anthropic",
  }),
  anthropicProxy
);

function handleCompatibilityRequest(req: Request, res: Response, next: any) {
  const type = req.params.type;
  const action = req.params.action;
  const alreadyInChatFormat = Boolean(req.body.messages);
  const compatModel = `claude-3-${type}-20240229`;
  req.log.info(
    { type, inputModel: req.body.model, compatModel, alreadyInChatFormat },
    "Handling Anthropic compatibility request"
  );

  if (action === "messages" || alreadyInChatFormat) {
    return sendErrorToClient({
      req,
      res,
      options: {
        title: "Unnecessary usage of compatibility endpoint",
        message: `Your client seems to already support the new Claude API format. This endpoint is intended for clients that do not yet support the new format.\nUse the normal \`/anthropic\` proxy endpoint instead.`,
        format: "unknown",
        statusCode: 400,
        reqId: req.id,
        obj: {
          requested_endpoint: "/anthropic/" + type,
          correct_endpoint: "/anthropic",
        },
      },
    });
  }

  req.body.model = compatModel;
  next();
}

function maybeReassignModel(req: Request) {
  const model = req.body.model;
  if (!model.startsWith("gpt-")) return;
  req.body.model = "claude-2.1";
}

export const anthropic = anthropicRouter;


code:
import { Request, RequestHandler, Response, Router } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { v4 } from "uuid";
import { config } from "../config";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  createPreprocessorMiddleware,
  signAwsRequest,
  finalizeSignedRequest,
  createOnProxyReqHandler,
} from "./middleware/request";
import {
  ProxyResHandlerWithBody,
  createOnProxyResHandler,
} from "./middleware/response";
import { transformAnthropicChatResponseToAnthropicText } from "./anthropic";
import { sendErrorToClient } from "./middleware/response/error-generator";

const LATEST_AWS_V2_MINOR_VERSION = "1";

let modelsCache: any = null;
let modelsCacheTime = 0;

const getModelsResponse = () => {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60) {
    return modelsCache;
  }

  if (!config.awsCredentials) return { object: "list", data: [] };

  // https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
  const variants = [
    "anthropic.claude-v2",
    "anthropic.claude-v2:1",
    "anthropic.claude-3-sonnet-20240229-v1:0",
  ];

  const models = variants.map((id) => ({
    id,
    object: "model",
    created: new Date().getTime(),
    owned_by: "anthropic",
    permission: [],
    root: "claude",
    parent: null,
  }));

  modelsCache = { object: "list", data: models };
  modelsCacheTime = new Date().getTime();

  return modelsCache;
};

const handleModelRequest: RequestHandler = (_req, res) => {
  res.status(200).json(getModelsResponse());
};

/** Only used for non-streaming requests. */
const awsResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.inboundApi === "openai") {
    req.log.info("Transforming AWS Claude response to OpenAI format");
    body = transformAwsTextResponseToOpenAI(body, req);
  }

  if (
    req.inboundApi === "anthropic-text" &&
    req.outboundApi === "anthropic-chat"
  ) {
    req.log.info("Transforming AWS Claude chat response to Text format");
    body = transformAnthropicChatResponseToAnthropicText(body);
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  // AWS does not confirm the model in the response, so we have to add it
  body.model = req.body.model;

  res.status(200).json(body);
};

/**
 * Transforms a model response from the Anthropic API to match those from the
 * OpenAI API, for users using Claude via the OpenAI-compatible endpoint. This
 * is only used for non-streaming requests as streaming requests are handled
 * on-the-fly.
 */
function transformAwsTextResponseToOpenAI(
  awsBody: Record<string, any>,
  req: Request
): Record<string, any> {
  const totalTokens = (req.promptTokens ?? 0) + (req.outputTokens ?? 0);
  return {
    id: "aws-" + v4(),
    object: "chat.completion",
    created: Date.now(),
    model: req.body.model,
    usage: {
      prompt_tokens: req.promptTokens,
      completion_tokens: req.outputTokens,
      total_tokens: totalTokens,
    },
    choices: [
      {
        message: {
          role: "assistant",
          content: awsBody.completion?.trim(),
        },
        finish_reason: awsBody.stop_reason,
        index: 0,
      },
    ],
  };
}

const awsProxy = createQueueMiddleware({
  beforeProxy: signAwsRequest,
  proxyMiddleware: createProxyMiddleware({
    target: "bad-target-will-be-rewritten",
    router: ({ signedRequest }) => {
      if (!signedRequest) throw new Error("Must sign request before proxying");
      return `${signedRequest.protocol}//${signedRequest.hostname}`;
    },
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    on: {
      proxyReq: createOnProxyReqHandler({ pipeline: [finalizeSignedRequest] }),
      proxyRes: createOnProxyResHandler([awsResponseHandler]),
      error: handleProxyError,
    },
  }),
});

const nativeTextPreprocessor = createPreprocessorMiddleware(
  { inApi: "anthropic-text", outApi: "anthropic-text", service: "aws" },
  { afterTransform: [maybeReassignModel] }
);

const textToChatPreprocessor = createPreprocessorMiddleware(
  { inApi: "anthropic-text", outApi: "anthropic-chat", service: "aws" },
  { afterTransform: [maybeReassignModel] }
);

/**
 * Routes text completion prompts to aws anthropic-chat if they need translation
 * (claude-3 based models do not support the old text completion endpoint).
 */
const awsTextCompletionRouter: RequestHandler = (req, res, next) => {
  if (req.body.model?.includes("claude-3")) {
    textToChatPreprocessor(req, res, next);
  } else {
    nativeTextPreprocessor(req, res, next);
  }
};

const awsRouter = Router();
awsRouter.get("/v1/models", handleModelRequest);
// Native(ish) Anthropic text completion endpoint.
awsRouter.post("/v1/complete", ipLimiter, awsTextCompletionRouter, awsProxy);
// Native Anthropic chat completion endpoint.
awsRouter.post(
  "/v1/messages",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "anthropic-chat", outApi: "anthropic-chat", service: "aws" },
    { afterTransform: [maybeReassignModel] }
  ),
  awsProxy
);
// Temporary force-Claude3 endpoint
awsRouter.post(
  "/v1/sonnet/:action(complete|messages)",
  ipLimiter,
  handleCompatibilityRequest,
  createPreprocessorMiddleware({
    inApi: "anthropic-text",
    outApi: "anthropic-chat",
    service: "aws",
  }),
  awsProxy
);

// OpenAI-to-AWS Anthropic compatibility endpoint.
awsRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "openai", outApi: "anthropic-text", service: "aws" },
    { afterTransform: [maybeReassignModel] }
  ),
  awsProxy
);

/**
 * Tries to deal with:
 * - frontends sending AWS model names even when they want to use the OpenAI-
 *   compatible endpoint
 * - frontends sending Anthropic model names that AWS doesn't recognize
 * - frontends sending OpenAI model names because they expect the proxy to
 *   translate them
 */
function maybeReassignModel(req: Request) {
  const model = req.body.model;

  // If client already specified an AWS Claude model ID, use it
  if (model.includes("anthropic.claude")) {
    return;
  }

  const pattern =
    /^(claude-)?(instant-)?(v)?(\d+)(\.(\d+))?(-\d+k)?(-sonnet-?|-opus-?)(\d*)/i;
  const match = model.match(pattern);

  // If there's no match, return the latest v2 model
  if (!match) {
    req.body.model = `anthropic.claude-v2:${LATEST_AWS_V2_MINOR_VERSION}`;
    return;
  }

  const instant = match[2];
  const major = match[4];
  const minor = match[6];

  if (instant) {
    req.body.model = "anthropic.claude-instant-v1";
    return;
  }

  // There's only one v1 model
  if (major === "1") {
    req.body.model = "anthropic.claude-v1";
    return;
  }

  // Try to map Anthropic API v2 models to AWS v2 models
  if (major === "2") {
    if (minor === "0") {
      req.body.model = "anthropic.claude-v2";
      return;
    }
    req.body.model = `anthropic.claude-v2:${LATEST_AWS_V2_MINOR_VERSION}`;
    return;
  }

  // AWS currently only supports one v3 model.
  const variant = match[8]; // sonnet or opus
  const variantVersion = match[9];
  if (major === "3") {
    req.body.model = "anthropic.claude-3-sonnet-20240229-v1:0";
    return;
  }

  // Fallback to latest v2 model
  req.body.model = `anthropic.claude-v2:${LATEST_AWS_V2_MINOR_VERSION}`;
  return;
}

export function handleCompatibilityRequest(
  req: Request,
  res: Response,
  next: any
) {
  const action = req.params.action;
  const alreadyInChatFormat = Boolean(req.body.messages);
  const compatModel = "anthropic.claude-3-sonnet-20240229-v1:0";
  req.log.info(
    { inputModel: req.body.model, compatModel, alreadyInChatFormat },
    "Handling AWS compatibility request"
  );

  if (action === "messages" || alreadyInChatFormat) {
    return sendErrorToClient({
      req,
      res,
      options: {
        title: "Unnecessary usage of compatibility endpoint",
        message: `Your client seems to already support the new Claude API format. This endpoint is intended for clients that do not yet support the new format.\nUse the normal \`/aws/claude\` proxy endpoint instead.`,
        format: "unknown",
        statusCode: 400,
        reqId: req.id,
        obj: {
          requested_endpoint: "/aws/claude/sonnet",
          correct_endpoint: "/aws/claude",
        },
      },
    });
  }

  req.body.model = compatModel;
  next();
}

export const aws = awsRouter;

repo_name: Cucumber/oai

license: unknown


code:
import { RequestHandler, Router } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { keyPool } from "../shared/key-management";
import {
  ModelFamily,
  AzureOpenAIModelFamily,
  getAzureOpenAIModelFamily,
} from "../shared/models";
import { logger } from "../logger";
import { KNOWN_OPENAI_MODELS } from "./openai";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addAzureKey,
  createOnProxyReqHandler,
  createPreprocessorMiddleware,
  finalizeSignedRequest,
} from "./middleware/request";
import {
  createOnProxyResHandler,
  ProxyResHandlerWithBody,
} from "./middleware/response";

let modelsCache: any = null;
let modelsCacheTime = 0;

function getModelsResponse() {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60) {
    return modelsCache;
  }

  let available = new Set<AzureOpenAIModelFamily>();
  for (const key of keyPool.list()) {
    if (key.isDisabled || key.service !== "azure") continue;
    key.modelFamilies.forEach((family) =>
      available.add(family as AzureOpenAIModelFamily)
    );
  }
  const allowed = new Set<ModelFamily>(config.allowedModelFamilies);
  available = new Set([...available].filter((x) => allowed.has(x)));

  const models = KNOWN_OPENAI_MODELS.map((id) => ({
    id,
    object: "model",
    created: new Date().getTime(),
    owned_by: "azure",
    permission: [
      {
        id: "modelperm-" + id,
        object: "model_permission",
        created: new Date().getTime(),
        organization: "*",
        group: null,
        is_blocking: false,
      },
    ],
    root: id,
    parent: null,
  })).filter((model) => available.has(getAzureOpenAIModelFamily(model.id)));

  modelsCache = { object: "list", data: models };
  modelsCacheTime = new Date().getTime();

  return modelsCache;
}

const handleModelRequest: RequestHandler = (_req, res) => {
  res.status(200).json(getModelsResponse());
};

const azureOpenaiResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  res.status(200).json(body);
};

const azureOpenAIProxy = createQueueMiddleware({
  beforeProxy: addAzureKey,
  proxyMiddleware: createProxyMiddleware({
    target: "will be set by router",
    router: (req) => {
      if (!req.signedRequest) throw new Error("signedRequest not set");
      const { hostname, path } = req.signedRequest;
      return `https://${hostname}${path}`;
    },
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    on: {
      proxyReq: createOnProxyReqHandler({ pipeline: [finalizeSignedRequest] }),
      proxyRes: createOnProxyResHandler([azureOpenaiResponseHandler]),
      error: handleProxyError,
    },
  }),
});

const azureOpenAIRouter = Router();
azureOpenAIRouter.get("/v1/models", handleModelRequest);
azureOpenAIRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai",
    outApi: "openai",
    service: "azure",
  }),
  azureOpenAIProxy
);

export const azure = azureOpenAIRouter;


code:
import { config } from "../config";
import { RequestHandler } from "express";

const BLOCKED_REFERERS = config.blockedOrigins?.split(",") || [];

/** Disallow requests from blocked origins and referers. */
export const checkOrigin: RequestHandler = (req, res, next) => {
  const blocks = BLOCKED_REFERERS || [];
  for (const block of blocks) {
    if (
      req.headers.origin?.includes(block) ||
      req.headers.referer?.includes(block)
    ) {
      req.log.warn(
        { origin: req.headers.origin, referer: req.headers.referer },
        "Blocked request from origin or referer"
      );

      // VenusAI requests incorrectly say they accept HTML despite immediately
      // trying to parse the response as JSON, so we check the body type instead
      const hasJsonBody =
        req.headers["content-type"]?.includes("application/json");
      if (!req.accepts("html") || hasJsonBody) {
        return res.status(403).json({
          error: { type: "blocked_origin", message: config.blockMessage },
        });
      } else {
        const destination = config.blockRedirect || "https://openai.com";
        return res.status(403).send(
          `

  <title>Redirecting</title>
  <meta http-equiv="refresh" content="3; url=${destination}" />

${config.blockMessage}

Please hold while you are redirected to a more suitable service.

`
        );
      }
    }
  }
  next();
};


code:
/**
 * Authenticates RisuAI.xyz users using a special x-risu-tk header provided by
 * RisuAI.xyz. This lets us rate limit and limit queue concurrency properly,
 * since otherwise RisuAI.xyz users share the same IP address and can't be
 * distinguished.
 * Contributors: @kwaroran
 */
import crypto from "crypto";
import { Request, Response, NextFunction } from "express";
import { logger } from "../logger";

const log = logger.child({ module: "check-risu-token" });

const RISUAI_PUBLIC_KEY = `
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArEXBmHQfy/YdNIu9lfNC
xHbVwb2aYx07pBEmqQJtvVEOISj80fASxg+cMJH+/0a/Z4gQgzUJl0HszRpMXAfu
wmRoetedyC/6CLraHke0Qad/AEHAKwG9A+NwsHRv/cDfP8euAr20cnOyVa79bZsl
1wlHYQQGo+ve+P/FXtjLGJ/KZYr479F5jkIRKZxPE8mRmkhAVS/u+18QM94BzfoI
0LlbwvvCHe18QSX6viDK+HsqhhyYDh+0FgGNJw6xKYLdExbQt77FSukH7NaJmVAs
kYuIJbnAGw5Oq0L6dXFW2DFwlcLz51kPVOmDc159FsQjyuPnta7NiZAANS8KM1CJ
pwIDAQAB`;
let IMPORTED_RISU_KEY: CryptoKey | null = null;

type RisuToken = { id: string; expiresIn: number };
type SignedToken = { data: RisuToken; sig: string };

(async () => {
  try {
    log.debug("Importing Risu public key");
    IMPORTED_RISU_KEY = await crypto.subtle.importKey(
      "spki",
      Buffer.from(RISUAI_PUBLIC_KEY.replace(/\s/g, ""), "base64"),
      { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" },
      true,
      ["verify"]
    );
    log.debug("Imported Risu public key");
  } catch (err) {
    log.warn({ error: err.message }, "Error importing Risu public key");
    IMPORTED_RISU_KEY = null;
  }
})();

export async function checkRisuToken(
  req: Request,
  _res: Response,
  next: NextFunction
) {
  let header = req.header("x-risu-tk") || null;
  if (!header || !IMPORTED_RISU_KEY) {
    return next();
  }

  try {
    const { valid, data } = await validCheck(header);

    if (!valid || !data) {
      req.log.warn(
        { token: header, data },
        "Invalid RisuAI token; using IP instead"
      );
    } else {
      req.log.info("RisuAI token validated");
      req.risuToken = String(data.id);
    }
  } catch (err) {
    req.log.warn(
      { error: err.message },
      "Error validating RisuAI token; using IP instead"
    );
  }

  next();
}

async function validCheck(header: string) {
  let tk: SignedToken;
  try {
    tk = JSON.parse(
      Buffer.from(decodeURIComponent(header), "base64").toString("utf-8")
    );
  } catch (err) {
    log.warn({ error: err.message }, "Provided unparseable RisuAI token");
    return { valid: false };
  }
  const data: RisuToken = tk.data;
  const sig = Buffer.from(tk.sig, "base64");

  if (data.expiresIn < Math.floor(Date.now() / 1000)) {
    log.warn({ token: header }, "Provided expired RisuAI token");
    return { valid: false };
  }

  const valid = await crypto.subtle.verify(
    { name: "RSASSA-PKCS1-v1_5" },
    IMPORTED_RISU_KEY!,
    sig,
    Buffer.from(JSON.stringify(data))
  );

  if (!valid) {
    log.warn({ token: header }, "RisuAI token failed signature check");
  }

  return { valid, data };
}


code:
import type { Request, RequestHandler } from "express";
import { config } from "../config";
import { authenticate, getUser } from "../shared/users/user-store";

const GATEKEEPER = config.gatekeeper;
const PROXY_KEY = config.proxyKey;
const ADMIN_KEY = config.adminKey;

function getProxyAuthorizationFromRequest(req: Request): string | undefined {
  // Anthropic's API uses x-api-key instead of Authorization.  Some clients will
  // pass the _proxy_ key in this header too, instead of providing it as a
  // Bearer token in the Authorization header.  So we need to check both.
  // Prefer the Authorization header if both are present.

  if (req.headers.authorization) {
    const token = req.headers.authorization?.slice("Bearer ".length);
    delete req.headers.authorization;
    return token;
  }

  if (req.headers["x-api-key"]) {
    const token = req.headers["x-api-key"]?.toString();
    delete req.headers["x-api-key"];
    return token;
  }

  return undefined;
}

export const gatekeeper: RequestHandler = (req, res, next) => {
  const token = getProxyAuthorizationFromRequest(req);

  // TODO: Generate anonymous users based on IP address for public or proxy_key
  // modes so that all middleware can assume a user of some sort is present.

  if (ADMIN_KEY && token === ADMIN_KEY) {
    return next();
  }

  if (GATEKEEPER === "none") {
    return next();
  }

  if (GATEKEEPER === "proxy_key" && token === PROXY_KEY) {
    return next();
  }

  if (GATEKEEPER === "user_token" && token) {
    // RisuAI users all come from a handful of aws lambda IPs so we cannot use
    // IP alone to distinguish between them and prevent usertoken sharing.
    // Risu sends a signed token in the request headers with an anonymous user
    // ID that we can instead use to associate requests with an individual.
    const ip = req.risuToken?.length ?
      `risu${req.risuToken}-${req.ip}` :
      req.ip;

    const { user, result } = authenticate(token, ip);

    switch (result) {
      case "success":
        req.user = user;
        return next();
      case "limited":
        return res.status(403).json({
          error: `Forbidden: no more IPs can authenticate with this token`,
        });
      case "disabled":
        const bannedUser = getUser(token);
        if (bannedUser?.disabledAt) {
          const reason = bannedUser.disabledReason || "Token disabled";
          return res.status(403).json({ error: `Forbidden: ${reason}` });
        }
    }
  }

  res.status(401).json({ error: "Unauthorized" });
};


code:
import { Request, RequestHandler, Router } from "express";
import { createProxyMiddleware } from "http-proxy-middleware";
import { v4 } from "uuid";
import { config } from "../config";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  createOnProxyReqHandler,
  createPreprocessorMiddleware,
  finalizeSignedRequest,
} from "./middleware/request";
import {
  createOnProxyResHandler,
  ProxyResHandlerWithBody,
} from "./middleware/response";
import { addGoogleAIKey } from "./middleware/request/preprocessors/add-google-ai-key";

let modelsCache: any = null;
let modelsCacheTime = 0;

// https://ai.google.dev/models/gemini
// TODO: list models https://ai.google.dev/tutorials/rest_quickstart#list_models

const getModelsResponse = () => {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60) {
    return modelsCache;
  }

  if (!config.googleAIKey) return { object: "list", data: [] };

  const googleAIVariants = ["gemini-pro", "gemini-1.0-pro", "gemini-1.5-pro"];

  const models = googleAIVariants.map((id) => ({
    id,
    object: "model",
    created: new Date().getTime(),
    owned_by: "google",
    permission: [],
    root: "google",
    parent: null,
  }));

  modelsCache = { object: "list", data: models };
  modelsCacheTime = new Date().getTime();

  return modelsCache;
};

const handleModelRequest: RequestHandler = (_req, res) => {
  res.status(200).json(getModelsResponse());
};

/** Only used for non-streaming requests. */
const googleAIResponseHandler: ProxyResHandlerWithBody = async (
  _proxyRes,
  req,
  res,
  body
) => {
  if (typeof body !== "object") {
    throw new Error("Expected body to be an object");
  }

  if (config.promptLogging) {
    const host = req.get("host");
    body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
  }

  if (req.inboundApi === "openai") {
    req.log.info("Transforming Google AI response to OpenAI format");
    body = transformGoogleAIResponse(body, req);
  }

  if (req.tokenizerInfo) {
    body.proxy_tokenizer = req.tokenizerInfo;
  }

  res.status(200).json(body);
};

function transformGoogleAIResponse(
  resBody: Record<string, any>,
  req: Request
): Record<string, any> {
  const totalTokens = (req.promptTokens ?? 0) + (req.outputTokens ?? 0);
  const parts = resBody.candidates[0].content?.parts ?? [{ text: "" }];
  const content = parts[0].text.replace(/^(.{0,50}?): /, () => "");
  return {
    id: "goo-" + v4(),
    object: "chat.completion",
    created: Date.now(),
    model: req.body.model,
    usage: {
      prompt_tokens: req.promptTokens,
      completion_tokens: req.outputTokens,
      total_tokens: totalTokens,
    },
    choices: [
      {
        message: { role: "assistant", content },
        finish_reason: resBody.candidates[0].finishReason,
        index: 0,
      },
    ],
  };
}

const googleAIProxy = createQueueMiddleware({
  beforeProxy: addGoogleAIKey,
  proxyMiddleware: createProxyMiddleware({
    target: "bad-target-will-be-rewritten",
    router: ({ signedRequest }) => {
      const { protocol, hostname, path } = signedRequest;
      return `${protocol}//${hostname}${path}`;
    },
    changeOrigin: true,
    selfHandleResponse: true,
    logger,
    on: {
      proxyReq: createOnProxyReqHandler({ pipeline: [finalizeSignedRequest] }),
      proxyRes: createOnProxyResHandler([googleAIResponseHandler]),
      error: handleProxyError,
    },
  }),
});

const googleAIRouter = Router();
googleAIRouter.get("/v1/models", handleModelRequest);
// OpenAI-to-Google AI compatibility endpoint.
googleAIRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai",
    outApi: "google-ai",
    service: "google-ai",
  }),
  googleAIProxy
);

export const googleAI = googleAIRouter;


code:
import { Request, Response } from "express";
import httpProxy from "http-proxy";
import { ZodError } from "zod";
import { generateErrorMessage } from "zod-error";
import { assertNever } from "../../shared/utils";
import { QuotaExceededError } from "./request/preprocessors/apply-quota-limits";
import { sendErrorToClient } from "./response/error-generator";
import { HttpError } from "../../shared/errors";

const OPENAI_CHAT_COMPLETION_ENDPOINT = "/v1/chat/completions";
const OPENAI_TEXT_COMPLETION_ENDPOINT = "/v1/completions";
const OPENAI_EMBEDDINGS_ENDPOINT = "/v1/embeddings";
const OPENAI_IMAGE_COMPLETION_ENDPOINT = "/v1/images/generations";
const ANTHROPIC_COMPLETION_ENDPOINT = "/v1/complete";
const ANTHROPIC_MESSAGES_ENDPOINT = "/v1/messages";
const ANTHROPIC_SONNET_COMPAT_ENDPOINT = "/v1/sonnet";
const ANTHROPIC_OPUS_COMPAT_ENDPOINT = "/v1/opus";

export function isTextGenerationRequest(req: Request) {
  return (
    req.method === "POST" &&
    [
      OPENAI_CHAT_COMPLETION_ENDPOINT,
      OPENAI_TEXT_COMPLETION_ENDPOINT,
      ANTHROPIC_COMPLETION_ENDPOINT,
      ANTHROPIC_MESSAGES_ENDPOINT,
      ANTHROPIC_SONNET_COMPAT_ENDPOINT,
      ANTHROPIC_OPUS_COMPAT_ENDPOINT,
    ].some((endpoint) => req.path.startsWith(endpoint))
  );
}

export function isImageGenerationRequest(req: Request) {
  return (
    req.method === "POST" &&
    req.path.startsWith(OPENAI_IMAGE_COMPLETION_ENDPOINT)
  );
}

export function isEmbeddingsRequest(req: Request) {
  return (
    req.method === "POST" && req.path.startsWith(OPENAI_EMBEDDINGS_ENDPOINT)
  );
}

export function sendProxyError(
  req: Request,
  res: Response,
  statusCode: number,
  statusMessage: string,
  errorPayload: Record<string, any>
) {
  const msg =
    statusCode === 500
      ? `The proxy encountered an error while trying to process your prompt.`
      : `The proxy encountered an error while trying to send your prompt to the upstream service.`;

  if (req.tokenizerInfo && typeof errorPayload.error === "object") {
    errorPayload.error.proxy_tokenizer = req.tokenizerInfo;
  }

  sendErrorToClient({
    options: {
      format: req.inboundApi,
      title: `Proxy error (HTTP ${statusCode} ${statusMessage})`,
      message: `${msg} Further technical details are provided below.`,
      obj: errorPayload,
      reqId: req.id,
      model: req.body?.model,
    },
    req,
    res,
  });
}

export const handleProxyError: httpProxy.ErrorCallback = (err, req, res) => {
  req.log.error(err, `Error during http-proxy-middleware request`);
  classifyErrorAndSend(err, req as Request, res as Response);
};

export const classifyErrorAndSend = (
  err: Error,
  req: Request,
  res: Response
) => {
  try {
    const { statusCode, statusMessage, userMessage, ...errorDetails } =
      classifyError(err);
    sendProxyError(req, res, statusCode, statusMessage, {
      error: { message: userMessage, ...errorDetails },
    });
  } catch (error) {
    req.log.error(error, `Error writing error response headers, giving up.`);
    res.end();
  }
};

function classifyError(err: Error): {
  /** HTTP status code returned to the client. */
  statusCode: number;
  /** HTTP status message returned to the client. */
  statusMessage: string;
  /** Message displayed to the user. */
  userMessage: string;
  /** Short error type, e.g. "proxy_validation_error". */
  type: string;
} & Record<string, any> {
  const defaultError = {
    statusCode: 500,
    statusMessage: "Internal Server Error",
    userMessage: `Reverse proxy error: ${err.message}`,
    type: "proxy_internal_error",
    stack: err.stack,
  };

  switch (err.constructor.name) {
    case "HttpError":
      if ((err as HttpError).status === 402) {
        return {
          statusCode: 402,
          statusMessage: "No Keys Available",
          userMessage: err.message,
          type: "proxy_no_keys_available",
        };
      } else return defaultError;
    case "ZodError":
      const userMessage = generateErrorMessage((err as ZodError).issues, {
        prefix: "Request validation failed. ",
        path: { enabled: true, label: null, type: "breadcrumbs" },
        code: { enabled: false },
        maxErrors: 3,
        transform: ({ issue, ...rest }) => {
          return `At '${rest.pathComponent}': ${issue.message}`;
        },
      });
      return {
        statusCode: 400,
        statusMessage: "Bad Request",
        userMessage,
        type: "proxy_validation_error",
      };
    case "ZoomerForbiddenError":
      // Mimics a ban notice from OpenAI, thrown when blockZoomerOrigins blocks
      // a request.
      return {
        statusCode: 403,
        statusMessage: "Forbidden",
        userMessage: `Your account has been disabled for violating our terms of service.`,
        type: "organization_account_disabled",
        code: "policy_violation",
      };
    case "ForbiddenError":
      return {
        statusCode: 403,
        statusMessage: "Forbidden",
        userMessage: `Request is not allowed. (${err.message})`,
        type: "proxy_forbidden",
      };
    case "QuotaExceededError":
      return {
        statusCode: 429,
        statusMessage: "Too Many Requests",
        userMessage: `You've exceeded your token quota for this model type.`,
        type: "proxy_quota_exceeded",
        info: (err as QuotaExceededError).quotaInfo,
      };
    case "Error":
      if ("code" in err) {
        switch (err.code) {
          case "ENOTFOUND":
            return {
              statusCode: 502,
              statusMessage: "Bad Gateway",
              userMessage: `Reverse proxy encountered a DNS error while trying to connect to the upstream service.`,
              type: "proxy_network_error",
              code: err.code,
            };
          case "ECONNREFUSED":
            return {
              statusCode: 502,
              statusMessage: "Bad Gateway",
              userMessage: `Reverse proxy couldn't connect to the upstream service.`,
              type: "proxy_network_error",
              code: err.code,
            };
          case "ECONNRESET":
            return {
              statusCode: 504,
              statusMessage: "Gateway Timeout",
              userMessage: `Reverse proxy timed out while waiting for the upstream service to respond.`,
              type: "proxy_network_error",
              code: err.code,
            };
        }
      }
      return defaultError;
    default:
      return defaultError;
  }
}

export function getCompletionFromBody(req: Request, body: Record<string, any>) {
  const format = req.outboundApi;
  switch (format) {
    case "openai":
    case "mistral-ai":
      // Can be null if the model wants to invoke tools rather than return a
      // completion.
      return body.choices[0].message.content || "";
    case "openai-text":
      return body.choices[0].text;
    case "anthropic-chat":
      if (!body.content) {
        req.log.error(
          { body: JSON.stringify(body) },
          "Received empty Anthropic chat completion"
        );
        return "";
      }
      return body.content
        .map(({ text, type }: { type: string; text: string }) =>
          type === "text" ? text : `[Unsupported content type: ${type}]`
        )
        .join("\n");
    case "anthropic-text":
      if (!body.completion) {
        req.log.error(
          { body: JSON.stringify(body) },
          "Received empty Anthropic text completion"
        );
        return "";
      }
      return body.completion.trim();
    case "google-ai":
      if ("choices" in body) {
        return body.choices[0].message.content;
      }
      return body.candidates[0].content.parts[0].text;
    case "openai-image":
      return body.data?.map((item: any) => item.url).join("\n");
    default:
      assertNever(format);
  }
}

export function getModelFromBody(req: Request, body: Record<string, any>) {
  const format = req.outboundApi;
  switch (format) {
    case "openai":
    case "openai-text":
    case "mistral-ai":
      return body.model;
    case "openai-image":
      return req.body.model;
    case "anthropic-chat":
    case "anthropic-text":
      // Anthropic confirms the model in the response, but AWS Claude doesn't.
      return body.model || req.body.model;
    case "google-ai":
      // Google doesn't confirm the model in the response.
      return req.body.model;
    default:
      assertNever(format);
  }
}


code:
import type { Request } from "express";
import type { ClientRequest } from "http";
import type { ProxyReqCallback } from "http-proxy";

export { createOnProxyReqHandler } from "./onproxyreq-factory";
export {
  createPreprocessorMiddleware,
  createEmbeddingsPreprocessorMiddleware,
} from "./preprocessor-factory";

// Express middleware (runs before http-proxy-middleware, can be async)
export { addAzureKey } from "./preprocessors/add-azure-key";
export { applyQuotaLimits } from "./preprocessors/apply-quota-limits";
export { validateContextSize } from "./preprocessors/validate-context-size";
export { countPromptTokens } from "./preprocessors/count-prompt-tokens";
export { languageFilter } from "./preprocessors/language-filter";
export { setApiFormat } from "./preprocessors/set-api-format";
export { signAwsRequest } from "./preprocessors/sign-aws-request";
export { transformOutboundPayload } from "./preprocessors/transform-outbound-payload";

// http-proxy-middleware callbacks (runs on onProxyReq, cannot be async)
export { addKey, addKeyForEmbeddingsRequest } from "./onproxyreq/add-key";
export { addAnthropicPreamble } from "./onproxyreq/add-anthropic-preamble";
export { blockZoomerOrigins } from "./onproxyreq/block-zoomer-origins";
export { checkModelFamily } from "./onproxyreq/check-model-family";
export { finalizeBody } from "./onproxyreq/finalize-body";
export { finalizeSignedRequest } from "./onproxyreq/finalize-signed-request";
export { stripHeaders } from "./onproxyreq/strip-headers";

/**
 * Middleware that runs prior to the request being handled by http-proxy-
 * middleware.
 *
 * Async functions can be used here, but you will not have access to the proxied
 * request/response objects, nor the data set by ProxyRequestMiddleware
 * functions as they have not yet been run.
 *
 * User will have been authenticated by the time this middleware runs, but your
 * request won't have been assigned an API key yet.
 *
 * Note that these functions only run once ever per request, even if the request
 * is automatically retried by the request queue middleware.
 */
export type RequestPreprocessor = (req: Request) => void | Promise<void>;

/**
 * Callbacks that run immediately before the request is sent to the API in
 * response to http-proxy-middleware's `proxyReq` event.
 *
 * Async functions cannot be used here as HPM's event emitter is not async and
 * will not wait for the promise to resolve before sending the request.
 *
 * Note that these functions may be run multiple times per request if the
 * first attempt is rate limited and the request is automatically retried by the
 * request queue middleware.
 */
export type HPMRequestCallback = ProxyReqCallback<ClientRequest, Request>;

export const forceModel = (model: string) => (req: Request) =>
  void (req.body.model = model);


code:
import {
  applyQuotaLimits,
  blockZoomerOrigins,
  checkModelFamily,
  HPMRequestCallback,
  stripHeaders,
} from "./index";

type ProxyReqHandlerFactoryOptions = { pipeline: HPMRequestCallback[] };

/**
 * Returns an http-proxy-middleware request handler that runs the given set of
 * onProxyReq callback functions in sequence.
 *
 * These will run each time a request is proxied, including on automatic retries
 * by the queue after encountering a rate limit.
 */
export const createOnProxyReqHandler = ({
  pipeline,
}: ProxyReqHandlerFactoryOptions): HPMRequestCallback => {
  const callbackPipeline = [
    checkModelFamily,
    applyQuotaLimits,
    blockZoomerOrigins,
    stripHeaders,
    ...pipeline,
  ];
  return (proxyReq, req, res, options) => {
    // The streaming flag must be set before any other onProxyReq handler runs,
    // as it may influence the behavior of subsequent handlers.
    // Image generation requests can't be streamed.
    // TODO: this flag is set in too many places
    req.isStreaming =
      req.isStreaming || req.body.stream === true || req.body.stream === "true";
    req.body.stream = req.isStreaming;

    try {
      for (const fn of callbackPipeline) {
        fn(proxyReq, req, res, options);
      }
    } catch (error) {
      proxyReq.destroy(error);
    }
  };
};


code:
import { AnthropicKey, Key } from "../../../../shared/key-management";
import { isTextGenerationRequest } from "../../common";
import { HPMRequestCallback } from "../index";

/**
 * Some keys require the prompt to start with `\n\nHuman:`. There is no way to
 * know this without trying to send the request and seeing if it fails. If a
 * key is marked as requiring a preamble, it will be added here.
 */
export const addAnthropicPreamble: HPMRequestCallback = (_proxyReq, req) => {
  if (
    !isTextGenerationRequest(req) ||
    req.key?.service !== "anthropic" ||
    req.outboundApi !== "anthropic-text"
  ) {
    return;
  }

  let preamble = "";
  let prompt = req.body.prompt;
  assertAnthropicKey(req.key);
  if (req.key.requiresPreamble && prompt) {
    preamble = prompt.startsWith("\n\nHuman:") ? "" : "\n\nHuman:";
    req.log.debug({ key: req.key.hash, preamble }, "Adding preamble to prompt");
  }
  req.body.prompt = preamble + prompt;
};

function assertAnthropicKey(key: Key): asserts key is AnthropicKey {
  if (key.service !== "anthropic") {
    throw new Error(`Expected an Anthropic key, got '${key.service}'`);
  }
}


code:
import { Key, OpenAIKey, keyPool } from "../../../../shared/key-management";
import { isEmbeddingsRequest } from "../../common";
import { HPMRequestCallback } from "../index";
import { assertNever } from "../../../../shared/utils";

/** Add a key that can service this request to the request object. */
export const addKey: HPMRequestCallback = (proxyReq, req) => {
  let assignedKey: Key;

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

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

  if (req.inboundApi === req.outboundApi) {
    assignedKey = keyPool.get(req.body.model, req.service);
  } else {
    switch (req.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.
      case "anthropic-chat":
      case "anthropic-text":
        assignedKey = keyPool.get("claude-v1", req.service);
        break;
      case "openai-text":
        assignedKey = keyPool.get("gpt-3.5-turbo-instruct", req.service);
        break;
      case "openai":
        throw new Error(
          "OpenAI Chat as an API translation target is not supported"
        );
      case "google-ai":
        throw new Error("add-key should not be used for this model.");
      case "mistral-ai":
        throw new Error("Mistral AI should never be translated");
      case "openai-image":
        assignedKey = keyPool.get("dall-e-3", req.service);
        break;
      default:
        assertNever(req.outboundApi);
    }
  }

  req.key = assignedKey;
  req.log.info(
    {
      key: assignedKey.hash,
      model: req.body?.model,
      fromApi: req.inboundApi,
      toApi: req.outboundApi,
    },
    "Assigned key to request"
  );

  // TODO: KeyProvider should assemble all necessary headers
  switch (assignedKey.service) {
    case "anthropic":
      proxyReq.setHeader("X-API-Key", assignedKey.key);
      break;
    case "openai":
      const key: OpenAIKey = assignedKey as OpenAIKey;
      if (key.organizationId) {
        proxyReq.setHeader("OpenAI-Organization", key.organizationId);
      }
      proxyReq.setHeader("Authorization", `Bearer ${assignedKey.key}`);
      break;
    case "mistral-ai":
      proxyReq.setHeader("Authorization", `Bearer ${assignedKey.key}`);
      break;
    case "azure":
      const azureKey = assignedKey.key;
      proxyReq.setHeader("api-key", azureKey);
      break;
    case "aws":
    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: HPMRequestCallback = (
  proxyReq,
  req
) => {
  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");
  }

  req.body = { input: req.body.input, model: "text-embedding-ada-002" };

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

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

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


code:
import { HPMRequestCallback } 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: HPMRequestCallback = (_proxyReq, 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 { HPMRequestCallback } from "../index";
import { config } from "../../../../config";
import { ForbiddenError } from "../../../../shared/errors";
import { getModelFamilyForRequest } from "../../../../shared/models";

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


code:
import { fixRequestBody } from "http-proxy-middleware";
import type { HPMRequestCallback } from "../index";

/** Finalize the rewritten request body. Must be the last rewriter. */
export const finalizeBody: HPMRequestCallback = (proxyReq, req) => {
  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 updatedBody = JSON.stringify(req.body);
    proxyReq.setHeader("Content-Length", Buffer.byteLength(updatedBody));
    (req as any).rawBody = Buffer.from(updatedBody);

    // body-parser and http-proxy-middleware don't play nice together
    fixRequestBody(proxyReq, req);
  }
};


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

/**
 * For AWS/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: HPMRequestCallback = (proxyReq, req) => {
  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.
  proxyReq.path = 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.
  proxyReq.getRawHeaderNames().forEach(proxyReq.removeHeader.bind(proxyReq));
  Object.entries(req.signedRequest.headers).forEach(([key, value]) => {
    proxyReq.setHeader(key, value);
  });

  // Don't use fixRequestBody here because it adds a content-length header.
  // Amazon doesn't want that and it breaks the signature.
  proxyReq.write(req.signedRequest.body);
};


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

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

  proxyReq.removeHeader("cf-connecting-ip");
  proxyReq.removeHeader("forwarded");
  proxyReq.removeHeader("true-client-ip");
  proxyReq.removeHeader("x-forwarded-for");
  proxyReq.removeHeader("x-real-ip");
};


code:
import { RequestHandler } from "express";
import { ZodIssue } from "zod";
import { initializeSseStream } from "../../../shared/streaming";
import { classifyErrorAndSend } from "../common";
import {
  RequestPreprocessor,
  validateContextSize,
  countPromptTokens,
  setApiFormat,
  transformOutboundPayload,
  languageFilter,
} 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 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 a preprocessor on every re-attempt, pass it to createQueueMiddleware.
 * It will run after these preprocessors, but before the request is sent to
 * http-proxy-middleware.
 */
export const createPreprocessorMiddleware = (
  apiFormat: Parameters<typeof setApiFormat>[0],
  { beforeTransform, afterTransform }: RequestPreprocessorOptions = {}
): RequestHandler => {
  const preprocessors: RequestPreprocessor[] = [
    setApiFormat(apiFormat),
    ...(beforeTransform ?? []),
    transformOutboundPayload,
    countPromptTokens,
    languageFilter,
    ...(afterTransform ?? []),
    validateContextSize,
  ];
  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 msg = error?.issues
        ?.map((issue: ZodIssue) => issue.message)
        .join("; ");
      req.log.info(msg, "Prompt validation failed.");
    } 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!" }],
      proxy_note:
        "This response was generated by the proxy's test message handler and did not go to the API.",
    });
  }
};

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

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


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

export const addAzureKey: RequestPreprocessor = (req) => {
  const apisValid = req.inboundApi === "openai" && req.outboundApi === "openai";
  const serviceValid = req.service === "azure";
  if (!apisValid || !serviceValid) {
    throw new Error("addAzureKey called on invalid request");
  }

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

  const model = req.body.model.startsWith("azure-")
    ? req.body.model
    : `azure-${req.body.model}`;

  req.key = keyPool.get(model, "azure");
  req.body.model = model;

  // Handles the sole Azure API deviation from the OpenAI spec (that I know of)
  const notNullOrUndefined = (x: any) => x !== null && x !== undefined;
  if ([req.body.logprobs, req.body.top_logprobs].some(notNullOrUndefined)) {
    // OpenAI wants logprobs: true/false and top_logprobs: number
    // Azure seems to just want to combine them into logprobs: number
    // if (typeof req.body.logprobs === "boolean") {
    //   req.body.logprobs = req.body.top_logprobs || undefined;
    //   delete req.body.top_logprobs
    // }

    // Temporarily just disabling logprobs for Azure because their model support
    // is random: `This model does not support the 'logprobs' parameter.`
    delete req.body.logprobs;
    delete req.body.top_logprobs;
  }

  req.log.info(
    { key: req.key.hash, model },
    "Assigned Azure OpenAI key to request"
  );

  const cred = req.key as AzureOpenAIKey;
  const { resourceName, deploymentId, apiKey } = getCredentialsFromKey(cred);

  req.signedRequest = {
    method: "POST",
    protocol: "https:",
    hostname: `${resourceName}.openai.azure.com`,
    path: `/openai/deployments/${deploymentId}/chat/completions?api-version=2023-09-01-preview`,
    headers: {
      ["host"]: `${resourceName}.openai.azure.com`,
      ["content-type"]: "application/json",
      ["api-key"]: apiKey,
    },
    body: JSON.stringify(req.body),
  };
};

function getCredentialsFromKey(key: AzureOpenAIKey) {
  const [resourceName, deploymentId, apiKey] = key.key.split(":");
  if (!resourceName || !deploymentId || !apiKey) {
    throw new Error("Assigned Azure OpenAI key is not in the correct format.");
  }
  return { resourceName, deploymentId, apiKey };
}


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

export const addGoogleAIKey: RequestPreprocessor = (req) => {
  const apisValid = req.inboundApi === "openai" && req.outboundApi === "google-ai";
  const serviceValid = req.service === "google-ai";
  if (!apisValid || !serviceValid) {
    throw new Error("addGoogleAIKey called on invalid request");
  }

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

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

  req.log.info(
    { key: req.key.hash, model },
    "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}

  req.isStreaming = req.isStreaming || req.body.stream;
  delete req.body.stream;

  req.signedRequest = {
    method: "POST",
    protocol: "https:",
    hostname: "generativelanguage.googleapis.com",
    path: `/v1beta/models/${model}:${req.isStreaming ? "streamGenerateContent" : "generateContent"}?key=${req.key.key}`,
    headers: {
      ["host"]: `generativelanguage.googleapis.com`,
      ["content-type"]: "application/json",
    },
    body: JSON.stringify(req.body),
  };
};


code:
import { hasAvailableQuota } from "../../../../shared/users/user-store";
import { isImageGenerationRequest, isTextGenerationRequest } from "../../common";
import { HPMRequestCallback } 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: HPMRequestCallback = (_proxyReq, 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";
import { countTokens } from "../../../../shared/tokenization";
import { assertNever } from "../../../../shared/utils";
import {
  AnthropicChatMessage,
  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_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;
      const prompt: AnthropicChatMessage[] = 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": {
      req.outputTokens = req.body.max_tokens;
      const prompt: MistralAIChatMessage[] = req.body.messages;
      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 { 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";

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

function getPromptFromRequest(req: Request) {
  const service = req.outboundApi;
  const body = req.body;
  switch (service) {
    case "anthropic-chat":
      return flattenAnthropicMessages(body.messages);
    case "anthropic-text":
      return body.prompt;
    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 "openai-text":
    case "openai-image":
      return body.prompt;
    case "google-ai":
      return body.prompt.text;
    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: {
  inApi: Request["inboundApi"];
  outApi: APIFormat;
  service: LLMService;
}): RequestPreprocessor => {
  return function configureRequestApiFormat(req) {
    req.inboundApi = api.inApi;
    req.outboundApi = api.outApi;
    req.service = api.service;
  };
};


code:
import express 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/anthropic";
import { keyPool } from "../../../../shared/key-management";
import { RequestPreprocessor } 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: RequestPreprocessor = async (req) => {
  const { model, stream } = req.body;
  req.key = keyPool.get(model, "aws");

  req.isStreaming = stream === true || stream === "true";

  // same as addAnthropicPreamble for non-AWS requests, but has to happen here
  if (req.outboundApi === "anthropic-text") {
    let preamble = req.body.prompt.startsWith("\n\nHuman:") ? "" : "\n\nHuman:";
    req.body.prompt = preamble + req.body.prompt;
  }

  // AWS uses mostly the same parameters as Anthropic, with a few removed params
  // and much stricter validation on unused parameters. Rather than treating it
  // as a separate schema we will use the anthropic ones and strip the unused
  // parameters.
  // TODO: This should happen in transform-outbound-payload.ts
  let strippedParams: Record<string, unknown>;
  if (req.outboundApi === "anthropic-chat") {
    strippedParams = AnthropicV1MessagesSchema.pick({
      messages: 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";
  } else {
    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);
  }

  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.
  req.headers["anthropic-version"] = "2023-06-01";

  // 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/${model}/invoke${stream ? "-with-response-stream" : ""}`,
    headers: {
      ["Host"]: host,
      ["content-type"]: "application/json",
    },
    body: JSON.stringify(strippedParams),
  });

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

  req.signedRequest = 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);
}


code:
import {
  anthropicTextToAnthropicChat,
  openAIToAnthropicText,
} from "../../../../shared/api-schemas/anthropic";
import { openAIToOpenAIText } from "../../../../shared/api-schemas/openai-text";
import { openAIToOpenAIImage } from "../../../../shared/api-schemas/openai-image";
import { openAIToGoogleAI } from "../../../../shared/api-schemas/google-ai";
import { fixMistralPrompt } from "../../../../shared/api-schemas/mistral-ai";
import { API_SCHEMA_VALIDATORS } from "../../../../shared/api-schemas";
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 sameService = req.inboundApi === req.outboundApi;
  const alreadyTransformed = req.retryCount > 0;
  const notTransformable =
    !isTextGenerationRequest(req) && !isImageGenerationRequest(req);

  if (alreadyTransformed || notTransformable) return;

  if (req.inboundApi === "mistral-ai") {
    const messages = req.body.messages;
    req.body.messages = fixMistralPrompt(messages);
    req.log.info(
      { old: messages.length, new: req.body.messages.length },
      "Fixed Mistral prompt"
    );
  }

  if (sameService) {
    const result = API_SCHEMA_VALIDATORS[req.inboundApi].safeParse(req.body);
    if (!result.success) {
      req.log.error(
        { issues: result.error.issues, body: req.body },
        "Request validation failed"
      );
      throw result.error;
    }
    req.body = result.data;
    return;
  }

  if (
    req.inboundApi === "anthropic-text" &&
    req.outboundApi === "anthropic-chat"
  ) {
    req.body = anthropicTextToAnthropicChat(req);
    return;
  }

  if (req.inboundApi === "openai" && req.outboundApi === "anthropic-text") {
    req.body = openAIToAnthropicText(req);
    return;
  }

  if (req.inboundApi === "openai" && req.outboundApi === "google-ai") {
    req.body = openAIToGoogleAI(req);
    return;
  }

  if (req.inboundApi === "openai" && req.outboundApi === "openai-text") {
    req.body = openAIToOpenAIText(req);
    return;
  }

  if (req.inboundApi === "openai" && req.outboundApi === "openai-image") {
    req.body = openAIToOpenAIImage(req);
    return;
  }

  throw new Error(
    `'${req.inboundApi}' -> '${req.outboundApi}' request proxying is not supported. Make sure your client is configured to use the correct API.`
  );
};


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;
const GOOGLE_AI_MAX_CONTEXT = 32000;
const MISTRAL_AI_MAX_CONTENT = 32768;

/**
 * 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":
      proxyMax = MISTRAL_AI_MAX_CONTENT;
      break;
    case "openai-image":
      return;
    default:
      assertNever(req.outboundApi);
  }
  proxyMax ||= Number.MAX_SAFE_INTEGER;

  let modelMax: number;
  if (model.match(/gpt-3.5-turbo-16k/)) {
    modelMax = 16384;
  } 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(/gpt-3.5-turbo/)) {
    modelMax = 4096;
  } 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-\d{3}$/)) {
    modelMax = GOOGLE_AI_MAX_CONTEXT;
  } else if (model.match(/^mistral-(tiny|small|medium)$/)) {
    modelMax = MISTRAL_AI_MAX_CONTENT;
  } else if (model.match(/^anthropic\.claude-3-sonnet/)) {
    modelMax = 200000;
  } else if (model.match(/^anthropic\.claude-v2:\d/)) {
    modelMax = 200000;
  } else if (model.match(/^anthropic\.claude/)) {
    // Not sure if AWS Claude has the same context limit as Anthropic Claude.
    modelMax = 100000;
  } 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 express from "express";
import { APIFormat } from "../../../shared/key-management";
import { assertNever } from "../../../shared/utils";
import { initializeSseStream } from "../../../shared/streaming";

function getMessageContent({
  title,
  message,
  obj,
}: {
  title: string;
  message: string;
  obj?: Record<string, any>;
}) {
  /*
  Constructs 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_tokenizer": {
       "tokenizer": "@anthropic-ai/tokenizer",
       "token_count": 6104,
       "tokenization_duration_ms": 4.0765,
       "prompt_tokens": 6104,
       "completion_tokens": 30,
       "max_model_tokens": 200000,
       "max_proxy_tokens": 9007199254740991
      }
    },
   "proxy_note": "The requested Claude model might not exist, or the key might not be provisioned for it."
  }
  ```
   */
  const note = obj?.proxy_note || obj?.error?.message || "";
  const friendlyMessage = note ? `${message}\n\n***\n\n*${note}*` : message;
  const details = JSON.parse(JSON.stringify(obj ?? {}));
  let stack = "";
  if (details.stack) {
    stack = `\n\nInclude this trace when reporting an issue.\n\`\`\`\n${details.stack}\n\`\`\``;
    delete details.stack;
  }
  return `\n\n**${title}**\n${friendlyMessage}${
    obj ? `\n\`\`\`\n${JSON.stringify(obj, null, 2)}\n\`\`\`\n${stack}` : ""
  }`;
}

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

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

export function sendErrorToClient({
  options,
  req,
  res,
}: {
  options: ErrorGeneratorOptions;
  req: express.Request;
  res: express.Response;
}) {
  const { format: inputFormat } = options;

  // This is an error thrown before we know the format of the request, so we
  // can't send a response in the format the client expects.
  const format =
    inputFormat === "unknown" ? tryInferFormat(req.body) : inputFormat;
  if (format === "unknown") {
    return res.status(options.statusCode || 400).json({
      error: options.message,
      details: options.obj,
    });
  }

  const completion = buildSpoofedCompletion({ ...options, format });
  const event = buildSpoofedSSE({ ...options, format });
  const isStreaming =
    req.isStreaming || req.body.stream === true || req.body.stream === "true";

  if (isStreaming) {
    if (!res.headersSent) {
      initializeSseStream(res);
    }
    res.write(event);
    res.write(`data: [DONE]\n\n`);
    res.end();
  } else {
    res.status(200).json(completion);
  }
}

/**
 * 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 "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":
      // TODO: Native Google AI non-streaming responses are not supported, this
      // is an untested guess at what the response should look like.
      return {
        id: "error-" + id,
        object: "chat.completion",
        created: Date.now(),
        model,
        candidates: [
          {
            content: { parts: [{ text: content }], role: "model" },
            finishReason: title,
            index: 0,
            tokenCount: null,
            safetyRatings: [],
          },
        ],
      };
    case "openai-image":
      throw new Error(
        `Spoofed completions not supported for ${format} requests`
      );
    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 "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":
      return JSON.stringify({
        candidates: [
          {
            content: { parts: [{ text: content }], role: "model" },
            finishReason: title,
            index: 0,
            tokenCount: null,
            safetyRatings: [],
          },
        ],
      });
    case "openai-image":
      throw new Error(`SSE not supported for ${format} requests`);
    default:
      assertNever(format);
  }

  if (format === "anthropic-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 { pipeline, Readable, Transform } from "stream";
import StreamArray from "stream-json/streamers/StreamArray";
import { StringDecoder } from "string_decoder";
import { promisify } from "util";
import { APIFormat, keyPool } from "../../../shared/key-management";
import {
  copySseResponseHeaders,
  initializeSseStream,

} from "../../../shared/streaming";
import type { logger } from "../../../logger";
import { enqueue } from "../../queue";
import { decodeResponseBody, RawResponseBodyHandler, RetryableError } from ".";
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 { buildSpoofedSSE } from "./error-generator";

const pipelineAsync = promisify(pipeline);

/**
 * `handleStreamedResponse` consumes and transforms a streamed response from the
 * upstream service, forwarding events to the client in their requested format.
 * 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.
 *
 * In the event of an error, the request's streaming flag is unset and the non-
 * streaming response handler is called instead.
 *
 * 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 { hash } = req.key!;
  if (!req.isStreaming) {
    throw new Error("handleStreamedResponse called for non-streaming request.");
  }

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

  req.log.debug(
    { headers: proxyRes.headers, key: hash },
    `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: proxyRes.headers["content-type"],
    api: req.outboundApi,
    logger: req.log,
  };

  // Decoder turns the raw response stream into a stream of events in some
  // format (text/event-stream, vnd.amazon.event-stream, streaming JSON, etc).
  const decoder = getDecoder({ ...streamOptions, input: proxyRes });
  // Adapter transforms the decoded events into server-sent events.
  const adapter = new SSEStreamAdapter(streamOptions);
  // Aggregator compiles all events into a single response object.
  const aggregator = new EventAggregator({ format: req.outboundApi });
  // Transformer converts server-sent events from one vendor's API message
  // format to another.
  const transformer = new SSEMessageTransformer({
    inputFormat: req.outboundApi,
    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 pipelineAsync(proxyRes, decoder, adapter, transformer);
    req.log.debug({ key: hash }, `Finished proxying SSE stream.`);
    res.end();
    return aggregator.getFinalResponse();
  } catch (err) {
    if (err instanceof RetryableError) {
      keyPool.markRateLimited(req.key!);
      req.log.warn(
        { key: req.key!.hash, retryCount: req.retryCount },
        `Re-enqueueing request due to retryable error during streaming response.`
      );
      req.retryCount++;
      await enqueue(req);
    } 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();
    }
    throw err;
  }
};

function getDecoder(options: {
  input: Readable;
  api: APIFormat;
  logger: typeof logger;
  contentType?: string;
}) {
  const { api, contentType, input, logger } = options;
  if (contentType?.includes("application/vnd.amazon.eventstream")) {
    return getAwsEventStreamDecoder({ input, logger });
  } else if (api === "google-ai") {
    return StreamArray.withParser();
  } else {
    // Passthrough stream, but 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();
      },
    });
  }
}