""" Token 使用统计服务 """ import logging from typing import Dict, Any, List from datetime import datetime, timedelta, date from sqlalchemy.orm import Session from sqlalchemy import func, and_, extract from app.models.token_usage import TokenUsage from app.core.config import settings logger = logging.getLogger(__name__) class TokenService: """Token 统计服务类""" def __init__(self, db: Session): self.db = db def get_today_usage(self) -> int: """获取今日 Token 使用量""" today = date.today() result = self.db.query(func.sum(TokenUsage.tokens_used)).filter( func.date(TokenUsage.usage_date) == today ).scalar() return result or 0 def get_yesterday_usage(self) -> int: """获取昨日 Token 使用量""" yesterday = date.today() - timedelta(days=1) result = self.db.query(func.sum(TokenUsage.tokens_used)).filter( func.date(TokenUsage.usage_date) == yesterday ).scalar() return result or 0 def get_month_usage(self, year: Optional[int] = None, month: Optional[int] = None) -> int: """获取当月 Token 使用量""" if not year or not month: now = datetime.now() year = now.year month = now.month result = self.db.query(func.sum(TokenUsage.tokens_used)).filter( and_( extract('year', TokenUsage.usage_date) == year, extract('month', TokenUsage.usage_date) == month ) ).scalar() return result or 0 def get_remaining_tokens(self) -> int: """获取今日剩余 Token""" today_used = self.get_today_usage() remaining = settings.MAX_TOKENS_PER_DAY - today_used return max(0, remaining) def get_daily_usages(self, days: int = 30) -> List[Dict[str, Any]]: """ 获取每日 Token 使用(用于图表) :param days: 天数 :return: 每日使用列表 """ end_date = date.today() start_date = end_date - timedelta(days=days - 1) results = self.db.query( func.date(TokenUsage.usage_date).label('usage_date'), func.sum(TokenUsage.tokens_used).label('total_tokens') ).filter( and_( func.date(TokenUsage.usage_date) >= start_date, func.date(TokenUsage.usage_date) <= end_date ) ).group_by( func.date(TokenUsage.usage_date) ).order_by( func.date(TokenUsage.usage_date) ).all() # 转换为字典列表 usage_dict = {str(row.usage_date): row.total_tokens for row in results} # 填充缺失的日期 daily_usages = [] current_date = start_date while current_date <= end_date: date_str = str(current_date) tokens = usage_dict.get(date_str, 0) daily_usages.append({ "date": date_str, "tokens": tokens }) current_date += timedelta(days=1) return daily_usages def get_monthly_usages(self, months: int = 12) -> List[Dict[str, Any]]: """ 获取每月 Token 使用(用于图表) :param months: 月数 :return: 每月使用列表 """ now = datetime.now() results = [] for i in range(months): # 计算月份 month_offset = months - 1 - i target_date = now - timedelta(days=30 * month_offset) year = target_date.year month = target_date.month # 查询该月的使用量 usage = self.get_month_usage(year, month) results.append({ "month": f"{year}-{month:02d}", "tokens": usage }) return results def get_user_token_usage( self, user_id: int, days: int = 30 ) -> List[Dict[str, Any]]: """ 获取指定用户的 Token 使用 :param user_id: 用户 ID :param days: 天数 :return: 每日使用列表 """ end_date = date.today() start_date = end_date - timedelta(days=days - 1) results = self.db.query( func.date(TokenUsage.usage_date).label('usage_date'), func.sum(TokenUsage.tokens_used).label('total_tokens') ).filter( and_( TokenUsage.virtual_user_id == user_id, func.date(TokenUsage.usage_date) >= start_date, func.date(TokenUsage.usage_date) <= end_date ) ).group_by( func.date(TokenUsage.usage_date) ).order_by( func.date(TokenUsage.usage_date) ).all() return [ {"date": str(row.usage_date), "tokens": row.total_tokens} for row in results ] def check_token_limit_exceeded(self) -> bool: """检查是否超出 Token 限额""" today_used = self.get_today_usage() return today_used >= settings.MAX_TOKENS_PER_DAY # 工厂函数 def get_token_service(db: Session) -> TokenService: """获取 Token 服务实例""" return TokenService(db)