Upload 2 files

app.py

@@ -129,16 +129,28 @@ class AuthManager:
         self._session: requests.Session = create_custom_session()
         self._logger: logging.Logger = logging.getLogger(__name__)
         self.model_status = {model: True for model in MODEL_INFO.keys()}
-        # 添加新的属性来跟踪认证请求
         self._last_auth_attempt = 0
         self._auth_attempts = 0
         self._auth_window_start = time.time()
         self._backoff_delay = AUTH_RETRY_DELAY
+        self._failed_attempts = 0  # 添加失败尝试计数
+        self._last_failure_time = 0  # 添加最后失败时间
+        self._is_valid = True  # 添加账号有效性标志
 
     def _should_attempt_auth(self) -> bool:
         """检查是否应该尝试认证请求"""
         current_time = time.time()
         
+        # 如果账号已被标记为无效，直接返回False
+        if not self._is_valid:
+            return False
+            
+        # 如果连续失败次数过多，标记账号为无效
+        if self._failed_attempts >= 5:
+            self._is_valid = False
+            self._logger.warning(f"Account {self._email} marked as invalid due to too many failures")
+            return False
+            
         # 检查是否在退避期内
         if current_time - self._last_auth_attempt < self._backoff_delay:
             return False
@@ -149,6 +161,7 @@ class AuthManager:
             self._auth_window_start = current_time
             self._auth_attempts = 0
             self._backoff_delay = AUTH_RETRY_DELAY
+            self._failed_attempts = 0  # 重置失败计数
             
         # 检查请求数量
         if self._auth_attempts >= AUTH_MAX_REQUESTS:
@@ -156,10 +169,39 @@ class AuthManager:
             
         return True
 
+    def _make_request(self, method: str, url: str, **kwargs) -> requests.Response:
+        """改进的请求方法，支持代理"""
+        global proxy_pool
+        max_retries = 3
+        
+        for attempt in range(max_retries):
+            try:
+                # 获取代理
+                proxy = proxy_pool.get_proxy() if proxy_pool else None
+                
+                if proxy:
+                    kwargs['proxies'] = proxy
+                    kwargs['timeout'] = 10  # 使用代理时设置较短的超时
+                
+                response = self._session.request(method, url, **kwargs)
+                response.raise_for_status()
+                return response
+                
+            except requests.RequestException as e:
+                if proxy:
+                    proxy_pool.remove_proxy(proxy)
+                
+                if attempt == max_retries - 1:
+                    self._logger.error(f"请求错误 ({method} {url}): {e}")
+                    raise
+                
+                time.sleep(1)  # 短暂等待后重试
+        
+        raise requests.RequestException("Max retries exceeded")
+
     def login(self) -> bool:
-        """改进的登录方法，包含速率限制和退避机制"""
+        """改进的登录方法，使用代理"""
         if not self._should_attempt_auth():
-            logger.warning(f"Rate limit reached for {self._email}, waiting {self._backoff_delay}s")
             return False
 
         try:
@@ -178,29 +220,55 @@ class AuthManager:
             
             if response.status_code == 429:
                 self._backoff_delay *= AUTH_BACKOFF_FACTOR
+                self._failed_attempts += 1
+                self._last_failure_time = time.time()
                 logger.warning(f"Rate limit hit, increasing backoff to {self._backoff_delay}s")
                 return False
                 
-            response.raise_for_status()
+            if response.status_code == 400:
+                self._failed_attempts += 1
+                self._last_failure_time = time.time()
+                if self._failed_attempts >= 5:
+                    self._is_valid = False
+                    logger.error(f"Account {self._email} marked as invalid due to authentication failures")
+                return False
+                
             self._user_info = response.json()
             self._refresh_token = self._user_info.get('refresh_token', '')
             self._access_token = self._user_info.get('access_token', '')
             self._token_expiry = time.time() + self._user_info.get('expires_in', 3600)
             
-            # 重置退避延迟
+            # 重置失败相关计数
+            self._failed_attempts = 0
             self._backoff_delay = AUTH_RETRY_DELAY
             self._log_values()
             return True
             
         except requests.RequestException as e:
             logger.error(f"\033[91m登录请求错误: {e}\033[0m")
+            self._failed_attempts += 1
+            self._last_failure_time = time.time()
             self._backoff_delay *= AUTH_BACKOFF_FACTOR
             return False
 
+    def is_valid(self) -> bool:
+        """检查账号是否有效"""
+        return self._is_valid
+
+    def reset_status(self) -> None:
+        """重置账号状态"""
+        self._is_valid = True
+        self._failed_attempts = 0
+        self._backoff_delay = AUTH_RETRY_DELAY
+        self._auth_attempts = 0
+        self._auth_window_start = time.time()
+
     def refresh_user_token(self) -> bool:
+        """改进的令牌刷新方法，使用代理"""
         url = f"{_API_BASE_URL}/auth/v1/token?grant_type=refresh_token"
         headers = self._get_headers(with_content_type=True)
         data = {"refresh_token": self._refresh_token}
