Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import json
 import uuid
+import re
 from datetime import datetime
 from flask import Flask, request, Response, jsonify
 import socketio
@@ -76,10 +77,11 @@ def normalize_content(content):
 
 def calculate_tokens(text):
     """
-    使用字符数近似计算 token 数量。
-    这里假设平均每 4 个字符代表 1 个 token。
+    计算输入文本的 token 数量。我们通过正则表达式来进行 tokenization，确保标点符号和单词都被正确计算为 token。
     """
-    return len(text) // 4  # 近似每 4 个字符为 1 个 token
+    # 使用正则表达式匹配单词和标点符号
+    tokens = re.findall(r"\w+|[^\w\s]", text)
+    return len(tokens)
 
 @app.route('/')
 def root():
@@ -122,7 +124,7 @@ def messages():
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
-        response_text = []  # 用于存储输出的文本块
+        response_text = []
 
         if not stream:
             # 处理 stream 为 false 的情况
@@ -142,7 +144,7 @@ def messages():
                     "model": model,  # 动态模型
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": 0},  # 动态 input_tokens
+                    "usage": {"input_tokens": input_tokens, "output_tokens": 0},  # 动态 input_tokens，output_tokens 初始化为 0
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -220,20 +222,15 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            # 合并所有的输出块
-            output_text = ''.join(response_text)
-            # 动态计算 output_tokens
-            output_tokens = calculate_tokens(output_text)
+            # 计算输出 token 数量
+            output_tokens = calculate_tokens(''.join(response_text))
 
-            # 生成 content_block_stop 事件
             yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            # 生成 message_delta 事件，包含 output_tokens
             yield create_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                "usage": {"output_tokens": output_tokens},  # 动态 output_tokens
+                "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},  # 动态 output_tokens
             })
-            # 生成 message_stop 事件
             yield create_event("message_stop", {"type": "message_stop"})  # 确保发送 message_stop 事件
 
         return Response(generate(), content_type='text/event-stream')
@@ -300,14 +297,12 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         # 等待响应完成
         response_event.wait(timeout=30)
 
-        # 合并所有的输出块
-        output_text = ''.join(response_text)
-        # 动态计算 output_tokens
-        output_tokens = calculate_tokens(output_text)
+        # 计算输出 token 数量
+        output_tokens = calculate_tokens(''.join(response_text))
 
         # 生成完整的响应
         full_response = {
-            "content": [{"text": output_text, "type": "text"}],  # 合并所有文本块
+            "content": [{"text": ''.join(response_text), "type": "text"}],  # 合并所有文本块
             "id": msg_id,
             "model": model,  # 动态模型
             "role": "assistant",