From 1a4b4401eebca2ba0b9b0ed9f34b9c6b5e935d8c Mon Sep 17 00:00:00 2001
From: Alan P <alanpog@gmail.com>
Date: Thu, 23 Mar 2023 17:51:51 -0400
Subject: [PATCH] include prompt in token count (#104)

---
 pages/api/chat.ts | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/pages/api/chat.ts b/pages/api/chat.ts
index f96e0e7..313cd00 100644
--- a/pages/api/chat.ts
+++ b/pages/api/chat.ts
@@ -7,7 +7,7 @@ import { init, Tiktoken } from "@dqbd/tiktoken/lite/init";
 import wasm from "../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module";
 
 export const config = {
-  runtime: "edge"
+  runtime: "edge",
 };
 
 const handler = async (req: Request): Promise<Response> => {
@@ -15,10 +15,22 @@ const handler = async (req: Request): Promise<Response> => {
     const { model, messages, key, prompt } = (await req.json()) as ChatBody;
 
     await init((imports) => WebAssembly.instantiate(wasm, imports));
-    const encoding = new Tiktoken(tiktokenModel.bpe_ranks, tiktokenModel.special_tokens, tiktokenModel.pat_str);
+    const encoding = new Tiktoken(
+      tiktokenModel.bpe_ranks,
+      tiktokenModel.special_tokens,
+      tiktokenModel.pat_str
+    );
 
     const tokenLimit = model.id === OpenAIModelID.GPT_4 ? 6000 : 3000;
-    let tokenCount = 0;
+
+    let promptToSend = prompt;
+    if (!promptToSend) {
+      promptToSend = DEFAULT_SYSTEM_PROMPT;
+    }
+
+    const prompt_tokens = encoding.encode(promptToSend);
+
+    let tokenCount = prompt_tokens.length;
     let messagesToSend: Message[] = [];
 
     for (let i = messages.length - 1; i >= 0; i--) {
@@ -34,11 +46,6 @@ const handler = async (req: Request): Promise<Response> => {
 
     encoding.free();
 
-    let promptToSend = prompt;
-    if (!promptToSend) {
-      promptToSend = DEFAULT_SYSTEM_PROMPT;
-    }
-
     const stream = await OpenAIStream(model, promptToSend, key, messagesToSend);
 
     return new Response(stream);