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**