From c9e9e0262ea5bbc99e715ad55a3f035c7d73c20d Mon Sep 17 00:00:00 2001
From: liuwei <liuwei@wdtrgf.com.cn>
Date: Mon, 10 Mar 2025 12:21:30 +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

---
 base/func_doubao.py              | 26 +++++++++++++++++++--
 message_summary/message_split.py | 39 --------------------------------
 2 files changed, 24 insertions(+), 41 deletions(-)
 delete mode 100644 message_summary/message_split.py

diff --git a/base/func_doubao.py b/base/func_doubao.py
index 0007acb..1f70ceb 100644
--- a/base/func_doubao.py
+++ b/base/func_doubao.py
@@ -89,13 +89,35 @@ class Doubao():
         return False
 
 
-# 解析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
diff --git a/message_summary/message_split.py b/message_summary/message_split.py
deleted file mode 100644
index 2455b75..0000000
--- a/message_summary/message_split.py
+++ /dev/null
@@ -1,39 +0,0 @@
-import json
-
-# 模拟返回的API数据
-data_string = '''{
-    "id":"chatcmpl-1739339885",
-    "object":"chat.completion",
-    "created":1739339885,
-    "model":"windsurf",
-    "choices":[
-        {
-            "index":0,
-            "message":{
-                "role":"assistant",
-                "content":"[聊天主题] 婚姻与创业话题讨论  \n☆☆参与者：wxid_atv7kzgxv9lg21, wxid_j1f8su26egm222, wxid_nn630aasb1nb22, wxid_q1ugj6gbjj3o12, wxid_ne15uczjd2mi21, wxid_ayf2rispnhwi21  \n时间点：2025-02-12 12:59:14 - 2025-02-12 13:56:02  \n内容摘要： 本次聊天涵盖了多个话题，包括使用道具、杭州新兴科技公司、医疗手术、杜蕾斯品牌、创业公司、以及婚姻中性行为的问题。wxid_j1f8su26egm222提到杭州出现了新的科技公司，而wxid_ayf2rispnhwi21则提出了关于婚姻和性行为的疑问，引发了群内的讨论。  \n☑ 点评： 讨论内容较为分散，涉及多个话题。关于杭州科技公司的讨论显示出对当地创业环境的关注，而关于婚姻和性行为的问题则反映了个人生活中的实际困惑。这些话题虽然不直接相关，但都反映了群成员对生活不同方面的关注。  \n· 趣味互动 wxid_ayf2rispnhwi21关于婚姻和性行为的提问引发了群内的笑声和调侃，体现了群成员之间轻松的互动氛围。  \n· 待跟进事项 无具体行动项。未来可以考虑深入讨论某一特定话题，如创业或婚姻问题。  \n· 结语 今日的讨论涉及多个生活领域，展现了群成员对不同话题的兴趣和幽默感。期待未来能有更多深入的讨论，鼓励大家继续积极参与。  \n注： 本次讨论话题较为多样，未深入展开某一单一主题。"
-            },
-            "finish_reason":"stop"
-        }
-    ],
-    "usage":{
-        "completion_tokens":1024,
-        "prompt_tokens":0,
-        "total_tokens":1024
-    }
-}'''
-
-# 解析JSON
-def extract_content(data_string):
-    try:
-        data = json.loads(data_string)
-        # 提取content字段
-        content = data["choices"][0]["message"].get("content", "")
-        return content
-    except json.JSONDecodeError:
-        print("Invalid JSON")
-        return None
-
-# 获取并打印content内容
-content = extract_content(data_string)
-print(f"Extracted Content: {content}")