code:
import { Router } from "express";
import multer from "multer";
import { z } from "zod";
import { config } from "../../config";
import { HttpError } from "../../shared/errors";
import * as userStore from "../../shared/users/user-store";
import { parseSort, sortBy, paginate } from "../../shared/utils";
import { keyPool } from "../../shared/key-management";
import { MODEL_FAMILIES } from "../../shared/models";
import { getTokenCostUsd, prettyTokens } from "../../shared/stats";
import {
  User,
  UserPartialSchema,
  UserSchema,
  UserTokenCounts,
} from "../../shared/users/schema";

const router = Router();

const upload = multer({
  storage: multer.memoryStorage(),
  fileFilter: (_req, file, cb) => {
    if (file.mimetype !== "application/json") {
      cb(new Error("Invalid file type"));
    } else {
      cb(null, true);
    }
  },
});

router.get("/create-user", (req, res) => {
  const recentUsers = userStore
    .getUsers()
    .sort(sortBy(["createdAt"], false))
    .slice(0, 5);
  res.render("admin_create-user", {
    recentUsers,
    newToken: !!req.query.created,
  });
});

router.post("/create-user", (req, res) => {
  const body = req.body;

  const base = z.object({ type: UserSchema.shape.type.default("normal") });
  const tempUser = base
    .extend({
      temporaryUserDuration: z.coerce
        .number()
        .int()
        .min(1)
        .max(10080 * 4),
    })
    .merge(
      MODEL_FAMILIES.reduce((schema, model) => {
        return schema.extend({
          [`temporaryUserQuota_${model}`]: z.coerce.number().int().min(0),
        });
      }, z.object({}))
    )
    .transform((data: any) => {
      const expiresAt = Date.now() + data.temporaryUserDuration * 60 * 1000;
      const tokenLimits = MODEL_FAMILIES.reduce((limits, model) => {
        limits[model] = data[`temporaryUserQuota_${model}`];
        return limits;
      }, {} as UserTokenCounts);
      return { ...data, expiresAt, tokenLimits };
    });

  const createSchema = body.type === "temporary" ? tempUser : base;
  const result = createSchema.safeParse(body);
  if (!result.success) {
    throw new HttpError(
      400,
      result.error.issues.flatMap((issue) => issue.message).join(", ")
    );
  }

  userStore.createUser({ ...result.data });
  return res.redirect(`/admin/manage/create-user?created=true`);
});

router.get("/view-user/:token", (req, res) => {
  const user = userStore.getUser(req.params.token);
  if (!user) throw new HttpError(404, "User not found");
  res.render("admin_view-user", { user });
});

router.get("/list-users", (req, res) => {
  const sort = parseSort(req.query.sort) || ["sumTokens", "createdAt"];
  const requestedPageSize =
    Number(req.query.perPage) || Number(req.cookies.perPage) || 20;
  const perPage = Math.max(1, Math.min(1000, requestedPageSize));
  const users = userStore
    .getUsers()
    .map((user) => {
      const sums = getSumsForUser(user);
      return { ...user, ...sums };
    })
    .sort(sortBy(sort, false));

  const page = Number(req.query.page) || 1;
  const { items, ...pagination } = paginate(users, page, perPage);

  return res.render("admin_list-users", {
    sort: sort.join(","),
    users: items,
    ...pagination,
  });
});

router.get("/import-users", (_req, res) => {
  res.render("admin_import-users");
});

router.post("/import-users", upload.single("users"), (req, res) => {
  if (!req.file) throw new HttpError(400, "No file uploaded");

  const data = JSON.parse(req.file.buffer.toString());
  const result = z.array(UserPartialSchema).safeParse(data.users);
  if (!result.success) throw new HttpError(400, result.error.toString());

  const upserts = result.data.map((user) => userStore.upsertUser(user));
  req.session.flash = {
    type: "success",
    message: `${upserts.length} users imported`,
  };
  res.redirect("/admin/manage/import-users");
});

router.get("/export-users", (_req, res) => {
  res.render("admin_export-users");
});

router.get("/export-users.json", (_req, res) => {
  const users = userStore.getUsers();
  res.setHeader("Content-Disposition", "attachment; filename=users.json");
  res.setHeader("Content-Type", "application/json");
  res.send(JSON.stringify({ users }, null, 2));
});

router.get("/", (_req, res) => {
  res.render("admin_index");
});

router.post("/edit-user/:token", (req, res) => {
  const result = UserPartialSchema.safeParse({
    ...req.body,
    token: req.params.token,
  });
  if (!result.success) {
    throw new HttpError(
      400,
      result.error.issues.flatMap((issue) => issue.message).join(", ")
    );
  }

  userStore.upsertUser(result.data);
  return res.status(200).json({ success: true });
});

router.post("/reactivate-user/:token", (req, res) => {
  const user = userStore.getUser(req.params.token);
  if (!user) throw new HttpError(404, "User not found");

  userStore.upsertUser({
    token: user.token,
    disabledAt: null,
    disabledReason: null,
  });
  return res.sendStatus(204);
});

router.post("/disable-user/:token", (req, res) => {
  const user = userStore.getUser(req.params.token);
  if (!user) throw new HttpError(404, "User not found");

  userStore.disableUser(req.params.token, req.body.reason);
  return res.sendStatus(204);
});

router.post("/refresh-user-quota", (req, res) => {
  const user = userStore.getUser(req.body.token);
  if (!user) throw new HttpError(404, "User not found");

  userStore.refreshQuota(user.token);
  req.session.flash = {
    type: "success",
    message: "User's quota was refreshed",
  };
  return res.redirect(`/admin/manage/view-user/${user.token}`);
});

router.post("/maintenance", (req, res) => {
  const action = req.body.action;
  let flash = { type: "", message: "" };
  switch (action) {
    case "recheck": {
      keyPool.recheck("openai");
      keyPool.recheck("anthropic");
      const size = keyPool
        .list()
        .filter((k) => k.service !== "google-palm").length;
      flash.type = "success";
      flash.message = `Scheduled recheck of ${size} keys for OpenAI and Anthropic.`;
      break;
    }
    case "resetQuotas": {
      const users = userStore.getUsers();
      users.forEach((user) => userStore.refreshQuota(user.token));
      const { claude, gpt4, turbo } = config.tokenQuota;
      flash.type = "success";
      flash.message = `All users' token quotas reset to ${turbo} (Turbo), ${gpt4} (GPT-4), ${claude} (Claude).`;
      break;
    }
    case "resetCounts": {
      const users = userStore.getUsers();
      users.forEach((user) => userStore.resetUsage(user.token));
      flash.type = "success";
      flash.message = `All users' token usage records reset.`;
      break;
    }
    default: {
      throw new HttpError(400, "Invalid action");
    }
  }

  req.session.flash = flash;

  return res.redirect(`/admin/manage`);
});

router.get("/download-stats", (_req, res) => {
  return res.render("admin_download-stats");
});

router.post("/generate-stats", (req, res) => {
  const body = req.body;

  const valid = z
    .object({
      anon: z.coerce.boolean().optional().default(false),
      sort: z.string().optional().default("prompts"),
      maxUsers: z.coerce
        .number()
        .int()
        .min(5)
        .max(1000)
        .optional()
        .default(1000),
      tableType: z.enum(["code", "markdown"]).optional().default("markdown"),
      format: z
        .string()
        .optional()
        .default("# Stats\n{{header}}\n{{stats}}\n{{time}}"),
    })
    .strict()
    .safeParse(body);

  if (!valid.success) {
    throw new HttpError(
      400,
      valid.error.issues.flatMap((issue) => issue.message).join(", ")
    );
  }

  const { anon, sort, format, maxUsers, tableType } = valid.data;
  const users = userStore.getUsers();

  let totalTokens = 0;
  let totalCost = 0;
  let totalPrompts = 0;
  let totalIps = 0;

  const lines = users
    .map((u) => {
      const sums = getSumsForUser(u);
      totalTokens += sums.sumTokens;
      totalCost += sums.sumCost;
      totalPrompts += u.promptCount;
      totalIps += u.ip.length;

      const getName = (u: User) => {
        const id = `...${u.token.slice(-5)}`;
        const banned = !!u.disabledAt;
        let nick = anon || !u.nickname ? "Anonymous" : u.nickname;

        if (tableType === "markdown") {
          nick = banned ? `~~${nick}~~` : nick;
          return `${nick.slice(0, 18)} | ${id}`;
        } else {
          // Strikethrough doesn't work within code blocks
          const dead = !!u.disabledAt ? "[dead] " : "";
          nick = `${dead}${nick}`;
          return `${nick.slice(0, 18).padEnd(18)} ${id}`.padEnd(27);
        }
      };

      const user = getName(u);
      const prompts = `${u.promptCount} proompts`.padEnd(14);
      const ips = `${u.ip.length} IPs`.padEnd(8);
      const tokens = `${sums.prettyUsage} tokens`.padEnd(30);
      const sortField = sort === "prompts" ? u.promptCount : sums.sumTokens;
      return { user, prompts, ips, tokens, sortField };
    })
    .sort((a, b) => b.sortField - a.sortField)
    .map(({ user, prompts, ips, tokens }, i) => {
      const pos = tableType === "markdown" ? (i + 1 + ".").padEnd(4) : "";
      return `${pos}${user} | ${prompts} | ${ips} | ${tokens}`;
    })
    .slice(0, maxUsers);

  const strTotalPrompts = `${totalPrompts} proompts`;
  const strTotalIps = `${totalIps} IPs`;
  const strTotalTokens = `${prettyTokens(totalTokens)} tokens`;
  const strTotalCost = `US$${totalCost.toFixed(2)} cost`;
  const header = `!!!Note ${users.length} users | ${strTotalPrompts} | ${strTotalIps} | ${strTotalTokens} | ${strTotalCost}`;
  const time = `\n-> *(as of ${new Date().toISOString()})* <-`;

  let table = [];
  table.push(lines.join("\n"));

  if (valid.data.tableType === "markdown") {
    table = ["User||Prompts|IPs|Usage", "---|---|---|---|---", ...table];
  } else {
    table = ["```text", ...table, "```"];
  }

  const result = format
    .replace("{{header}}", header)
    .replace("{{stats}}", table.join("\n"))
    .replace("{{time}}", time);

  res.setHeader(
    "Content-Disposition",
    `attachment; filename=proxy-stats-${new Date().toISOString()}.md`
  );
  res.setHeader("Content-Type", "text/markdown");
  res.send(result);
});

