From 0c8938dea5420b24f364aefadb0053a208f8fcce Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Wed, 8 Nov 2023 19:40:38 +0800
Subject: [PATCH] better cound token

---
 src/chatgpt.ts | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/chatgpt.ts b/src/chatgpt.ts
index b213ad6..74db715 100644
--- a/src/chatgpt.ts
+++ b/src/chatgpt.ts
@@ -47,20 +47,32 @@ 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 | MessageDetail[]
-): number {
-  const text =
-    typeof content === "string"
-      ? content
-      : content.map((c) => c?.text).join(" ");
+
+function calculate_token_length_from_text(text: string): number {
   const totalCount = text.length;
   const chineseCount = text.match(/[\u00ff-\uffff]|\S+/g)?.length ?? 0;
   const englishCount = totalCount - chineseCount;
   const tokenLength = englishCount / 4 + (chineseCount * 4) / 3;
   return ~~tokenLength;
 }
+// https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them
+export function calculate_token_length(
+  content: string | MessageDetail[]
+): number {
+  if (typeof content === "string") {
+    return calculate_token_length_from_text(content);
+  }
+  let tokens = 0;
+  for (const m of content) {
+    if (m.type === "text") {
+      tokens += calculate_token_length_from_text(m.text ?? "");
+    }
+    if (m.type === "image_url") {
+      tokens += m.image_url?.detail === "high" ? 65 * 4 : 65;
+    }
+  }
+  return 0;
+}
 
 class Chat {
   OPENAI_API_KEY: string;