Update app.py

app.py

@@ -7,6 +7,7 @@ import socketio
 import requests
 import logging
 from threading import Event
+import tiktoken  # 引入 tiktoken 库
 
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -74,21 +75,14 @@ def normalize_content(content):
         # 如果是其他类型，返回空字符串
         return ""
 
-def calculate_tokens_via_unicode(text):
+def calculate_tokens_via_tiktoken(text, model="gpt-3.5-turbo"):
     """
-    通过 Unicode 码点来计算文本的 token 数量。
-    每个中文字符、标点符号等都被视为一个 token。
-    对于英文，我们可以继续使用空格分词。
+    使用 tiktoken 库根据 GPT 模型计算 token 数量。
+    Claude 模型与 GPT 模型的 token 计算机制类似，因此可以使用 tiktoken。
     """
-    tokens = 0
-    for char in text:
-        # 如果是中文字符（Unicode 范围：\u4e00 - \u9fff），每个字符视为一个 token
-        if '\u4e00' <= char <= '\u9fff':
-            tokens += 1
-        else:
-            # 对于非中文字符，简单按空格分词
-            tokens += len(char.split())
-    return tokens
+    encoding = tiktoken.encoding_for_model(model)  # 获取模型的编码器
+    tokens = encoding.encode(text)  # 对文本进行 tokenization
+    return len(tokens)
 
 @app.route('/')
 def root():
@@ -126,8 +120,8 @@ def messages():
         # 使用 normalize_content 递归处理 msg['content']
         previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
 
-        # 动态计算输入的 token 数量
-        input_tokens = calculate_tokens_via_unicode(previous_messages)
+        # 动态计算输入的 token 数量，使用 tiktoken 进行 tokenization
+        input_tokens = calculate_tokens_via_tiktoken(previous_messages, model="gpt-3.5-turbo")
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
@@ -229,8 +223,8 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            # 动态计算输出的 token 数量
-            output_tokens = calculate_tokens_via_unicode(''.join(response_text))
+            # 动态计算输出的 token 数量，使用 tiktoken 进行 tokenization
+            output_tokens = calculate_tokens_via_tiktoken(''.join(response_text), model="gpt-3.5-turbo")
 
             yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
             yield create_event("message_delta", {
@@ -304,8 +298,8 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         # 等待响应完成
         response_event.wait(timeout=30)
 
-        # 动态计算输出的 token 数量
-        output_tokens = calculate_tokens_via_unicode(''.join(response_text))
+        # 动态计算输出的 token 数量，使用 tiktoken 进行 tokenization
+        output_tokens = calculate_tokens_via_tiktoken(''.join(response_text), model="gpt-3.5-turbo")
 
         # 生成完整的响应
         full_response = {