SiliconflowApiKeyPool/key_manager.py at main · Timmmmmms/SiliconflowApiKeyPool · GitHub

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
# 密钥管理模块
import time
from key_storage import KeyStorage
from balance_checker import BalanceChecker

class KeyManager:
    """
    密钥管理类，用于协调密钥存储和余额查询
    提供更高级的功能
    """

    def __init__(self, key_storage=None, balance_checker=None):
        """
        初始化密钥管理器

        Args:
            key_storage (KeyStorage, optional): 密钥存储对象
            balance_checker (BalanceChecker, optional): 余额查询对象
        """
        self.key_storage = key_storage or KeyStorage()
        self.balance_checker = balance_checker or BalanceChecker(self.key_storage)

    def add_key(self, key, check_balance=True):
        """
        添加一个新的API密钥

        Args:
            key (str): API密钥
            check_balance (bool): 是否立即检查余额

        Returns:
            tuple: (是否成功, 信息)
        """
        # 先添加到存储中
        if not self.key_storage.add_key(key):
            return (False, "密钥已存在")

        # 如果需要立即检查余额
        if check_balance:
            success, result = self.balance_checker.check_balance(key)
            if success:
                self.key_storage.update_key_status(key, result, time.time())
                return (True, f"添加成功，当前余额: {result}")
            else:
                self.key_storage.mark_key_invalid(key, result, time.time())
                return (False, f"添加失败，原因: {result}")

        return (True, "添加成功，等待余额检查")

    def batch_add_keys(self, keys, check_balance=True):
        """
        批量添加API密钥

        Args:
            keys (list): API密钥列表
            check_balance (bool): 是否立即检查余额

        Returns:
            tuple: (成功数量, 失败数量)
        """
        success_count = 0
        fail_count = 0

        for key in keys:
            success, _ = self.add_key(key, check_balance)
            if success:
                success_count += 1
            else:
                fail_count += 1

            # 添加延迟以避免API请求过于频繁
            if check_balance:
                time.sleep(0.5)

        return (success_count, fail_count)

    def get_available_key(self, min_balance=0.1):
        """
        获取一个可用的API密钥

        Args:
            min_balance (float): 最小余额要求

        Returns:
            str or None: 可用的API密钥，如果没有则返回None
        """
        valid_keys = self.key_storage.get_all_valid_keys()

        # 按余额从高到低排序
        sorted_keys = sorted(
            valid_keys.items(),
            key=lambda x: x[1].get("balance") or 0,
            reverse=True
        )

        for key, info in sorted_keys:
            if info.get("balance") and float(info["balance"]) >= min_balance:
                return key

        return None

    def update_all_keys_status(self):
        """
        更新所有密钥的状态

        Returns:
            tuple: (更新成功的密钥数, 更新失败的密钥数, 恢复的无效密钥数)
        """
        # 更新有效密钥
        success_count, fail_count = self.balance_checker.update_all_keys()

        # 检查部分无效密钥
        recovered_count = self.balance_checker.check_invalid_keys()

        return (success_count, fail_count, recovered_count)

    def get_statistics(self):
        """
        获取密钥池统计信息

        Returns:
            dict: 统计信息
        """
        valid_keys = self.key_storage.get_all_valid_keys()
        invalid_keys = self.key_storage.get_all_invalid_keys()

        # 计算总余额
        total_balance = sum(float(info.get("balance") or 0) for info in valid_keys.values())

        # 统计无效原因
        invalid_reasons = {}
        for info in invalid_keys.values():
            reason = info.get("reason", "未知原因")
            invalid_reasons[reason] = invalid_reasons.get(reason, 0) + 1

        return {
            "valid_count": len(valid_keys),
            "invalid_count": len(invalid_keys),
            "total_balance": total_balance,
            "invalid_reasons": invalid_reasons
        }

    def remove_key(self, key):
        """
        删除指定的密钥

        Args:
            key (str): 要删除的API密钥

        Returns:
            bool: 删除是否成功
        """
        return self.key_storage.remove_key(key)