code:
import crypto from "crypto";
import dotenv from "dotenv";
import type firebase from "firebase-admin";
import path from "path";
import pino from "pino";
import type { LLMService, ModelFamily } from "./shared/models";
import { MODEL_FAMILIES } from "./shared/models";

dotenv.config();

const startupLogger = pino({ level: "debug" }).child({ module: "startup" });
const isDev = process.env.NODE_ENV !== "production";

export const DATA_DIR = path.join(__dirname, "..", "data");
export const USER_ASSETS_DIR = path.join(DATA_DIR, "user-files");

type Config = {
  /** The port the proxy server will listen on. */
  port: number;
  /** The network interface the proxy server will listen on. */
  bindAddress: string;
  /** Comma-delimited list of OpenAI API keys. */
  openaiKey?: string;
  /** Comma-delimited list of Anthropic API keys. */
  anthropicKey?: string;
  /**
   * Comma-delimited list of Google AI API keys. Note that these are not the
   * same as the GCP keys/credentials used for Vertex AI; the models are the
   * same but the APIs are different. Vertex is the GCP product for enterprise.
   **/
  googleAIKey?: string;
  /**
   * Comma-delimited list of Mistral AI API keys.
   */
  mistralAIKey?: string;
  /**
   * Comma-delimited list of AWS credentials. Each credential item should be a
   * colon-delimited list of access key, secret key, and AWS region.
   *
   * The credentials must have access to the actions `bedrock:InvokeModel` and
   * `bedrock:InvokeModelWithResponseStream`. You must also have already
   * provisioned the necessary models in your AWS account, on the specific
   * regions specified for each credential. Models are region-specific.
   *
   * @example `AWS_CREDENTIALS=access_key_1:secret_key_1:us-east-1,access_key_2:secret_key_2:us-west-2`
   */
  awsCredentials?: string;
  /**
   * Comma-delimited list of GCP credentials. Each credential item should be a
   * colon-delimited list of access key, secret key, and GCP region.
   *
   * @example `GCP_CREDENTIALS=project1:1@1.com:us-east5:-----BEGIN PRIVATE KEY-----xxx-----END PRIVATE KEY-----,project2:2@2.com:us-east5:-----BEGIN PRIVATE KEY-----xxx-----END PRIVATE KEY-----`
   */
  gcpCredentials?: string;
  /**
   * Comma-delimited list of Azure OpenAI credentials. Each credential item
   * should be a colon-delimited list of Azure resource name, deployment ID, and
   * API key.
   *
   * The resource name is the subdomain in your Azure OpenAI deployment's URL,
   * e.g. `https://resource-name.openai.azure.com
   *
   * @example `AZURE_CREDENTIALS=resource_name_1:deployment_id_1:api_key_1,resource_name_2:deployment_id_2:api_key_2`
   */
  azureCredentials?: string;
  /**
   * The proxy key to require for requests. Only applicable if the user
   * management mode is set to 'proxy_key', and required if so.
   */
  proxyKey?: string;
  /**
   * The admin key used to access the /admin API or UI. Required if the user
   * management mode is set to 'user_token'.
   */
  adminKey?: string;
  /**
   * The password required to view the service info/status page. If not set, the
   * info page will be publicly accessible.
   */
  serviceInfoPassword?: string;
  /**
   * Which user management mode to use.
   * - `none`: No user management. Proxy is open to all requests with basic
   *   abuse protection.
   * - `proxy_key`: A specific proxy key must be provided in the Authorization
   *   header to use the proxy.
   * - `user_token`: Users must be created via by admins and provide their
   *   personal access token in the Authorization header to use the proxy.
   *   Configure this function and add users via the admin API or UI.
   */
  gatekeeper: "none" | "proxy_key" | "user_token";
  /**
   * Persistence layer to use for user management.
   * - `memory`: Users are stored in memory and are lost on restart (default)
   * - `firebase_rtdb`: Users are stored in a Firebase Realtime Database;
   *   requires `firebaseKey` and `firebaseRtdbUrl` to be set.
   */
  gatekeeperStore: "memory" | "firebase_rtdb";
  /** URL of the Firebase Realtime Database if using the Firebase RTDB store. */
  firebaseRtdbUrl?: string;
  /**
   * Base64-encoded Firebase service account key if using the Firebase RTDB
   * store. Note that you should encode the *entire* JSON key file, not just the
   * `private_key` field inside it.
   */
  firebaseKey?: string;
  /**
   * Maximum number of IPs allowed per user token.
   * Users with the manually-assigned `special` role are exempt from this limit.
   * - Defaults to 0, which means that users are not IP-limited.
   */
  maxIpsPerUser: number;
  /**
   * Whether a user token should be automatically disabled if it exceeds the
   * `maxIpsPerUser` limit, or if only connections from new IPs are be rejected.
   */
  maxIpsAutoBan: boolean;
  /**
   * Which captcha verification mode to use. Requires `user_token` gatekeeper.
   * Allows users to automatically obtain a token by solving a captcha.
   * - `none`: No captcha verification; tokens are issued manually.
   * - `proof_of_work`: Users must solve an Argon2 proof of work to obtain a
   *    temporary usertoken valid for a limited period.
   */
  captchaMode: "none" | "proof_of_work";
  /**
   * Duration (in hours) for which a PoW-issued temporary user token is valid.
   */
  powTokenHours: number;
  /**
   * The maximum number of IPs from which a single temporary user token can be
   * used. Upon reaching the limit, the `maxIpsAutoBan` behavior is triggered.
   */
  powTokenMaxIps: number;
  /**
   * Difficulty level for the proof-of-work challenge.
   * - `low`: 200 iterations
   * - `medium`: 900 iterations
   * - `high`: 1900 iterations
   * - `extreme`: 4000 iterations
   * - `number`: A custom number of iterations to use.
   *
   * Difficulty level only affects the number of iterations used in the PoW,
   * not the complexity of the hash itself. Therefore, the average time-to-solve
   * will scale linearly with the number of iterations.
   *
   * Refer to docs/proof-of-work.md for guidance and hashrate benchmarks.
   */
  powDifficultyLevel: "low" | "medium" | "high" | "extreme" | number;
  /**
   * Duration (in minutes) before a PoW challenge expires. Users' browsers must
   * solve the challenge within this time frame or it will be rejected. Should
   * be kept somewhat low to prevent abusive clients from working on many
   * challenges in parallel, but you may need to increase this value for higher
   * difficulty levels or older devices will not be able to solve the challenge
   * in time.
   *
   * Defaults to 30 minutes.
   */
  powChallengeTimeout: number;
  /**
   * Duration (in hours) before expired temporary user tokens are purged from
   * the user database. Users can refresh expired tokens by solving a faster PoW
   * challenge as long as the original token has not been purged. Once purged,
   * the user must solve a full PoW challenge to obtain a new token.
   *
   * Defaults to 48 hours. At 0, tokens are purged immediately upon expiry.
   */
  powTokenPurgeHours: number;
  /**
   * Maximum number of active temporary user tokens that can be associated with
   * a single IP address. Note that this may impact users sending requests from
   * hosted AI chat clients such as Agnaistic or RisuAI, as they may share IPs.
   *
   * When the limit is reached, the oldest token with the same IP will be
   * expired. At 0, no limit is enforced. Defaults to 0.
   */
  // powMaxTokensPerIp: number;
  /** Per-user limit for requests per minute to text and chat models. */
  textModelRateLimit: number;
  /** Per-user limit for requests per minute to image generation models. */
  imageModelRateLimit: number;
  /**
   * For OpenAI, the maximum number of context tokens (prompt + max output) a
   * user can request before their request is rejected.
   * Context limits can help prevent excessive spend.
   * - Defaults to 0, which means no limit beyond OpenAI's stated maximums.
   */
  maxContextTokensOpenAI: number;
  /**
   * For Anthropic, the maximum number of context tokens a user can request.
   * Claude context limits can prevent requests from tying up concurrency slots
   * for too long, which can lengthen queue times for other users.
   * - Defaults to 0, which means no limit beyond Anthropic's stated maximums.
   */
  maxContextTokensAnthropic: number;
  /** For OpenAI, the maximum number of sampled tokens a user can request. */
  maxOutputTokensOpenAI: number;
  /** For Anthropic, the maximum number of sampled tokens a user can request. */
  maxOutputTokensAnthropic: number;
  /** Whether requests containing the following phrases should be rejected. */
  rejectPhrases: string[];
  /** Message to return when rejecting requests. */
  rejectMessage: string;
  /** Verbosity level of diagnostic logging. */
  logLevel: "trace" | "debug" | "info" | "warn" | "error";
  /**
   * Whether to allow the usage of AWS credentials which could be logging users'
   * model invocations. By default, such keys are treated as if they were
   * disabled because users may not be aware that their usage is being logged.
   *
   * Some credentials do not have the policy attached that allows the proxy to
   * confirm logging status, in which case the proxy assumes that logging could
   * be enabled and will refuse to use the key. If you still want to use such a
   * key and can't attach the policy, you can set this to true.
   */
  allowAwsLogging?: boolean;
  /**
   * Path to the SQLite database file for storing data such as event logs. By
   * default, the database will be stored at `data/database.sqlite`.
   *
   * Ensure target is writable by the server process, and be careful not to
   * select a path that is served publicly. The default path is safe.
   */
  sqliteDataPath?: string;
  /**
   * Whether to log events, such as generated completions, to the database.
   * Events are associated with IP+user token pairs. If user_token mode is
   * disabled, no events will be logged.
   *
   * Currently there is no pruning mechanism for the events table, so it will
   * grow indefinitely. You may want to periodically prune the table manually.
   */
  eventLogging?: boolean;
  /**
   * When hashing prompt histories, how many messages to trim from the end.
   * If zero, only the full prompt hash will be stored.
   * If greater than zero, for each number N, a hash of the prompt with the
   * last N messages removed will be stored.
   *
   * Experimental function, config may change in future versions.
   */
  eventLoggingTrim?: number;
  /** Whether prompts and responses should be logged to persistent storage. */
  promptLogging?: boolean;
  /** Which prompt logging backend to use. */
  promptLoggingBackend?: "google_sheets" | "file";
  /** Prefix for prompt logging files when using the file backend. */
  promptLoggingFilePrefix?: string;
  /** Base64-encoded Google Sheets API key. */
  googleSheetsKey?: string;
  /** Google Sheets spreadsheet ID. */
  googleSheetsSpreadsheetId?: string;
  /** Whether to periodically check keys for usage and validity. */
  checkKeys: boolean;
  /** Whether to publicly show total token costs on the info page. */
  showTokenCosts: boolean;
  /**
   * Comma-separated list of origins to block. Requests matching any of these
   * origins or referers will be rejected.
   * - Partial matches are allowed, so `reddit` will match `www.reddit.com`.
   * - Include only the hostname, not the protocol or path, e.g:
   *  `reddit.com,9gag.com,gaiaonline.com`
   */
  blockedOrigins?: string;
  /** Message to return when rejecting requests from blocked origins. */
  blockMessage?: string;
  /** Destination URL to redirect blocked requests to, for non-JSON requests. */
  blockRedirect?: string;
  /** Which model families to allow requests for. Applies only to OpenAI. */
  allowedModelFamilies: ModelFamily[];
  /**
   * The number of (LLM) tokens a user can consume before requests are rejected.
   * Limits include both prompt and response tokens. `special` users are exempt.
   * - Defaults to 0, which means no limit.
   * - Changes are not automatically applied to existing users. Use the
   * admin API or UI to update existing users, or use the QUOTA_REFRESH_PERIOD
   * setting to periodically set all users' quotas to these values.
   */
  tokenQuota: { [key in ModelFamily]: number };
  /**
   * The period over which to enforce token quotas. Quotas will be fully reset
   * at the start of each period, server time. Unused quota does not roll over.
   * You can also provide a cron expression for a custom schedule. If not set,
   * quotas will never automatically refresh.
   * - Defaults to unset, which means quotas will never automatically refresh.
   */
  quotaRefreshPeriod?: "hourly" | "daily" | string;
  /** Whether to allow users to change their own nicknames via the UI. */
  allowNicknameChanges: boolean;
  /** Whether to show recent DALL-E image generations on the homepage. */
  showRecentImages: boolean;
  /**
   * If true, cookies will be set without the `Secure` attribute, allowing
   * the admin UI to used over HTTP.
   */
  useInsecureCookies: boolean;
  /**
   * Whether to use a more minimal public Service Info page with static content.
   * Disables all stats pertaining to traffic, prompt/token usage, and queues.
   * The full info page will appear if you have signed in as an admin using the
   * configured ADMIN_KEY and go to /admin/service-info.
   **/
  staticServiceInfo?: boolean;
  /**
   * Trusted proxy hops. If you are deploying the server behind a reverse proxy
   * (Nginx, Cloudflare Tunnel, AWS WAF, etc.) the IP address of incoming
   * requests will be the IP address of the proxy, not the actual user.
   *
   * Depending on your hosting configuration, there may be multiple proxies/load
   * balancers between your server and the user. Each one will append the
   * incoming IP address to the `X-Forwarded-For` header. The user's real IP
   * address will be the first one in the list, assuming the header has not been
   * tampered with. Setting this value correctly ensures that the server doesn't
   * trust values in `X-Forwarded-For` not added by trusted proxies.
   *
   * In order for the server to determine the user's real IP address, you need
   * to tell it how many proxies are between the user and the server so it can
   * select the correct IP address from the `X-Forwarded-For` header.
   *
   * *WARNING:* If you set it incorrectly, the proxy will either record the
   * wrong IP address, or it will be possible for users to spoof their IP
   * addresses and bypass rate limiting. Check the request logs to see what
   * incoming X-Forwarded-For values look like.
   *
   * Examples:
   *  - X-Forwarded-For: "34.1.1.1, 172.1.1.1, 10.1.1.1" => trustedProxies: 3
   *  - X-Forwarded-For: "34.1.1.1" => trustedProxies: 1
   *  - no X-Forwarded-For header => trustedProxies: 0 (the actual IP of the incoming request will be used)
   *
   * As of 2024/01/08:
   * For HuggingFace or Cloudflare Tunnel, use 1.
   * For Render, use 3.
   * For deployments not behind a load balancer, use 0.
   *
   * You should double check against your actual request logs to be sure.
   *
   * Defaults to 1, as most deployments are on HuggingFace or Cloudflare Tunnel.
   */
  trustedProxies?: number;
  /**
   * Whether to allow OpenAI tool usage.  The proxy doesn't impelment any
   * support for tools/function calling but can pass requests and responses as
   * is. Note that the proxy also cannot accurately track quota usage for
   * requests involving tools, so you must opt in to this feature at your own
   * risk.
   */
  allowOpenAIToolUsage?: boolean;
  /**
   * Which services will accept prompts containing images, for use with
   * multimodal models. Users with `special` role are exempt from this
   * restriction.
   *
   * Do not enable this feature for untrusted users, as malicious users could
   * send images which violate your provider's terms of service or local laws.
   *
   * Defaults to no services, meaning image prompts are disabled. Use a comma-
   * separated list. Available services are:
   * openai,anthropic,google-ai,mistral-ai,aws,gcp,azure
   */
  allowedVisionServices: LLMService[];
  /**
   * Allows overriding the default proxy endpoint route. Defaults to /proxy.
   * A leading slash is required.
   */
  proxyEndpointRoute: string;
  /**
   * If set, only requests from these IP addresses will be permitted to use the
   * admin API and UI. Provide a comma-separated list of IP addresses or CIDR
   * ranges. If not set, the admin API and UI will be open to all requests.
   */
  adminWhitelist: string[];
  /**
   * If set, requests from these IP addresses will be blocked from using the
   * application. Provide a comma-separated list of IP addresses or CIDR ranges.
   * If not set, no IP addresses will be blocked.
   *
   * Takes precedence over the adminWhitelist.
   */
  ipBlacklist: string[];
  /**
   * If set, pushes requests further back into the queue according to their
   * token costs by factor*tokens*milliseconds (or more intuitively
   * factor*thousands_of_tokens*seconds).
   * Accepts floats.
   */
  tokensPunishmentFactor: number;
  /**
   * Configuration for HTTP requests made by the proxy to other servers, such
   * as when checking keys or forwarding users' requests to external services.
   * If not set, all requests will be made using the default agent.
   *
   * If set, the proxy may make requests to other servers using the specified
   * settings. This is useful if you wish to route users' requests through
   * another proxy or VPN, or if you have multiple network interfaces and want
   * to use a specific one for outgoing requests.
   */
  httpAgent?: {
    /**
     * The name of the network interface to use. The first external IPv4 address
     * belonging to this interface will be used for outgoing requests.
     */
    interface?: string;
    /**
     * The URL of a proxy server to use. Supports SOCKS4, SOCKS5, HTTP, and
     * HTTPS. If not set, the proxy will be made using the default agent.
     * - SOCKS4: `socks4://some-socks-proxy.com:9050`
     * - SOCKS5: `socks5://username:password@some-socks-proxy.com:9050`
     * - HTTP: `http://proxy-server-over-tcp.com:3128`
     * - HTTPS: `https://proxy-server-over-tls.com:3129`
     *
     * **Note:** If your proxy server issues a certificate, you may need to set
     * `NODE_EXTRA_CA_CERTS` to the path to your certificate, otherwise this
     * application will reject TLS connections.
     */
    proxyUrl?: string;
  };
};

