Update app.py

app.py

@@ -7,9 +7,7 @@ import socketio
 import requests
 import logging
 from threading import Event
-
-# 如果使用 GPT 模型的 tokenization，可以引入 tiktoken
-# import tiktoken  # 如果需要使用 GPT 的 token 化库
+import tiktoken  # 使用 tiktoken 库进行 tokenization
 
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -77,14 +75,14 @@ def normalize_content(content):
         # 如果是其他类型，返回空字符串
         return ""
 
-def calculate_input_tokens(text):
+def calculate_tokens(text):
     """
-    计算输入文本的 token 数量。
-    这里我们简单地通过空格分词来模拟 token 计数。
-    如果使用 GPT 模型，可以使用 tiktoken 库进行 tokenization。
+    使用 tiktoken 库来计算输入和输出文本的 token 数量。
+    Claude 模型可能有不同的 tokenization 规则，但 tiktoken 是一个很好的近似工具。
     """
-    # 使用简单的空格分词计数
-    tokens = text.split()
+    # 使用 tiktoken 的 GPT-3.5 模型进行 tokenization
+    encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")  # 使用 gpt-3.5-turbo 作为近似模型
+    tokens = encoding.encode(text)
     return len(tokens)
 
 @app.route('/')
@@ -124,7 +122,7 @@ def messages():
         previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
 
         # 动态计算输入的 token 数量
-        input_tokens = calculate_input_tokens(previous_messages)
+        input_tokens = calculate_tokens(previous_messages)
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
@@ -148,7 +146,7 @@ def messages():
                     "model": model,  # 动态模型
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": 1},  # 动态 input_tokens
+                    "usage": {"input_tokens": input_tokens, "output_tokens": 0},  # 动态 input_tokens
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -226,11 +224,15 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
+            output_text = ''.join(response_text)
+            # 动态计算 output_tokens
+            output_tokens = calculate_tokens(output_text)
+
             yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
             yield create_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                "usage": {"output_tokens": len(''.join(response_text))},
+                "usage": {"output_tokens": output_tokens},  # 动态 output_tokens
             })
             yield create_event("message_stop", {"type": "message_stop"})  # 确保发送 message_stop 事件
 
@@ -298,9 +300,13 @@ 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)
+
         # 生成完整的响应
         full_response = {
-            "content": [{"text": ''.join(response_text), "type": "text"}],  # 合并所有文本块
+            "content": [{"text": output_text, "type": "text"}],  # 合并所有文本块
             "id": msg_id,
             "model": model,  # 动态模型
             "role": "assistant",
@@ -309,7 +315,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
             "type": "message",
             "usage": {
                 "input_tokens": input_tokens,  # 动态 input_tokens
-                "output_tokens": len(''.join(response_text)),
+                "output_tokens": output_tokens,  # 动态 output_tokens
             },
         }
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')