+        
         try:
             response = self._make_request('POST', url, headers=headers, json=data)
             self._user_info = response.json()
@@ -285,16 +353,6 @@ class AuthManager:
             headers['Authorization'] = f'Bearer {self._access_token}'
         return headers
 
-    def _make_request(self, method: str, url: str, **kwargs) -> requests.Response:
-        """发送HTTP请求并处理异常。"""
-        try:
-            response = self._session.request(method, url, **kwargs)
-            response.raise_for_status()
-            return response
-        except requests.RequestException as e:
-            self._logger.error(f"请求错误 ({method} {url}): {e}")
-            raise
-
     def is_model_available(self, model):
         return self.model_status.get(model, True)
 
@@ -324,7 +382,9 @@ class MultiAuthManager:
         start_index = self.current_index
         for _ in range(len(self.auth_managers)):
             auth_manager = self.auth_managers[self.current_index]
-            if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
+            if (auth_manager.is_valid() and 
+                auth_manager.is_model_available(model) and 
+                auth_manager._should_attempt_auth()):
                 return auth_manager
             self.current_index = (self.current_index + 1) % len(self.auth_managers)
             if self.current_index == start_index:
@@ -338,12 +398,13 @@ class MultiAuthManager:
                 return auth_manager
         return None
 
-    def reset_all_model_status(self):
+    def reset_all_accounts(self):
+        """重置所有账号状态"""
         for auth_manager in self.auth_managers:
-            auth_manager.reset_model_status()
+            auth_manager.reset_status()
 
 def require_auth(func: Callable) -> Callable:
-    """���饰器，确保在调用API之前有有效的token。"""
+    """饰器，确保在调用API之前有有效的token。"""
     @wraps(func)
     def wrapper(self, *args, **kwargs):
         if not self.ensure_valid_token():
@@ -788,6 +849,7 @@ def make_request(payload, auth_manager, model_id):
     global multi_auth_manager
     max_retries = 3
     retry_delay = 1
+    proxy_retries = 2  # 每个账号最多尝试2次代理
     
     logger.info(f"尝试发送请求，模型：{model_id}")
     
@@ -819,53 +881,64 @@ def make_request(payload, auth_manager, model_id):
         logger.info(f"尝试使用账号 {auth_manager._email}")
 
         for attempt in range(max_retries):
-            try:
-                url = get_notdiamond_url()
-                headers = get_notdiamond_headers(auth_manager)
-                
-                # 获取代理
-                proxy = proxy_pool.get_proxy()
-                
-                response = executor.submit(
-                    requests.post, 
-                    url, 
-                    headers=headers, 
-                    json=payload, 
-                    stream=True,
-                    proxies=proxy if proxy else None,
-                    timeout=30  # 添加超时设置
-                ).result()
-                
-                if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
-                    logger.info(f"请求成功，使用账号 {auth_manager._email}")
-                    return response
-                
-                # 如果请求失败且使用了代理，移除该代理
-                if proxy:
-                    proxy_pool.remove_proxy(proxy)
-                
-                headers_cache.clear()
-                
-                if response.status_code == 401:
-                    logger.info(f"Token expired for account {auth_manager._email}, attempting refresh")
-                    if auth_manager.ensure_valid_token():
+            for proxy_attempt in range(proxy_retries):
+                try:
+                    url = get_notdiamond_url()
+                    headers = get_notdiamond_headers(auth_manager)
+                    
+                    # 获取代理
+                    proxy = proxy_pool.get_proxy()
+                    if not proxy and proxy_attempt == 0:
+                        # 如果第一次获取代理失败，等待一下再试
+                        time.sleep(1)
                         continue
