Upload app.py

app.py

@@ -273,7 +273,7 @@ class AuthManager:
             return ""
 
     def _get_headers(self, with_content_type: bool = False) -> Dict[str, str]:
-        """成请求头。"""
+        """生成请求头。"""
         headers = {
             'apikey': self._fetch_apikey(),
             'user-agent': _USER_AGENT
@@ -309,53 +309,48 @@ class MultiAuthManager:
         self.current_index = 0
         self._last_rotation = time.time()
         self._rotation_interval = 300  # 5分钟轮转间隔
-        self._last_successful_index = 0  # 记录上一次成功的账号索引
-        self._last_success_date = datetime.now().date()  # 记录上一次成功的日期
-
-    def _should_rotate(self) -> bool:
-        """检查是否应该轮转到下一个账号"""
-        current_date = datetime.now().date()
-        # 如果是新的一天且上一个成功的账号失败了，才进行完整轮询
-        return current_date > self._last_success_date
+        self.last_successful_index = 0  # 新增：记录上一次成功的账号索引
+        self.daily_rotation_triggered = False  # 新增：标记是否触发过每日轮询
 
     def get_next_auth_manager(self, model):
-        """改进的账号选择逻辑，优先使用上次成功的账号"""
-        current_date = datetime.now().date()
+        """改进的账号选择逻辑，从上一次成功的账号开始"""
+        # 如果当前账号可用，直接返回
+        current_auth = self.auth_managers[self.last_successful_index]
+        if current_auth.is_model_available(model) and current_auth._should_attempt_auth():
+            return current_auth
+
+        # 如果当前账号不可用，且未触发每日轮询，则等待到第二天
+        if not self.daily_rotation_triggered:
+            return None
+
+        # 进行账号轮询
+        start_index = (self.last_successful_index + 1) % len(self.auth_managers)
+        current_index = start_index
         
-        # 首先尝试上次成功的账号
-        auth_manager = self.auth_managers[self._last_successful_index]
-        if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
-            return auth_manager
-            
-        # 如果是新的一天，或者上次成功的账号不可用，才进行完整轮询
-        if self._should_rotate():
-            start_index = (self._last_successful_index + 1) % len(self.auth_managers)
-            current_index = start_index
-            
-            while True:
-                auth_manager = self.auth_managers[current_index]
-                if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
-                    self._last_successful_index = current_index
-                    return auth_manager
-                    
-                current_index = (current_index + 1) % len(self.auth_managers)
-                if current_index == start_index:
-                    break
-                    
+        while True:
+            auth_manager = self.auth_managers[current_index]
+            if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
+                self.last_successful_index = current_index
+                return auth_manager
+                
+            current_index = (current_index + 1) % len(self.auth_managers)
+            if current_index == start_index:
+                break
+                
         return None
 
-    def update_last_success(self, auth_manager):
-        """更新最后一次成功使用的账号信息"""
+    def set_successful_account(self, auth_manager):
+        """记录成功使用的账号索引"""
         for i, manager in enumerate(self.auth_managers):
             if manager._email == auth_manager._email:
-                self._last_successful_index = i
-                self._last_success_date = datetime.now().date()
+                self.last_successful_index = i
                 break
 
     def ensure_valid_token(self, model):
-        auth_manager = self.get_next_auth_manager(model)
-        if auth_manager and auth_manager.ensure_valid_token():
-            return auth_manager
+        for _ in range(len(self.auth_managers)):
+            auth_manager = self.get_next_auth_manager(model)
+            if auth_manager and auth_manager.ensure_valid_token():
+                return auth_manager
         return None
 
     def reset_all_model_status(self):
@@ -800,7 +795,7 @@ def build_payload(request_data, model_id):
     return payload
 
 def make_request(payload, auth_manager, model_id):
-    """发送请求并处理可能的认证刷新和模型定错误。"""
+    """发送请求并处理可能的认证刷新和模型特定错误。"""
     global multi_auth_manager
     max_retries = 3
     retry_delay = 1
@@ -826,6 +821,9 @@ def make_request(payload, auth_manager, model_id):
     while len(tried_accounts) < len(multi_auth_manager.auth_managers):
         auth_manager = multi_auth_manager.get_next_auth_manager(model_id)
         if not auth_manager:
+            if not multi_auth_manager.daily_rotation_triggered:
+                logger.info("当前账号不可用，等待至次日重试")
+                raise Exception("当前账号不可用，请等待至次日重试")
             break
             
         # 如果这个账号已经尝试过，继续下一个
@@ -849,8 +847,7 @@ def make_request(payload, auth_manager, model_id):
                 
                 if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
                     logger.info(f"请求成功，使用账号 {auth_manager._email}")
-                    # 更新最后一次成功使用的账号信息
-                    multi_auth_manager.update_last_success(auth_manager)
+                    multi_auth_manager.set_successful_account(auth_manager)  # 记录成功的账号
                     return response
                 
                 headers_cache.clear()
@@ -886,41 +883,43 @@ def make_request(payload, auth_manager, model_id):
     raise Exception("所有账号均不可用，且注册新账号失败")
 
 def health_check():
-    """改进的健康检查函数，每60秒只检查一个账号"""
-    check_index = 0  # 用于跟踪当前要检查的账号索引
+    """改进的健康检查函数，每次只检查一个账号"""
+    check_index = 0
+    last_check_time = time.time()
+    last_date = datetime.now().date()
     
     while True:
         try:
-            if multi_auth_manager and multi_auth_manager.auth_managers:
+            if multi_auth_manager:
                 current_time = time.time()
-                current_time_local = time.localtime()
-                
-                # 获取当前要检查的账号
-                auth_manager = multi_auth_manager.auth_managers[check_index]
-                email = auth_manager._email
-                
-                # 检查当前账号
-                if auth_manager._should_attempt_auth():
-                    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}")
-                else:
-                    logger.info(f"Skipping health check for {email} due to rate limiting")
-                
-                # 更新索引，准备下一次检查下一个账号
-                check_index = (check_index + 1) % len(multi_auth_manager.auth_managers)
+                current_date = datetime.now().date()
                 
-                # 每天凌晨重置所有账号的模型使用状态
-                if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
+                # 检查是否是新的一天
+                if current_date > last_date:
+                    multi_auth_manager.daily_rotation_triggered = True
                     multi_auth_manager.reset_all_model_status()
-                    logger.info("Reset model status for all accounts")
+                    logger.info("New day started - Reset model status for all accounts")
+                    last_date = current_date
+                
+                # 每60秒检查一个账号
+                if current_time - last_check_time >= 60:
+                    auth_manager = multi_auth_manager.auth_managers[check_index]
+                    
+                    if auth_manager._should_attempt_auth():
+                        if not auth_manager.ensure_valid_token():
+                            logger.warning(f"Auth token validation failed during health check for {auth_manager._email}")
+                            auth_manager.clear_auth()
+                        else:
+                            logger.info(f"Health check passed for {auth_manager._email}")
+                    
+                    # 更新检查索引和时间
+                    check_index = (check_index + 1) % len(multi_auth_manager.auth_managers)
+                    last_check_time = current_time
                     
         except Exception as e:
             logger.error(f"Health check error: {e}")
             
-        sleep(60)  # 每60秒检查一个账号
+        sleep(1)  # 降低CPU使用率
 
 # 为了兼容 Flask CLI 和 Gunicorn，修改启动逻辑
 if __name__ != "__main__":