function getSumsForUser(user: User) {
  const sums = MODEL_FAMILIES.reduce(
    (s, model) => {
      const tokens = user.tokenCounts[model] ?? 0;
      s.sumTokens += tokens;
      s.sumCost += getTokenCostUsd(model, tokens);
      return s;
    },
    { sumTokens: 0, sumCost: 0, prettyUsage: "" }
  );
  sums.prettyUsage = `${prettyTokens(sums.sumTokens)} ($${sums.sumCost.toFixed(
    2
  )})`;
  return sums;
}

export { router as usersWebRouter };


code:
<%- include("partials/shared_header", { title: "Create User - OAI Reverse Proxy Admin" }) %>

Create User Token

User token types:

Normal - Standard users.

Special - Exempt from token quotas and MAX_IPS_PER_USER enforcement.

Temporary - Disabled after a specified duration. Quotas never refresh.

<form action="/admin/manage/create-user" method="post">
  <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
  Type
  <select name="type">
    <option value="normal">Normal</option>
    <option value="special">Special</option>
    <option value="temporary">Temporary</option>
  </select>
  <input type="submit" value="Create" />
  <fieldset id="temporaryUserOptions" style="display: none">
    <legend>Temporary User Options</legend>

        Temporary users will be disabled after the specified duration, and their records will be deleted 72 hours after that.
        These options apply only to new
        temporary users; existing ones use whatever options were in effect when they were created.

      Access duration (in minutes)
      <input type="number" name="temporaryUserDuration" id="temporaryUserDuration" value="60" class="full-width" />

      6 hours:360
      12 hours:720
      1 day:1440
      1 week:10080

Token Quotas

Temporary users' quotas are never refreshed.

      <% Object.entries(quota).forEach(function([model, tokens]) { %>
      <label class="quota-label" for="temporaryUserQuota_<%= model %>"><%= model %>
      <input
        type="number"
        name="temporaryUserQuota_<%= model %>"
        id="temporaryUserQuota_<%= model %>"
        value="0"
        data-fieldtype="tokenquota"
        data-default="<%= tokens %>" />
      <% }) %>

  </fieldset>
</form>
<% if (newToken) { %>

Just created <%= recentUsers[0].token %>.

<% } %>

Recent Tokens

  <% recentUsers.forEach(function(user) { %>

<a href="/admin/manage/view-user/<%= user.token %>"><%= user.token %>

  <% }) %>

<%- include("partials/admin-footer") %>

repo_name: FASTER/oai

license: unknown


code:
<%- include("partials/shared_header", { title: "Download Stats - OAI Reverse Proxy Admin" }) %>

Download Stats

  Download usage statistics to a Markdown document. You can paste this into a service like Rentry.org to share it.

Options

  <form id="statsForm" action="/admin/manage/generate-stats" method="post"
    style="display: flex; flex-direction: column;">
    <input id="_csrf" type="hidden" name="_csrf" value="<%= csrfToken %>" />

      Anonymize
      <input id="anon" type="checkbox" name="anon" value="true" />

      Sort
      <select id="sort" name="sort">
        <option value="tokens" selected>By Token Count</option>
        <option value="prompts">By Prompt Count</option>
      </select>

      Max Users
      <input id="maxUsers" type="number" name="maxUsers" value="1000" />

      Table Type
      <select id="tableType" name="tableType">
        <option value="markdown" selected>Markdown Table</option>
        <option value="code">Code Block</option>
      </select>

      Custom Format

{{header}}

{{stats}}

{{time}}

      <textarea id="format" name="format" rows="10" cols="50" placeholder="{{stats}}">
# Stats
{{header}}
{{stats}}
{{time}}
    </textarea>

      <button type="submit">Download</button>
      <button id="copyButton" type="button">Copy to Clipboard</button>

  </form>

<%- include("partials/admin-footer") %>


code:
<%- include("partials/shared_header", { title: "Error" }) %>

⚠ Error <%= status %>: <%= message %>

<%= stack %>

      Go Back | Go Home


code:
<%- include("partials/shared_header", { title: "Export Users - OAI Reverse Proxy Admin" }) %>

Export Users

  Export users to JSON. The JSON will be an array of objects under the key
  users. You can use this JSON to import users later.

<button onclick="exportUsers()">Export</button>
<%- include("partials/admin-footer") %>


code:
<%- include("partials/shared_header", { title: "Import Users - OAI Reverse Proxy Admin" }) %>

Import Users

  Import users from JSON. The JSON should be an array of objects under the key
  users. Each object should have the following fields:

token (required): a unique identifier for the user

nickname (optional): a nickname for the user, max 80 chars

ip (optional): IP addresses the user has connected from

    type (optional): either normal or
    special

    promptCount (optional): the number of times the user has sent a
    prompt

    tokenCounts (optional): the number of tokens the user has
    consumed. This should be an object with keys turbo,
    gpt4, and claude.

    tokenLimits (optional): the number of tokens the user can
    consume. This should be an object with keys turbo,
    gpt4, and claude.

    createdAt (optional): the timestamp when the user was created

    disabledAt (optional): the timestamp when the user was disabled

    disabledReason (optional): the reason the user was disabled

  If a user with the same token already exists, the existing user will be
  updated with the new values.

<form action="/admin/manage/import-users?_csrf=<%= csrfToken %>" method="post" enctype="multipart/form-data">
  <input type="file" name="users" />
  <input type="submit" value="Import" />
</form>
</form>
<%- include("partials/admin-footer") %>


code:
<%- include("partials/shared_header", { title: "OAI Reverse Proxy Admin" }) %>

OAI Reverse Proxy Admin

<% if (!persistenceEnabled) { %>

  ⚠ Users will be lost when the server restarts because persistence is not configured.

Be sure to export your users and import them again after restarting the server if you want to keep them.

  See the

    user management documentation
  to learn how to set up persistence.

<% } %>

Users

List Users

Create User

Import Users

Export Users

Download Rentry Stats

Maintenance

<form id="maintenanceForm" action="/admin/manage/maintenance" method="post">
  <input id="_csrf" type="hidden" name="_csrf" value="<%= csrfToken %>" />
  <input id="hiddenAction" type="hidden" name="action" value="" />

    <fieldset>
      <legend>Key Recheck</legend>
      <button id="recheck-keys" type="button" onclick="submitForm('recheck')">Force Key Recheck</button>
      Triggers a recheck of all keys without restarting the server.
    </fieldset>
    <% if (quotasEnabled) { %>
    <fieldset>
      <legend>Bulk Quota Management</legend>

        <button id="refresh-quotas" type="button" onclick="submitForm('resetQuotas')">Refresh All Quotas</button>
        Resets all users' quotas to the values set in the TOKEN_QUOTA_* environment variables.

        <button id="clear-token-counts" type="button" onclick="submitForm('resetCounts')">Clear All Token Counts</button>
        Resets all users' token records to zero.

    </fieldset>
    <% } %>

</form>

<%- include("partials/admin-footer") %>


code:
<%- include("partials/shared_header", { title: "Users - OAI Reverse Proxy Admin" }) %>