// To change configs, create a file called .env in the root directory.
// See .env.example for an example.
export const config: Config = {
  port: getEnvWithDefault("PORT", 7860),
  bindAddress: getEnvWithDefault("BIND_ADDRESS", "0.0.0.0"),
  openaiKey: getEnvWithDefault("OPENAI_KEY", ""),
  anthropicKey: getEnvWithDefault("ANTHROPIC_KEY", ""),
  googleAIKey: getEnvWithDefault("GOOGLE_AI_KEY", ""),
  mistralAIKey: getEnvWithDefault("MISTRAL_AI_KEY", ""),
  awsCredentials: getEnvWithDefault("AWS_CREDENTIALS", ""),
  gcpCredentials: getEnvWithDefault("GCP_CREDENTIALS", ""),
  azureCredentials: getEnvWithDefault("AZURE_CREDENTIALS", ""),
  proxyKey: getEnvWithDefault("PROXY_KEY", ""),
  adminKey: getEnvWithDefault("ADMIN_KEY", ""),
  serviceInfoPassword: getEnvWithDefault("SERVICE_INFO_PASSWORD", ""),
  sqliteDataPath: getEnvWithDefault(
    "SQLITE_DATA_PATH",
    path.join(DATA_DIR, "database.sqlite")
  ),
  eventLogging: getEnvWithDefault("EVENT_LOGGING", false),
  eventLoggingTrim: getEnvWithDefault("EVENT_LOGGING_TRIM", 5),
  gatekeeper: getEnvWithDefault("GATEKEEPER", "none"),
  gatekeeperStore: getEnvWithDefault("GATEKEEPER_STORE", "memory"),
  maxIpsPerUser: getEnvWithDefault("MAX_IPS_PER_USER", 0),
  maxIpsAutoBan: getEnvWithDefault("MAX_IPS_AUTO_BAN", false),
  captchaMode: getEnvWithDefault("CAPTCHA_MODE", "none"),
  powTokenHours: getEnvWithDefault("POW_TOKEN_HOURS", 24),
  powTokenMaxIps: getEnvWithDefault("POW_TOKEN_MAX_IPS", 2),
  powDifficultyLevel: getEnvWithDefault("POW_DIFFICULTY_LEVEL", "low"),
  powChallengeTimeout: getEnvWithDefault("POW_CHALLENGE_TIMEOUT", 30),
  powTokenPurgeHours: getEnvWithDefault("POW_TOKEN_PURGE_HOURS", 48),
  firebaseRtdbUrl: getEnvWithDefault("FIREBASE_RTDB_URL", undefined),
  firebaseKey: getEnvWithDefault("FIREBASE_KEY", undefined),
  textModelRateLimit: getEnvWithDefault("TEXT_MODEL_RATE_LIMIT", 4),
  imageModelRateLimit: getEnvWithDefault("IMAGE_MODEL_RATE_LIMIT", 4),
  maxContextTokensOpenAI: getEnvWithDefault("MAX_CONTEXT_TOKENS_OPENAI", 32768),
  maxContextTokensAnthropic: getEnvWithDefault(
    "MAX_CONTEXT_TOKENS_ANTHROPIC",
    32768
  ),
  maxOutputTokensOpenAI: getEnvWithDefault(
    ["MAX_OUTPUT_TOKENS_OPENAI", "MAX_OUTPUT_TOKENS"],
    1024
  ),
  maxOutputTokensAnthropic: getEnvWithDefault(
    ["MAX_OUTPUT_TOKENS_ANTHROPIC", "MAX_OUTPUT_TOKENS"],
    1024
  ),
  allowedModelFamilies: getEnvWithDefault(
    "ALLOWED_MODEL_FAMILIES",
    getDefaultModelFamilies()
  ),
  rejectPhrases: parseCsv(getEnvWithDefault("REJECT_PHRASES", "")),
  rejectMessage: getEnvWithDefault(
    "REJECT_MESSAGE",
    "This content violates /aicg/'s acceptable use policy."
  ),
  logLevel: getEnvWithDefault("LOG_LEVEL", "info"),
  checkKeys: getEnvWithDefault("CHECK_KEYS", !isDev),
  showTokenCosts: getEnvWithDefault("SHOW_TOKEN_COSTS", false),
  allowAwsLogging: getEnvWithDefault("ALLOW_AWS_LOGGING", false),
  promptLogging: getEnvWithDefault("PROMPT_LOGGING", false),
  promptLoggingBackend: getEnvWithDefault("PROMPT_LOGGING_BACKEND", undefined),
  promptLoggingFilePrefix: getEnvWithDefault(
    "PROMPT_LOGGING_FILE_PREFIX",
    "prompt-logs"
  ),
  googleSheetsKey: getEnvWithDefault("GOOGLE_SHEETS_KEY", undefined),
  googleSheetsSpreadsheetId: getEnvWithDefault(
    "GOOGLE_SHEETS_SPREADSHEET_ID",
    undefined
  ),
  blockedOrigins: getEnvWithDefault("BLOCKED_ORIGINS", undefined),
  blockMessage: getEnvWithDefault(
    "BLOCK_MESSAGE",
    "You must be over the age of majority in your country to use this service."
  ),
  blockRedirect: getEnvWithDefault("BLOCK_REDIRECT", "https://www.9gag.com"),
  tokenQuota: MODEL_FAMILIES.reduce(
    (acc, family: ModelFamily) => {
      acc[family] = getEnvWithDefault(
        `TOKEN_QUOTA_${family.toUpperCase().replace(/-/g, "_")}`,
        0
      ) as number;
      return acc;
    },
    {} as { [key in ModelFamily]: number }
  ),
  quotaRefreshPeriod: getEnvWithDefault("QUOTA_REFRESH_PERIOD", undefined),
  allowNicknameChanges: getEnvWithDefault("ALLOW_NICKNAME_CHANGES", true),
  showRecentImages: getEnvWithDefault("SHOW_RECENT_IMAGES", true),
  useInsecureCookies: getEnvWithDefault("USE_INSECURE_COOKIES", isDev),
  staticServiceInfo: getEnvWithDefault("STATIC_SERVICE_INFO", false),
  trustedProxies: getEnvWithDefault("TRUSTED_PROXIES", 1),
  allowOpenAIToolUsage: getEnvWithDefault("ALLOW_OPENAI_TOOL_USAGE", false),
  allowedVisionServices: parseCsv(
    getEnvWithDefault("ALLOWED_VISION_SERVICES", "")
  ) as LLMService[],
  proxyEndpointRoute: getEnvWithDefault("PROXY_ENDPOINT_ROUTE", "/proxy"),
  adminWhitelist: parseCsv(
    getEnvWithDefault("ADMIN_WHITELIST", "0.0.0.0/0,::/0")
  ),
  ipBlacklist: parseCsv(getEnvWithDefault("IP_BLACKLIST", "")),
  tokensPunishmentFactor: getEnvWithDefault("TOKENS_PUNISHMENT_FACTOR", 0.0),
  httpAgent: {
    interface: getEnvWithDefault("HTTP_AGENT_INTERFACE", undefined),
    proxyUrl: getEnvWithDefault("HTTP_AGENT_PROXY_URL", undefined),
  },
} as const;

