""" 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()