import time import random import uuid import json import os import logging import threading from typing import Dict, List, Optional, Any, Union, Tuple, Callable # 初始化日志 logger = logging.getLogger("sora-api.key_manager") class KeyManager: def __init__(self, storage_file: str = "api_keys.json"): """ 初始化密钥管理器 Args: storage_file: 密钥存储文件路径 """ self.keys = [] # 密钥列表 self.storage_file = storage_file self.usage_stats = {} # 使用统计 self._lock = threading.RLock() # 添加可重入锁以支持并发访问 self._working_keys = {} # 新增：记录正在工作中的密钥 {key_value: task_id} self._load_keys() def _load_keys(self) -> None: """从环境变量或文件加载密钥""" keys_loaded = False # 先尝试从环境变量加载 api_keys_str = os.getenv("API_KEYS", "") if api_keys_str: try: env_data = json.loads(api_keys_str) self._process_keys_data(env_data) if len(self.keys) > 0: logger.info(f"已从环境变量加载 {len(self.keys)} 个密钥") keys_loaded = True else: logger.warning("环境变量API_KEYS存在但未包含有效密钥") except json.JSONDecodeError as e: logger.error(f"解析环境变量API keys失败: {str(e)}") # 如果环境变量未设置、解析失败或未加载到密钥，从文件加载 if not keys_loaded: try: if os.path.exists(self.storage_file): logger.info(f"尝试从文件加载密钥: {self.storage_file}") with open(self.storage_file, 'r', encoding='utf-8') as f: data = json.load(f) keys_before = len(self.keys) self._process_keys_data(data) keys_loaded = len(self.keys) > keys_before logger.info(f"已从文件加载 {len(self.keys) - keys_before} 个密钥") else: logger.warning(f"密钥文件不存在: {self.storage_file}") except Exception as e: logger.error(f"加载密钥失败: {str(e)}") if len(self.keys) == 0: logger.warning("未能从环境变量或文件加载任何密钥") def _process_keys_data(self, data): """处理不同格式的密钥数据""" # 处理不同的数据格式 if isinstance(data, list): # 旧版格式：直接是密钥列表 raw_keys = data self.keys = [] self.usage_stats = {} # 为每个密钥创建完整的记录 for key_info in raw_keys: if isinstance(key_info, dict): key_value = key_info.get("key") if not key_value: logger.warning(f"忽略无效密钥配置: {key_info}") continue # 确保有ID key_id = key_info.get("id") or str(uuid.uuid4()) # 构建完整的密钥记录 key_record = { "id": key_id, "name": key_info.get("name", ""), "key": key_value, "weight": key_info.get("weight", 1), "max_rpm": key_info.get("max_rpm", 60), "requests": 0, "last_reset": time.time(), "available": key_info.get("is_enabled", True), "is_enabled": key_info.get("is_enabled", True), "created_at": key_info.get("created_at", time.time()), "last_used": key_info.get("last_used"), "notes": key_info.get("notes") } self.keys.append(key_record) # 初始化使用统计 self.usage_stats[key_id] = { "total_requests": 0, "successful_requests": 0, "failed_requests": 0, "daily_usage": {}, "average_response_time": 0 } elif isinstance(key_info, str): # 如果是字符串，直接作为密钥值 key_id = str(uuid.uuid4()) self.keys.append({ "id": key_id, "name": "", "key": key_info, "weight": 1, "max_rpm": 60, "requests": 0, "last_reset": time.time(), "available": True, "is_enabled": True, "created_at": time.time(), "last_used": None, "notes": None }) # 初始化使用统计 self.usage_stats[key_id] = { "total_requests": 0, "successful_requests": 0, "failed_requests": 0, "daily_usage": {}, "average_response_time": 0 } else: # 新版格式：包含keys和usage_stats的字典 self.keys = data.get('keys', []) self.usage_stats = data.get('usage_stats', {}) def _save_keys(self) -> None: """保存密钥到文件""" try: with open(self.storage_file, 'w', encoding='utf-8') as f: json.dump({ 'keys': self.keys, 'usage_stats': self.usage_stats }, f, ensure_ascii=False, indent=2) # 同时更新Config中的API_KEYS try: from .config import Config Config.API_KEYS = self.keys except (ImportError, AttributeError): logger.debug("无法更新Config中的API_KEYS") except Exception as e: logger.error(f"保存密钥失败: {str(e)}") def add_key(self, key_value: str, name: str = "", weight: int = 1, rate_limit: int = 60, is_enabled: bool = True, notes: str = None) -> Dict[str, Any]: """ 添加密钥 Args: key_value: 密钥值 name: 密钥名称 weight: 权重 rate_limit: 速率限制(每分钟请求数) is_enabled: 是否启用 notes: 备注 Returns: 添加的密钥信息 """ with self._lock: # 使用锁保护添加过程 # 检查密钥是否已存在 for key in self.keys: if key.get("key") == key_value: return key key_id = str(uuid.uuid4()) new_key = { "id": key_id, "name": name, "key": key_value, "weight": weight, "max_rpm": rate_limit, "requests": 0, "last_reset": time.time(), "available": is_enabled, "is_enabled": is_enabled, "created_at": time.time(), "last_used": None, "notes": notes } self.keys.append(new_key) # 初始化使用统计 self.usage_stats[key_id] = { "total_requests": 0, "successful_requests": 0, "failed_requests": 0, "daily_usage": {}, "average_response_time": 0 } self._save_keys() logger.info(f"已添加密钥: {name or key_id}") return new_key def get_all_keys(self) -> List[Dict[str, Any]]: """获取所有密钥信息（已隐藏完整密钥值）""" with self._lock: # 使用锁保护读取过程 result = [] for key in self.keys: key_copy = key.copy() if "key" in key_copy: # 只显示密钥前6位和后4位 full_key = key_copy["key"] if len(full_key) > 10: key_copy["key"] = full_key[:6] + "..." + full_key[-4:] # 增加临时禁用信息的处理 if key_copy.get("temp_disabled_until"): temp_disabled_until = key_copy["temp_disabled_until"] # 确保temp_disabled_until是时间戳格式 if isinstance(temp_disabled_until, (int, float)): # 转换为可读格式，但保留原始时间戳，让前端可以自行处理 disabled_until_date = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(temp_disabled_until)) key_copy["temp_disabled_until_formatted"] = disabled_until_date key_copy["temp_disabled_remaining"] = int(temp_disabled_until - time.time()) result.append(key_copy) return result def get_key_by_id(self, key_id: str) -> Optional[Dict[str, Any]]: """根据ID获取密钥信息""" with self._lock: # 使用锁保护读取过程 for key in self.keys: if key.get("id") == key_id: return key return None def update_key(self, key_id: str, **kwargs) -> Optional[Dict[str, Any]]: """ 更新密钥信息 Args: key_id: 密钥ID **kwargs: 要更新的字段 Returns: 更新后的密钥信息，未找到则返回None """ with self._lock: # 使用锁保护更新过程 for key in self.keys: if key.get("id") == key_id: # 更新提供的字段 for field, value in kwargs.items(): if value is not None: if field == "is_enabled": key["available"] = value # 同步更新available字段 key[field] = value self._save_keys() logger.info(f"已更新密钥: {key.get('name') or key_id}") return key logger.warning(f"未找到密钥: {key_id}") return None def delete_key(self, key_id: str) -> bool: """ 删除密钥 Args: key_id: 密钥ID Returns: 是否成功删除 """ with self._lock: # 使用锁保护删除过程 original_length = len(self.keys) self.keys = [key for key in self.keys if key.get("id") != key_id] # 如果成功删除，保存密钥 if len(self.keys) < original_length: self._save_keys() return True return False def batch_import_keys(self, keys_data: List[Dict[str, Any]]) -> Dict[str, int]: """ 批量导入密钥 Args: keys_data: 密钥数据列表，每个元素为包含密钥信息的字典 Returns: 导入结果统计 """ with self._lock: # 使用锁保护导入过程 imported_count = 0 skipped_count = 0 # 获取现有密钥值 existing_keys = {key.get("key") for key in self.keys} for key_data in keys_data: key_value = key_data.get("key") if not key_value: continue # 检查密钥是否已存在 if key_value in existing_keys: skipped_count += 1 continue # 添加新密钥 key_id = str(uuid.uuid4()) new_key = { "id": key_id, "name": key_data.get("name", ""), "key": key_value, "weight": key_data.get("weight", 1), "max_rpm": key_data.get("rate_limit", 60), "requests": 0, "last_reset": time.time(), "available": key_data.get("enabled", True), "is_enabled": key_data.get("enabled", True), "created_at": time.time(), "last_used": None, "notes": key_data.get("notes") } self.keys.append(new_key) existing_keys.add(key_value) # 添加到已存在集合中 # 初始化使用统计 self.usage_stats[key_id] = { "total_requests": 0, "successful_requests": 0, "failed_requests": 0, "daily_usage": {}, "average_response_time": 0 } imported_count += 1 # 保存密钥 if imported_count > 0: self._save_keys() return { "imported": imported_count, "skipped": skipped_count } def get_key(self) -> Optional[str]: """获取下一个可用的密钥""" with self._lock: # 使用锁保护整个获取密钥过程 if not self.keys: logger.warning("没有可用的密钥") return None # 重置计数器（如果需要） current_time = time.time() temporary_disabled_updated = False for key in self.keys: # 检查是否有被临时禁用的密钥需要重新启用 if key.get("temp_disabled_until") and current_time > key.get("temp_disabled_until"): key["is_enabled"] = True key["available"] = True key["temp_disabled_until"] = None temporary_disabled_updated = True logger.info(f"密钥 {key.get('name') or key.get('id')} 的临时禁用已解除") if current_time - key["last_reset"] >= 60: key["requests"] = 0 key["last_reset"] = current_time if not key.get("temp_disabled_until"): # 只有未被临时禁用的密钥才会被重新激活 key["available"] = key.get("is_enabled", True) # 如果有任何临时禁用的密钥被更新，保存变更 if temporary_disabled_updated: self._save_keys() # 筛选可用的密钥，排除工作中的密钥 available_keys = [] for k in self.keys: key_value = k.get("key", "") clean_key = key_value.replace("Bearer ", "") if key_value.startswith("Bearer ") else key_value # 检查此密钥是否在工作中 is_working = clean_key in self._working_keys if k.get("available", False) and not is_working: available_keys.append(k) if not available_keys: logger.warning("没有可用的密钥（所有密钥都达到速率限制、被禁用或正在工作中）") return None # 根据权重选择密钥 weights = [k.get("weight", 1) for k in available_keys] selected_idx = random.choices(range(len(available_keys)), weights=weights, k=1)[0] selected_key = available_keys[selected_idx] # 更新使用统计 selected_key["requests"] += 1 selected_key["last_used"] = current_time # 检查是否达到速率限制 if selected_key["requests"] >= selected_key.get("max_rpm", 60): selected_key["available"] = False # 保存数据 - 并发环境下调整为每次都保存，避免状态不一致 # 原来是随机保存（10%的概率） self._save_keys() # 确保返回的密钥包含"Bearer "前缀 key_value = selected_key["key"] if not key_value.startswith("Bearer "): key_value = f"Bearer {key_value}" return key_value def record_request_result(self, key: str, success: bool, response_time: float = 0) -> None: """ 记录请求结果 Args: key: 密钥值 success: 请求是否成功 response_time: 响应时间(秒) """ if not key: logger.warning("记录请求结果失败：密钥为空") return with self._lock: # 使用锁保护记录过程 # 去掉可能的Bearer前缀 key_for_search = key.replace("Bearer ", "") if key.startswith("Bearer ") else key # 查找对应的密钥ID key_id = None key_info = None for k in self.keys: stored_key = k.get("key", "").replace("Bearer ", "") if k.get("key", "").startswith("Bearer ") else k.get("key", "") if stored_key == key_for_search: key_id = k.get("id") key_info = k break if not key_id: logger.warning(f"记录请求结果失败：未找到密钥 {key_for_search[:6]}...") return # 初始化usage_stats如果该密钥还没有统计数据 if key_id not in self.usage_stats: self.usage_stats[key_id] = { "total_requests": 0, "successful_requests": 0, "failed_requests": 0, "daily_usage": {}, "average_response_time": 0 } # 记录请求结果 stats = self.usage_stats[key_id] stats["total_requests"] += 1 if success: stats["successful_requests"] += 1 else: stats["failed_requests"] += 1 # 记录响应时间 if response_time > 0: if stats["average_response_time"] == 0: stats["average_response_time"] = response_time else: # 使用加权平均 old_avg = stats["average_response_time"] total = stats["total_requests"] # 避免 total 为0或1时产生问题，尽管前面 total_requests 已经增加了 if total > 0: stats["average_response_time"] = ((old_avg * (total - 1)) + response_time) / total else: # 理论上不应该发生，因为 total_requests 已经增加了 stats["average_response_time"] = response_time # 记录每日使用情况 today = time.strftime("%Y-%m-%d") if today not in stats["daily_usage"]: stats["daily_usage"][today] = {"successful": 0, "failed": 0} # 初始化每日统计 # 根据成功与否更新每日统计 if success: stats["daily_usage"][today]["successful"] += 1 else: stats["daily_usage"][today]["failed"] += 1 # 保留最近30天的数据 if len(stats["daily_usage"]) > 30: # 获取所有日期并排序，然后删除最早的 sorted_dates = sorted(stats["daily_usage"].keys()) if sorted_dates: # 确保列表不为空 oldest_date = sorted_dates[0] del stats["daily_usage"][oldest_date] # 更新密钥的最后使用时间 if key_info and "last_used" in key_info: key_info["last_used"] = time.time() # 并发环境下每次都保存，确保统计准确性 self._save_keys() def get_usage_stats(self) -> Dict[str, Any]: """获取使用统计信息""" with self._lock: # 使用锁保护读取过程 total_keys = len(self.keys) active_keys = sum(1 for k in self.keys if k.get("is_enabled", False)) available_keys = sum(1 for k in self.keys if k.get("available", False)) total_requests = sum(stats.get("total_requests", 0) for stats in self.usage_stats.values()) successful_requests = sum(stats.get("successful_requests", 0) for stats in self.usage_stats.values()) # 计算成功率 success_rate = (successful_requests / total_requests * 100) if total_requests > 0 else 0 # 计算每个密钥的平均响应时间 avg_response_times = [stats.get("average_response_time", 0) for stats in self.usage_stats.values() if stats.get("average_response_time", 0) > 0] overall_avg_response_time = sum(avg_response_times) / len(avg_response_times) if avg_response_times else 0 # 获取过去7天的使用情况 past_7_days = {} for key_id, stats in self.usage_stats.items(): daily_usage = stats.get("daily_usage", {}) for date, count_data in daily_usage.items(): if date not in past_7_days: past_7_days[date] = {"successful": 0, "failed": 0} # 正确处理字典类型的count_data past_7_days[date]["successful"] += count_data.get("successful", 0) past_7_days[date]["failed"] += count_data.get("failed", 0) # 只保留最近7天 dates = sorted(past_7_days.keys(), reverse=True)[:7] past_7_days = {date: past_7_days[date] for date in dates} return { "total_keys": total_keys, "active_keys": active_keys, "available_keys": available_keys, "total_requests": total_requests, "successful_requests": successful_requests, "failed_requests": total_requests - successful_requests, "success_rate": success_rate, "average_response_time": overall_avg_response_time, "past_7_days": past_7_days } def mark_key_as_working(self, key: str, task_id: str) -> None: """ 将密钥标记为工作中状态 Args: key: API密钥值（可能包含Bearer前缀） task_id: 关联的任务ID """ with self._lock: clean_key = key.replace("Bearer ", "") if key.startswith("Bearer ") else key self._working_keys[clean_key] = task_id logger.debug(f"密钥已标记为工作中，关联任务ID: {task_id}") def release_key(self, key: str) -> None: """ 释放工作中的密钥 Args: key: API密钥值（可能包含Bearer前缀） """ with self._lock: clean_key = key.replace("Bearer ", "") if key.startswith("Bearer ") else key if clean_key in self._working_keys: del self._working_keys[clean_key] logger.debug(f"密钥已释放") def is_key_working(self, key: str) -> bool: """ 检查密钥是否正在工作中 Args: key: API密钥值（可能包含Bearer前缀） Returns: bool: 是否在工作中 """ with self._lock: clean_key = key.replace("Bearer ", "") if key.startswith("Bearer ") else key return clean_key in self._working_keys def mark_key_invalid(self, key: str) -> Optional[str]: """ 将指定的密钥标记为无效（临时禁用而不是永久禁用），并返回一个新的可用密钥 Args: key: API密钥值（可能包含Bearer前缀） Returns: Optional[str]: 新的可用密钥，如果没有可用密钥则返回None """ # 调用临时禁用方法，设置24小时禁用时间 return self.mark_key_temp_disabled(key, hours=24.0) def mark_key_temp_disabled(self, key: str, hours: float = 12.0) -> Optional[str]: """ 将指定的密钥临时禁用指定小时数，并返回一个新的可用密钥 Args: key: API密钥值（可能包含Bearer前缀） hours: 禁用小时数 Returns: Optional[str]: 新的可用密钥，如果没有可用密钥则返回None """ with self._lock: # 使用锁保护临时禁用过程 # 去掉可能的Bearer前缀 key_for_search = key.replace("Bearer ", "") if key.startswith("Bearer ") else key # 检查是否是因为密钥在工作中导致的错误 if key_for_search in self._working_keys: logger.warning(f"尝试禁用正在工作中的密钥（任务ID: {self._working_keys[key_for_search]}），跳过禁用操作") # 获取一个新密钥返回，但不禁用当前密钥 new_key = self.get_key() if new_key: logger.info(f"已返回新密钥，但未禁用工作中的密钥") return new_key else: logger.warning("没有可用的备用密钥") return None # 查找对应的密钥 key_found = False disabled_key_id = None for key_info in self.keys: stored_key = key_info.get("key", "").replace("Bearer ", "") if key_info.get("key", "").startswith("Bearer ") else key_info.get("key", "") if stored_key == key_for_search: # 标记密钥为临时禁用 disabled_until = time.time() + (hours * 3600) # 当前时间加上禁用小时数 key_info["available"] = False key_info["temp_disabled_until"] = disabled_until key_info["notes"] = (key_info.get("notes") or "") + f"\n[自动] 在 {time.strftime('%Y-%m-%d %H:%M:%S')} 被临时禁用{hours}小时" key_found = True disabled_key_id = key_info.get("id") logger.warning(f"密钥 {key_info.get('name') or key_info.get('id')} 被临时禁用{hours}小时") break if key_found: # 保存更改 self._save_keys() # 获取新的密钥，排除已禁用的 new_key = self.get_key() if new_key: logger.info(f"已自动切换到新的密钥") return new_key else: logger.warning("没有可用的备用密钥") return None else: logger.warning(f"未找到要临时禁用的密钥") return None def retry_request(self, original_key: str, request_func: Callable, max_retries: int = 1, max_key_switches: int = 3) -> Tuple[bool, Any, str]: """ 出错时自动重试请求，并在需要时切换密钥 Args: original_key: 原始API密钥（可能包含Bearer前缀） request_func: 执行请求的函数，接受一个参数(密钥)并返回(成功标志, 结果) max_retries: 使用同一密钥的最大重试次数 max_key_switches: 最大密钥切换次数 Returns: Tuple[bool, Any, str]: (是否成功, 请求结果, 使用的密钥) """ current_key = original_key current_key_switches = 0 # 首先用原始密钥尝试 for attempt in range(max_retries + 1): # +1是因为第一次不算重试 try: success, result = request_func(current_key) # 成功的请求不应该导致密钥被禁用 if success: # 记录请求成功，避免不必要的密钥禁用 with self._lock: self.record_request_result(current_key, True) return True, result, current_key logger.warning(f"请求失败(尝试 {attempt+1}/{max_retries+1}): {result}") except Exception as e: logger.error(f"请求异常(尝试 {attempt+1}/{max_retries+1}): {str(e)}") # 如果这不是最后一次尝试，等待一秒后重试 if attempt < max_retries: time.sleep(1) # 如果原始密钥的所有重试都失败，尝试切换密钥 tried_keys = set([current_key.replace("Bearer ", "") if current_key.startswith("Bearer ") else current_key]) while current_key_switches < max_key_switches: # 获取新的密钥 with self._lock: new_key = self.get_key() if not new_key: logger.warning("没有更多可用的密钥") break # 确保不使用已经尝试过的密钥 clean_new_key = new_key.replace("Bearer ", "") if new_key.startswith("Bearer ") else new_key if clean_new_key in tried_keys: continue tried_keys.add(clean_new_key) current_key = new_key current_key_switches += 1 logger.info(f"切换到新密钥 (切换 {current_key_switches}/{max_key_switches})") # 用新密钥尝试 for attempt in range(max_retries + 1): try: success, result = request_func(current_key) if success: # 记录请求成功 with self._lock: self.record_request_result(current_key, True) return True, result, current_key logger.warning(f"使用新密钥请求失败(尝试 {attempt+1}/{max_retries+1}): {result}") except Exception as e: logger.error(f"使用新密钥请求异常(尝试 {attempt+1}/{max_retries+1}): {str(e)}") # 如果这不是最后一次尝试，等待一秒后重试 if attempt < max_retries: time.sleep(1) # 所有尝试都失败，临时禁用原始密钥 # 但是在并发环境下，这可能是因为网络或服务问题，而非密钥问题 # 增加额外检查以减少不必要的密钥禁用 should_disable = True # 在临时禁用前，确认是否是密钥问题而非服务或网络问题 # 此处可以添加额外逻辑来判断是否应该禁用密钥 if should_disable: logger.error(f"所有重试和密钥切换尝试都失败，临时禁用原始密钥") with self._lock: self.mark_key_temp_disabled(original_key, hours=6.0) # 减少禁用时间，避免资源浪费 else: logger.warning(f"所有重试和密钥切换尝试都失败，但可能是服务问题而非密钥问题，不禁用密钥") # 返回最后一次尝试的结果 return False, result, current_key # 创建全局密钥管理器实例 storage_file = os.getenv("KEYS_STORAGE_FILE", "api_keys.json") # 如果提供了绝对路径则直接使用，否则使用相对路径 if not os.path.isabs(storage_file): base_dir = os.getenv("BASE_DIR", os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) storage_file = os.path.join(base_dir, storage_file) key_manager = KeyManager(storage_file=storage_file) logger.info(f"初始化全局密钥管理器，存储文件: {storage_file}")