function generateSigningKey() {
  if (process.env.COOKIE_SECRET !== undefined) {
    // legacy, replaced by SIGNING_KEY
    return process.env.COOKIE_SECRET;
  } else if (process.env.SIGNING_KEY !== undefined) {
    return process.env.SIGNING_KEY;
  }

  const secrets = [
    config.adminKey,
    config.openaiKey,
    config.anthropicKey,
    config.googleAIKey,
    config.mistralAIKey,
    config.awsCredentials,
    config.gcpCredentials,
    config.azureCredentials,
  ];
  if (secrets.filter((s) => s).length === 0) {
    startupLogger.warn(
      "No SIGNING_KEY or secrets are set. All sessions, cookies, and proofs of work will be invalidated on restart."
    );
    return crypto.randomBytes(32).toString("hex");
  }

  startupLogger.info("No SIGNING_KEY set; one will be generated from secrets.");
  startupLogger.info(
    "It's recommended to set SIGNING_KEY explicitly to ensure users' sessions and cookies always persist across restarts."
  );
  const seed = secrets.map((s) => s || "n/a").join("");
  return crypto.createHash("sha256").update(seed).digest("hex");
}

const signingKey = generateSigningKey();
export const SECRET_SIGNING_KEY = signingKey;

export async function assertConfigIsValid() {
  if (process.env.MODEL_RATE_LIMIT !== undefined) {
    const limit =
      parseInt(process.env.MODEL_RATE_LIMIT, 10) || config.textModelRateLimit;

    config.textModelRateLimit = limit;
    config.imageModelRateLimit = Math.max(Math.floor(limit / 2), 1);

    startupLogger.warn(
      { textLimit: limit, imageLimit: config.imageModelRateLimit },
      "MODEL_RATE_LIMIT is deprecated. Use TEXT_MODEL_RATE_LIMIT and IMAGE_MODEL_RATE_LIMIT instead."
    );
  }

  if (process.env.ALLOW_IMAGE_PROMPTS === "true") {
    const hasAllowedServices = config.allowedVisionServices.length > 0;
    if (!hasAllowedServices) {
      config.allowedVisionServices = ["openai", "anthropic"];
      startupLogger.warn(
        { allowedVisionServices: config.allowedVisionServices },
        "ALLOW_IMAGE_PROMPTS is deprecated. Use ALLOWED_VISION_SERVICES instead."
      );
    }
  }

  if (config.promptLogging && !config.promptLoggingBackend) {
    throw new Error(
      "Prompt logging is enabled but no backend is configured. Set PROMPT_LOGGING_BACKEND to 'google_sheets' or 'file'."
    );
  }

  if (!["none", "proxy_key", "user_token"].includes(config.gatekeeper)) {
    throw new Error(
      `Invalid gatekeeper mode: ${config.gatekeeper}. Must be one of: none, proxy_key, user_token.`
    );
  }

  if (config.gatekeeper === "user_token" && !config.adminKey) {
    throw new Error(
      "`user_token` gatekeeper mode requires an `ADMIN_KEY` to be set."
    );
  }

  if (
    config.captchaMode === "proof_of_work" &&
    config.gatekeeper !== "user_token"
  ) {
    throw new Error(
      "Captcha mode 'proof_of_work' requires gatekeeper mode 'user_token'."
    );
  }

  if (config.captchaMode === "proof_of_work") {
    const val = config.powDifficultyLevel;
    const isDifficulty =
      typeof val === "string" &&
      ["low", "medium", "high", "extreme"].includes(val);
    const isIterations =
      typeof val === "number" && Number.isInteger(val) && val > 0;
    if (!isDifficulty && !isIterations) {
      throw new Error(
        "Invalid POW_DIFFICULTY_LEVEL. Must be one of: low, medium, high, extreme, or a positive integer."
      );
    }
  }

  if (config.gatekeeper === "proxy_key" && !config.proxyKey) {
    throw new Error(
      "`proxy_key` gatekeeper mode requires a `PROXY_KEY` to be set."
    );
  }

  if (
    config.gatekeeperStore === "firebase_rtdb" &&
    (!config.firebaseKey || !config.firebaseRtdbUrl)
  ) {
    throw new Error(
      "Firebase RTDB store requires `FIREBASE_KEY` and `FIREBASE_RTDB_URL` to be set."
    );
  }

  if (Object.values(config.httpAgent || {}).filter(Boolean).length === 0) {
    delete config.httpAgent;
  } else if (config.httpAgent) {
    if (config.httpAgent.interface && config.httpAgent.proxyUrl) {
      throw new Error(
        "Cannot set both `HTTP_AGENT_INTERFACE` and `HTTP_AGENT_PROXY_URL`."
      );
    }
  }

  // Ensure forks which add new secret-like config keys don't unwittingly expose
  // them to users.
  for (const key of getKeys(config)) {
    const maybeSensitive = ["key", "credentials", "secret", "password"].some(
      (sensitive) =>
        key.toLowerCase().includes(sensitive) && !["checkKeys"].includes(key)
    );
    const secured = new Set([...SENSITIVE_KEYS, ...OMITTED_KEYS]);
    if (maybeSensitive && !secured.has(key))
      throw new Error(
        `Config key "${key}" may be sensitive but is exposed. Add it to SENSITIVE_KEYS or OMITTED_KEYS.`
      );
  }
}

