assist/src/analytics/token_monitor.py

# -*- coding: utf-8 -*-
"""
Token消耗监控模块
监控AI调用的Token使用情况和成本
"""

import json
import logging
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime, timedelta
from dataclasses import dataclass
from collections import defaultdict
from ..core.database import db_manager
from ..core.models import Conversation
from ..core.redis_manager import redis_manager
from ..config.config import Config

logger = logging.getLogger(__name__)

@dataclass
class TokenUsage:
    """Token使用记录"""
    timestamp: datetime
    user_id: str
    work_order_id: Optional[int]
    model_name: str
    input_tokens: int
    output_tokens: int
    total_tokens: int
    cost: float
    response_time: float
    success: bool
    error_message: Optional[str] = None

class TokenMonitor:
    """Token消耗监控器"""
    
    def __init__(self):
        # 使用统一的Redis管理器
        
        # Token价格配置（每1000个token的价格，单位：元）
        self.token_prices = {
            "qwen-plus-latest": {
                "input": 0.002,   # 输入token价格
                "output": 0.006   # 输出token价格
            },
            "qwen-turbo": {
                "input": 0.0008,
                "output": 0.002
            },
            "qwen-max": {
                "input": 0.02,
                "output": 0.06
            }
        }
        
        # 监控阈值
        self.thresholds = {
            "daily_cost_limit": 100.0,      # 每日成本限制（元）
            "hourly_cost_limit": 20.0,      # 每小时成本限制（元）
            "token_limit_per_request": 10000,  # 单次请求token限制
            "error_rate_threshold": 0.1     # 错误率阈值
        }
    
    def _get_redis_client(self):
        """获取Redis客户端"""
        return redis_manager.get_connection()
    
    def record_token_usage(
        self,
        user_id: str,
        work_order_id: Optional[int],
        model_name: str,
        input_tokens: int,
        output_tokens: int,
        response_time: float,
        success: bool = True,
        error_message: Optional[str] = None
    ) -> TokenUsage:
        """记录Token使用情况"""
        try:
            total_tokens = input_tokens + output_tokens
            
            # 计算成本
            cost = self._calculate_cost(model_name, input_tokens, output_tokens)
            
            # 创建使用记录
            usage = TokenUsage(
                timestamp=datetime.now(),
                user_id=user_id,
                work_order_id=work_order_id,
                model_name=model_name,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                cost=cost,
                response_time=response_time,
                success=success,
                error_message=error_message
            )
            
            # 保存到Redis
            self._save_to_redis(usage)
            
            # 检查阈值
            self._check_thresholds(usage)
            
            logger.info(f"Token使用记录: {total_tokens} tokens, 成本: {cost:.4f}元")
            return usage
            
        except Exception as e:
            logger.error(f"记录Token使用失败: {e}")
            return None
    
    def _calculate_cost(self, model_name: str, input_tokens: int, output_tokens: int) -> float:
        """计算Token成本"""
        if model_name not in self.token_prices:
            model_name = "qwen-plus-latest"  # 默认模型
        
        prices = self.token_prices[model_name]
        input_cost = (input_tokens / 1000) * prices["input"]
        output_cost = (output_tokens / 1000) * prices["output"]
        
        return input_cost + output_cost
    
    def _save_to_redis(self, usage: TokenUsage):
        """保存到Redis"""
        redis_client = self._get_redis_client()

        if not redis_client:
            return
        
        try:
            # 保存到时间序列
            timestamp = usage.timestamp.timestamp()
            usage_data = {
                "user_id": usage.user_id,
                "work_order_id": usage.work_order_id,
                "model_name": usage.model_name,
                "input_tokens": usage.input_tokens,
                "output_tokens": usage.output_tokens,
                "total_tokens": usage.total_tokens,
                "cost": usage.cost,
                "response_time": usage.response_time,
                "success": usage.success,
                "error_message": usage.error_message
            }
            
            # 保存到多个键
            redis_client.zadd(
                "token_usage:daily",
                {json.dumps(usage_data, ensure_ascii=False): timestamp}
            )
            
            redis_client.zadd(
                f"token_usage:user:{usage.user_id}",
                {json.dumps(usage_data, ensure_ascii=False): timestamp}
            )
            
            if usage.work_order_id:
                redis_client.zadd(
                    f"token_usage:work_order:{usage.work_order_id}",
                    {json.dumps(usage_data, ensure_ascii=False): timestamp}
                )
            
            # 设置过期时间（保留30天）
            redis_client.expire("token_usage:daily", 30 * 24 * 3600)
            
        except Exception as e:
            logger.error(f"保存Token使用到Redis失败: {e}")
    
    def _check_thresholds(self, usage: TokenUsage):
        """检查阈值并触发预警"""
        try:
            # 检查单次请求token限制
            if usage.total_tokens > self.thresholds["token_limit_per_request"]:
                self._trigger_alert(
                    "high_token_usage",
                    f"单次请求Token使用过多: {usage.total_tokens}",
                    "warning"
                )
            
            # 检查今日成本
            daily_cost = self.get_daily_cost(usage.timestamp.date())
            if daily_cost > self.thresholds["daily_cost_limit"]:
                self._trigger_alert(
                    "daily_cost_exceeded",
                    f"今日成本超限: {daily_cost:.2f}元",
                    "critical"
                )
            
            # 检查每小时成本
            hourly_cost = self.get_hourly_cost(usage.timestamp)
            if hourly_cost > self.thresholds["hourly_cost_limit"]:
                self._trigger_alert(
                    "hourly_cost_exceeded",
                    f"每小时成本超限: {hourly_cost:.2f}元",
                    "warning"
                )
            
        except Exception as e:
            logger.error(f"检查阈值失败: {e}")
    
    def _trigger_alert(self, alert_type: str, message: str, severity: str):
        """触发预警"""
        try:
            from ..core.models import Alert
            
            with db_manager.get_session() as session:
                alert = Alert(
                    rule_name=f"Token监控_{alert_type}",
                    alert_type=alert_type,
                    level=severity,
                    severity=severity,
                    message=message,
                    is_active=True,
                    created_at=datetime.now()
                )
                session.add(alert)
                session.commit()
                
            logger.warning(f"Token监控预警: {message}")
            
        except Exception as e:
            logger.error(f"触发Token监控预警失败: {e}")
    
    def get_daily_cost(self, date: datetime.date) -> float:
        """获取指定日期的成本"""
        try:
            redis_client = self._get_redis_client()

            if not redis_client:
                return 0.0
            
            start_time = datetime.combine(date, datetime.min.time()).timestamp()
            end_time = datetime.combine(date, datetime.max.time()).timestamp()
            
            # 从Redis获取当日数据
            usage_records = redis_client.zrangebyscore(
                "token_usage:daily",
                start_time,
                end_time,
                withscores=True
            )
            
            total_cost = 0.0
            for record_data, _ in usage_records:
                try:
                    record = json.loads(record_data)
                    total_cost += record.get("cost", 0)
                except json.JSONDecodeError:
                    continue
            
            return total_cost
            
        except Exception as e:
            logger.error(f"获取日成本失败: {e}")
            return 0.0
    
    def get_hourly_cost(self, timestamp: datetime) -> float:
        """获取指定小时的成本"""
        try:
            redis_client = self._get_redis_client()

            if not redis_client:
                return 0.0
            
            # 获取当前小时的数据
            hour_start = timestamp.replace(minute=0, second=0, microsecond=0)
            hour_end = hour_start + timedelta(hours=1)
            
            start_time = hour_start.timestamp()
            end_time = hour_end.timestamp()
            
            usage_records = redis_client.zrangebyscore(
                "token_usage:daily",
                start_time,
                end_time,
                withscores=True
            )
            
            total_cost = 0.0
            for record_data, _ in usage_records:
                try:
                    record = json.loads(record_data)
                    total_cost += record.get("cost", 0)
                except json.JSONDecodeError:
                    continue
            
            return total_cost
            
        except Exception as e:
            logger.error(f"获取小时成本失败: {e}")
            return 0.0
    
    def get_user_token_stats(self, user_id: str, days: int = 7) -> Dict[str, Any]:
        """获取用户Token使用统计"""
        try:
            redis_client = self._get_redis_client()

            if not redis_client:
                return {}
            
            end_time = datetime.now().timestamp()
            start_time = (datetime.now() - timedelta(days=days)).timestamp()
            
            usage_records = redis_client.zrangebyscore(
                f"token_usage:user:{user_id}",
                start_time,
                end_time,
                withscores=True
            )
            
            stats = {
                "total_tokens": 0,
                "total_cost": 0.0,
                "total_requests": 0,
                "successful_requests": 0,
                "failed_requests": 0,
                "avg_response_time": 0.0,
                "model_usage": defaultdict(int),
                "daily_usage": defaultdict(lambda: {"tokens": 0, "cost": 0})
            }
            
            response_times = []
            
            for record_data, timestamp in usage_records:
                try:
                    record = json.loads(record_data)
                    
                    stats["total_tokens"] += record.get("total_tokens", 0)
                    stats["total_cost"] += record.get("cost", 0)
                    stats["total_requests"] += 1
                    
                    if record.get("success", True):
                        stats["successful_requests"] += 1
                    else:
                        stats["failed_requests"] += 1
                    
                    model_name = record.get("model_name", "unknown")
                    stats["model_usage"][model_name] += 1
                    
                    if record.get("response_time"):
                        response_times.append(record["response_time"])
                    
                    # 按日期统计
                    date_str = datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d")
                    stats["daily_usage"][date_str]["tokens"] += record.get("total_tokens", 0)
                    stats["daily_usage"][date_str]["cost"] += record.get("cost", 0)
                    
                except json.JSONDecodeError:
                    continue
            
            # 计算平均响应时间
            if response_times:
                stats["avg_response_time"] = sum(response_times) / len(response_times)
            
            # 计算成功率
            if stats["total_requests"] > 0:
                stats["success_rate"] = stats["successful_requests"] / stats["total_requests"]
            else:
                stats["success_rate"] = 0
            
            return dict(stats)
            
        except Exception as e:
            logger.error(f"获取用户Token统计失败: {e}")
            return {}
    
    def get_system_token_stats(self, days: int = 7) -> Dict[str, Any]:
        """获取系统Token使用统计"""
        try:
            redis_client = self._get_redis_client()

            if not redis_client:
                return {}
            
            end_time = datetime.now().timestamp()
            start_time = (datetime.now() - timedelta(days=days)).timestamp()
            
            usage_records = redis_client.zrangebyscore(
                "token_usage:daily",
                start_time,
                end_time,
                withscores=True
            )
            
            stats = {
                "total_tokens": 0,
                "total_cost": 0.0,
                "total_requests": 0,
                "successful_requests": 0,
                "failed_requests": 0,
                "unique_users": set(),
                "model_usage": defaultdict(int),
                "daily_usage": defaultdict(lambda: {"tokens": 0, "cost": 0, "requests": 0})
            }
            
            for record_data, timestamp in usage_records:
                try:
                    record = json.loads(record_data)
                    
                    stats["total_tokens"] += record.get("total_tokens", 0)
                    stats["total_cost"] += record.get("cost", 0)
                    stats["total_requests"] += 1
                    
                    if record.get("success", True):
                        stats["successful_requests"] += 1
                    else:
                        stats["failed_requests"] += 1
                    
                    stats["unique_users"].add(record.get("user_id", ""))
                    
                    model_name = record.get("model_name", "unknown")
                    stats["model_usage"][model_name] += 1
                    
                    # 按日期统计
                    date_str = datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d")
                    stats["daily_usage"][date_str]["tokens"] += record.get("total_tokens", 0)
                    stats["daily_usage"][date_str]["cost"] += record.get("cost", 0)
                    stats["daily_usage"][date_str]["requests"] += 1
                    
                except json.JSONDecodeError:
                    continue
            
            # 计算成功率
            if stats["total_requests"] > 0:
                stats["success_rate"] = stats["successful_requests"] / stats["total_requests"]
            else:
                stats["success_rate"] = 0
            
            stats["unique_users"] = len(stats["unique_users"])
            
            return dict(stats)
            
        except Exception as e:
            logger.error(f"获取系统Token统计失败: {e}")
            return {}
    
    def get_cost_trend(self, days: int = 30) -> List[Dict[str, Any]]:
        """获取成本趋势"""
        try:
            trend_data = []
            
            for i in range(days):
                date = datetime.now().date() - timedelta(days=i)
                daily_cost = self.get_daily_cost(date)
                
                trend_data.append({
                    "date": date.isoformat(),
                    "cost": daily_cost
                })
            
            return list(reversed(trend_data))
            
        except Exception as e:
            logger.error(f"获取成本趋势失败: {e}")
            return []
    
    def cleanup_old_data(self, days: int = 30) -> int:
        """清理旧数据"""
        try:
            redis_client = self._get_redis_client()

            if not redis_client:
                return 0
            
            cutoff_time = (datetime.now() - timedelta(days=days)).timestamp()
            
            # 清理每日数据
            removed_count = redis_client.zremrangebyscore(
                "token_usage:daily",
                0,
                cutoff_time
            )
            
            # 清理用户数据
            user_keys = redis_client.keys("token_usage:user:*")
            for key in user_keys:
                redis_client.zremrangebyscore(key, 0, cutoff_time)
            
            # 清理工单数据
            work_order_keys = redis_client.keys("token_usage:work_order:*")
            for key in work_order_keys:
                redis_client.zremrangebyscore(key, 0, cutoff_time)
            
            logger.info(f"清理Token监控数据成功: 数量={removed_count}")
            return removed_count
            
        except Exception as e:
            logger.error(f"清理Token监控数据失败: {e}")
            return 0