User Token List

    <% if (users.length === 0) { %>

No users found.

    <% } else { %>
    <input type="checkbox" id="toggle-nicknames" onchange="toggleNicknames()" />
    Show Nicknames

          User
          <th <% if (sort.includes("ip")) { %>class="active"<% } %> >IPs
          <th <% if (sort.includes("promptCount")) { %>class="active"<% } %> >Prompts
          <th <% if (sort.includes("sumCost")) { %>class="active"<% } %> >Usage
          Type
          <th <% if (sort.includes("createdAt")) { %>class="active"<% } %> >Created (UTC)
          <th <% if (sort.includes("lastUsedAt")) { %>class="active"<% } %> >Last Used (UTC)
          Banned?

        <% users.forEach(function(user){ %>

            <a href="/admin/manage/view-user/<%= user.token %>">
            <%= user.token %>
            <% if (user.nickname) { %>
            <%= user.nickname %>
            <% } else { %>
            <%= "..." + user.token.slice(-5) %>
            <% } %>

          <%= user.ip.length %>
          <%= user.promptCount %>
          <%= user.prettyUsage %>
          <%= user.type %>
          <%= user.createdAt %>
          <%= user.lastUsedAt ?? "never" %>

            <% if (user.disabledAt) { %>
            <a title="Unban" href="#" class="unban" data-token="<%= user.token %>">🔄
            <% } else { %>
            <a title="Ban" href="#" class="ban" data-token="<%= user.token %>">🚫
            <% } %>
          <%= user.disabledAt ? "Yes" : "No" %> <%= user.disabledReason ? `(${user.disabledReason})` : "" %>

        <% }); %>

      <% if (page > 1) { %>

<a href="/admin/manage/list-users?sort=<%= sort %>&page=<%= page - 1 %>">«

      <% } %> <% for (var i = 1; i <= pageCount; i++) { %>
      <li <% if (i === page) { %>class="active"<% } %>><a href="/admin/manage/list-users?sort=<%= sort %>&page=<%= i %>"><%= i %>

      <% } %> <% if (page < pageCount) { %>

<a href="/admin/manage/list-users?sort=<%= sort %>&page=<%= page + 1 %>">»

      <% } %>

Showing <%= page * pageSize - pageSize + 1 %> to <%= users.length + page * pageSize - pageSize %> of <%= totalCount %> users.

    <%- include("partials/shared_pagination") %>
    <% } %>

<%- include("partials/admin-ban-xhr-script") %>

<%- include("partials/admin-footer") %>


code:
<%- include("partials/shared_header", { title: "Login" }) %>

Login

    <form action="/admin/login" method="post">
      <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
      Admin Key
      <input type="password" name="token" />
      <input type="submit" value="Login" />
    </form>


code:
<%- include("partials/shared_header", { title: "View User - OAI Reverse Proxy Admin" }) %>

View User

      Key
      Value

      Token
      <%- user.token %>

      Nickname
      <%- user.nickname ?? "none" %>

        <a title="Edit" id="edit-nickname" href="#" data-field="nickname" data-token="<%= user.token %>">✏

      Type
      <%- user.type %>

        <a title="Edit" id="edit-type" href="#" data-field="type" data-token="<%= user.token %>">✏

      Prompts
      <%- user.promptCount %>

      Created At
      <%- user.createdAt %>

      Last Used At
      <%- user.lastUsedAt || "never" %>

      Disabled At
      <%- user.disabledAt %>

        <% if (user.disabledAt) { %>
        <a title="Unban" href="#" class="unban" data-token="<%= user.token %>">🔄
        <% } else { %>
        <a title="Ban" href="#" class="ban" data-token="<%= user.token %>">🚫
        <% } %>

      Disabled Reason
      <%- user.disabledReason %>
      <% if (user.disabledAt) { %>

        <a title="Edit" id="edit-disabledReason" href="#" data-field="disabledReason"
          data-token="<%= user.token %>">✏

      <% } %>

      IP Address Limit
      <%- (user.maxIps ?? maxIps) || "Unlimited" %>

        <a title="Edit" id="edit-maxIps" href="#" data-field="maxIps" data-token="<%= user.token %>">✏

      IPs
      <%- include("partials/shared_user_ip_list", { user, shouldRedact: false }) %>

      Admin Note 🔒

      <%- user.adminNote ?? "none" %>

        <a title="Edit" id="edit-adminNote" href="#" data-field="adminNote" data-token="<%= user.token %>">✏

    <% if (user.type === "temporary") { %>

      Expires At
      <%- user.expiresAt %>

    <% } %>

<form style="display:none" id="current-values">
  <input type="hidden" name="token" value="<%- user.token %>" />
  <% ["nickname", "type", "disabledAt", "disabledReason", "maxIps", "adminNote"].forEach(function (key) { %>
  <input type="hidden" name="<%- key %>" value="<%- user[key] %>" />
  <% }); %>
</form>

Quota Information

<% if (quotasEnabled) { %>
<form action="/admin/manage/refresh-user-quota" method="POST">
  <input type="hidden" name="token" value="<%- user.token %>" />
  <input type="hidden" name="_csrf" value="<%- csrfToken %>" />
  <button type="submit" class="btn btn-primary">Refresh Quotas for User</button>
</form>
<% } %> <%- include("partials/shared_quota-info", { quota, user }) %>

Back to User List

<%- include("partials/admin-ban-xhr-script") %> <%- include("partials/admin-footer") %>


code:
<script>
  document.querySelectorAll("td.actions a.ban").forEach(function (a) {
    a.addEventListener("click", function (e) {
      e.preventDefault();
      var token = a.getAttribute("data-token");
      if (confirm("Are you sure you want to ban this user?")) {
        let reason = prompt("Reason for ban:");
        fetch("/admin/manage/disable-user/" + token, {
          method: "POST",
          credentials: "same-origin",
          body: JSON.stringify({ reason, _csrf: document.querySelector("meta[name=csrf-token]").getAttribute("content") }),
          headers: { "Content-Type": "application/json" },
        }).then(() => window.location.reload());
      }
    });
  });

  document.querySelectorAll("td.actions a.unban").forEach(function (a) {
    a.addEventListener("click", function (e) {
      e.preventDefault();
      var token = a.getAttribute("data-token");
      if (confirm("Are you sure you want to unban this user?")) {
        fetch("/admin/manage/reactivate-user/" + token, {
          method: "POST",
          credentials: "same-origin",
          body: JSON.stringify({ _csrf: document.querySelector("meta[name=csrf-token]").getAttribute("content") }),
          headers: { "Content-Type": "application/json" },
        }).then(() => window.location.reload());
      }
    });
  });
</script>


code:
import dotenv from "dotenv";
import type firebase from "firebase-admin";
import pino from "pino";
import type { ModelFamily } from "./shared/models";
dotenv.config();

// Can't import the usual logger here because it itself needs the config.
const startupLogger = pino({ level: "debug" }).child({ module: "startup" });

const isDev = process.env.NODE_ENV !== "production";

type PromptLoggingBackend = "google_sheets";