/**
 * Config keys that are masked on the info page, but not hidden as their
 * presence may be relevant to the user due to privacy implications.
 */
export const SENSITIVE_KEYS: (keyof Config)[] = [
  "googleSheetsSpreadsheetId",
  "httpAgent",
];

/**
 * Config keys that are not displayed on the info page at all, generally because
 * they are not relevant to the user or can be inferred from other config.
 */
export const OMITTED_KEYS = [
  "port",
  "bindAddress",
  "logLevel",
  "openaiKey",
  "anthropicKey",
  "googleAIKey",
  "mistralAIKey",
  "awsCredentials",
  "gcpCredentials",
  "azureCredentials",
  "proxyKey",
  "adminKey",
  "serviceInfoPassword",
  "rejectPhrases",
  "rejectMessage",
  "showTokenCosts",
  "promptLoggingFilePrefix",
  "googleSheetsKey",
  "firebaseKey",
  "firebaseRtdbUrl",
  "sqliteDataPath",
  "eventLogging",
  "eventLoggingTrim",
  "gatekeeperStore",
  "maxIpsPerUser",
  "blockedOrigins",
  "blockMessage",
  "blockRedirect",
  "allowNicknameChanges",
  "showRecentImages",
  "useInsecureCookies",
  "staticServiceInfo",
  "checkKeys",
  "allowedModelFamilies",
  "trustedProxies",
  "proxyEndpointRoute",
  "adminWhitelist",
  "ipBlacklist",
  "powTokenPurgeHours",
] satisfies (keyof Config)[];
type OmitKeys = (typeof OMITTED_KEYS)[number];

type Printable<T> = {
  [P in keyof T as Exclude<P, OmitKeys>]: T[P] extends object
    ? Printable<T[P]>
    : string;
};
type PublicConfig = Printable<Config>;

const getKeys = Object.keys as <T extends object>(obj: T) => Array<keyof T>;

export function listConfig(obj: Config = config) {
  const result: Record<string, unknown> = {};
  for (const key of getKeys(obj)) {
    const value = obj[key]?.toString() || "";

    const shouldMask = SENSITIVE_KEYS.includes(key);
    const shouldOmit =
      OMITTED_KEYS.includes(key as OmitKeys) ||
      value === "" ||
      value === "undefined";

    if (shouldOmit) {
      continue;
    }

    const validKey = key as keyof Printable<Config>;

    if (value && shouldMask) {
      result[validKey] = "********";
    } else {
      result[validKey] = value;
    }

    if (typeof obj[key] === "object" && !Array.isArray(obj[key])) {
      result[key] = listConfig(obj[key] as unknown as Config);
    }
  }
  return result as PublicConfig;
}

/**
 * Tries to get a config value from one or more environment variables (in
 * order), falling back to a default value if none are set.
 */
function getEnvWithDefault<T>(env: string | string[], defaultValue: T): T {
  const value = Array.isArray(env)
    ? env.map((e) => process.env[e]).find((v) => v !== undefined)
    : process.env[env];
  if (value === undefined) {
    return defaultValue;
  }
  try {
    if (
      [
        "OPENAI_KEY",
        "ANTHROPIC_KEY",
        "GOOGLE_AI_KEY",
        "AWS_CREDENTIALS",
        "GCP_CREDENTIALS",
        "AZURE_CREDENTIALS",
      ].includes(String(env))
    ) {
      return value as unknown as T;
    }

    // Intended to be used for comma-delimited lists
    if (Array.isArray(defaultValue)) {
      return value.split(",").map((v) => v.trim()) as T;
    }

    return JSON.parse(value) as T;
  } catch (err) {
    return value as unknown as T;
  }
}

