From c8d74fe8fdc23c504fc826c4704f61e7a0680a7c Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Fri, 2 Jun 2023 01:59:12 -0500
Subject: [PATCH] includes tokenizer debug info on responses

---
 src/proxy/anthropic.ts         | 8 +++++++-
 src/proxy/middleware/common.ts | 8 ++++++--
 src/proxy/openai.ts            | 5 +++++
 src/types/custom.d.ts          | 3 +++
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/proxy/anthropic.ts b/src/proxy/anthropic.ts
index 334a41b..c76a584 100644
--- a/src/proxy/anthropic.ts
+++ b/src/proxy/anthropic.ts
@@ -108,10 +108,16 @@ const anthropicResponseHandler: ProxyResHandlerWithBody = async (
     body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
   }
 
-  if (!req.originalUrl.includes("/v1/complete") && !req.originalUrl.includes("/complete")) {
+  if (req.inboundApi === "openai") {
     req.log.info("Transforming Anthropic response to OpenAI format");
     body = transformAnthropicResponse(body);
   }
+
+  // TODO: Remove once tokenization is stable
+  if (req.debug) {
+    body.proxy_tokenizer_debug_info = req.debug;
+  }
+
   res.status(200).json(body);
 };
 
diff --git a/src/proxy/middleware/common.ts b/src/proxy/middleware/common.ts
index 7c70655..dda5324 100644
--- a/src/proxy/middleware/common.ts
+++ b/src/proxy/middleware/common.ts
@@ -45,6 +45,9 @@ export function writeErrorResponse(
     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);
   }
 }
@@ -57,7 +60,8 @@ export const handleProxyError: httpProxy.ErrorCallback = (err, req, res) => {
 export const handleInternalError = (
   err: Error,
   req: Request,
-  res: Response
+  res: Response,
+  errorType: string = "proxy_internal_error"
 ) => {
   try {
     const isZod = err instanceof ZodError;
@@ -86,7 +90,7 @@ export const handleInternalError = (
     } else {
       writeErrorResponse(req, res, 500, {
         error: {
-          type: "proxy_rewriter_error",
+          type: errorType,
           proxy_note: `Reverse proxy encountered an error before it could reach the upstream API.`,
           message: err.message,
           stack: err.stack,
diff --git a/src/proxy/openai.ts b/src/proxy/openai.ts
index 1c8832e..11380bd 100644
--- a/src/proxy/openai.ts
+++ b/src/proxy/openai.ts
@@ -125,6 +125,11 @@ const openaiResponseHandler: ProxyResHandlerWithBody = async (
     body.proxy_note = `Prompts are logged on this proxy instance. See ${host} for more information.`;
   }
 
+  // TODO: Remove once tokenization is stable
+  if (req.debug) {
+    body.proxy_tokenizer_debug_info = req.debug;
+  }
+
   res.status(200).json(body);
 };
 
diff --git a/src/types/custom.d.ts b/src/types/custom.d.ts
index e81bd1b..fe8acf1 100644
--- a/src/types/custom.d.ts
+++ b/src/types/custom.d.ts
@@ -18,6 +18,9 @@ declare global {
       onAborted?: () => void;
       proceed: () => void;
       heartbeatInterval?: NodeJS.Timeout;
+      promptTokens?: number;
+      // TODO: remove later
+      debug: Record<string, any>;
     }
   }
 }