type Config = {
  /** The port the proxy server will listen on. */
  port: number;
  /** Comma-delimited list of OpenAI API keys. */
  openaiKey?: string;
  /** Comma-delimited list of Anthropic API keys. */
  anthropicKey?: string;
  /** Comma-delimited list of Google PaLM API keys. */
  googlePalmKey?: 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;
  /**
   * 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 per user, after which their token is disabled.
   * 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;
  /** Per-IP limit for requests per minute to OpenAI's completions endpoint. */
  modelRateLimit: 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 disallowed characters should be rejected. */
  rejectDisallowed?: boolean;
  /** Message to return when rejecting requests. */
  rejectMessage?: string;
  /** Verbosity level of diagnostic logging. */
  logLevel: "trace" | "debug" | "info" | "warn" | "error";
  /** Whether prompts and responses should be logged to persistent storage. */
  promptLogging?: boolean;
  /** Which prompt logging backend to use. */
  promptLoggingBackend?: PromptLoggingBackend;
  /** 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;
  /** Desination 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;
};

// 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),
  openaiKey: getEnvWithDefault("OPENAI_KEY", ""),
  anthropicKey: getEnvWithDefault("ANTHROPIC_KEY", ""),
  googlePalmKey: getEnvWithDefault("GOOGLE_PALM_KEY", ""),
  proxyKey: getEnvWithDefault("PROXY_KEY", ""),
  adminKey: getEnvWithDefault("ADMIN_KEY", ""),
  gatekeeper: getEnvWithDefault("GATEKEEPER", "none"),
  gatekeeperStore: getEnvWithDefault("GATEKEEPER_STORE", "memory"),
  maxIpsPerUser: getEnvWithDefault("MAX_IPS_PER_USER", 0),
  firebaseRtdbUrl: getEnvWithDefault("FIREBASE_RTDB_URL", undefined),
  firebaseKey: getEnvWithDefault("FIREBASE_KEY", undefined),
  modelRateLimit: getEnvWithDefault("MODEL_RATE_LIMIT", 4),
  maxContextTokensOpenAI: getEnvWithDefault("MAX_CONTEXT_TOKENS_OPENAI", 0),
  maxContextTokensAnthropic: getEnvWithDefault(
    "MAX_CONTEXT_TOKENS_ANTHROPIC",
    0
  ),
  maxOutputTokensOpenAI: getEnvWithDefault(
    ["MAX_OUTPUT_TOKENS_OPENAI", "MAX_OUTPUT_TOKENS"],
    300
  ),
  maxOutputTokensAnthropic: getEnvWithDefault(
    ["MAX_OUTPUT_TOKENS_ANTHROPIC", "MAX_OUTPUT_TOKENS"],
    400
  ),
  allowedModelFamilies: getEnvWithDefault("ALLOWED_MODEL_FAMILIES", [
    "turbo",
    "gpt4",
    "gpt4-32k",
    "claude",
  ]),
  rejectDisallowed: getEnvWithDefault("REJECT_DISALLOWED", false),
  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),
  promptLogging: getEnvWithDefault("PROMPT_LOGGING", false),
  promptLoggingBackend: getEnvWithDefault("PROMPT_LOGGING_BACKEND", undefined),
  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: {
    turbo: getEnvWithDefault("TOKEN_QUOTA_TURBO", 0),
    gpt4: getEnvWithDefault("TOKEN_QUOTA_GPT4", 0),
    "gpt4-32k": getEnvWithDefault("TOKEN_QUOTA_GPT4_32K", 0),
    claude: getEnvWithDefault("TOKEN_QUOTA_CLAUDE", 0),
    bison: getEnvWithDefault("TOKEN_QUOTA_BISON", 0),
  },
  quotaRefreshPeriod: getEnvWithDefault("QUOTA_REFRESH_PERIOD", undefined),
  allowNicknameChanges: getEnvWithDefault("ALLOW_NICKNAME_CHANGES", true),
} as const;

function generateCookieSecret() {
  if (process.env.COOKIE_SECRET !== undefined) {
    return process.env.COOKIE_SECRET;
  }

  const seed = "" + config.adminKey + config.openaiKey + config.anthropicKey;
  const crypto = require("crypto");
  return crypto.createHash("sha256").update(seed).digest("hex");
}

export const COOKIE_SECRET = generateCookieSecret();

export async function assertConfigIsValid() {
  if (process.env.TURBO_ONLY === "true") {
    startupLogger.warn(
      "TURBO_ONLY is deprecated. Use ALLOWED_MODEL_FAMILIES=turbo instead."
    );
    config.allowedModelFamilies = config.allowedModelFamilies.filter(
      (f) => !f.includes("gpt4")
    );
  }

  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.gatekeeper === "proxy_key" && !config.proxyKey) {
    throw new Error(
      "`proxy_key` gatekeeper mode requires a `PROXY_KEY` to be set."
    );
  }

  if (config.gatekeeper !== "proxy_key" && config.proxyKey) {
    throw new Error(
      "`PROXY_KEY` is set, but gatekeeper mode is not `proxy_key`. Make sure to set `GATEKEEPER=proxy_key`."
    );
  }

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

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

  await maybeInitializeFirebase();
}

/**
 * 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"];

/**
 * 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: (keyof Config)[] = [
  "port",
  "logLevel",
  "openaiKey",
  "anthropicKey",
  "googlePalmKey",
  "proxyKey",
  "adminKey",
  "checkKeys",
  "showTokenCosts",
  "googleSheetsKey",
  "firebaseKey",
  "firebaseRtdbUrl",
  "gatekeeperStore",
  "maxIpsPerUser",
  "blockedOrigins",
  "blockMessage",
  "blockRedirect",
  "allowNicknameChanges",
];

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

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

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

    if (shouldOmit) {
      continue;
    }

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

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

/**
 * 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_PALM_KEY"].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;
  }
}

let firebaseApp: firebase.app.App | undefined;

async function maybeInitializeFirebase() {
  if (!config.gatekeeperStore.startsWith("firebase")) {
    return;
  }

  const firebase = await import("firebase-admin");
  const firebaseKey = Buffer.from(config.firebaseKey!, "base64").toString();
  const app = firebase.initializeApp({
    credential: firebase.credential.cert(JSON.parse(firebaseKey)),
    databaseURL: config.firebaseRtdbUrl,
  });

  await app.database().ref("connection-test").set(Date.now());

  firebaseApp = app;
}

export function getFirebaseApp(): firebase.app.App {
  if (!firebaseApp) {
    throw new Error("Firebase app not initialized.");
  }
  return firebaseApp;
}


code:
import fs from "fs";
import { Request, Response } from "express";
import showdown from "showdown";
import { config, listConfig } from "./config";
import {
  AnthropicKey,
  GooglePalmKey,
  OpenAIKey,
  keyPool,
} from "./shared/key-management";
import { ModelFamily, OpenAIModelFamily } from "./shared/models";
import { getUniqueIps } from "./proxy/rate-limit";
import { getEstimatedWaitTime, getQueueLength } from "./proxy/queue";
import { getTokenCostUsd, prettyTokens } from "./shared/stats";
import { assertNever } from "./shared/utils";

const INFO_PAGE_TTL = 2000;
let infoPageHtml: string | undefined;
let infoPageLastUpdated = 0;

type KeyPoolKey = ReturnType<typeof keyPool.list>[0];
const keyIsOpenAIKey = (k: KeyPoolKey): k is OpenAIKey =>
  k.service === "openai";
const keyIsAnthropicKey = (k: KeyPoolKey): k is AnthropicKey =>
  k.service === "anthropic";
const keyIsGooglePalmKey = (k: KeyPoolKey): k is GooglePalmKey =>
  k.service === "google-palm";

type ModelAggregates = {
  active: number;
  trial?: number;
  revoked?: number;
  overQuota?: number;
  pozzed?: number;
  queued: number;
  queueTime: string;
  tokens: number;
};
type ModelAggregateKey = `${ModelFamily}__${keyof ModelAggregates}`;
type ServiceAggregates = {
  status?: string;
  openaiKeys?: number;
  openaiOrgs?: number;
  anthropicKeys?: number;
  palmKeys?: number;
  proompts: number;
  tokens: number;
  tokenCost: number;
  openAiUncheckedKeys?: number;
  anthropicUncheckedKeys?: number;
} & {
  [modelFamily in ModelFamily]?: ModelAggregates;
};

const modelStats = new Map<ModelAggregateKey, number>();
const serviceStats = new Map<keyof ServiceAggregates, number>();

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

  // Sometimes huggingface doesn't send the host header and makes us guess.
  const baseUrl =
    process.env.SPACE_ID && !req.get("host")?.includes("hf.space")
      ? getExternalUrlForHuggingfaceSpaceId(process.env.SPACE_ID)
      : req.protocol + "://" + req.get("host");

  res.send(cacheInfoPageHtml(baseUrl));
};

function getCostString(cost: number) {
  if (!config.showTokenCosts) return "";
  return ` ($${cost.toFixed(2)})`;
}

function cacheInfoPageHtml(baseUrl: string) {
  const keys = keyPool.list();

  modelStats.clear();
  serviceStats.clear();
  keys.forEach(addKeyToAggregates);

  const openaiKeys = serviceStats.get("openaiKeys") || 0;
  const anthropicKeys = serviceStats.get("anthropicKeys") || 0;
  const palmKeys = serviceStats.get("palmKeys") || 0;
  const proompts = serviceStats.get("proompts") || 0;
  const tokens = serviceStats.get("tokens") || 0;
  const tokenCost = serviceStats.get("tokenCost") || 0;

  const info = {
    uptime: Math.floor(process.uptime()),
    endpoints: {
      ...(openaiKeys ? { openai: baseUrl + "/proxy/openai" } : {}),
      ...(openaiKeys
        ? { ["openai2"]: baseUrl + "/proxy/openai/turbo-instruct" }
        : {}),
      ...(anthropicKeys ? { anthropic: baseUrl + "/proxy/anthropic" } : {}),
      ...(palmKeys ? { "google-palm": baseUrl + "/proxy/google-palm" } : {}),
    },
    proompts,
    tookens: `${prettyTokens(tokens)}${getCostString(tokenCost)}`,
    ...(config.modelRateLimit ? { proomptersNow: getUniqueIps() } : {}),
    openaiKeys,
    anthropicKeys,
    palmKeys,
    ...(openaiKeys ? getOpenAIInfo() : {}),
    ...(anthropicKeys ? getAnthropicInfo() : {}),
    ...(palmKeys ? { "palm-bison": getPalmInfo() } : {}),
    config: listConfig(),
    build: process.env.BUILD_INFO || "dev",
  };

  const title = getServerTitle();
  const headerHtml = buildInfoPageHeader(new showdown.Converter(), title);

  const pageBody = `<!DOCTYPE html>

    <meta charset="utf-8" />
    <meta name="robots" content="noindex" />
    <title>${title}</title>

    ${headerHtml}

Service Info

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

    ${getSelfServiceLinks()}

`;

  infoPageHtml = pageBody;
  infoPageLastUpdated = Date.now();

  return pageBody;
}

function getUniqueOpenAIOrgs(keys: KeyPoolKey[]) {
  const orgIds = new Set(
    keys.filter((k) => k.service === "openai").map((k: any) => k.organizationId)
  );
  return orgIds.size;
}

function increment<T extends keyof ServiceAggregates | ModelAggregateKey>(
  map: Map<T, number>,
  key: T,
  delta = 1
) {
  map.set(key, (map.get(key) || 0) + delta);
}

function addKeyToAggregates(k: KeyPoolKey) {
  increment(serviceStats, "proompts", k.promptCount);
  increment(serviceStats, "openaiKeys", k.service === "openai" ? 1 : 0);
  increment(serviceStats, "anthropicKeys", k.service === "anthropic" ? 1 : 0);
  increment(serviceStats, "palmKeys", k.service === "google-palm" ? 1 : 0);

  let sumTokens = 0;
  let sumCost = 0;
  let family: ModelFamily;
  const families = k.modelFamilies.filter((f) =>
    config.allowedModelFamilies.includes(f)
  );

  switch (k.service) {
    case "openai":
    case "openai-text":
      if (!keyIsOpenAIKey(k)) throw new Error("Invalid key type");
      increment(
        serviceStats,
        "openAiUncheckedKeys",
        Boolean(k.lastChecked) ? 0 : 1
      );

      // Technically this would not account for keys that have tokens recorded
      // on models they aren't provisioned for, but that would be strange
      k.modelFamilies.forEach((f) => {
        const tokens = k[`${f}Tokens`];
        sumTokens += tokens;
        sumCost += getTokenCostUsd(f, tokens);
        increment(modelStats, `${f}__tokens`, tokens);
      });

      if (families.includes("gpt4-32k")) {
        family = "gpt4-32k";
      } else if (families.includes("gpt4")) {
        family = "gpt4";
      } else {
        family = "turbo";
      }
      break;
    case "anthropic":
      if (!keyIsAnthropicKey(k)) throw new Error("Invalid key type");
      family = "claude";
      sumTokens += k.claudeTokens;
      sumCost += getTokenCostUsd(family, k.claudeTokens);
      increment(modelStats, `${family}__tokens`, k.claudeTokens);
      increment(modelStats, `${family}__pozzed`, k.isPozzed ? 1 : 0);
      increment(
        serviceStats,
        "anthropicUncheckedKeys",
        Boolean(k.lastChecked) ? 0 : 1
      );
      break;
    case "google-palm":
      if (!keyIsGooglePalmKey(k)) throw new Error("Invalid key type");
      family = "bison";
      sumTokens += k.bisonTokens;
      sumCost += getTokenCostUsd(family, k.bisonTokens);
      increment(modelStats, `${family}__tokens`, k.bisonTokens);
      break;
    default:
      assertNever(k.service);
  }

  increment(serviceStats, "tokens", sumTokens);
  increment(serviceStats, "tokenCost", sumCost);
  increment(modelStats, `${family}__active`, k.isDisabled ? 0 : 1);
  increment(modelStats, `${family}__trial`, k.isTrial ? 1 : 0);
  if ("isRevoked" in k) {
    increment(modelStats, `${family}__revoked`, k.isRevoked ? 1 : 0);
  }
  if ("isOverQuota" in k) {
    increment(modelStats, `${family}__overQuota`, k.isOverQuota ? 1 : 0);
  }
}

function getOpenAIInfo() {
  const info: { status?: string; openaiKeys?: number; openaiOrgs?: number } & {
    [modelFamily in OpenAIModelFamily]?: {
      usage?: string;
      activeKeys: number;
      trialKeys?: number;
      revokedKeys?: number;
      overQuotaKeys?: number;
      proomptersInQueue?: number;
      estimatedQueueTime?: string;
    };
  } = {};

  const allowedFamilies = new Set(config.allowedModelFamilies);
  let families = new Set<OpenAIModelFamily>();
  const keys = keyPool.list().filter((k) => {
    const isOpenAI = keyIsOpenAIKey(k);
    if (isOpenAI) k.modelFamilies.forEach((f) => families.add(f));
    return isOpenAI;
  }) as Omit<OpenAIKey, "key">[];
  families = new Set([...families].filter((f) => allowedFamilies.has(f)));

  if (config.checkKeys) {
    const unchecked = serviceStats.get("openAiUncheckedKeys") || 0;
    if (unchecked > 0) {
      info.status = `Checking ${unchecked} keys...`;
    }
    info.openaiKeys = keys.length;
    info.openaiOrgs = getUniqueOpenAIOrgs(keys);

    families.forEach((f) => {
      const tokens = modelStats.get(`${f}__tokens`) || 0;
      const cost = getTokenCostUsd(f, tokens);

      info[f] = {
        usage: `${prettyTokens(tokens)} tokens${getCostString(cost)}`,
        activeKeys: modelStats.get(`${f}__active`) || 0,
        trialKeys: modelStats.get(`${f}__trial`) || 0,
        revokedKeys: modelStats.get(`${f}__revoked`) || 0,
        overQuotaKeys: modelStats.get(`${f}__overQuota`) || 0,
      };
    });
  } else {
    info.status = "Key checking is disabled.";
    info.turbo = { activeKeys: keys.filter((k) => !k.isDisabled).length };
    info.gpt4 = {
      activeKeys: keys.filter(
        (k) => !k.isDisabled && k.modelFamilies.includes("gpt4")
      ).length,
    };
  }

  families.forEach((f) => {
    if (info[f]) {
      const { estimatedQueueTime, proomptersInQueue } = getQueueInformation(f);
      info[f]!.proomptersInQueue = proomptersInQueue;
      info[f]!.estimatedQueueTime = estimatedQueueTime;
    }
  });

  return info;
}

function getAnthropicInfo() {
  const claudeInfo: Partial<ModelAggregates> = {
    active: modelStats.get("claude__active") || 0,
    pozzed: modelStats.get("claude__pozzed") || 0,
  };

  const queue = getQueueInformation("claude");
  claudeInfo.queued = queue.proomptersInQueue;
  claudeInfo.queueTime = queue.estimatedQueueTime;

  const tokens = modelStats.get("claude__tokens") || 0;
  const cost = getTokenCostUsd("claude", tokens);

  const unchecked =
    (config.checkKeys && serviceStats.get("anthropicUncheckedKeys")) || 0;

  return {
    claude: {
      usage: `${prettyTokens(tokens)} tokens${getCostString(cost)}`,
      ...(unchecked > 0 ? { status: `Checking ${unchecked} keys...` } : {}),
      activeKeys: claudeInfo.active,
      ...(config.checkKeys ? { pozzedKeys: claudeInfo.pozzed } : {}),
      proomptersInQueue: claudeInfo.queued,
      estimatedQueueTime: claudeInfo.queueTime,
    },
  };
}

function getPalmInfo() {
  const bisonInfo: Partial<ModelAggregates> = {
    active: modelStats.get("bison__active") || 0,
  };

  const queue = getQueueInformation("bison");
  bisonInfo.queued = queue.proomptersInQueue;
  bisonInfo.queueTime = queue.estimatedQueueTime;

  const tokens = modelStats.get("bison__tokens") || 0;
  const cost = getTokenCostUsd("bison", tokens);

  return {
    usage: `${prettyTokens(tokens)} tokens${getCostString(cost)}`,
    activeKeys: bisonInfo.active,
    proomptersInQueue: bisonInfo.queued,
    estimatedQueueTime: bisonInfo.queueTime,
  };
}

const customGreeting = fs.existsSync("greeting.md")
  ? fs.readFileSync("greeting.md", "utf8")
  : null;

/**
 * If the server operator provides a `greeting.md` file, it will be included in
 * the rendered info page.
 **/
function buildInfoPageHeader(converter: showdown.Converter, title: string) {
  // TODO: use some templating engine instead of this mess
  let infoBody = `
# ${title}`;
  if (config.promptLogging) {
    infoBody += `\n## Prompt logging is enabled!
The server operator has enabled prompt logging. The prompts you send to this proxy and the AI responses you receive may be saved.

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/prompt-logging/index.ts).

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

  const waits: string[] = [];
  infoBody += `\n## Estimated Wait Times\nIf the AI is busy, your prompt will processed when a slot frees up.`;

  if (config.openaiKey) {
    // TODO: un-fuck this
    const keys = keyPool.list().filter((k) => k.service === "openai");

    const turboWait = getQueueInformation("turbo").estimatedQueueTime;
    waits.push(`**Turbo:** ${turboWait}`);

    const gpt4Wait = getQueueInformation("gpt4").estimatedQueueTime;
    const hasGpt4 = keys.some((k) => k.modelFamilies.includes("gpt4"));
    const allowedGpt4 = config.allowedModelFamilies.includes("gpt4");
    if (hasGpt4 && allowedGpt4) {
      waits.push(`**GPT-4:** ${gpt4Wait}`);
    }

    const gpt432kWait = getQueueInformation("gpt4-32k").estimatedQueueTime;
    const hasGpt432k = keys.some((k) => k.modelFamilies.includes("gpt4-32k"));
    const allowedGpt432k = config.allowedModelFamilies.includes("gpt4-32k");
    if (hasGpt432k && allowedGpt432k) {
      waits.push(`**GPT-4-32k:** ${gpt432kWait}`);
    }
  }

  if (config.anthropicKey) {
    const claudeWait = getQueueInformation("claude").estimatedQueueTime;
    waits.push(`**Claude:** ${claudeWait}`);
  }
  infoBody += "\n\n" + waits.join(" / ");

  if (customGreeting) {
    infoBody += `\n## Server Greeting\n${customGreeting}`;
  }
  return converter.makeHtml(infoBody);
}

function getSelfServiceLinks() {
  if (config.gatekeeper !== "user_token") return "";
  return `<footer style="font-size: 0.8em;"><hr /><a target="_blank" href="/user/lookup">Check your user token info</a></footer>`;
}

/** Returns queue time in seconds, or minutes + seconds if over 60 seconds. */
function getQueueInformation(partition: ModelFamily) {
  const waitMs = getEstimatedWaitTime(partition);
  const waitTime =
    waitMs < 60000
      ? `${Math.round(waitMs / 1000)}sec`
      : `${Math.round(waitMs / 60000)}min, ${Math.round(
          (waitMs % 60000) / 1000
        )}sec`;
  return {
    proomptersInQueue: getQueueLength(partition),
    estimatedQueueTime: waitMs > 2000 ? waitTime : "no wait",
  };
}

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 getExternalUrlForHuggingfaceSpaceId(spaceId: string) {
  // Huggingface broke their amazon elb config and no longer sends the
  // x-forwarded-host header. This is a workaround.
  try {
    const [username, spacename] = spaceId.split("/");
    return `https://${username}-${spacename.replace(/_/g, "-")}.hf.space`;
  } catch (e) {
    return "";
  }
}


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

export const logger = pino({
  level: config.logLevel,
});


code:
import { Request, RequestHandler, Router } from "express";
import * as http from "http";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { logger } from "../logger";
import { createQueueMiddleware } from "./queue";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addKey,
  applyQuotaLimits,
  addAnthropicPreamble,
  blockZoomerOrigins,
  createPreprocessorMiddleware,
  finalizeBody,
  languageFilter,
  removeOriginHeaders,
} from "./middleware/request";
import {
  ProxyResHandlerWithBody,
  createOnProxyResHandler,
} from "./middleware/response";

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 is 100k by default it seems
    "claude-2.0",
  ];

  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 rewriteAnthropicRequest = (
  proxyReq: http.ClientRequest,
  req: Request,
  res: http.ServerResponse
) => {
  const rewriterPipeline = [
    applyQuotaLimits,
    addKey,
    addAnthropicPreamble,
    languageFilter,
    blockZoomerOrigins,
    removeOriginHeaders,
    finalizeBody,
  ];

  try {
    for (const rewriter of rewriterPipeline) {
      rewriter(proxyReq, req, res, {});
    }
  } catch (error) {
    req.log.error(error, "Error while executing proxy rewriter");
    proxyReq.destroy(error as Error);
  }
};

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

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

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

  // TODO: Remove once tokenization is stable
  if (req.debug) {
    body.proxy_tokenizer_debug_info = req.debug;
  }

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

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

