Spaces:
Running
Running
File size: 33,359 Bytes
b064311 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 |
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}") |