Token based and model conditional limits (#36)

* use tiktoken for api limit

* model conditional char limits on frontend

* adjust for completion tokens

---------

Co-authored-by: Alan Pogrebinschi <alanpog@gmail.com>

components/Chat/Chat.tsx

@@ -18,7 +18,17 @@ interface Props {
   onModelChange: (conversation: Conversation, model: OpenAIModel) => void;
 }
 
-export const Chat: FC<Props> = ({ conversation, models, messageIsStreaming, modelError, messageError, loading, lightMode, onSend, onModelChange }) => {
+export const Chat: FC<Props> = ({
+  conversation,
+  models,
+  messageIsStreaming,
+  modelError,
+  messageError,
+  loading,
+  lightMode,
+  onSend,
+  onModelChange,
+}) => {
   const [currentMessage, setCurrentMessage] = useState<Message>();
 
   const messagesEndRef = useRef<HTMLDivElement>(null);
@@ -36,8 +46,13 @@ export const Chat: FC<Props> = ({ conversation, models, messageIsStreaming, mode
       {modelError ? (
         <div className="flex flex-col justify-center mx-auto h-full w-[300px] sm:w-[500px] space-y-6">
           <div className="text-center text-red-500">Error fetching models.</div>
-          <div className="text-center text-red-500">Make sure your OpenAI API key is set in the bottom left of the sidebar or in a .env.local file and refresh.</div>
+          <div className="text-center text-red-500">
-          <div className="text-center text-red-500">If you completed this step, OpenAI may be experiencing issues.</div>
+            Make sure your OpenAI API key is set in the bottom left of the
+            sidebar or in a .env.local file and refresh.
+          </div>
+          <div className="text-center text-red-500">
+            If you completed this step, OpenAI may be experiencing issues.
+          </div>
         </div>
       ) : (
         <>
@@ -48,15 +63,21 @@ export const Chat: FC<Props> = ({ conversation, models, messageIsStreaming, mode
                   <ModelSelect
                     model={conversation.model}
                     models={models}
-                    onModelChange={(model) => onModelChange(conversation, model)}
+                    onModelChange={(model) =>
+                      onModelChange(conversation, model)
+                    }
                   />
                 </div>
 
-                <div className="text-4xl text-center text-neutral-600 dark:text-neutral-200 pt-[160px] sm:pt-[280px]">{models.length === 0 ? "Loading..." : "Chatbot UI"}</div>
+                <div className="text-4xl text-center text-neutral-600 dark:text-neutral-200 pt-[160px] sm:pt-[280px]">
+                  {models.length === 0 ? "Loading..." : "Chatbot UI"}
+                </div>
               </>
             ) : (
               <>
-                <div className="flex justify-center py-2 text-neutral-500 bg-neutral-100 dark:bg-[#444654] dark:text-neutral-200 text-sm border border-b-neutral-300 dark:border-none">Model: {conversation.model.name}</div>
+                <div className="flex justify-center py-2 text-neutral-500 bg-neutral-100 dark:bg-[#444654] dark:text-neutral-200 text-sm border border-b-neutral-300 dark:border-none">
+                  Model: {conversation.model.name}
+                </div>
 
                 {conversation.messages.map((message, index) => (
                   <ChatMessage
@@ -91,6 +112,7 @@ export const Chat: FC<Props> = ({ conversation, models, messageIsStreaming, mode
                 setCurrentMessage(message);
                 onSend(message, false);
               }}
+              model={conversation.model}
             />
           )}
         </>

components/Chat/ChatInput.tsx

@@ -1,13 +1,14 @@
-import { Message } from "@/types";
+import { Message, OpenAIModel, OpenAIModelID } from "@/types";
 import { IconSend } from "@tabler/icons-react";
 import { FC, KeyboardEvent, useEffect, useRef, useState } from "react";
 
 interface Props {
   messageIsStreaming: boolean;
   onSend: (message: Message) => void;
+  model: OpenAIModel;
 }
 
-export const ChatInput: FC<Props> = ({ onSend, messageIsStreaming }) => {
+export const ChatInput: FC<Props> = ({ onSend, messageIsStreaming, model }) => {
   const [content, setContent] = useState<string>();
   const [isTyping, setIsTyping] = useState<boolean>(false);
 
@@ -15,8 +16,10 @@ export const ChatInput: FC<Props> = ({ onSend, messageIsStreaming }) => {
 
   const handleChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
     const value = e.target.value;
-    if (value.length > 4000) {
+    const maxLength = model.id === OpenAIModelID.GPT_3_5 ? 12000 : 24000;
-      alert("Message limit is 4000 characters");
+
+    if (value.length > maxLength) {
+      alert(`Message limit is ${maxLength} characters`);
       return;
     }
 
@@ -42,8 +45,10 @@ export const ChatInput: FC<Props> = ({ onSend, messageIsStreaming }) => {
   };
 
   const isMobile = () => {
-    const userAgent = typeof window.navigator === "undefined" ? "" : navigator.userAgent;
+    const userAgent =
-    const mobileRegex = /Android|webOS|iPhone|iPad|iPod|BlackBerry|IEMobile|Opera Mini|Mobile|mobile|CriOS/i;
+      typeof window.navigator === "undefined" ? "" : navigator.userAgent;
+    const mobileRegex =
+      /Android|webOS|iPhone|iPad|iPod|BlackBerry|IEMobile|Opera Mini|Mobile|mobile|CriOS/i;
     return mobileRegex.test(userAgent);
   };
 
@@ -72,7 +77,7 @@ export const ChatInput: FC<Props> = ({ onSend, messageIsStreaming }) => {
           resize: "none",
           bottom: `${textareaRef?.current?.scrollHeight}px`,
           maxHeight: "400px",
-          overflow: "auto"
+          overflow: "auto",
         }}
         placeholder="Type a message..."
         value={content}

next.config.js

@@ -1,6 +1,15 @@
 /** @type {import('next').NextConfig} */
 const nextConfig = {
   reactStrictMode: true,
-}
 
-module.exports = nextConfig
+  webpack(config, { isServer, dev }) {
+    config.experiments = {
+      asyncWebAssembly: true,
+      layers: true,
+    };
+
+    return config;
+  },
+};
+
+module.exports = nextConfig;

package-lock.json

@@ -8,6 +8,7 @@
       "name": "ai-chatbot-starter",
       "version": "0.1.0",
       "dependencies": {
+        "@dqbd/tiktoken": "^1.0.2",
         "@tabler/icons-react": "^2.9.0",
         "@types/node": "18.15.0",
         "@types/react": "18.0.28",
@@ -43,6 +44,11 @@
         "node": ">=6.9.0"
       }
     },
+    "node_modules/@dqbd/tiktoken": {
+      "version": "1.0.2",
+      "resolved": "https://registry.npmjs.org/@dqbd/tiktoken/-/tiktoken-1.0.2.tgz",
+      "integrity": "sha512-AjGTBRWsMoVmVeN55NLyupyM8TNamOUBl6tj5t/leLDVup3CFGO9tVagNL1jf3GyZLkWZSTmYVbPQ/M2LEcNzw=="
+    },
     "node_modules/@eslint-community/eslint-utils": {
       "version": "4.3.0",
       "resolved": "https://registry.npmjs.org/@eslint-community/eslint-utils/-/eslint-utils-4.3.0.tgz",
@@ -5543,6 +5549,11 @@
         "regenerator-runtime": "^0.13.11"
       }
     },
+    "@dqbd/tiktoken": {
+      "version": "1.0.2",
+      "resolved": "https://registry.npmjs.org/@dqbd/tiktoken/-/tiktoken-1.0.2.tgz",
+      "integrity": "sha512-AjGTBRWsMoVmVeN55NLyupyM8TNamOUBl6tj5t/leLDVup3CFGO9tVagNL1jf3GyZLkWZSTmYVbPQ/M2LEcNzw=="
+    },
     "@eslint-community/eslint-utils": {
       "version": "4.3.0",
       "resolved": "https://registry.npmjs.org/@eslint-community/eslint-utils/-/eslint-utils-4.3.0.tgz",

package.json

@@ -9,6 +9,7 @@
     "lint": "next lint"
   },
   "dependencies": {
+    "@dqbd/tiktoken": "^1.0.2",
     "@tabler/icons-react": "^2.9.0",
     "@types/node": "18.15.0",
     "@types/react": "18.0.28",

pages/api/chat.ts

@@ -1,5 +1,9 @@
-import { Message, OpenAIModel } from "@/types";
+import { Message, OpenAIModel, OpenAIModelID } from "@/types";
 import { OpenAIStream } from "@/utils/server";
+import tiktokenModel from "@dqbd/tiktoken/encoders/cl100k_base.json";
+import { init, Tiktoken } from "@dqbd/tiktoken/lite/init";
+// @ts-expect-error
+import wasm from "../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module";
 
 export const config = {
   runtime: "edge"
@@ -13,19 +17,26 @@ const handler = async (req: Request): Promise<Response> => {
       key: string;
     };
 
-    const charLimit = 12000;
+    await init((imports) => WebAssembly.instantiate(wasm, imports));
-    let charCount = 0;
+    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 messagesToSend: Message[] = [];
 
     for (let i = messages.length - 1; i >= 0; i--) {
       const message = messages[i];
-      if (charCount + message.content.length > charLimit) {
+      const tokens = encoding.encode(message.content);
+
+      if (tokenCount + tokens.length > tokenLimit) {
         break;
       }
-      charCount += message.content.length;
+      tokenCount += tokens.length;
       messagesToSend = [message, ...messagesToSend];
     }
 
+    encoding.free();
+
     const stream = await OpenAIStream(model, key, messagesToSend);
 
     return new Response(stream);