const anthropicProxy = createQueueMiddleware(
  createProxyMiddleware({
    target: "https://api.anthropic.com",
    changeOrigin: true,
    on: {
      proxyReq: rewriteAnthropicRequest,
      proxyRes: createOnProxyResHandler([anthropicResponseHandler]),
      error: handleProxyError,
    },
    selfHandleResponse: true,
    logger,
    pathRewrite: {
      // Send OpenAI-compat requests to the real Anthropic endpoint.
      "^/v1/chat/completions": "/v1/complete",
    },
  })
);

const anthropicRouter = Router();
// Fix paths because clients don't consistently use the /v1 prefix.
anthropicRouter.use((req, _res, next) => {
  if (!req.path.startsWith("/v1/")) {
    req.url = `/v1${req.url}`;
  }
  next();
});
anthropicRouter.get("/v1/models", handleModelRequest);
anthropicRouter.post(
  "/v1/complete",
  ipLimiter,
  createPreprocessorMiddleware({ inApi: "anthropic", outApi: "anthropic" }),
  anthropicProxy
);
// OpenAI-to-Anthropic compatibility endpoint.
anthropicRouter.post(
  "/v1/chat/completions",
  ipLimiter,
  createPreprocessorMiddleware({ inApi: "openai", outApi: "anthropic" }),
  anthropicProxy
);
// Redirect browser requests to the homepage.
anthropicRouter.get("*", (req, res, next) => {
  const isBrowser = req.headers["user-agent"]?.includes("Mozilla");
  if (isBrowser) {
    res.redirect("/");
  } else {
    next();
  }
});