function parseCsv(val: string): string[] {
  if (!val) return [];

  const regex = /(".*?"|[^",]+)(?=\s*,|\s*$)/g;
  const matches = val.match(regex) || [];
  return matches.map((item) => item.replace(/^"|"$/g, "").trim());
}

function getDefaultModelFamilies(): ModelFamily[] {
  return MODEL_FAMILIES.filter(
    (f) => !f.includes("dall-e") && !f.includes("o1")
  ) as ModelFamily[];
}


code:
/** This whole module kinda sucks */
import fs from "fs";
import express, { Router, Request, Response } from "express";
import showdown from "showdown";
import { config } from "./config";
import { buildInfo, ServiceInfo } from "./service-info";
import { getLastNImages } from "./shared/file-storage/image-history";
import { keyPool } from "./shared/key-management";
import { MODEL_FAMILY_SERVICE, ModelFamily } from "./shared/models";
import { withSession } from "./shared/with-session";
import { checkCsrfToken, injectCsrfToken } from "./shared/inject-csrf";

const INFO_PAGE_TTL = 2000;
const MODEL_FAMILY_FRIENDLY_NAME: { [f in ModelFamily]: string } = {
  turbo: "GPT-4o Mini / 3.5 Turbo",
  gpt4: "GPT-4",
  "gpt4-32k": "GPT-4 32k",
  "gpt4-turbo": "GPT-4 Turbo",
  gpt4o: "GPT-4o",
  o1: "OpenAI o1",
  "o1-mini": "OpenAI o1 mini",
  "dall-e": "DALL-E",
  claude: "Claude (Sonnet)",
  "claude-opus": "Claude (Opus)",
  "gemini-flash": "Gemini Flash",
  "gemini-pro": "Gemini Pro",
  "gemini-ultra": "Gemini Ultra",
  "mistral-tiny": "Mistral 7B",
  "mistral-small": "Mistral Nemo",
  "mistral-medium": "Mistral Medium",
  "mistral-large": "Mistral Large",
  "aws-claude": "AWS Claude (Sonnet)",
  "aws-claude-opus": "AWS Claude (Opus)",
  "aws-mistral-tiny": "AWS Mistral 7B",
  "aws-mistral-small": "AWS Mistral Nemo",
  "aws-mistral-medium": "AWS Mistral Medium",
  "aws-mistral-large": "AWS Mistral Large",
  "gcp-claude": "GCP Claude (Sonnet)",
  "gcp-claude-opus": "GCP Claude (Opus)",
  "azure-turbo": "Azure GPT-3.5 Turbo",
  "azure-gpt4": "Azure GPT-4",
  "azure-gpt4-32k": "Azure GPT-4 32k",
  "azure-gpt4-turbo": "Azure GPT-4 Turbo",
  "azure-gpt4o": "Azure GPT-4o",
  "azure-o1": "Azure o1",
  "azure-o1-mini": "Azure o1 mini",
  "azure-dall-e": "Azure DALL-E",
};

const converter = new showdown.Converter();
const customGreeting = fs.existsSync("greeting.md")
  ? `<div id="servergreeting">${fs.readFileSync("greeting.md", "utf8")}</div>`
  : "";
let infoPageHtml: string | undefined;
let infoPageLastUpdated = 0;

export const handleInfoPage = (req: Request, res: Response) => {
  if (infoPageLastUpdated + INFO_PAGE_TTL > Date.now()) {
    return res.send(infoPageHtml);
  }

  const baseUrl =
    process.env.SPACE_ID && !req.get("host")?.includes("hf.space")
      ? getExternalUrlForHuggingfaceSpaceId(process.env.SPACE_ID)
      : req.protocol + "://" + req.get("host");

  const info = buildInfo(baseUrl + config.proxyEndpointRoute);
  infoPageHtml = renderPage(info);
  infoPageLastUpdated = Date.now();

  res.send(infoPageHtml);
};

export function renderPage(info: ServiceInfo) {
  const title = getServerTitle();
  const headerHtml = buildInfoPageHeader(info);

  return `<!doctype html>

    <meta charset="utf-8" />
    <meta name="robots" content="noindex" />
    <title>${title}</title>
    <link rel="stylesheet" href="/res/css/reset.css" media="screen" />
    <link rel="stylesheet" href="/res/css/sakura.css" media="screen" />
    <link rel="stylesheet" href="/res/css/sakura-dark.css" media="screen and (prefers-color-scheme: dark)" />

    ${headerHtml}

    ${getSelfServiceLinks()}

Service Info

${JSON.stringify(info, null, 2)}

`;
}

/**
 * If the server operator provides a `greeting.md` file, it will be included in
 * the rendered info page.
 **/
function buildInfoPageHeader(info: ServiceInfo) {
  const title = getServerTitle();
  // TODO: use some templating engine instead of this mess
  let infoBody = `# ${title}`;
  if (config.promptLogging) {
    infoBody += `\n## Prompt Logging Enabled
This proxy keeps full logs of all prompts and AI responses. Prompt logs are anonymous and do not contain IP addresses or timestamps.

[You can see the type of data logged here, along with the rest of the code.](https://gitgud.io/khanon/oai-reverse-proxy/-/blob/main/src/shared/prompt-logging/index.ts).

**If you are uncomfortable with this, don't send prompts to this proxy!**`;
  }

  if (config.staticServiceInfo) {
    return converter.makeHtml(infoBody + customGreeting);
  }

  const waits: string[] = [];

  for (const modelFamily of config.allowedModelFamilies) {
    const service = MODEL_FAMILY_SERVICE[modelFamily];

    const hasKeys = keyPool.list().some((k) => {
      return k.service === service && k.modelFamilies.includes(modelFamily);
    });

    const wait = info[modelFamily]?.estimatedQueueTime;
    if (hasKeys && wait) {
      waits.push(
        `**${MODEL_FAMILY_FRIENDLY_NAME[modelFamily] || modelFamily}**: ${wait}`
      );
    }
  }

  infoBody += "\n\n" + waits.join(" / ");

  infoBody += customGreeting;

  infoBody += buildRecentImageSection();

  return converter.makeHtml(infoBody);
}

function getSelfServiceLinks() {
  if (config.gatekeeper !== "user_token") return "";

  const links = [["Check your user token", "/user/lookup"]];
  if (config.captchaMode !== "none") {
    links.unshift(["Request a user token", "/user/captcha"]);
  }

  return `
${links
    .map(([text, link]) => `<a href="${link}">${text}</a>`)
    .join(" | ")}
`;
}

function getServerTitle() {
  // Use manually set title if available
  if (process.env.SERVER_TITLE) {
    return process.env.SERVER_TITLE;
  }

  // Huggingface
  if (process.env.SPACE_ID) {
    return `${process.env.SPACE_AUTHOR_NAME} / ${process.env.SPACE_TITLE}`;
  }

  // Render
  if (process.env.RENDER) {
    return `Render / ${process.env.RENDER_SERVICE_NAME}`;
  }

  return "OAI Reverse Proxy";
}

function buildRecentImageSection() {
  const dalleModels: ModelFamily[] = ["azure-dall-e", "dall-e"];
  if (
    !config.showRecentImages ||
    dalleModels.every((f) => !config.allowedModelFamilies.includes(f))
  ) {
    return "";
  }

  let html = `<h2>Recent DALL-E Generations</h2>`;
  const recentImages = getLastNImages(12).reverse();
  if (recentImages.length === 0) {
    html += `<p>No images yet.</p>`;
    return html;
  }

  html += `<div style="display: flex; flex-wrap: wrap;" id="recent-images">`;
  for (const { url, prompt } of recentImages) {
    const thumbUrl = url.replace(/\.png$/, "_t.jpg");
    const escapedPrompt = escapeHtml(prompt);
    html += `

`;
  }
  html += `</div>`;
  html += `<p style="clear: both; text-align: center;"><a href="/user/image-history">View all recent images</a></p>`;

  return html;
}

function escapeHtml(unsafe: string) {
  return unsafe
    .replace(/&/g, "&")
    .replace(/</g, "<")
    .replace(/>/g, ">")
    .replace(/"/g, """)
    .replace(/'/g, "'")
    .replace(/\[/g, "[")
    .replace(/]/g, "]");
}

function getExternalUrlForHuggingfaceSpaceId(spaceId: string) {
  try {
    const [username, spacename] = spaceId.split("/");
    return `https://${username}-${spacename.replace(/_/g, "-")}.hf.space`;
  } catch (e) {
    return "";
  }
}

function checkIfUnlocked(
  req: Request,
  res: Response,
  next: express.NextFunction
) {
  if (config.serviceInfoPassword?.length && !req.session?.unlocked) {
    return res.redirect("/unlock-info");
  }
  next();
}

const infoPageRouter = Router();
if (config.serviceInfoPassword?.length) {
  infoPageRouter.use(
    express.json({ limit: "1mb" }),
    express.urlencoded({ extended: true, limit: "1mb" })
  );
  infoPageRouter.use(withSession);
  infoPageRouter.use(injectCsrfToken, checkCsrfToken);
  infoPageRouter.post("/unlock-info", (req, res) => {
    if (req.body.password !== config.serviceInfoPassword) {
      return res.status(403).send("Incorrect password");
    }
    req.session!.unlocked = true;
    res.redirect("/");
  });
  infoPageRouter.get("/unlock-info", (_req, res) => {
    if (_req.session?.unlocked) return res.redirect("/");

    res.send(`
      <form method="post" action="/unlock-info">

Unlock Service Info

        <input type="hidden" name="_csrf" value="${res.locals.csrfToken}" />
        <input type="password" name="password" placeholder="Password" />
        <button type="submit">Unlock</button>
      </form>
    `);
  });
  infoPageRouter.use(checkIfUnlocked);
}
infoPageRouter.get("/", handleInfoPage);
infoPageRouter.get("/status", (req, res) => {
  res.json(buildInfo(req.protocol + "://" + req.get("host"), false));
});
export { infoPageRouter };

license: unknown


code:
import pino from "pino";
import { config } from "./config";

const transport =
  process.env.NODE_ENV === "production"
    ? undefined
    : {
        target: "pino-pretty",
        options: {
          singleLine: true,
          messageFormat: "{if module}\x1b[90m[{module}] \x1b[39m{end}{msg}",
          ignore: "module",
        },
      };

export const logger = pino({
  level: config.logLevel,
  base: { pid: process.pid, module: "server" },
  transport,
});


code:
import { NextFunction, Request, Response } from "express";

export function addV1(req: Request, res: Response, next: NextFunction) {
  // Clients don't consistently use the /v1 prefix so we'll add it for them.
  if (!req.path.startsWith("/v1/") && !req.path.startsWith("/v1beta/")) {
    req.url = `/v1${req.url}`;
  }
  next();
}


code:
import { Request, RequestHandler, Router } from "express";
import { config } from "../config";
import { ipLimiter } from "./rate-limit";
import {
  addKey,
  createPreprocessorMiddleware,
  finalizeBody,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";
import { ProxyReqManager } from "./middleware/request/proxy-req-manager";

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-haiku-20240307",
    "claude-3-5-haiku-20241022",
    "claude-3-opus-20240229",
    "claude-3-opus-latest",
    "claude-3-sonnet-20240229",
    "claude-3-5-sonnet-20240620",
    "claude-3-5-sonnet-20241022",
    "claude-3-5-sonnet-latest",
  ];

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

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

  let newBody = body;
  switch (`${req.inboundApi}<-${req.outboundApi}`) {
    case "openai<-anthropic-text":
      req.log.info("Transforming Anthropic Text back to OpenAI format");
      newBody = transformAnthropicTextResponseToOpenAI(body, req);
      break;
    case "openai<-anthropic-chat":
      req.log.info("Transforming Anthropic Chat back to OpenAI format");
      newBody = transformAnthropicChatResponseToOpenAI(body);
      break;
    case "anthropic-text<-anthropic-chat":
      req.log.info("Transforming Anthropic Chat back to Anthropic chat format");
      newBody = transformAnthropicChatResponseToAnthropicText(body);
      break;
  }

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

function flattenChatResponse(
  content: { type: string; text: string }[]
): string {
  return content
    .map((part: { type: string; text: string }) =>
      part.type === "text" ? part.text : ""
    )
    .join("\n");
}

export function transformAnthropicChatResponseToAnthropicText(
  anthropicBody: Record<string, any>
): Record<string, any> {
  return {
    type: "completion",
    id: "ant-" + anthropicBody.id,
    completion: flattenChatResponse(anthropicBody.content),
    stop_reason: anthropicBody.stop_reason,
    stop: anthropicBody.stop_sequence,
    model: anthropicBody.model,
    usage: anthropicBody.usage,
  };
}

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

export function transformAnthropicChatResponseToOpenAI(
  anthropicBody: Record<string, any>
): Record<string, any> {
  return {
    id: "ant-" + anthropicBody.id,
    object: "chat.completion",
    created: Date.now(),
    model: anthropicBody.model,
    usage: anthropicBody.usage,
    choices: [
      {
        message: {
          role: "assistant",
          content: flattenChatResponse(anthropicBody.content),
        },
        finish_reason: anthropicBody.stop_reason,
        index: 0,
      },
    ],
  };
}

/**
 * If a client using the OpenAI compatibility endpoint requests an actual OpenAI
 * model, reassigns it to Sonnet.
 */
function maybeReassignModel(req: Request) {
  const model = req.body.model;
  if (model.includes("claude")) return; // use whatever model the user requested
  req.body.model = "claude-3-5-sonnet-latest";
}

/**
 * If client requests more than 4096 output tokens the request must have a
 * particular version header.
 * https://docs.anthropic.com/en/release-notes/api#july-15th-2024
 */
function setAnthropicBetaHeader(req: Request) {
  const { max_tokens_to_sample } = req.body;
  if (max_tokens_to_sample > 4096) {
    req.headers["anthropic-beta"] = "max-tokens-3-5-sonnet-2024-07-15";
  }
}

function selectUpstreamPath(manager: ProxyReqManager) {
  const req = manager.request;
  const pathname = req.url.split("?")[0];
  req.log.debug({ pathname }, "Anthropic path filter");
  const isText = req.outboundApi === "anthropic-text";
  const isChat = req.outboundApi === "anthropic-chat";
  if (isChat && pathname === "/v1/complete") {
    manager.setPath("/v1/messages");
  }
  if (isText && pathname === "/v1/chat/completions") {
    manager.setPath("/v1/complete");
  }
  if (isChat && pathname === "/v1/chat/completions") {
    manager.setPath("/v1/messages");
  }
  if (isChat && ["sonnet", "opus"].includes(req.params.type)) {
    manager.setPath("/v1/messages");
  }
}

const anthropicProxy = createQueuedProxyMiddleware({
  target: "https://api.anthropic.com",
  mutations: [selectUpstreamPath, addKey, finalizeBody],
  blockingResponseHandler: anthropicBlockingResponseHandler,
});

const nativeAnthropicChatPreprocessor = createPreprocessorMiddleware(
  { inApi: "anthropic-chat", outApi: "anthropic-chat", service: "anthropic" },
  { afterTransform: [setAnthropicBetaHeader] }
);

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 preprocessAnthropicTextRequest: RequestHandler = (req, res, next) => {
  if (req.body.model?.startsWith("claude-3")) {
    textToChatPreprocessor(req, res, next);
  } else {
    nativeTextPreprocessor(req, res, next);
  }
};

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

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

/**
 * Routes an OpenAI prompt to either the legacy Claude text completion endpoint
 * or the new Claude chat completion endpoint, based on the requested model.
 */
const preprocessOpenAICompatRequest: RequestHandler = (req, res, next) => {
  maybeReassignModel(req);
  if (req.body.model?.includes("claude-3")) {
    oaiToChatPreprocessor(req, res, next);
  } else {
    oaiToTextPreprocessor(req, res, next);
  }
};

const anthropicRouter = Router();
anthropicRouter.get("/v1/models", handleModelRequest);
// Native Anthropic chat completion endpoint.
anthropicRouter.post(
  "/v1/messages",
  ipLimiter,
  nativeAnthropicChatPreprocessor,
  anthropicProxy
);
// Anthropic text completion endpoint. Translates to Anthropic chat completion
// if the requested model is a Claude 3 model.
anthropicRouter.post(
  "/v1/complete",
  ipLimiter,
  preprocessAnthropicTextRequest,
  anthropicProxy
);
// OpenAI-to-Anthropic compatibility endpoint. Accepts an OpenAI chat completion
// request and transforms/routes it to the appropriate Anthropic format and
// endpoint based on the requested model.
anthropicRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  preprocessOpenAICompatRequest,
  anthropicProxy
);

export const anthropic = anthropicRouter;


code:
import { Request, RequestHandler, Router } from "express";
import { v4 } from "uuid";
import {
  transformAnthropicChatResponseToAnthropicText,
  transformAnthropicChatResponseToOpenAI,
} from "./anthropic";
import { ipLimiter } from "./rate-limit";
import {
  createPreprocessorMiddleware,
  finalizeSignedRequest,
  signAwsRequest,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

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

  let newBody = body;
  switch (`${req.inboundApi}<-${req.outboundApi}`) {
    case "openai<-anthropic-text":
      req.log.info("Transforming Anthropic Text back to OpenAI format");
      newBody = transformAwsTextResponseToOpenAI(body, req);
      break;
    case "openai<-anthropic-chat":
      req.log.info("Transforming AWS Anthropic Chat back to OpenAI format");
      newBody = transformAnthropicChatResponseToOpenAI(body);
      break;
    case "anthropic-text<-anthropic-chat":
      req.log.info("Transforming AWS Anthropic Chat back to Text format");
      newBody = transformAnthropicChatResponseToAnthropicText(body);
      break;
  }

  // AWS does not always confirm the model in the response, so we have to add it
  if (!newBody.model && req.body.model) {
    newBody.model = req.body.model;
  }

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

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 awsClaudeProxy = createQueuedProxyMiddleware({
  target: ({ signedRequest }) => {
    if (!signedRequest) throw new Error("Must sign request before proxying");
    return `${signedRequest.protocol}//${signedRequest.hostname}`;
  },
  mutations: [signAwsRequest, finalizeSignedRequest],
  blockingResponseHandler: awsBlockingResponseHandler,
});

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 preprocessAwsTextRequest: RequestHandler = (req, res, next) => {
  if (req.body.model?.includes("claude-3")) {
    textToChatPreprocessor(req, res, next);
  } else {
    nativeTextPreprocessor(req, res, next);
  }
};

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

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

/**
 * Routes an OpenAI prompt to either the legacy Claude text completion endpoint
 * or the new Claude chat completion endpoint, based on the requested model.
 */
const preprocessOpenAICompatRequest: RequestHandler = (req, res, next) => {
  if (req.body.model?.includes("claude-3")) {
    oaiToAwsChatPreprocessor(req, res, next);
  } else {
    oaiToAwsTextPreprocessor(req, res, next);
  }
};

const awsClaudeRouter = Router();
// Native(ish) Anthropic text completion endpoint.
awsClaudeRouter.post(
  "/v1/complete",
  ipLimiter,
  preprocessAwsTextRequest,
  awsClaudeProxy
);
// Native Anthropic chat completion endpoint.
awsClaudeRouter.post(
  "/v1/messages",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "anthropic-chat", outApi: "anthropic-chat", service: "aws" },
    { afterTransform: [maybeReassignModel] }
  ),
  awsClaudeProxy
);

// OpenAI-to-AWS Anthropic compatibility endpoint.
awsClaudeRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  preprocessOpenAICompatRequest,
  awsClaudeProxy
);