-                
-                if response.status_code == 403:
-                    logger.warning(f"Model {model_id} usage limit reached for account {auth_manager._email}")
-                    auth_manager.set_model_unavailable(model_id)
-                    break
-                
-                logger.error(f"Request failed with status {response.status_code} for account {auth_manager._email}")
-                
-            except Exception as e:
-                logger.error(f"Request attempt {attempt + 1} failed for account {auth_manager._email}: {e}")
-                # 如果是代理导致的错误，移除该代理
-                if proxy:
-                    proxy_pool.remove_proxy(proxy)
-                if attempt < max_retries - 1:
-                    time.sleep(retry_delay)
-                continue
+                    
+                    response = executor.submit(
+                        requests.post, 
+                        url, 
+                        headers=headers, 
+                        json=payload, 
+                        stream=True,
+                        proxies=proxy if proxy else None,
+                        timeout=30
+                    ).result()
+                    
+                    if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
+                        logger.info(f"请求成功，使用账号 {auth_manager._email}")
+                        return response
+                    
+                    # 如果请求失败且使用了代理，移除该代理
+                    if proxy:
+                        proxy_pool.remove_proxy(proxy)
+                        if proxy_attempt < proxy_retries - 1:
+                            # 如果还有代理重试机会，继续下一次尝试
+                            continue
+                    
+                    headers_cache.clear()
+                    
+                    if response.status_code == 401:
+                        logger.info(f"Token expired for account {auth_manager._email}, attempting refresh")
+                        if auth_manager.ensure_valid_token():
+                            break  # 跳出代理重试循环，使用新token重试
+                    
+                    if response.status_code == 403:
+                        logger.warning(f"Model {model_id} usage limit reached for account {auth_manager._email}")
+                        auth_manager.set_model_unavailable(model_id)
+                        break  # 跳出代理重试循环，尝试下一个账号
+                    
+                    logger.error(f"Request failed with status {response.status_code} for account {auth_manager._email}")
+                    break  # 跳出代理重试循环
+                    
+                except Exception as e:
+                    logger.error(f"Request attempt {attempt + 1} failed for account {auth_manager._email}: {e}")
+                    if proxy:
+                        proxy_pool.remove_proxy(proxy)
+                        if proxy_attempt < proxy_retries - 1:
+                            continue
+                    break  # 跳出代理重试循环
+            
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
 
     # 所有账号都尝试过且失败后，尝试注册新账号
     if len(tried_accounts) == len(multi_auth_manager.auth_managers):
@@ -879,43 +952,66 @@ def make_request(payload, auth_manager, model_id):
     raise Exception("所有账号均不可用，且注册新账号失败")
 
 def health_check():
-    """改进的健康检查函数"""
-    last_check_time = {}  # 用于跟踪每个账号的最后检查时间
+    """改进的健康检查函数，使用线程池控制并发"""
+    last_check_time = {}
+    check_interval = 60  # 每个账号的检查间隔（秒）
+    max_concurrent_checks = 3  # 最大并发检查数量
     
-    while True:
-        try:
-            if multi_auth_manager:
-                current_time = time.time()
-                
-                for auth_manager in multi_auth_manager.auth_managers:
-                    email = auth_manager._email
+    with ThreadPoolExecutor(max_workers=max_concurrent_checks) as executor:
+        while True:
+            try:
+                if multi_auth_manager:
+                    current_time = time.time()
+                    check_futures = []
                     
-                    # 检查是否需要进行健康检查
-                    if email not in last_check_time or \
-                       current_time - last_check_time[email] >= AUTH_CHECK_INTERVAL:
+                    for auth_manager in multi_auth_manager.auth_managers:
+                        email = auth_manager._email
                         
-                        if not auth_manager._should_attempt_auth():
-                            logger.info(f"Skipping health check for {email} due to rate limiting")
-                            continue
+                        # 检查是否需要进行健康检查
+                        if email not in last_check_time or \
+                           current_time - last_check_time[email] >= check_interval:
                             
-                        if not auth_manager.ensure_valid_token():
-                            logger.warning(f"Auth token validation failed during health check for {email}")
-                            auth_manager.clear_auth()
-                        else:
-                            logger.info(f"Health check passed for {email}")
+                            if not auth_manager._should_attempt_auth():
+                                logger.info(f"Skipping health check for {email} due to rate limiting")
+                                continue
                             
-                        last_check_time[email] = current_time
-                
-                # 每天重置所有账号的模型使用状态
-                current_time_local = time.localtime()
-                if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
-                    multi_auth_manager.reset_all_model_status()
-                    logger.info("Reset model status for all accounts")
+                            # 提交健康检查任务
+                            future = executor.submit(perform_health_check, auth_manager)
+                            check_futures.append((email, future))
+                            last_check_time[email] = current_time
                     
-        except Exception as e:
-            logger.error(f"Health check error: {e}")
+                    # 等待所有检查完成
+                    for email, future in check_futures:
+                        try:
+                            result = future.result(timeout=30)
+                            if result:
+                                logger.info(f"Health check passed for {email}")
+                            else:
+                                logger.warning(f"Health check failed for {email}")
+                        except Exception as e:
+                            logger.error(f"Health check error for {email}: {e}")
+                    
+                    # 每天重置所有账号的模型使用状态
+                    current_time_local = time.localtime()
+                    if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
+                        multi_auth_manager.reset_all_accounts()
+                        logger.info("Reset model status for all accounts")
+                    
+            except Exception as e:
+                logger.error(f"Health check error: {e}")
             
-        sleep(60)  # 主循环每分钟运行一次
+            time.sleep(10)  # 主循环每10秒运行一次
+
+def perform_health_check(auth_manager):
+    """执行单个账号的健康检查"""
+    try:
+        if auth_manager.ensure_valid_token():
+            return True
+        auth_manager.clear_auth()
+        return False
+    except Exception as e:
+        logger.error(f"Health check error for {auth_manager._email}: {e}")
+        return False
 
 # 为了兼容 Flask CLI 和 Gunicorn，修改启动逻辑
 if __name__ != "__main__":