# -*- coding: utf-8 -*- """ 后台任务调度器模块。 使用 APScheduler 设置和管理周期性任务，例如： - 清理旧日志文件。 - 每日重置 API Key 的使用计数 (RPD, TPD)。 - 定期生成并记录使用情况报告。 - 定期刷新 Key 分数缓存。 - 定期清理内存数据库中的旧上下文记录 (如果使用内存数据库)。 """ import logging # 导入日志模块 import asyncio # 导入 asyncio 库 from typing import TYPE_CHECKING, List, Dict, Any # 导入类型提示 from apscheduler.schedulers.asyncio import AsyncIOScheduler # 导入 APScheduler 的异步调度器 from apscheduler.executors.asyncio import AsyncIOExecutor # 导入 APScheduler 的异步执行器 # 从其他模块导入必要的组件 from app import config # 导入应用配置模块 from app.config import ( # 导入具体的配置项 USAGE_REPORT_INTERVAL_MINUTES, # 使用情况报告的间隔时间（分钟） MEMORY_CONTEXT_CLEANUP_INTERVAL_SECONDS, # 内存上下文清理任务的运行间隔（秒） CACHE_REFRESH_INTERVAL_SECONDS, # Key 分数缓存的刷新间隔（秒） REPORT_LOG_LEVEL_INT # 报告日志的级别 (整数形式) ) # 导入 ContextStore 类用于类型提示 from app.core.context.store import ContextStore # 导入数据库工具模块中的内存数据库标志 from app.core.database.utils import IS_MEMORY_DB # (新路径) # 导入日志配置模块中的日志清理函数 from app.handlers.log_config import cleanup_old_logs # (路径不变) # 导入报告模块中的任务函数 from app.core.reporting.daily_reset import reset_daily_counts # 每日计数重置函数 (新路径) from app.core.reporting.reporter import report_usage # 使用情况报告生成函数 (新路径) # 导入 Key 检查器模块中的分数刷新函数 (注意：下划线前缀表示内部使用) from app.core.keys.checker import _refresh_all_key_scores # (新路径) # 条件导入 APIKeyManager 用于类型提示，避免循环导入 if TYPE_CHECKING: from app.core.keys.manager import APIKeyManager # (新路径) logger = logging.getLogger('my_logger') # 获取日志记录器实例 # --- ANSI 转义码定义 (用于彩色日志输出) --- COLOR_TITLE = "\033[1;94m" # 亮蓝色 COLOR_SEPARATOR = "\033[0;90m" # 亮黑色/深灰色 COLOR_SECTION_HEADER = "\033[1;96m" # 亮青色 COLOR_POSITIVE = "\033[1;92m" # 亮绿色 COLOR_WARNING = "\033[1;93m" # 亮黄色 COLOR_ERROR = "\033[1;91m" # 亮红色 COLOR_INFO = "\033[0;37m" # 白色 COLOR_RESET = "\033[0m" # 重置颜色和样式 SEPARATOR_LINE = f"{COLOR_SEPARATOR}{'=' * 60}{COLOR_RESET}" # 定义分隔符行样式 # --- 创建调度器实例 --- # 使用 AsyncIOScheduler 适合在 asyncio 应用中使用 scheduler = AsyncIOScheduler( executors={ 'default': {'type': 'threadpool', 'max_workers': 20}, # 配置默认线程池执行器 'asyncio': AsyncIOExecutor() # 添加 asyncio 执行器，用于运行异步任务 } ) # --- 同步包装器 (用于从调度器调用异步函数) --- # run_cleanup_memory_context 同步包装器已不再需要，将直接调度异步方法。 # --- 报告格式化辅助函数 (用于 log_usage_report) --- # 这些函数将 report_usage 返回的字典数据格式化为带颜色的文本行，以便在日志中清晰显示。 def _format_key_usage_summary(models_summary: List[Dict[str, Any]]) -> List[str]: """(内部辅助函数) 格式化报告中的 Key 使用情况聚合部分。""" lines = [f"\n{SEPARATOR_LINE}\n{COLOR_SECTION_HEADER} Key 使用情况聚合 {COLOR_RESET}\n{SEPARATOR_LINE}"] # 添加章节标题和分隔符 if not models_summary: # 如果没有模型数据 lines.append(f" {COLOR_WARNING}暂无 Key 使用数据。{COLOR_RESET}") # 添加提示信息 else: # 遍历每个模型的摘要信息 for model_summary in models_summary: # 添加模型名称、Key 数量、总 RPD、总 TPD 输入 lines.append(f" {COLOR_POSITIVE}模型: {model_summary.get('model_name', 'N/A')} (Key 数量: {model_summary.get('key_count', 0)}){COLOR_RESET}") lines.append(f" 今日总 RPD: {model_summary.get('total_rpd_today', 0):,} | 今日总 TPD 输入: {model_summary.get('total_tpd_input_today', 0):,}") lines.append(" 状态分布:") # 添加状态分布小标题 status_distribution = model_summary.get("status_distribution", []) # 获取状态分布列表 if not status_distribution: # 如果没有状态分布数据 lines.append(f" {COLOR_WARNING}无状态数据。{COLOR_RESET}") # 添加提示信息 else: # 遍历每个状态及其对应的 Key 数量 for status_info in status_distribution: lines.append(f" - 数量: {status_info.get('count', 0)}, 状态: {status_info.get('status', 'N/A')}") return lines # 返回格式化后的文本行列表 def _format_overall_stats(report_data: Dict[str, Any]) -> List[str]: """(内部辅助函数) 格式化报告中的总体统计、预测、缓存使用情况和 Key 筛选跟踪部分。""" # 从报告数据中提取相关部分 overall_stats = report_data.get("overall_stats", {}) cache_stats = report_data.get("cache_stats", {}) key_selection_stats = report_data.get("key_selection_stats", {}) # 获取 Key 筛选统计数据 lines = [f"\n{SEPARATOR_LINE}\n{COLOR_SECTION_HEADER} 总体统计与预测 {COLOR_RESET}\n{SEPARATOR_LINE}"] # 添加章节标题 # 添加总体统计信息 lines.append(f" 活跃 Key 数量: {overall_stats.get('active_keys_count', 0)}") lines.append(f" 启动时无效 Key 数量: {overall_stats.get('invalid_keys_at_startup', 0)}") # 添加 RPD 容量估算 lines.append(" RPD 容量估算:") rpd_capacity_estimations = overall_stats.get("rpd_capacity_estimation", []) if not rpd_capacity_estimations: lines.append(f" {COLOR_WARNING}无 RPD 容量估算数据。{COLOR_RESET}") else: for estimation in rpd_capacity_estimations: if "capacity" in estimation: lines.append(f" - 基于模型 {estimation.get('based_on', 'N/A')}: 限制 {estimation.get('limit', 'N/A')}, 容量 {estimation.get('capacity', 0):,}") else: lines.append(f" - 基于模型 {estimation.get('based_on', 'N/A')}: {estimation.get('message', 'N/A')}") # 添加 TPD 输入容量估算 lines.append(" TPD 输入容量估算:") tpd_input_capacity_estimations = overall_stats.get("tpd_input_capacity_estimation", []) if not tpd_input_capacity_estimations: lines.append(f" {COLOR_WARNING}无 TPD 输入容量估算数据。{COLOR_RESET}") else: for estimation in tpd_input_capacity_estimations: if "capacity" in estimation: lines.append(f" - 基于模型 {estimation.get('based_on', 'N/A')}: 限制 {estimation.get('limit', 'N/A')}, 容量 {estimation.get('capacity', 0):,}") else: lines.append(f" - 基于模型 {estimation.get('based_on', 'N/A')}: {estimation.get('message', 'N/A')}") # 添加今日用量和预估 lines.append(f" 今日已用 RPD: {overall_stats.get('current_rpd_today', 0):,}") lines.append(f" 今日已用 TPD 输入: {overall_stats.get('current_tpd_input_today', 0):,}") lines.append(f" 预估今日总 RPD: {overall_stats.get('estimated_rpd_today', 'N/A')} (基于已过去 {overall_stats.get('estimation_fraction_of_day', 0):.1%} 时间)") lines.append(f" 预估今日总 TPD 输入: {overall_stats.get('estimated_tpd_input_today', 'N/A')} (基于已过去 {overall_stats.get('estimation_fraction_of_day', 0):.1%} 时间)") # 添加历史平均 RPD historical_average_usage = overall_stats.get("historical_average_usage", {}) lines.append(f" 历史平均每日 RPD ({historical_average_usage.get('days_included', 0)} 天): {historical_average_usage.get('avg_daily_rpd', 'N/A')}") # 添加缓存使用情况统计 lines.append(f"\n{COLOR_SECTION_HEADER} 缓存使用情况 {COLOR_RESET}") # 添加子标题 lines.append(f" 缓存命中次数: {cache_stats.get('hit_count', 0):,}") lines.append(f" 缓存未命中次数: {cache_stats.get('miss_count', 0):,}") lines.append(f" 节省的总 Token 数: {cache_stats.get('total_tokens_saved', 0):,}") lines.append(f" 缓存命中率: {cache_stats.get('hit_rate', 'N/A')}") # 添加命中率 # 添加 Key 筛选原因统计 lines.append(f"\n{COLOR_SECTION_HEADER} Key 筛选原因统计 {COLOR_RESET}") # 添加子标题 total_by_reason = key_selection_stats.get('total_by_reason', {}) # 获取按原因统计的总数 if not total_by_reason: # 如果没有数据 lines.append(f" {COLOR_WARNING}无 Key 筛选原因数据。{COLOR_RESET}") # 添加提示 else: lines.append(" 总计 (按原因):") # 添加小标题 # 按次数降序显示原因 for reason, count in sorted(total_by_reason.items(), key=lambda item: item[1], reverse=True): lines.append(f" - {reason}: {count:,} 次") # 可以选择性地添加 details_by_key 的格式化输出，但可能会很长 # details_by_key = key_selection_stats.get('details_by_key', {}) # if details_by_key: # lines.append("\n 详情 (按 Key 和原因):") # for key, reasons in sorted(details_by_key.items()): # lines.append(f" - Key '{key}':") # for reason, count in sorted(reasons.items(), key=lambda item: item[1], reverse=True): # lines.append(f" - {reason}: {count:,} 次") return lines # 返回格式化后的文本行列表 def _format_key_suggestion(suggestion: str) -> List[str]: """(内部辅助函数) 格式化报告中的 Key 数量建议部分，根据建议内容添加颜色。""" lines = [f"\n{SEPARATOR_LINE}\n{COLOR_SECTION_HEADER} Key 数量建议 {COLOR_RESET}\n{SEPARATOR_LINE}"] # 添加章节标题 # 根据建议内容判断颜色 if "强烈建议" in suggestion or "接近或达到上限" in suggestion: lines.append(f" {COLOR_ERROR}{suggestion}{COLOR_RESET}") # 紧急建议用红色 elif "建议增加" in suggestion or "用量较高" in suggestion: lines.append(f" {COLOR_WARNING}{suggestion}{COLOR_RESET}") # 普通增加建议用黄色 elif "可以考虑减少" in suggestion: lines.append(f" {COLOR_INFO}{suggestion}{COLOR_RESET}") # 减少建议用白色/默认色 else: # 保持当前数量的建议 lines.append(f" {COLOR_POSITIVE}{suggestion}{COLOR_RESET}") # 保持建议用绿色 return lines # 返回格式化后的文本行列表 def _format_top_ips(top_ips: Dict[str, Any]) -> List[str]: """(内部辅助函数) 格式化报告中的 Top 5 IP 地址统计部分。""" lines = [f"\n{SEPARATOR_LINE}\n{COLOR_SECTION_HEADER} Top 5 IP 地址统计 {COLOR_RESET}\n{SEPARATOR_LINE}"] # 添加章节标题 # --- Top 请求 IP --- lines.append(f"\n {COLOR_INFO}Top 请求 IP:{COLOR_RESET}") # 添加子标题 # 遍历时间段：今天、本周、本月 for period in ["today", "week", "month"]: lines.append(f" {period.capitalize()}:") # 添加时间段标题 ip_list = top_ips.get("requests", {}).get(period, []) # 获取对应时间段的 IP 列表 if not ip_list: # 如果列表为空 lines.append(f" {COLOR_WARNING}无数据。{COLOR_RESET}") # 添加提示 else: # 如果有数据 # 遍历 Top IP 信息 for ip_info in ip_list: lines.append(f" - {ip_info.get('ip', 'N/A')}: {ip_info.get('count', 0):,} 次请求") # 添加 IP 和请求次数 # --- Top Token IP (输入) --- lines.append(f"\n {COLOR_INFO}Top Token IP (输入):{COLOR_RESET}") # 添加子标题 # 遍历时间段 for period in ["today", "week", "month"]: lines.append(f" {period.capitalize()}:") # 添加时间段标题 ip_list = top_ips.get("tokens", {}).get(period, []) # 获取对应时间段的 IP 列表 if not ip_list: # 如果列表为空 lines.append(f" {COLOR_WARNING}无数据。{COLOR_RESET}") # 添加提示 else: # 如果有数据 # 遍历 Top IP 信息 for ip_info in ip_list: lines.append(f" - {ip_info.get('ip', 'N/A')}: {ip_info.get('tokens', 0):,} Tokens") # 添加 IP 和 Token 数量 return lines # 返回格式化后的文本行列表 # --- 周期性报告日志记录函数 --- def log_usage_report(key_manager: 'APIKeyManager'): """ (同步函数) 生成使用情况报告并将其记录到日志。 此函数由 APScheduler 定期调用。 Args: key_manager (APIKeyManager): APIKeyManager 的实例。 """ logger.info("尝试执行 log_usage_report 任务...") # 新增：确认任务是否被调用 logger.info("正在记录周期性使用情况报告到终端...") # 记录开始日志 try: # 1. 调用 report_usage 函数获取包含所有统计数据的字典 # report_usage 内部会获取锁并复制共享数据 report_data = report_usage(key_manager) # 2. (已移至 report_usage 内部) 获取缓存和 Key 筛选统计数据 # from app.core import tracking # cache_stats = {...} # key_selection_stats = {...} # report_data["cache_stats"] = cache_stats # report_data["key_selection_stats"] = key_selection_stats # 3. 使用辅助函数将报告数据字典格式化为带颜色的文本行列表 report_lines = [f"{COLOR_TITLE}--- API 使用情况报告 ({report_data.get('timestamp', 'N/A')}) ---{COLOR_RESET}"] # 报告标题 report_lines.extend(_format_key_usage_summary(report_data.get("key_usage_summary", {}).get("models", []))) # 格式化 Key 使用摘要 report_lines.extend(_format_overall_stats(report_data)) # 格式化总体统计、缓存、筛选 report_lines.extend(_format_key_suggestion(report_data.get("key_suggestion", "无建议。"))) # 格式化 Key 建议 report_lines.extend(_format_top_ips(report_data.get("top_ips", {}))) # 格式化 Top IP # 4. 将所有报告行合并为一个字符串 full_report_string = "\n".join(report_lines) # 5. 使用配置中定义的日志级别记录完整的报告字符串 logger.log(REPORT_LOG_LEVEL_INT, full_report_string) logger.info("周期性使用情况报告已记录到终端。") # 记录完成日志 except Exception as e: # 捕获报告生成或记录过程中发生的异常 logger.error(f"记录周期性使用情况报告到终端失败: {e}", exc_info=True) # 记录错误日志 # --- 调度器设置和启动 --- def _add_memory_context_cleanup_job(context_store_manager: ContextStore): """ (内部辅助函数) 向调度器添加用于清理内存上下文的定时任务。 仅在 CONTEXT_STORAGE_MODE 为 'memory' 时添加。 Args: context_store_manager (ContextStore): ContextStore 的实例。 """ if config.CONTEXT_STORAGE_MODE == "memory": # 检查是否为内存上下文模式 cleanup_interval = config.MEMORY_CONTEXT_CLEANUP_INTERVAL_SECONDS if cleanup_interval <= 0: logger.info("内存上下文清理任务的间隔配置为非正数，任务将不被添加。") return scheduler.add_job( context_store_manager.perform_memory_cleanup, # 直接调用异步方法 'interval', # 触发器类型：间隔触发 seconds=cleanup_interval, # 触发间隔（秒） id='memory_context_cleanup', # 作业的唯一 ID name='内存上下文清理 (ContextStore)', # 作业的名称 replace_existing=True, # 如果已存在同 ID 作业，则替换 executor='asyncio' # 确保使用 asyncio 执行器 ) logger.info(f"内存上下文清理任务 (ContextStore) 已添加，运行间隔: {cleanup_interval} 秒。") else: logger.info("非内存上下文存储模式，跳过添加内存上下文清理任务 (ContextStore)。") def setup_scheduler(key_manager: 'APIKeyManager', context_store_manager: ContextStore): """ 设置 APScheduler，添加所有需要的后台定时任务。 Args: key_manager (APIKeyManager): APIKeyManager 的实例。 context_store_manager (ContextStore): ContextStore 的实例。 """ logger.info("正在设置后台任务调度器...") # 记录开始设置日志 # --- 添加日志清理任务 --- # 使用 cron 触发器，每天凌晨 3:00 执行 # args=[30] 表示保留最近 30 天的日志 scheduler.add_job(cleanup_old_logs, 'cron', hour=3, minute=0, args=[30], id='log_cleanup', name='日志清理', replace_existing=True) # --- 添加每日计数重置任务 --- # 使用 cron 触发器，在太平洋时间 (PT) 的午夜 00:00 执行 # 指定 executor='asyncio' 来确保协程被正确执行 scheduler.add_job(reset_daily_counts, 'cron', hour=0, minute=0, timezone='America/Los_Angeles', id='daily_reset', name='每日限制重置', replace_existing=True, executor='asyncio') # --- 添加周期性使用报告任务 --- # 使用 interval 触发器，按配置的分钟数间隔执行 scheduler.add_job(log_usage_report, 'interval', minutes=USAGE_REPORT_INTERVAL_MINUTES, args=[key_manager], id='usage_report', name='使用报告', replace_existing=True) # --- 添加 Key 分数缓存更新任务 --- # 使用 interval 触发器，按配置的秒数间隔执行 # 注意：_refresh_all_key_scores 是一个异步函数，需要指定使用 asyncio 执行器 scheduler.add_job(_refresh_all_key_scores, 'interval', seconds=CACHE_REFRESH_INTERVAL_SECONDS, args=[key_manager], id='key_score_update', name='Key 得分更新', replace_existing=True, executor='asyncio') # --- 添加内存数据库上下文清理任务 (如果需要) --- _add_memory_context_cleanup_job(context_store_manager) # 记录已调度的任务名称 job_names = [job.name for job in scheduler.get_jobs()] logger.info(f"后台任务已调度: {', '.join(job_names)}") # 记录调度完成日志 def start_scheduler(): """ 启动后台任务调度器。 如果调度器已在运行，则不执行任何操作。 """ try: if not scheduler.running: # 检查调度器是否已在运行 scheduler.start() # 启动调度器 logger.info("后台调度器已启动。") # 记录启动日志 else: logger.info("后台调度器已在运行。") # 记录已运行日志 except Exception as e: # 捕获启动过程中可能发生的异常 logger.error(f"启动后台调度器失败: {e}", exc_info=True) # 记录错误日志 def shutdown_scheduler(): """ 关闭后台任务调度器。 通常在应用关闭时调用，以确保任务被优雅地停止。 """ if scheduler.running: # 检查调度器是否在运行 logger.info("正在关闭后台调度器...") # 记录关闭日志 scheduler.shutdown() # 关闭调度器 logger.info("后台调度器已关闭。") # 记录关闭完成日志