/**
 * 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
 *
 * If client sends AWS model ID it will be used verbatim. Otherwise, various
 * strategies are used to try to map a non-AWS model name to AWS model ID.
 */
function maybeReassignModel(req: Request) {
  const model = req.body.model;

  // If it looks like an AWS model, use it as-is
  if (model.includes("anthropic.claude")) {
    return;
  }

  // Anthropic model names can look like:
  // - claude-v1
  // - claude-2.1
  // - claude-3-5-sonnet-20240620
  // - claude-3-opus-latest
  const pattern =
    /^(claude-)?(instant-)?(v)?(\d+)([.-](\d))?(-\d+k)?(-sonnet-|-opus-|-haiku-)?(latest|\d*)/i;
  const match = model.match(pattern);

  if (!match) {
    throw new Error(`Provided model name (${model}) doesn't resemble a Claude model ID.`);
  }

  const [_, _cl, instant, _v, major, _sep, minor, _ctx, rawName, rev] = match;

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

  const ver = minor ? `${major}.${minor}` : major;
  const name = rawName?.match(/([a-z]+)/)?.[1] || "";

  switch (ver) {
    case "1":
    case "1.0":
      req.body.model = "anthropic.claude-v1";
      return;
    case "2":
    case "2.0":
      req.body.model = "anthropic.claude-v2";
      return;
    case "2.1":
      req.body.model = "anthropic.claude-v2:1";
      return;
    case "3":
    case "3.0":
      // there is only one snapshot for all Claude 3 models so there is no need
      // to check the revision
      switch (name) {
        case "sonnet":
          req.body.model = "anthropic.claude-3-sonnet-20240229-v1:0";
          return;
        case "haiku":
          req.body.model = "anthropic.claude-3-haiku-20240307-v1:0";
          return;
        case "opus":
          req.body.model = "anthropic.claude-3-opus-20240229-v1:0";
          return;
      }
      break;
    case "3.5":
      switch (name) {
        case "sonnet":
          switch (rev) {
            case "20241022":
            case "latest":
              req.body.model = "anthropic.claude-3-5-sonnet-20241022-v2:0";
              return;
            case "20240620":
              req.body.model = "anthropic.claude-3-5-sonnet-20240620-v1:0";
              return;
          }
          break;
        case "haiku":
          switch (rev) {
            case "20241022":
            case "latest":
              req.body.model = "anthropic.claude-3-5-haiku-20241022-v1:0";
              return;
          }
        case "opus":
          // Add after model id is announced never
          break;
      }
  }

  throw new Error(`Provided model name (${model}) could not be mapped to a known AWS Claude model ID.`);
}

export const awsClaude = awsClaudeRouter;


code:
import { Request, Router } from "express";
import {
  detectMistralInputApi,
  transformMistralTextToMistralChat,
} from "./mistral-ai";
import { ipLimiter } from "./rate-limit";
import { ProxyResHandlerWithBody } from "./middleware/response";
import {
  createPreprocessorMiddleware,
  finalizeSignedRequest,
  signAwsRequest,
} from "./middleware/request";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

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

  let newBody = body;
  if (req.inboundApi === "mistral-ai" && req.outboundApi === "mistral-text") {
    newBody = transformMistralTextToMistralChat(body);
  }
  // AWS does not always confirm the model in the response, so we have to add it
  if (!newBody.model && req.body.model) {
    newBody.model = req.body.model;
  }

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

const awsMistralProxy = createQueuedProxyMiddleware({
  target: ({ signedRequest }) => {
    if (!signedRequest) throw new Error("Must sign request before proxying");
    return `${signedRequest.protocol}//${signedRequest.hostname}`;
  },
  mutations: [signAwsRequest,finalizeSignedRequest],
  blockingResponseHandler: awsMistralBlockingResponseHandler,
});

function maybeReassignModel(req: Request) {
  const model = req.body.model;

  // If it looks like an AWS model, use it as-is
  if (model.startsWith("mistral.")) {
    return;
  }
  // Mistral 7B Instruct
  else if (model.includes("7b")) {
    req.body.model = "mistral.mistral-7b-instruct-v0:2";
  }
  // Mistral 8x7B Instruct
  else if (model.includes("8x7b")) {
    req.body.model = "mistral.mixtral-8x7b-instruct-v0:1";
  }
  // Mistral Large (Feb 2024)
  else if (model.includes("large-2402")) {
    req.body.model = "mistral.mistral-large-2402-v1:0";
  }
  // Mistral Large 2 (July 2024)
  else if (model.includes("large")) {
    req.body.model = "mistral.mistral-large-2407-v1:0";
  }
  // Mistral Small (Feb 2024)
  else if (model.includes("small")) {
    req.body.model = "mistral.mistral-small-2402-v1:0";
  } else {
    throw new Error(
      `Can't map '${model}' to a supported AWS model ID; make sure you are requesting a Mistral model supported by Amazon Bedrock`
    );
  }
}

const nativeMistralChatPreprocessor = createPreprocessorMiddleware(
  { inApi: "mistral-ai", outApi: "mistral-ai", service: "aws" },
  {
    beforeTransform: [detectMistralInputApi],
    afterTransform: [maybeReassignModel],
  }
);

const awsMistralRouter = Router();
awsMistralRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  nativeMistralChatPreprocessor,
  awsMistralProxy
);

