From dcf1b5528a13e542b56cefb905a2454298d31da2 Mon Sep 17 00:00:00 2001
From: liuwei <liuwei@wdtrgf.com.cn>
Date: Mon, 10 Mar 2025 11:51:27 +0800
Subject: [PATCH] =?UTF-8?q?=E6=80=BB=E7=BB=93=E5=8A=A0=E5=85=A5=EF=BC=8Cto?=
 =?UTF-8?q?kens=E6=98=BE=E7=A4=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 message_summary/message_summary_4o.py | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/message_summary/message_summary_4o.py b/message_summary/message_summary_4o.py
index 908396e..0d5c2f4 100644
--- a/message_summary/message_summary_4o.py
+++ b/message_summary/message_summary_4o.py
@@ -2,13 +2,35 @@ import requests
 import json
 
 
-# 解析JSON
 def extract_content(data_string):
     try:
         data = json.loads(data_string)
+
         # 提取content字段
         content = data["choices"][0]["message"].get("content", "")
-        return content
+
+        # 提取tokens相关内容，加入容错处理
+        tokens_usage = data.get("usage", {})
+
+        # 确保tokens_usage是字典类型
+        if isinstance(tokens_usage, dict):
+            prompt_tokens = tokens_usage.get("prompt_tokens", 0)
+            completion_tokens = tokens_usage.get("completion_tokens", 0)
+            total_tokens = tokens_usage.get("total_tokens", 0)
+        else:
+            prompt_tokens = completion_tokens = total_tokens = 0
+
+        # 如果tokens信息为空，提供默认值或提示
+        if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
+            tokens_info = "\n\n【tokens】暂无数据"
+        else:
+            tokens_info = (f"\n\n【tokens】输入: {prompt_tokens} 生成: {completion_tokens} 总: {total_tokens}")
+
+        # 将tokens信息添加到content后面，返回为字符串
+        content_with_tokens = content + tokens_info
+
+        return content_with_tokens
+
     except json.JSONDecodeError:
         print("Invalid JSON")
         return None
@@ -138,7 +160,6 @@ def message_summary(content):
     print(response.text)
     return extract_content(response.text)
 
-
 #  备用prompt ，找时间测试效果
 # prompt = f"""
 # 📊 **群聊话题统计与总结 Prompt**