From 7c34379ecbab814707f6156af339d0777afeb983 Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Tue, 28 Mar 2023 21:12:34 +0800
Subject: [PATCH] calculate token and forget some message

---
 src/chatbox.tsx |  5 +++++
 src/chatgpt.ts  | 15 +++++++++------
 src/message.tsx |  5 +++++
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/chatbox.tsx b/src/chatbox.tsx
index 240f831..98b47bd 100644
--- a/src/chatbox.tsx
+++ b/src/chatbox.tsx
@@ -102,6 +102,8 @@ export default function ChatBOX(props: {
     client.apiEndpoint = chatStore.apiEndpoint;
     client.sysMessageContent = chatStore.systemMessageContent;
     client.messages = chatStore.history.slice(chatStore.postBeginIndex);
+    // try forget message before sending request
+    client.forgetSomeMessages();
     try {
       setShowGenerating(true);
       const response = await client._fetch(chatStore.streamMode);
@@ -140,6 +142,9 @@ export default function ChatBOX(props: {
       return;
     }
     chatStore.history.push({ role: "user", content: inputMsg.trim() });
+    // manually calculate token length
+    chatStore.totalTokens += client.calculate_token_length(inputMsg.trim());
+    client.total_tokens += client.calculate_token_length(inputMsg.trim());
     setChatStore({ ...chatStore });
     setInputMsg("");
     await complete();
diff --git a/src/chatgpt.ts b/src/chatgpt.ts
index 589fba5..02c918e 100644
--- a/src/chatgpt.ts
+++ b/src/chatgpt.ts
@@ -25,6 +25,14 @@ export interface FetchResponse {
     index: number | undefined;
   }[];
 }
+// https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them
+export function calculate_token_length(content: string): number {
+  const totalCount = content.length;
+  const chineseCount = content.match(/[\u00ff-\uffff]|\S+/g)?.length ?? 0;
+  const englishCount = totalCount - chineseCount;
+  const tokenLength = englishCount / 4 + (chineseCount * 4) / 3;
+  return ~~tokenLength;
+}
 
 class Chat {
   OPENAI_API_KEY: string;
@@ -114,13 +122,8 @@ class Chat {
     return this._fetch(true);
   }
 
-  // https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them
   calculate_token_length(content: string): number {
-    const totalCount = content.length;
-    const chineseCount = content.match(/[\u00ff-\uffff]|\S+/g)?.length ?? 0;
-    const englishCount = totalCount - chineseCount;
-    const tokenLength = englishCount / 4 + (chineseCount * 4) / 3;
-    return ~~tokenLength;
+    return calculate_token_length(content);
   }
 
   user(...messages: string[]) {
diff --git a/src/message.tsx b/src/message.tsx
index da66afc..0ef04aa 100644
--- a/src/message.tsx
+++ b/src/message.tsx
@@ -1,4 +1,5 @@
 import { ChatStore } from "./app";
+import { calculate_token_length } from "./chatgpt";
 
 interface Props {
   messageIndex: number;
@@ -24,6 +25,10 @@ export default function Message(props: Props) {
         ) {
           chatStore.history.splice(messageIndex, 1);
           chatStore.postBeginIndex = Math.max(chatStore.postBeginIndex - 1, 0);
+          chatStore.totalTokens = Math.max(
+            0,
+            chatStore.totalTokens - calculate_token_length(chat.content)
+          );
           setChatStore({ ...chatStore });
         }
       }}