export const awsMistral = awsMistralRouter;


code:
/* Shared code between AWS Claude and AWS Mistral endpoints. */

import { Request, Response, Router } from "express";
import { config } from "../config";
import { addV1 } from "./add-v1";
import { awsClaude } from "./aws-claude";
import { awsMistral } from "./aws-mistral";
import { AwsBedrockKey, keyPool } from "../shared/key-management";

const awsRouter = Router();
awsRouter.get(["/:vendor?/v1/models", "/:vendor?/models"], handleModelsRequest);
awsRouter.use("/claude", addV1, awsClaude);
awsRouter.use("/mistral", addV1, awsMistral);

const MODELS_CACHE_TTL = 10000;
let modelsCache: Record<string, any> = {};
let modelsCacheTime: Record<string, number> = {};
function handleModelsRequest(req: Request, res: Response) {
  if (!config.awsCredentials) return { object: "list", data: [] };

  const vendor = req.params.vendor?.length
    ? req.params.vendor === "claude"
      ? "anthropic"
      : req.params.vendor
    : "all";

  const cacheTime = modelsCacheTime[vendor] || 0;
  if (new Date().getTime() - cacheTime < MODELS_CACHE_TTL) {
    return res.json(modelsCache[vendor]);
  }

  const availableModelIds = new Set<string>();
  for (const key of keyPool.list()) {
    if (key.isDisabled || key.service !== "aws") continue;
    (key as AwsBedrockKey).modelIds.forEach((id) => availableModelIds.add(id));
  }

  // https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
  const models = [
    "anthropic.claude-v2",
    "anthropic.claude-v2:1",
    "anthropic.claude-3-haiku-20240307-v1:0",
    "anthropic.claude-3-5-haiku-20241022-v1:0",
    "anthropic.claude-3-sonnet-20240229-v1:0",
    "anthropic.claude-3-5-sonnet-20240620-v1:0",
    "anthropic.claude-3-5-sonnet-20241022-v2:0",
    "anthropic.claude-3-opus-20240229-v1:0",
    "mistral.mistral-7b-instruct-v0:2",
    "mistral.mixtral-8x7b-instruct-v0:1",
    "mistral.mistral-large-2402-v1:0",
    "mistral.mistral-large-2407-v1:0",
    "mistral.mistral-small-2402-v1:0",
  ]
    .filter((id) => availableModelIds.has(id))
    .map((id) => {
      const vendor = id.match(/^(.*)\./)?.[1];
      return {
        id,
        object: "model",
        created: new Date().getTime(),
        owned_by: vendor,
        permission: [],
        root: vendor,
        parent: null,
      };
    });

  modelsCache[vendor] = {
    object: "list",
    data: models.filter((m) => vendor === "all" || m.root === vendor),
  };
  modelsCacheTime[vendor] = new Date().getTime();

  return res.json(modelsCache[vendor]);
}

export const aws = awsRouter;


code:
import { RequestHandler, Router } from "express";
import { config } from "../config";
import { generateModelList } from "./openai";
import { ipLimiter } from "./rate-limit";
import {
  addAzureKey,
  createPreprocessorMiddleware,
  finalizeSignedRequest,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

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

const handleModelRequest: RequestHandler = (_req, res) => {
  if (new Date().getTime() - modelsCacheTime < 1000 * 60) {
    return res.status(200).json(modelsCache);
  }

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

  const result = generateModelList("azure");

  modelsCache = { object: "list", data: result };
  modelsCacheTime = new Date().getTime();
  res.status(200).json(modelsCache);
};

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

  res.status(200).json({ ...body, proxy: body.proxy });
};

const azureOpenAIProxy = createQueuedProxyMiddleware({
  target: ({ signedRequest }) => {
    if (!signedRequest) throw new Error("Must sign request before proxying");
    const { hostname, protocol } = signedRequest;
    return `${protocol}//${hostname}`;
  },
  mutations: [addAzureKey, finalizeSignedRequest],
  blockingResponseHandler: azureOpenaiResponseHandler,
});

const azureOpenAIRouter = Router();
azureOpenAIRouter.get("/v1/models", handleModelRequest);
azureOpenAIRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai",
    outApi: "openai",
    service: "azure",
  }),
  azureOpenAIProxy
);
azureOpenAIRouter.post(
  "/v1/images/generations",
  ipLimiter,
  createPreprocessorMiddleware({
    inApi: "openai-image",
    outApi: "openai-image",
    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, Response, RequestHandler } from "express";
import { config } from "../config";
import { authenticate, getUser } from "../shared/users/user-store";
import { sendErrorToClient } from "./middleware/response/error-generator";

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.
  // Google AI uses a key querystring parameter.

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

  if (req.query.key) {
    const token = req.query.key?.toString();
    delete req.query.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 sendError(
          req,
          res,
          403,
          `Forbidden: no more IP addresses allowed for this user token`,
          { currentIp: ip, maxIps: user?.maxIps }
        );
      case "disabled":
        const bannedUser = getUser(token);
        if (bannedUser?.disabledAt) {
          const reason = bannedUser.disabledReason || "User token disabled";
          return sendError(req, res, 403, `Forbidden: ${reason}`);
        }
    }
  }

  sendError(req, res, 401, "Unauthorized");
};

function sendError(
  req: Request,
  res: Response,
  status: number,
  message: string,
  data: any = {}
) {
  const isPost = req.method === "POST";
  const hasBody = isPost && req.body;
  const hasModel = hasBody && req.body.model;

  if (!hasModel) {
    return res.status(status).json({ error: message });
  }

  sendErrorToClient({
    req,
    res,
    options: {
      title: `Proxy gatekeeper error (HTTP ${status})`,
      message,
      format: "unknown",
      statusCode: status,
      reqId: req.id,
      obj: data,
    },
  });
}


code:
import { Request, RequestHandler, Router } from "express";
import { config } from "../config";
import { transformAnthropicChatResponseToOpenAI } from "./anthropic";
import { ipLimiter } from "./rate-limit";
import {
  createPreprocessorMiddleware,
  finalizeSignedRequest,
  signGcpRequest,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

const LATEST_GCP_SONNET_MINOR_VERSION = "20240229";

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

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

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

  // https://docs.anthropic.com/en/docs/about-claude/models
  const variants = [
    "claude-3-haiku@20240307",
    "claude-3-5-haiku@20241022",
    "claude-3-sonnet@20240229",
    "claude-3-5-sonnet@20240620",
    "claude-3-5-sonnet-v2@20241022",
    "claude-3-opus@20240229",
  ];

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

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

  let newBody = body;
  switch (`${req.inboundApi}<-${req.outboundApi}`) {
    case "openai<-anthropic-chat":
      req.log.info("Transforming Anthropic Chat back to OpenAI format");
      newBody = transformAnthropicChatResponseToOpenAI(body);
      break;
  }

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

const gcpProxy = createQueuedProxyMiddleware({
  target: ({ signedRequest }) => {
    if (!signedRequest) throw new Error("Must sign request before proxying");
    return `${signedRequest.protocol}//${signedRequest.hostname}`;
  },
  mutations: [signGcpRequest, finalizeSignedRequest],
  blockingResponseHandler: gcpBlockingResponseHandler,
});

const oaiToChatPreprocessor = createPreprocessorMiddleware(
  { inApi: "openai", outApi: "anthropic-chat", service: "gcp" },
  { afterTransform: [maybeReassignModel] }
);

/**
 * Routes an OpenAI prompt to either the legacy Claude text completion endpoint
 * or the new Claude chat completion endpoint, based on the requested model.
 */
const preprocessOpenAICompatRequest: RequestHandler = (req, res, next) => {
  oaiToChatPreprocessor(req, res, next);
};

const gcpRouter = Router();
gcpRouter.get("/v1/models", handleModelRequest);
// Native Anthropic chat completion endpoint.
gcpRouter.post(
  "/v1/messages",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "anthropic-chat", outApi: "anthropic-chat", service: "gcp" },
    { afterTransform: [maybeReassignModel] }
  ),
  gcpProxy
);

// OpenAI-to-GCP Anthropic compatibility endpoint.
gcpRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  preprocessOpenAICompatRequest,
  gcpProxy
);

/**
 * Tries to deal with:
 * - frontends sending GCP model names even when they want to use the OpenAI-
 *   compatible endpoint
 * - frontends sending Anthropic model names that GCP doesn't recognize
 * - frontends sending OpenAI model names because they expect the proxy to
 *   translate them
 *
 * If client sends GCP model ID it will be used verbatim. Otherwise, various
 * strategies are used to try to map a non-GCP model name to GCP model ID.
 */
function maybeReassignModel(req: Request) {
  const model = req.body.model;

  // If it looks like an GCP model, use it as-is
  // if (model.includes("anthropic.claude")) {
  if (model.startsWith("claude-") && model.includes("@")) {
    return;
  }

  // Anthropic model names can look like:
  // - claude-v1
  // - claude-2.1
  // - claude-3-5-sonnet-20240620-v1:0
  const pattern =
    /^(claude-)?(instant-)?(v)?(\d+)([.-](\d{1}))?(-\d+k)?(-sonnet-|-opus-|-haiku-)?(\d*)/i;
  const match = model.match(pattern);

  // If there's no match, fallback to Claude3 Sonnet as it is most likely to be
  // available on GCP.
  if (!match) {
    req.body.model = `claude-3-sonnet@${LATEST_GCP_SONNET_MINOR_VERSION}`;
    return;
  }

  const [_, _cl, instant, _v, major, _sep, minor, _ctx, name, rev] = match;

  // TODO: rework this to function similarly to aws-claude.ts maybeReassignModel
  const ver = minor ? `${major}.${minor}` : major;
  switch (ver) {
    case "3":
    case "3.0":
      if (name.includes("opus")) {
        req.body.model = "claude-3-opus@20240229";
      } else if (name.includes("haiku")) {
        req.body.model = "claude-3-haiku@20240307";
      } else {
        req.body.model = "claude-3-sonnet@20240229";
      }
      return;
    case "3.5":
      switch (name) {
        case "sonnet":
          switch (rev) {
            case "20241022":
            case "latest":
              req.body.model = "claude-3-5-sonnet-v2@20241022";
              return;
            case "20240620":
              req.body.model = "claude-3-5-sonnet@20240620";
              return;
          }
          break;
        case "haiku":
          req.body.model = "claude-3-5-haiku@20241022";
          return;
        case "opus":
          // Add after model ids are announced late 2024
          break;
      }
  }

  // Fallback to Claude3 Sonnet
  req.body.model = `claude-3-sonnet@${LATEST_GCP_SONNET_MINOR_VERSION}`;
  return;
}