export const anthropic = anthropicRouter;


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: Uint8Array; 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) {
      req.log.warn(
        { token: header, data },
        "Invalid RisuAI token; using IP instead"
      );
    } else {
      req.log.info("RisuAI token validated");
      req.risuToken = header;
    }
  } 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, data: "[unparseable]" };
  }
  const data: RisuToken = tk.data;
  const sig = Buffer.from(tk.sig, "base64");

  if (data.expiresIn < Math.floor(Date.now() / 1000)) {
    return { valid: false };
  }

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

  return { valid, data };
}


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

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

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

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

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

  return undefined;
}

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

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

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

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

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

  if (GATEKEEPER === "user_token" && token) {
    const user = authenticate(token, req.ip);
    if (user) {
      req.user = user;
      return next();
    } else {
      const maybeBannedUser = getUser(token);
      if (maybeBannedUser?.disabledAt) {
        return res.status(403).json({
          error: `Forbidden: ${
            maybeBannedUser.disabledReason || "Token disabled"
          }`,
        });
      }
    }
  }

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


code:
/* Pretends to be a KoboldAI API endpoint and translates incoming Kobold
requests to OpenAI API equivalents. */

import { Request, Response, Router } from "express";
import http from "http";
import { createProxyMiddleware } from "http-proxy-middleware";
import { config } from "../config";
import { logger } from "../logger";
import { ipLimiter } from "./rate-limit";
import { handleProxyError } from "./middleware/common";
import {
  addKey,
  createPreprocessorMiddleware,
  finalizeBody,
  languageFilter,
  transformKoboldPayload,
} from "./middleware/request";
import {
  createOnProxyResHandler,
  ProxyResHandlerWithBody,
} from "./middleware/response";

export const handleModelRequest = (_req: Request, res: Response) => {
  res.status(200).json({ result: "Connected to OpenAI reverse proxy" });
};

export const handleSoftPromptsRequest = (_req: Request, res: Response) => {
  res.status(200).json({ soft_prompts_list: [] });
};

const rewriteRequest = (
  proxyReq: http.ClientRequest,
  req: Request,
  res: Response
) => {
  req.body.stream = false;
  const rewriterPipeline = [
    addKey,
    transformKoboldPayload,
    languageFilter,
    finalizeBody,
  ];

  try {
    for (const rewriter of rewriterPipeline) {
      rewriter(proxyReq, req, res, {});
    }
  } catch (error) {
    logger.error(error, "Error while executing proxy rewriter");
    proxyReq.destroy(error as Error);
  }
};

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

  const koboldResponse = {
    results: [{ text: body.choices[0].message.content }],
    model: body.model,
    ...(config.promptLogging && {
      proxy_note: `Prompt logging is enabled on this proxy instance. See ${req.get(
        "host"
      )} for more information.`,
    }),
  };

  res.send(JSON.stringify(koboldResponse));
};

const koboldOaiProxy = createProxyMiddleware({
  target: "https://api.openai.com",
  changeOrigin: true,
  pathRewrite: {
    "^/api/v1/generate": "/v1/chat/completions",
  },
  on: {
    proxyReq: rewriteRequest,
    proxyRes: createOnProxyResHandler([koboldResponseHandler]),
    error: handleProxyError,
  },
  selfHandleResponse: true,
  logger,
});

const koboldRouter = Router();
koboldRouter.use((req, res) => {
  logger.warn(`Unhandled kobold request: ${req.method} ${req.path}`);
  res.status(404).json({ error: "Not found" });
});

export const kobold = koboldRouter;


code:
import { Request, Response } from "express";
import httpProxy from "http-proxy";
import { ZodError } from "zod";
import { APIFormat } from "../../shared/key-management";
import { assertNever } from "../../shared/utils";
import { QuotaExceededError } from "./request/apply-quota-limits";

