From 76fdd047e7a9427dee18785d1cf60cc0e0999554 Mon Sep 17 00:00:00 2001
From: Yidadaa <yidadaa@qq.com>
Date: Thu, 15 Jun 2023 00:14:38 +0800
Subject: [PATCH] feat: new token count function

---
 app/store/chat.ts  |  3 ++-
 app/utils/token.ts | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 app/utils/token.ts

diff --git a/app/store/chat.ts b/app/store/chat.ts
index 3c9953e5..f56d9b9d 100644
--- a/app/store/chat.ts
+++ b/app/store/chat.ts
@@ -11,6 +11,7 @@ import { StoreKey } from "../constant";
 import { api, RequestMessage } from "../client/api";
 import { ChatControllerPool } from "../client/controller";
 import { prettyObject } from "../utils/format";
+import { estimateTokenLength } from "../utils/token";
 
 export type ChatMessage = RequestMessage & {
   date: string;
@@ -102,7 +103,7 @@ interface ChatStore {
 }
 
 function countMessages(msgs: ChatMessage[]) {
-  return msgs.reduce((pre, cur) => pre + cur.content.length, 0);
+  return msgs.reduce((pre, cur) => pre + estimateTokenLength(cur.content), 0);
 }
 
 export const useChatStore = create<ChatStore>()(
diff --git a/app/utils/token.ts b/app/utils/token.ts
new file mode 100644
index 00000000..ec8139b2
--- /dev/null
+++ b/app/utils/token.ts
@@ -0,0 +1,22 @@
+export function estimateTokenLength(input: string): number {
+  let tokenLength = 0;
+
+  for (let i = 0; i < input.length; i++) {
+    const charCode = input.charCodeAt(i);
+
+    if (charCode < 128) {
+      // ASCII character
+      if (charCode <= 122 && charCode >= 65) {
+        // a-Z
+        tokenLength += 0.25;
+      } else {
+        tokenLength += 0.5;
+      }
+    } else {
+      // Unicode character
+      tokenLength += 1.5;
+    }
+  }
+
+  return tokenLength;
+}