428 lines
15 KiB
Python
428 lines
15 KiB
Python
"""
|
||
Token 使用统计服务
|
||
|
||
负责记录、查询和分析 LLM 调用的 token 使用情况
|
||
数据持久化到 data/token_usage 目录，按月份组织
|
||
采用双层存储：
|
||
1. JSONL 文件 - 详细记录（按月存储）
|
||
2. 索引文件 - 快速聚合统计（按 API URL、日期等维度）
|
||
"""
|
||
import json
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional, Any
|
||
from datetime import datetime
|
||
from collections import defaultdict
|
||
|
||
try:
|
||
from backend.models.internal import TokenUsageRecord, TokenUsageStatus
|
||
from backend.core.config import settings
|
||
except ImportError:
|
||
from models.internal import TokenUsageRecord, TokenUsageStatus
|
||
from core.config import settings
|
||
|
||
|
||
class TokenUsageService:
|
||
"""
|
||
Token 使用统计服务
|
||
|
||
功能：
|
||
- 记录每次 LLM 调用的 token 使用情况
|
||
- 按月份、日期、角色、聊天、API URL 维度统计
|
||
- 支持中断和失败标记
|
||
- 数据持久化到文件系统（JSONL + 索引）
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.token_usage_dir = settings.DATA_PATH / "token_usage"
|
||
self.token_usage_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# ✅ 索引文件目录 - 用于快速聚合查询
|
||
self.index_dir = self.token_usage_dir / "indexes"
|
||
self.index_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
def _get_month_file(self, year: int, month: int) -> Path:
|
||
"""获取指定月份的统计文件路径"""
|
||
month_dir = self.token_usage_dir / f"{year}"
|
||
month_dir.mkdir(parents=True, exist_ok=True)
|
||
return month_dir / f"{month:02d}.jsonl"
|
||
|
||
def _load_month_records(self, year: int, month: int) -> List[TokenUsageRecord]:
|
||
"""加载指定月份的所有记录"""
|
||
file_path = self._get_month_file(year, month)
|
||
|
||
if not file_path.exists():
|
||
return []
|
||
|
||
records = []
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
for line in f:
|
||
if line.strip():
|
||
data = json.loads(line)
|
||
records.append(TokenUsageRecord(**data))
|
||
except Exception as e:
|
||
print(f"[TokenUsage] 加载记录失败: {e}")
|
||
|
||
return records
|
||
|
||
def _save_record(self, record: TokenUsageRecord):
|
||
"""保存单条记录到对应的月份文件"""
|
||
dt = datetime.fromtimestamp(record.timestamp)
|
||
file_path = self._get_month_file(dt.year, dt.month)
|
||
|
||
try:
|
||
with open(file_path, 'a', encoding='utf-8') as f:
|
||
f.write(json.dumps(record.model_dump(), ensure_ascii=False) + '\n')
|
||
|
||
# ✅ 同时更新索引文件（用于快速查询）
|
||
self._update_indexes(record)
|
||
except Exception as e:
|
||
print(f"[TokenUsage] 保存记录失败: {e}")
|
||
|
||
def _update_indexes(self, record: TokenUsageRecord):
|
||
"""
|
||
更新索引文件 - 实现高效的按维度聚合查询
|
||
|
||
索引结构：
|
||
- indexes/api_urls.json - 按 API URL 聚合
|
||
- indexes/daily/{year}-{month}.json - 按日聚合
|
||
"""
|
||
dt = datetime.fromtimestamp(record.timestamp)
|
||
|
||
# 1. 更新 API URL 索引
|
||
if record.apiUrl:
|
||
api_url_index = self.index_dir / "api_urls.json"
|
||
self._update_api_url_index(api_url_index, record)
|
||
|
||
# 2. 更新每日索引
|
||
daily_index = self.index_dir / "daily" / f"{dt.year}-{dt.month:02d}.json"
|
||
daily_index.parent.mkdir(parents=True, exist_ok=True)
|
||
self._update_daily_index(daily_index, record)
|
||
|
||
def _update_api_url_index(self, index_file: Path, record: TokenUsageRecord):
|
||
"""更新 API URL 索引文件"""
|
||
index_data = {}
|
||
|
||
# 加载现有索引
|
||
if index_file.exists():
|
||
try:
|
||
with open(index_file, 'r', encoding='utf-8') as f:
|
||
index_data = json.load(f)
|
||
except:
|
||
index_data = {}
|
||
|
||
# 更新统计
|
||
api_url = record.apiUrl
|
||
if api_url not in index_data:
|
||
index_data[api_url] = {
|
||
"totalPromptTokens": 0,
|
||
"totalCompletionTokens": 0,
|
||
"totalTokens": 0,
|
||
"count": 0,
|
||
"firstUsed": record.timestamp,
|
||
"lastUsed": record.timestamp
|
||
}
|
||
|
||
stats = index_data[api_url]
|
||
stats["totalPromptTokens"] += record.promptTokens
|
||
stats["totalCompletionTokens"] += record.completionTokens
|
||
stats["totalTokens"] += record.totalTokens
|
||
stats["count"] += 1
|
||
stats["lastUsed"] = max(stats["lastUsed"], record.timestamp)
|
||
stats["firstUsed"] = min(stats["firstUsed"], record.timestamp)
|
||
|
||
# 保存索引
|
||
with open(index_file, 'w', encoding='utf-8') as f:
|
||
json.dump(index_data, f, ensure_ascii=False, indent=2)
|
||
|
||
def _update_daily_index(self, index_file: Path, record: TokenUsageRecord):
|
||
"""更新每日索引文件"""
|
||
dt = datetime.fromtimestamp(record.timestamp)
|
||
day_key = f"{dt.year}-{dt.month:02d}-{dt.day:02d}"
|
||
|
||
index_data = {}
|
||
|
||
# 加载现有索引
|
||
if index_file.exists():
|
||
try:
|
||
with open(index_file, 'r', encoding='utf-8') as f:
|
||
index_data = json.load(f)
|
||
except:
|
||
index_data = {}
|
||
|
||
# 更新统计
|
||
if day_key not in index_data:
|
||
index_data[day_key] = {
|
||
"promptTokens": 0,
|
||
"completionTokens": 0,
|
||
"totalTokens": 0,
|
||
"count": 0
|
||
}
|
||
|
||
stats = index_data[day_key]
|
||
stats["promptTokens"] += record.promptTokens
|
||
stats["completionTokens"] += record.completionTokens
|
||
stats["totalTokens"] += record.totalTokens
|
||
stats["count"] += 1
|
||
|
||
# 保存索引
|
||
with open(index_file, 'w', encoding='utf-8') as f:
|
||
json.dump(index_data, f, ensure_ascii=False, indent=2)
|
||
|
||
async def record_usage(
|
||
self,
|
||
chat_id: str,
|
||
role_name: str,
|
||
chat_name: str,
|
||
prompt_tokens: int,
|
||
completion_tokens: int,
|
||
total_tokens: int,
|
||
status: TokenUsageStatus = TokenUsageStatus.COMPLETED,
|
||
message_id: Optional[str] = None,
|
||
floor: Optional[int] = None,
|
||
error_message: Optional[str] = None,
|
||
duration: Optional[float] = None,
|
||
model: Optional[str] = None,
|
||
api_provider: Optional[str] = None,
|
||
api_url: Optional[str] = None
|
||
) -> TokenUsageRecord:
|
||
"""
|
||
记录一次 LLM 调用的 token 使用情况
|
||
|
||
Args:
|
||
chat_id: 聊天ID
|
||
role_name: 角色名称
|
||
chat_name: 聊天名称
|
||
prompt_tokens: 输入 token 数
|
||
completion_tokens: 输出 token 数
|
||
total_tokens: 总 token 数
|
||
status: 请求状态
|
||
message_id: 关联的消息ID
|
||
floor: 楼层号
|
||
error_message: 错误信息
|
||
duration: 请求耗时
|
||
model: 使用的模型
|
||
api_provider: API 提供商
|
||
api_url: API URL地址
|
||
|
||
Returns:
|
||
TokenUsageRecord: 创建的记录
|
||
"""
|
||
record = TokenUsageRecord(
|
||
id=str(uuid.uuid4()),
|
||
chatId=chat_id,
|
||
roleName=role_name,
|
||
chatName=chat_name,
|
||
messageId=message_id,
|
||
floor=floor,
|
||
promptTokens=prompt_tokens,
|
||
completionTokens=completion_tokens,
|
||
totalTokens=total_tokens,
|
||
status=status,
|
||
errorMessage=error_message,
|
||
duration=duration,
|
||
model=model,
|
||
apiProvider=api_provider,
|
||
apiUrl=api_url
|
||
)
|
||
|
||
self._save_record(record)
|
||
return record
|
||
|
||
async def get_stats_by_month(
|
||
self,
|
||
year: int,
|
||
month: int,
|
||
role_name: Optional[str] = None,
|
||
chat_name: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
获取指定月份的统计数据
|
||
|
||
Args:
|
||
year: 年份
|
||
month: 月份
|
||
role_name: 角色名称（可选，用于过滤）
|
||
chat_name: 聊天名称（可选，用于过滤）
|
||
|
||
Returns:
|
||
统计数据字典
|
||
"""
|
||
records = self._load_month_records(year, month)
|
||
|
||
# 过滤
|
||
if role_name:
|
||
records = [r for r in records if r.roleName == role_name]
|
||
if chat_name:
|
||
records = [r for r in records if r.chatName == chat_name]
|
||
|
||
# 统计
|
||
total_prompt = sum(r.promptTokens for r in records)
|
||
total_completion = sum(r.completionTokens for r in records)
|
||
total_tokens = sum(r.totalTokens for r in records)
|
||
|
||
completed_count = sum(1 for r in records if r.status == TokenUsageStatus.COMPLETED)
|
||
interrupted_count = sum(1 for r in records if r.status == TokenUsageStatus.INTERRUPTED)
|
||
failed_count = sum(1 for r in records if r.status == TokenUsageStatus.FAILED)
|
||
|
||
# 按日期分组
|
||
daily_stats = defaultdict(lambda: {
|
||
"promptTokens": 0,
|
||
"completionTokens": 0,
|
||
"totalTokens": 0,
|
||
"count": 0
|
||
})
|
||
|
||
for r in records:
|
||
dt = datetime.fromtimestamp(r.timestamp)
|
||
day_key = f"{dt.year}-{dt.month:02d}-{dt.day:02d}"
|
||
daily_stats[day_key]["promptTokens"] += r.promptTokens
|
||
daily_stats[day_key]["completionTokens"] += r.completionTokens
|
||
daily_stats[day_key]["totalTokens"] += r.totalTokens
|
||
daily_stats[day_key]["count"] += 1
|
||
|
||
# 按角色分组
|
||
role_stats = defaultdict(lambda: {
|
||
"promptTokens": 0,
|
||
"completionTokens": 0,
|
||
"totalTokens": 0,
|
||
"count": 0
|
||
})
|
||
|
||
for r in records:
|
||
role_stats[r.roleName]["promptTokens"] += r.promptTokens
|
||
role_stats[r.roleName]["completionTokens"] += r.completionTokens
|
||
role_stats[r.roleName]["totalTokens"] += r.totalTokens
|
||
role_stats[r.roleName]["count"] += 1
|
||
|
||
# 按聊天分组
|
||
chat_stats = defaultdict(lambda: {
|
||
"promptTokens": 0,
|
||
"completionTokens": 0,
|
||
"totalTokens": 0,
|
||
"count": 0
|
||
})
|
||
|
||
for r in records:
|
||
chat_key = f"{r.roleName}/{r.chatName}"
|
||
chat_stats[chat_key]["promptTokens"] += r.promptTokens
|
||
chat_stats[chat_key]["completionTokens"] += r.completionTokens
|
||
chat_stats[chat_key]["totalTokens"] += r.totalTokens
|
||
chat_stats[chat_key]["count"] += 1
|
||
|
||
# ✅ 按 API URL 分组
|
||
api_url_stats = defaultdict(lambda: {
|
||
"promptTokens": 0,
|
||
"completionTokens": 0,
|
||
"totalTokens": 0,
|
||
"count": 0
|
||
})
|
||
|
||
for r in records:
|
||
if r.apiUrl:
|
||
api_url_stats[r.apiUrl]["promptTokens"] += r.promptTokens
|
||
api_url_stats[r.apiUrl]["completionTokens"] += r.completionTokens
|
||
api_url_stats[r.apiUrl]["totalTokens"] += r.totalTokens
|
||
api_url_stats[r.apiUrl]["count"] += 1
|
||
|
||
return {
|
||
"year": year,
|
||
"month": month,
|
||
"totalRecords": len(records),
|
||
"totalPromptTokens": total_prompt,
|
||
"totalCompletionTokens": total_completion,
|
||
"totalTokens": total_tokens,
|
||
"completedCount": completed_count,
|
||
"interruptedCount": interrupted_count,
|
||
"failedCount": failed_count,
|
||
"dailyStats": dict(daily_stats),
|
||
"roleStats": dict(role_stats),
|
||
"chatStats": dict(chat_stats),
|
||
"apiUrlStats": dict(api_url_stats), # ✅ 新增
|
||
"records": [r.model_dump() for r in records[:100]] # 最近100条记录
|
||
}
|
||
|
||
async def list_months(self) -> List[Dict[str, int]]:
|
||
"""列出所有有数据的月份"""
|
||
months = []
|
||
|
||
if not self.token_usage_dir.exists():
|
||
return months
|
||
|
||
for year_dir in sorted(self.token_usage_dir.iterdir()):
|
||
if year_dir.is_dir() and year_dir.name.isdigit():
|
||
year = int(year_dir.name)
|
||
for month_file in sorted(year_dir.glob("*.jsonl")):
|
||
month = int(month_file.stem)
|
||
months.append({"year": year, "month": month})
|
||
|
||
return months
|
||
|
||
async def get_api_url_stats(self) -> Dict[str, Any]:
|
||
"""
|
||
✅ 获取按 API URL 分组的统计数据（从索引文件快速读取）
|
||
|
||
Returns:
|
||
{api_url: {totalPromptTokens, totalCompletionTokens, totalTokens, count, firstUsed, lastUsed}}
|
||
"""
|
||
api_url_index = self.index_dir / "api_urls.json"
|
||
|
||
if not api_url_index.exists():
|
||
return {}
|
||
|
||
try:
|
||
with open(api_url_index, 'r', encoding='utf-8') as f:
|
||
return json.load(f)
|
||
except Exception as e:
|
||
print(f"[TokenUsage] 读取 API URL 索引失败: {e}")
|
||
return {}
|
||
|
||
async def get_daily_stats(self, year: int, month: int) -> Dict[str, Any]:
|
||
"""
|
||
✅ 获取指定月份的每日统计数据（从索引文件快速读取）
|
||
|
||
Args:
|
||
year: 年份
|
||
month: 月份
|
||
|
||
Returns:
|
||
{day_key: {promptTokens, completionTokens, totalTokens, count}}
|
||
"""
|
||
daily_index = self.index_dir / "daily" / f"{year}-{month:02d}.json"
|
||
|
||
if not daily_index.exists():
|
||
return {}
|
||
|
||
try:
|
||
with open(daily_index, 'r', encoding='utf-8') as f:
|
||
return json.load(f)
|
||
except Exception as e:
|
||
print(f"[TokenUsage] 读取每日索引失败: {e}")
|
||
return {}
|
||
|
||
async def get_available_roles(self, year: int, month: int) -> List[str]:
|
||
"""获取指定月份有数据的角色列表"""
|
||
records = self._load_month_records(year, month)
|
||
roles = set(r.roleName for r in records)
|
||
return sorted(list(roles))
|
||
|
||
async def get_available_chats(
|
||
self,
|
||
year: int,
|
||
month: int,
|
||
role_name: Optional[str] = None
|
||
) -> List[str]:
|
||
"""获取指定月份有数据的聊天列表"""
|
||
records = self._load_month_records(year, month)
|
||
|
||
if role_name:
|
||
records = [r for r in records if r.roleName == role_name]
|
||
|
||
chats = set(f"{r.roleName}/{r.chatName}" for r in records)
|
||
return sorted(list(chats))
|
||
|
||
|
||
# 全局实例
|
||
token_usage_service = TokenUsageService()
|