const OPENAI_CHAT_COMPLETION_ENDPOINT = "/v1/chat/completions";
const OPENAI_TEXT_COMPLETION_ENDPOINT = "/v1/completions";
const ANTHROPIC_COMPLETION_ENDPOINT = "/v1/complete";

/** Returns true if we're making a request to a completion endpoint. */
export function isCompletionRequest(req: Request) {
  // 99% sure this function is not needed anymore
  return (
    req.method === "POST" &&
    [
      OPENAI_CHAT_COMPLETION_ENDPOINT,
      OPENAI_TEXT_COMPLETION_ENDPOINT,
      ANTHROPIC_COMPLETION_ENDPOINT,
    ].some((endpoint) => req.path.startsWith(endpoint))
  );
}

export function writeErrorResponse(
  req: Request,
  res: Response,
  statusCode: number,
  errorPayload: Record<string, any>
) {
  const errorSource = errorPayload.error?.type?.startsWith("proxy")
    ? "proxy"
    : "upstream";

  // If we're mid-SSE stream, send a data event with the error payload and end
  // the stream. Otherwise just send a normal error response.
  if (
    res.headersSent ||
    res.getHeader("content-type") === "text/event-stream"
  ) {
    const errorContent =
      statusCode === 403
        ? JSON.stringify(errorPayload)
        : JSON.stringify(errorPayload, null, 2);

    const msg = buildFakeSseMessage(
      `${errorSource} error (${statusCode})`,
      errorContent,
      req
    );
    res.write(msg);
    res.write(`data: [DONE]\n\n`);
    res.end();
  } else {
    if (req.debug) {
      errorPayload.error.proxy_tokenizer_debug_info = req.debug;
    }
    res.status(statusCode).json(errorPayload);
  }
}

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

export const handleInternalError = (
  err: Error,
  req: Request,
  res: Response
) => {
  try {
    if (err instanceof ZodError) {
      writeErrorResponse(req, res, 400, {
        error: {
          type: "proxy_validation_error",
          proxy_note: `Reverse proxy couldn't validate your request when trying to transform it. Your client may be sending invalid data.`,
          issues: err.issues,
          stack: err.stack,
          message: err.message,
        },
      });
    } else if (err.name === "ForbiddenError") {
      // Spoofs a vaguely threatening OpenAI error message. Only invoked by the
      // block-zoomers rewriter to scare off tiktokers.
      writeErrorResponse(req, res, 403, {
        error: {
          type: "organization_account_disabled",
          code: "policy_violation",
          param: null,
          message: err.message,
        },
      });
    } else if (err instanceof QuotaExceededError) {
      writeErrorResponse(req, res, 429, {
        error: {
          type: "proxy_quota_exceeded",
          code: "quota_exceeded",
          message: `You've exceeded your token quota for this model type.`,
          info: err.quotaInfo,
          stack: err.stack,
        },
      });
    } else {
      writeErrorResponse(req, res, 500, {
        error: {
          type: "proxy_internal_error",
          proxy_note: `Reverse proxy encountered an error before it could reach the upstream API.`,
          message: err.message,
          stack: err.stack,
        },
      });
    }
  } catch (e) {
    req.log.error(
      { error: e },
      `Error writing error response headers, giving up.`
    );
  }
};

export function buildFakeSseMessage(
  type: string,
  string: string,
  req: Request
) {
  let fakeEvent;
  const useBackticks = !type.includes("403");
  const msgContent = useBackticks
    ? `\`\`\`\n[${type}: ${string}]\n\`\`\`\n`
    : `[${type}: ${string}]`;

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

export function getCompletionForService({
  service,
  body,
  req,
}: {
  service: APIFormat;
  body: Record<string, any>;
  req?: Request;
}): { completion: string; model: string } {
  switch (service) {
    case "openai":
      return { completion: body.choices[0].message.content, model: body.model };
    case "openai-text":
      return { completion: body.choices[0].text, model: body.model };
    case "anthropic":
      return { completion: body.completion.trim(), model: body.model };
    case "google-palm":
      return { completion: body.candidates[0].output, model: req?.body.model };
    default:
      assertNever(service);
  }
}


code:
import { AnthropicKey, Key } from "../../../shared/key-management";
import { isCompletionRequest } from "../common";
import { ProxyRequestMiddleware } from ".";

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

  assertAnthropicKey(req.key);

  if (req.key.requiresPreamble) {
    let prompt = req.body.prompt;
    const preamble = prompt.startsWith("\n\nHuman:") ? "" : "\n\nHuman:";
    req.log.debug({ key: req.key.hash, preamble }, "Prompt requres preamble");
    prompt = preamble + prompt;

    // Adds `Assistant:` to the end of the prompt if the turn closest to the
    // end is from the `Human:` persona.
    const humanIndex = prompt.lastIndexOf("\n\nHuman:");
    const assistantIndex = prompt.lastIndexOf("\n\nAssistant:");
    const shouldAddAssistant = humanIndex > assistantIndex;
    req.log.debug(
      {
        key: req.key.hash,
        shouldAdd: shouldAddAssistant,
        hIndex: humanIndex,
        aIndex: assistantIndex,
      },
      "Possibly adding Assistant: to prompt"
    );
    if (shouldAddAssistant) prompt += "\n\nAssistant:";
    req.body.prompt = prompt;
  }
};

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


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

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

  if (!isCompletionRequest(req)) {
    // Horrible, horrible hack to stop the proxy from complaining about clients
    // not sending a model when they are requesting the list of models (which
    // requires a key, but obviously not a model).
    // TODO: shouldn't even proxy /models to the upstream API, just fake it
    // using the models our key pool has available.
    req.body.model = "gpt-3.5-turbo";
  }

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

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

  // TODO: use separate middleware to deal with stream flags
  req.isStreaming = req.body.stream === true || req.body.stream === "true";
  req.body.stream = req.isStreaming;

  if (req.inboundApi === req.outboundApi) {
    assignedKey = keyPool.get(req.body.model);
  } else {
    switch (req.outboundApi) {
      // If we are translating between API formats we may need to select a model
      // for the user, because the provided model is for the inbound API.
      case "anthropic":
        assignedKey = keyPool.get("claude-v1");
        break;
      case "google-palm":
        assignedKey = keyPool.get("text-bison-001");
        delete req.body.stream;
        break;
      case "openai-text":
        assignedKey = keyPool.get("gpt-3.5-turbo-instruct");
        break;
      case "openai":
        throw new Error(
          "OpenAI Chat as an API translation target is not supported"
        );
      default:
        assertNever(req.outboundApi);
    }
  }

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

  // TODO: KeyProvider should assemble all necessary headers
  switch (assignedKey.service) {
    case "anthropic":
      proxyReq.setHeader("X-API-Key", assignedKey.key);
      break;
    case "openai":
    case "openai-text":
      const key: OpenAIKey = assignedKey as OpenAIKey;
      if (key.organizationId) {
        proxyReq.setHeader("OpenAI-Organization", key.organizationId);
      }
      proxyReq.setHeader("Authorization", `Bearer ${assignedKey.key}`);
      break;
    case "google-palm":
      const originalPath = proxyReq.path;
      proxyReq.path = originalPath.replace(
        /(\?.*)?$/,
        `?key=${assignedKey.key}`
      );
      break;
    default:
      assertNever(assignedKey.service);
  }
};


code:
import { hasAvailableQuota } from "../../../shared/users/user-store";
import { isCompletionRequest } from "../common";
import { ProxyRequestMiddleware } from ".";

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