export const gcp = gcpRouter;


code:
import { Request, RequestHandler, Router } from "express";
import { v4 } from "uuid";
import { GoogleAIKey, keyPool } from "../shared/key-management";
import { config } from "../config";
import { ipLimiter } from "./rate-limit";
import {
  createPreprocessorMiddleware,
  finalizeSignedRequest,
} from "./middleware/request";
import { ProxyResHandlerWithBody } from "./middleware/response";
import { addGoogleAIKey } from "./middleware/request/mutators/add-google-ai-key";
import { createQueuedProxyMiddleware } from "./middleware/request/proxy-middleware-factory";

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 keys = keyPool
    .list()
    .filter((k) => k.service === "google-ai") as GoogleAIKey[];
  if (keys.length === 0) {
    modelsCache = { object: "list", data: [] };
    modelsCacheTime = new Date().getTime();
    return modelsCache;
  }

  const modelIds = Array.from(
    new Set(keys.map((k) => k.modelIds).flat())
  ).filter((id) => id.startsWith("models/gemini"));
  const models = modelIds.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());
};

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

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

  res.status(200).json({ ...newBody, proxy: body.proxy });
};

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 = createQueuedProxyMiddleware({
  target: ({ signedRequest }) => {
    if (!signedRequest) throw new Error("Must sign request before proxying");
    const { protocol, hostname} = signedRequest;
    return `${protocol}//${hostname}`;
  },
  mutations: [addGoogleAIKey, finalizeSignedRequest],
  blockingResponseHandler: googleAIBlockingResponseHandler,
});

const googleAIRouter = Router();
googleAIRouter.get("/v1/models", handleModelRequest);

// Native Google AI chat completion endpoint
googleAIRouter.post(
  "/v1beta/models/:modelId:(generateContent|streamGenerateContent)",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "google-ai", outApi: "google-ai", service: "google-ai" },
    { beforeTransform: [maybeReassignModel], afterTransform: [setStreamFlag] }
  ),
  googleAIProxy
);

// OpenAI-to-Google AI compatibility endpoint.
googleAIRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware(
    { inApi: "openai", outApi: "google-ai", service: "google-ai" },
    { afterTransform: [maybeReassignModel] }
  ),
  googleAIProxy
);

function setStreamFlag(req: Request) {
  const isStreaming = req.url.includes("streamGenerateContent");
  if (isStreaming) {
    req.body.stream = true;
    req.isStreaming = true;
  } else {
    req.body.stream = false;
    req.isStreaming = false;
  }
}

/**
 * Replaces requests for non-Google AI models with gemini-1.5-pro-latest.
 * Also strips models/ from the beginning of the model IDs.
 **/
function maybeReassignModel(req: Request) {
  // Ensure model is on body as a lot of middleware will expect it.
  const model = req.body.model || req.url.split("/").pop()?.split(":").shift();
  if (!model) {
    throw new Error("You must specify a model with your request.");
  }
  req.body.model = model;

  const requested = model;
  if (requested.startsWith("models/")) {
    req.body.model = requested.slice("models/".length);
  }

  if (requested.includes("gemini")) {
    return;
  }

  req.log.info({ requested }, "Reassigning model to gemini-1.5-pro-latest");
  req.body.model = "gemini-1.5-pro-latest";
}

export const googleAI = googleAIRouter;


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

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";
const GOOGLE_AI_COMPLETION_ENDPOINT = "/v1beta/models";

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,
      GOOGLE_AI_COMPLETION_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 API.`;

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

/**
 * Handles errors thrown during preparation of a proxy request (before it is
 * sent to the upstream API), typically due to validation, quota, or other
 * pre-flight checks. Depending on the error class, this function will send an
 * appropriate error response to the client, streaming it if necessary.
 */
export const classifyErrorAndSend = (
  err: Error,
  req: Request,
  res: Response | Socket
) => {
  if (res instanceof Socket) {
    // We should always have an Express response object here, but http-proxy's
    // ErrorCallback type says it could be just a Socket.
    req.log.error(err, "Caught error while proxying request to target but cannot send error response to client.");
    return res.destroy();
  }
  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":
      const statusCode = (err as HttpError).status;
      return {
        statusCode,
        statusMessage: `HTTP ${statusCode} ${http.STATUS_CODES[statusCode]}`,
        userMessage: `Reverse proxy error: ${err.message}`,
        type: "proxy_http_error",
      };
    case "BadRequestError":
      return {
        statusCode: 400,
        statusMessage: "Bad Request",
        userMessage: `Request is not valid. (${err.message})`,
        type: "proxy_bad_request",
      };
    case "NotFoundError":
      return {
        statusCode: 404,
        statusMessage: "Not Found",
        userMessage: `Requested resource not found. (${err.message})`,
        type: "proxy_not_found",
      };
    case "PaymentRequiredError":
      return {
        statusCode: 402,
        statusMessage: "No Keys Available",
        userMessage: err.message,
        type: "proxy_no_keys_available",
      };
    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":
      // Few possible values:
      // - choices[0].message.content
      // - choices[0].message with no content if model is invoking a tool
      return body.choices?.[0]?.message?.content || "";
    case "mistral-text":
      return body.outputs?.[0]?.text || "";
    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;
      }
      const text = body.candidates[0].content?.parts?.[0]?.text;
      if (!text) {
        req.log.warn(
          { body: JSON.stringify(body) },
          "Received empty Google AI text completion"
        );
        return "";
      }
      return text;
    case "openai-image":
      return body.data?.map((item: any) => item.url).join("\n");
    default:
      assertNever(format);
  }
}

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


code:
import type { Request } from "express";

import { ProxyReqManager } from "./proxy-req-manager";
export {
  createPreprocessorMiddleware,
  createEmbeddingsPreprocessorMiddleware,
} from "./preprocessor-factory";

// Preprocessors (runs before request is queued, usually body transformation/validation)
export { applyQuotaLimits } from "./preprocessors/apply-quota-limits";
export { blockZoomerOrigins } from "./preprocessors/block-zoomer-origins";
export { countPromptTokens } from "./preprocessors/count-prompt-tokens";
export { languageFilter } from "./preprocessors/language-filter";
export { setApiFormat } from "./preprocessors/set-api-format";
export { transformOutboundPayload } from "./preprocessors/transform-outbound-payload";
export { validateContextSize } from "./preprocessors/validate-context-size";
export { validateModelFamily } from "./preprocessors/validate-model-family";
export { validateVision } from "./preprocessors/validate-vision";

// Proxy request mutators (runs every time request is dequeued, before proxying, usually for auth/signing)
export { addKey, addKeyForEmbeddingsRequest } from "./mutators/add-key";
export { addAzureKey } from "./mutators/add-azure-key";
export { finalizeBody } from "./mutators/finalize-body";
export { finalizeSignedRequest } from "./mutators/finalize-signed-request";
export { signAwsRequest } from "./mutators/sign-aws-request";
export { signGcpRequest } from "./mutators/sign-vertex-ai-request";
export { stripHeaders } from "./mutators/strip-headers";

/**
 * Middleware that runs prior to the request being queued or handled by
 * http-proxy-middleware. You will not have access to the proxied
 * request/response objects since they have not yet been sent to the API.
 *
 * User will have been authenticated by the proxy's gatekeeper, but the request
 * won't have been assigned an upstream 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>;

/**
 * Middleware that runs immediately before the request is proxied to the
 * upstream API, after dequeueing the request from the request queue.
 *
 * Because these middleware may be run multiple times per request if a retryable
 * error occurs and the request put back in the queue, they must be idempotent.
 * A change manager is provided to allow the middleware to make changes to the
 * request which can be automatically reverted.
 */
export type ProxyReqMutator = (
  changeManager: ProxyReqManager
) => void | Promise<void>;


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

export const addAzureKey: ProxyReqMutator = async (manager) => {
  const req = manager.request;
  const validAPIs: APIFormat[] = ["openai", "openai-image"];
  const apisValid = [req.outboundApi, req.inboundApi].every((api) =>
    validAPIs.includes(api)
  );
  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}`;
  // TODO: untracked mutation to body, I think this should just be a
  // RequestPreprocessor because we don't need to do it every dequeue.
  req.body.model = model;

  const key = keyPool.get(model, "azure");
  manager.setKey(key);

  // Handles the sole Azure API deviation from the OpenAI spec (that I know of)
  // TODO: this should also probably be a RequestPreprocessor
  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: key.hash, model },
    "Assigned Azure OpenAI key to request"
  );

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

  const operation =
    req.outboundApi === "openai" ? "/chat/completions" : "/images/generations";
  const apiVersion =
    req.outboundApi === "openai" ? "2023-09-01-preview" : "2024-02-15-preview";

  manager.setSignedRequest({
    method: "POST",
    protocol: "https:",
    hostname: `${resourceName}.openai.azure.com`,
    path: `/openai/deployments/${deploymentId}${operation}?api-version=${apiVersion}`,
    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 };
}