Update app.py

app.py

@@ -8,6 +8,9 @@ import requests
 import logging
 from threading import Event
 
+# 如果使用 GPT 模型的 tokenization，可以引入 tiktoken
+# import tiktoken  # 如果需要使用 GPT 的 token 化库
+
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
 
@@ -74,6 +77,16 @@ def normalize_content(content):
         # 如果是其他类型，返回空字符串
         return ""
 
+def calculate_input_tokens(text):
+    """
+    计算输入文本的 token 数量。
+    这里我们简单地通过空格分词来模拟 token 计数。
+    如果使用 GPT 模型，可以使用 tiktoken 库进行 tokenization。
+    """
+    # 使用简单的空格分词计数
+    tokens = text.split()
+    return len(tokens)
+
 @app.route('/')
 def root():
     log_request(request.remote_addr, request.path, 200)
@@ -110,13 +123,16 @@ 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_input_tokens(previous_messages)
+
         msg_id = str(uuid.uuid4())
         response_event = Event()
         response_text = []
 
         if not stream:
             # 处理 stream 为 false 的情况
-            return handle_non_stream(previous_messages, msg_id, model)
+            return handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
         # 记录日志：此时请求上下文仍然有效
         log_request(request.remote_addr, request.path, 200)
@@ -132,7 +148,7 @@ def messages():
                     "model": model,  # 动态模型
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": 8, "output_tokens": 1},
+                    "usage": {"input_tokens": input_tokens, "output_tokens": 1},  # 动态 input_tokens
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -225,7 +241,7 @@ def messages():
         log_request(request.remote_addr, request.path, 400)
         return jsonify({"error": str(e)}), 400
 
-def handle_non_stream(previous_messages, msg_id, model):
+def handle_non_stream(previous_messages, msg_id, model, input_tokens):
     """
     处理 stream 为 false 的情况，返回完整的响应。
     """
@@ -292,7 +308,7 @@ def handle_non_stream(previous_messages, msg_id, model):
             "stop_sequence": None,
             "type": "message",
             "usage": {
-                "input_tokens": 8,
+                "input_tokens": input_tokens,  # 动态 input_tokens
                 "output_tokens": len(''.join(response_text)),
             },
         }