Upload app.py

app.py

@@ -311,41 +311,56 @@ class MultiAuthManager:
         self._rotation_interval = 300  # 5分钟轮转间隔
         self.last_successful_index = 0  # 记录上次成功的账号索引
         self.last_success_date = datetime.now().date()  # 记录上次成功的日期
+        self.model_usage = {}  # 记录每个模型最后使用的成功账号
 
     def ensure_valid_token(self, model):
-        """改进的token验证方法，优先使用上次成功的账号"""
-        # 首先尝试使用上次成功的账号
-        auth_manager = self.auth_managers[self.last_successful_index]
-        if auth_manager.is_model_available(model) and auth_manager.ensure_valid_token():
-            return auth_manager
-
-        # 如果上次成功的账号不可用，才尝试其他账号
-        start_index = (self.last_successful_index + 1) % len(self.auth_managers)
-        current = start_index
+        """改进的token验证方法，优先使用上次成功的模型特定账号"""
+        # 首先尝试使用该模型上次成功的账号
+        if model in self.model_usage:
+            last_index = self.model_usage[model]
+            auth_manager = self.auth_managers[last_index]
+            if auth_manager.is_model_available(model) and auth_manager.ensure_valid_token():
+                return auth_manager
+
+        # 如果该模型没有成功记录或上次的账号不可用，从当前位置开始尝试
+        start_index = self.current_index
+        tried_count = 0
         
-        while current != self.last_successful_index:
-            auth_manager = self.auth_managers[current]
+        while tried_count < len(self.auth_managers):
+            auth_manager = self.auth_managers[self.current_index]
             if auth_manager.is_model_available(model) and auth_manager.ensure_valid_token():
-                self.last_successful_index = current  # 更新最后成功的账号索引
-                self.last_success_date = datetime.now().date()
+                # 更新该模型的成功账号记录
+                self.model_usage[model] = self.current_index
                 return auth_manager
-            current = (current + 1) % len(self.auth_managers)
             
-        # 如果所有账号都不可用，检查是否是新的一天
+            self.current_index = (self.current_index + 1) % len(self.auth_managers)
+            tried_count += 1
+            
+        # 如果是新的一天，重置所有状态
         if datetime.now().date() > self.last_success_date:
             self.reset_all_model_status()
             self.last_success_date = datetime.now().date()
+            self.model_usage.clear()  # 清除所有模型的使用记录
+            
             # 重新尝试第一个账号
             auth_manager = self.auth_managers[0]
             if auth_manager.ensure_valid_token():
-                self.last_successful_index = 0
+                self.current_index = 0
+                self.model_usage[model] = 0
                 return auth_manager
                 
         return None
 
+    def update_model_success(self, model, index):
+        """更新模型成功使用的账号索引"""
+        self.model_usage[model] = index
+        self.last_success_date = datetime.now().date()
+
     def reset_all_model_status(self):
+        """重置所有账号的状态和模型使用记录"""
         for auth_manager in self.auth_managers:
             auth_manager.reset_model_status()
+        self.model_usage.clear()
 
 def require_auth(func: Callable) -> Callable:
     """装饰器，确保在调用API之前有有效的token。"""
@@ -552,7 +567,7 @@ def handle_non_stream_response(response, model, prompt_tokens):
         raise
 
 def generate_stream_response(response, model, prompt_tokens):
-    """生��流式 HTTP 响应。"""
+    """生流式 HTTP 响应。"""
     total_completion_tokens = 0
     
     for chunk in stream_notdiamond_response(response, model):
@@ -774,7 +789,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
@@ -799,12 +814,10 @@ def make_request(payload, auth_manager, model_id):
     tried_accounts = set()
     
     while len(tried_accounts) < len(multi_auth_manager.auth_managers):
-        # 使用 ensure_valid_token 替代 get_next_auth_manager
         auth_manager = multi_auth_manager.ensure_valid_token(model_id)
         if not auth_manager:
             break
             
-        # 如果这个账号已经尝试过，继续下一个
         if auth_manager._email in tried_accounts:
             continue
             
@@ -816,7 +829,6 @@ def make_request(payload, auth_manager, model_id):
                 url = get_notdiamond_url()
                 headers = get_notdiamond_headers(auth_manager)
                 
-                # 使用超时参数
                 response = executor.submit(
                     requests.post, 
                     url, 
@@ -829,8 +841,8 @@ 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}")
                     current_index = multi_auth_manager.auth_managers.index(auth_manager)
-                    multi_auth_manager.last_successful_index = current_index  # 直接更新索引
-                    multi_auth_manager.last_success_date = datetime.now().date()  # 更新日期
+                    # 更新模型特定的成功记录
+                    multi_auth_manager.update_model_success(model_id, current_index)
                     return response
                 
                 headers_cache.clear()