refactors request rewriters

src/config.ts

@@ -11,7 +11,7 @@ type Config = {
   /** Per-IP limit for requests per minute to OpenAI's completions endpoint. */
   modelRateLimit: number; // TODO
   /** Max number of tokens to generate. Requests which specify a higher value will be rewritten to use this value. */
-  maxOutputTokens: number; // TODO
+  maxOutputTokens: number;
   /** Logging threshold. */
   logLevel?: "debug" | "info" | "warn" | "error";
   /** Whether prompts and responses should be logged. */

src/proxy/common.ts

@@ -5,6 +5,7 @@ import { logger } from "../logger";
 import { keys } from "../keys";
 
 /** Handle and rewrite response to proxied requests to OpenAI */
+// TODO: This is a mess, fix it
 export const handleResponse = (
   proxyRes: http.IncomingMessage,
   req: Request,

src/proxy/openai.ts

@@ -1,36 +1,34 @@
 import { Request, Router } from "express";
 import * as http from "http";
-import { createProxyMiddleware, fixRequestBody } from "http-proxy-middleware";
+import { createProxyMiddleware } from "http-proxy-middleware";
 import { logger } from "../logger";
-import { Key, keys } from "../keys";
 import { handleResponse, onError } from "./common";
+import {
+  addKey,
+  disableStream,
+  finalizeBody,
+  limitOutputTokens,
+} from "./rewriters";
 
-/**
- * Modifies the request body to add a randomly selected API key.
- */
-const rewriteRequest = (proxyReq: http.ClientRequest, req: Request) => {
-  let key: Key;
+const rewriteRequest = (
+  proxyReq: http.ClientRequest,
+  req: Request,
+  res: http.ServerResponse
+) => {
+  const rewriterPipeline = [
+    addKey,
+    disableStream,
+    limitOutputTokens,
+    finalizeBody,
+  ];
 
   try {
-    key = keys.get(req.body?.model || "gpt-3.5")!;
-  } catch (err) {
-    proxyReq.destroy(err as any);
-    return;
-  }
-
-  req.key = key;
-  proxyReq.setHeader("Authorization", `Bearer ${key.key}`);
-
-  if (req.method === "POST" && req.body) {
-    if (req.body?.stream) {
-      req.body.stream = false;
-      const updatedBody = JSON.stringify(req.body);
-      proxyReq.setHeader("Content-Length", Buffer.byteLength(updatedBody));
-      (req as any).rawBody = Buffer.from(updatedBody);
+    for (const rewriter of rewriterPipeline) {
+      rewriter(proxyReq, req, res, {});
     }
-
-    // body-parser and http-proxy-middleware don't play nice together
-    fixRequestBody(proxyReq, req);
+  } catch (error) {
+    logger.error(error, "Error while executing proxy rewriter");
+    proxyReq.destroy(error as Error);
   }
 };
 

src/proxy/rewriters/add-key.ts

@@ -0,0 +1,10 @@
+import type { ExpressHttpProxyReqCallback } from ".";
+import { Key, keys } from "../../keys";
+
+/** Add an OpenAI key from the pool to the request. */
+export const addKey: ExpressHttpProxyReqCallback = (proxyReq, req) => {
+  let assignedKey: Key;
+  assignedKey = keys.get(req.body?.model || "gpt-3.5")!;
+  req.key = assignedKey;
+  proxyReq.setHeader("Authorization", `Bearer ${assignedKey.key}`);
+};

src/proxy/rewriters/disable-stream.ts

@@ -0,0 +1,8 @@
+import type { ExpressHttpProxyReqCallback } from ".";
+
+/** Disable token streaming as the proxy middleware doesn't support it. */
+export const disableStream: ExpressHttpProxyReqCallback = (_proxyReq, req) => {
+  if (req.method === "POST" && req.body && req.body.stream) {
+    req.body.stream = false;
+  }
+};

src/proxy/rewriters/finalize-body.ts

@@ -0,0 +1,14 @@
+import { fixRequestBody } from "http-proxy-middleware";
+import type { ExpressHttpProxyReqCallback } from ".";
+
+/** Finalize the rewritten request body. Must be the last rewriter. */
+export const finalizeBody: ExpressHttpProxyReqCallback = (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);
+  }
+};

src/proxy/rewriters/index.ts

@@ -0,0 +1,13 @@
+import type { Request } from "express";
+import type { ClientRequest } from "http";
+import type { ProxyReqCallback } from "http-proxy";
+
+export { addKey } from "./add-key";
+export { disableStream } from "./disable-stream";
+export { limitOutputTokens } from "./limit-output-tokens";
+export { finalizeBody } from "./finalize-body";
+
+export type ExpressHttpProxyReqCallback = ProxyReqCallback<
+  ClientRequest,
+  Request
+>;

src/proxy/rewriters/limit-output-tokens.ts

@@ -0,0 +1,21 @@
+import { config } from "../../config";
+import type { ExpressHttpProxyReqCallback } from ".";
+import { logger } from "../../logger";
+
+const MAX_TOKENS = config.maxOutputTokens;
+
+/** Enforce a maximum number of tokens requested from OpenAI. */
+export const limitOutputTokens: ExpressHttpProxyReqCallback = (
+  _proxyReq,
+  req
+) => {
+  if (req.method === "POST" && req.body?.max_tokens) {
+    const originalTokens = req.body.max_tokens;
+    req.body.max_tokens = Math.min(req.body.max_tokens, MAX_TOKENS);
+    if (originalTokens !== req.body.max_tokens) {
+      logger.warn(
+        `Limiting max_tokens from ${originalTokens} to ${req.body.max_tokens}`
+      );
+    }
+  }
+};