export const applyQuotaLimits: ProxyRequestMiddleware = (_proxyReq, req) => {
  if (!isCompletionRequest(req) || !req.user) {
    return;
  }

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


code:
import { isCompletionRequest } from "../common";
import { ProxyRequestMiddleware } from ".";

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

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

/**
 * Blocks requests from Janitor AI users with a fake, scary error message so I
 * stop getting emails asking for tech support.
 */
export const blockZoomerOrigins: ProxyRequestMiddleware = (_proxyReq, req) => {
  if (!isCompletionRequest(req)) {
    return;
  }

  const origin = req.headers.origin || req.headers.referer;
  if (origin && DISALLOWED_ORIGIN_SUBSTRINGS.some((s) => origin.includes(s))) {
    // Venus-derivatives send a test prompt to check if the proxy is working.
    // We don't want to block that just yet.
    if (req.body.messages[0]?.content === "Just say TEST") {
      return;
    }

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


code:
import { Request } from "express";
import { z } from "zod";
import { config } from "../../../config";
import { OpenAIPromptMessage, countTokens } from "../../../shared/tokenization";
import { RequestPreprocessor } from ".";
import { assertNever } from "../../../shared/utils";

const CLAUDE_MAX_CONTEXT = config.maxContextTokensAnthropic;
const OPENAI_MAX_CONTEXT = config.maxContextTokensOpenAI;
const BISON_MAX_CONTEXT = 8100;

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

  switch (service) {
    case "openai": {
      req.outputTokens = req.body.max_tokens;
      const prompt: OpenAIPromptMessage[] = req.body.messages;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "openai-text": {
      req.outputTokens = req.body.max_tokens;
      const prompt: string = req.body.prompt;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "anthropic": {
      req.outputTokens = req.body.max_tokens_to_sample;
      const prompt: string = req.body.prompt;
      result = await countTokens({ req, prompt, service });
      break;
    }
    case "google-palm": {
      req.outputTokens = req.body.maxOutputTokens;
      const prompt: string = req.body.prompt.text;
      result = await countTokens({ req, prompt, service });
      break;
    }
    default:
      assertNever(service);
  }

  req.promptTokens = result.token_count;

  // TODO: Remove once token counting is stable
  req.log.debug({ result: result }, "Counted prompt tokens.");
  req.debug = req.debug ?? {};
  req.debug = { ...req.debug, ...result };

  maybeTranslateOpenAIModel(req);
  validateContextSize(req);
};

function validateContextSize(req: Request) {
  assertRequestHasTokenCounts(req);
  const promptTokens = req.promptTokens;
  const outputTokens = req.outputTokens;
  const contextTokens = promptTokens + outputTokens;
  const model = req.body.model;

  let proxyMax: number;
  switch (req.outboundApi) {
    case "openai":
    case "openai-text":
      proxyMax = OPENAI_MAX_CONTEXT;
      break;
    case "anthropic":
      proxyMax = CLAUDE_MAX_CONTEXT;
      break;
    case "google-palm":
      proxyMax = BISON_MAX_CONTEXT;
      break;
    default:
      assertNever(req.outboundApi);
  }
  proxyMax ||= Number.MAX_SAFE_INTEGER;

  let modelMax = 0;
  if (model.match(/gpt-3.5-turbo-16k/)) {
    modelMax = 16384;
  } else if (model.match(/gpt-3.5-turbo/)) {
    modelMax = 4096;
  } else if (model.match(/gpt-4-32k/)) {
    modelMax = 32768;
  } else if (model.match(/gpt-4/)) {
    modelMax = 8192;
  } else if (model.match(/claude-(?:instant-)?v1(?:\.\d)?(?:-100k)/)) {
    modelMax = 100000;
  } else if (model.match(/claude-(?:instant-)?v1(?:\.\d)?$/)) {
    modelMax = 9000;
  } else if (model.match(/claude-2/)) {
    modelMax = 100000;
  } else if (model.match(/^text-bison-\d{3}$/)) {
    modelMax = BISON_MAX_CONTEXT;
  } else {
    // Don't really want to throw here because I don't want to have to update
    // this ASAP every time a new model is released.
    req.log.warn({ model }, "Unknown model, using 100k token limit.");
    modelMax = 100000;
  }

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

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

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

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

/**
 * For OpenAI-to-Anthropic requests, users can't specify the model, so we need
 * to pick one based on the final context size. Ideally this would happen in
 * the `transformOutboundPayload` preprocessor, but we don't have the context
 * size at that point (and need a transformed body to calculate it).
 */
function maybeTranslateOpenAIModel(req: Request) {
  if (req.inboundApi !== "openai" || req.outboundApi !== "anthropic") {
    return;
  }

  const bigModel = process.env.CLAUDE_BIG_MODEL || "claude-v1-100k";
  const contextSize = req.promptTokens! + req.outputTokens!;

  if (contextSize > 8500) {
    req.log.debug(
      { model: bigModel, contextSize },
      "Using Claude 100k model for OpenAI-to-Anthropic request"
    );
    req.body.model = bigModel;
  }
  // Small model is the default already set in `transformOutboundPayload`
}


code:
import { fixRequestBody } from "http-proxy-middleware";
import type { ProxyRequestMiddleware } from ".";

/** Finalize the rewritten request body. Must be the last rewriter. */
export const finalizeBody: ProxyRequestMiddleware = (proxyReq, req) => {
  if (["POST", "PUT", "PATCH"].includes(req.method ?? "") && req.body) {
    const updatedBody = JSON.stringify(req.body);
    proxyReq.setHeader("Content-Length", Buffer.byteLength(updatedBody));
    (req as any).rawBody = Buffer.from(updatedBody);

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


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

// Express middleware (runs before http-proxy-middleware, can be async)
export { applyQuotaLimits } from "./apply-quota-limits";
export { createPreprocessorMiddleware } from "./preprocess";
export { checkContextSize } from "./check-context-size";
export { setApiFormat } from "./set-api-format";
export { transformOutboundPayload } from "./transform-outbound-payload";

// HPM middleware (runs on onProxyReq, cannot be async)
export { addKey } from "./add-key";
export { addAnthropicPreamble } from "./add-anthropic-preamble";
export { blockZoomerOrigins } from "./block-zoomer-origins";
export { finalizeBody } from "./finalize-body";
export { languageFilter } from "./language-filter";
export { limitCompletions } from "./limit-completions";
export { removeOriginHeaders } from "./remove-origin-headers";
export { transformKoboldPayload } from "./transform-kobold-payload";

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

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


code:
import { Request } from "express";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { assertNever } from "../../../shared/utils";
import { isCompletionRequest } from "../common";
import { ProxyRequestMiddleware } from ".";

const DISALLOWED_REGEX =
  /[\u2E80-\u2E99\u2E9B-\u2EF3\u2F00-\u2FD5\u3005\u3007\u3021-\u3029\u3038-\u303B\u3400-\u4DB5\u4E00-\u9FD5\uF900-\uFA6D\uFA70-\uFAD9]/;

// Our shitty free-tier VMs will fall over if we test every single character in
// each 15k character request ten times a second. So we'll just sample 20% of
// the characters and hope that's enough.
const containsDisallowedCharacters = (text: string) => {
  const sampleSize = Math.ceil(text.length * 0.2);
  const sample = text
    .split("")
    .sort(() => 0.5 - Math.random())
    .slice(0, sampleSize)
    .join("");
  return DISALLOWED_REGEX.test(sample);
};

/** Block requests containing too many disallowed characters. */
export const languageFilter: ProxyRequestMiddleware = (_proxyReq, req) => {
  if (!config.rejectDisallowed) {
    return;
  }

  if (isCompletionRequest(req)) {
    const combinedText = getPromptFromRequest(req);
    if (containsDisallowedCharacters(combinedText)) {
      logger.warn(`Blocked request containing bad characters`);
      _proxyReq.destroy(new Error(config.rejectMessage));
    }
  }
};

function getPromptFromRequest(req: Request) {
  const service = req.outboundApi;
  const body = req.body;
  switch (service) {
    case "anthropic":
      return body.prompt;
    case "openai":
      return body.messages
        .map((m: { content: string }) => m.content)
        .join("\n");
    case "openai-text":
      return body.prompt;
    case "google-palm":
      return body.prompt.text;
    default:
      assertNever(service);
  }
}


code:
import { isCompletionRequest } from "../common";
import { ProxyRequestMiddleware } from ".";

/**
 * Don't allow multiple completions to be requested to prevent abuse.
 * OpenAI-only, Anthropic provides no such parameter.
 **/
export const limitCompletions: ProxyRequestMiddleware = (_proxyReq, req) => {
  if (isCompletionRequest(req) && req.outboundApi === "openai") {
    const originalN = req.body?.n || 1;
    req.body.n = 1;
    if (originalN !== req.body.n) {
      req.log.warn(`Limiting completion choices from ${originalN} to 1`);
    }
  }
};


code:
import { RequestHandler } from "express";
import { handleInternalError } from "../common";
import {
  RequestPreprocessor,
  checkContextSize,
  setApiFormat,
  transformOutboundPayload,
} from ".";

/**
 * Returns a middleware function that processes the request body into the given
 * API format, and then sequentially runs the given additional preprocessors.
 */
export const createPreprocessorMiddleware = (
  apiFormat: Parameters<typeof setApiFormat>[0],
  additionalPreprocessors?: RequestPreprocessor[]
): RequestHandler => {
  const preprocessors: RequestPreprocessor[] = [
    setApiFormat(apiFormat),
    ...(additionalPreprocessors ?? []),
    transformOutboundPayload,
    checkContextSize,
  ];

  return async function executePreprocessors(req, res, next) {
    try {
      for (const preprocessor of preprocessors) {
        await preprocessor(req);
      }
      next();
    } catch (error) {
      req.log.error(error, "Error while executing request preprocessor");
      handleInternalError(error as Error, req, res);
    }
  };
};


code:
import { ProxyRequestMiddleware } from ".";

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


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

export const setApiFormat = (api: {
  inApi: Request["inboundApi"];
  outApi: APIFormat;
}): RequestPreprocessor => {
  return (req) => {
    req.inboundApi = api.inApi;
    req.outboundApi = api.outApi;
  };
};