balance_checker.py

# 余额查询模块
import requests
import time
import concurrent.futures
from threading import Lock
from key_storage import KeyStorage

class BalanceChecker:
    """
    余额查询类，用于检查API密钥的余额
    """
    
    def __init__(self, key_storage, max_workers=10, request_interval=0.1):
        """
        初始化余额查询器
        
        Args:
            key_storage (KeyStorage): 密钥存储对象
            max_workers (int): 最大工作线程数
            request_interval (float): 请求间隔时间(秒)
        """
        self.key_storage = key_storage
        self.api_url = "https://api.siliconflow.cn/v1/user/info"
        self.max_workers = max_workers
        self.request_interval = request_interval
        self.lock = Lock()  # 用于线程安全操作
    
    def check_balance(self, key):
        """
        检查单个密钥的余额
        
        Args:
            key (str): API密钥
        
        Returns:
            tuple: (是否成功, 余额或错误信息)
        """
        headers = {
            "Authorization": f"Bearer {key}"
        }
        
        try:
            response = requests.get(self.api_url, headers=headers)
            data = response.json()
            
            # 检查API响应
            if data.get("status") and data.get("code") == 20000:
                balance = data.get("data", {}).get("balance")
                if balance is not None:
                    return (True, float(balance))
            
            # 如果响应不符合预期
            error_msg = data.get("message", "未知错误")
            return (False, error_msg)
            
        except Exception as e:
            return (False, str(e))
    
    def _process_key_result(self, key, success, result):
        """
        处理单个密钥的检查结果
        
        Args:
            key (str): API密钥
            success (bool): 检查是否成功
            result: 检查结果
            
        Returns:
            bool: 是否成功
        """
        current_time = time.time()
        
        with self.lock:  # 确保线程安全
            if success:
                # 更新密钥状态
                self.key_storage.update_key_status(key, result, current_time)
                return True
            else:
                # 标记密钥为无效
                self.key_storage.mark_key_invalid(key, result, current_time)
                return False
    
    def _check_key_with_delay(self, key):
        """
        检查单个密钥并添加延迟
        
        Args:
            key (str): API密钥
            
        Returns:
            tuple: (密钥, 是否成功, 结果)
        """
        success, result = self.check_balance(key)
        time.sleep(self.request_interval)  # 添加延迟以避免API请求过于频繁
        return (key, success, result)
    
    def update_all_keys(self):
        """
        更新所有密钥的余额信息 (多线程版本)
        
        Returns:
            tuple: (更新成功的密钥数, 更新失败的密钥数)
        """
        valid_keys = self.key_storage.get_all_valid_keys()
        keys_list = list(valid_keys.keys())
        
        success_count = 0
        fail_count = 0
        
        # 使用线程池并行处理所有密钥
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交所有任务
            future_to_key = {executor.submit(self._check_key_with_delay, key): key for key in keys_list}
            
            # 处理结果
            for future in concurrent.futures.as_completed(future_to_key):
                key, success, result = future.result()
                if self._process_key_result(key, success, result):
                    success_count += 1
                else:
                    fail_count += 1
        
        return (success_count, fail_count)
    
    def _process_invalid_key_result(self, key, success, result):
        """
        处理无效密钥的检查结果
        
        Args:
            key (str): API密钥
            success (bool): 检查是否成功
            result: 检查结果
            
        Returns:
            bool: 是否恢复有效
        """
        current_time = time.time()
        
        with self.lock:  # 确保线程安全
            if success and float(result) > 0:
                # 如果密钥恢复有效
                self.key_storage.update_key_status(key, result, current_time)
                return True
            else:
                # 更新检查时间
                if success:
                    reason = "余额不足"
                else:
                    reason = result
                self.key_storage.mark_key_invalid(key, reason, current_time)
                return False
    
    def check_invalid_keys(self, max_keys=10):
        """
        检查一部分无效密钥，看是否可以恢复 (多线程版本)
        
        Args:
            max_keys (int): 最大检查的无效密钥数量
        
        Returns:
            int: 恢复的密钥数量
        """
        invalid_keys = self.key_storage.get_all_invalid_keys()
        
        # 按最后检查时间排序，优先检查最久未检查的密钥
        sorted_keys = sorted(
            invalid_keys.items(),
            key=lambda x: x[1].get("last_check") or 0
        )
        
        # 限制检查数量
        keys_to_check = [key for key, _ in sorted_keys[:max_keys]]
        recovered_count = 0
        
        # 使用线程池并行处理选定的无效密钥
        with concurrent.futures.ThreadPoolExecutor(max_workers=min(self.max_workers, len(keys_to_check))) as executor:
            # 提交所有任务
            future_to_key = {executor.submit(self._check_key_with_delay, key): key for key in keys_to_check}
            
            # 处理结果
            for future in concurrent.futures.as_completed(future_to_key):
                key, success, result = future.result()
                if self._process_invalid_key_result(key, success, result):
                    recovered_count += 1
        
        return recovered_count