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}")