dan92 commited on
Commit
6c55cb7
·
verified ·
1 Parent(s): 773eac5

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +64 -43
app.py CHANGED
@@ -273,7 +273,7 @@ class AuthManager:
273
  return ""
274
 
275
  def _get_headers(self, with_content_type: bool = False) -> Dict[str, str]:
276
- """生成请求头。"""
277
  headers = {
278
  'apikey': self._fetch_apikey(),
279
  'user-agent': _USER_AGENT
@@ -309,32 +309,53 @@ class MultiAuthManager:
309
  self.current_index = 0
310
  self._last_rotation = time.time()
311
  self._rotation_interval = 300 # 5分钟轮转间隔
 
 
312
 
313
  def _should_rotate(self) -> bool:
314
  """检查是否应该轮转到下一个账号"""
315
- return time.time() - self._last_rotation >= self._rotation_interval
 
 
316
 
317
  def get_next_auth_manager(self, model):
318
- """改进的账号选择逻辑"""
 
 
 
 
 
 
 
 
319
  if self._should_rotate():
320
- self.current_index = (self.current_index + 1) % len(self.auth_managers)
321
- self._last_rotation = time.time()
322
-
323
- start_index = self.current_index
324
- for _ in range(len(self.auth_managers)):
325
- auth_manager = self.auth_managers[self.current_index]
326
- if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
327
- return auth_manager
328
- self.current_index = (self.current_index + 1) % len(self.auth_managers)
329
- if self.current_index == start_index:
330
- break
 
 
331
  return None
332
 
 
 
 
 
 
 
 
 
333
  def ensure_valid_token(self, model):
334
- for _ in range(len(self.auth_managers)):
335
- auth_manager = self.get_next_auth_manager(model)
336
- if auth_manager and auth_manager.ensure_valid_token():
337
- return auth_manager
338
  return None
339
 
340
  def reset_all_model_status(self):
@@ -779,7 +800,7 @@ def build_payload(request_data, model_id):
779
  return payload
780
 
781
  def make_request(payload, auth_manager, model_id):
782
- """发送请求并处理可能的认证刷新和模型特定错误。"""
783
  global multi_auth_manager
784
  max_retries = 3
785
  retry_delay = 1
@@ -828,6 +849,8 @@ def make_request(payload, auth_manager, model_id):
828
 
829
  if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
830
  logger.info(f"请求成功,使用账号 {auth_manager._email}")
 
 
831
  return response
832
 
833
  headers_cache.clear()
@@ -863,35 +886,33 @@ def make_request(payload, auth_manager, model_id):
863
  raise Exception("所有账号均不可用,且注册新账号失败")
864
 
865
  def health_check():
866
- """改进的健康检查函数"""
867
- last_check_time = {} # 用于跟踪每个账号的最后检查时间
868
 
869
  while True:
870
  try:
871
- if multi_auth_manager:
872
  current_time = time.time()
 
873
 
874
- for auth_manager in multi_auth_manager.auth_managers:
875
- email = auth_manager._email
876
-
877
- # 检查是否需要进行健康检查
878
- if email not in last_check_time or \
879
- current_time - last_check_time[email] >= AUTH_CHECK_INTERVAL:
880
-
881
- if not auth_manager._should_attempt_auth():
882
- logger.info(f"Skipping health check for {email} due to rate limiting")
883
- continue
884
-
885
- if not auth_manager.ensure_valid_token():
886
- logger.warning(f"Auth token validation failed during health check for {email}")
887
- auth_manager.clear_auth()
888
- else:
889
- logger.info(f"Health check passed for {email}")
890
-
891
- last_check_time[email] = current_time
892
 
893
- # 每天重置所有账号的模型使用状态
894
- current_time_local = time.localtime()
 
 
 
 
 
 
 
 
 
 
 
 
895
  if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
896
  multi_auth_manager.reset_all_model_status()
897
  logger.info("Reset model status for all accounts")
@@ -899,7 +920,7 @@ def health_check():
899
  except Exception as e:
900
  logger.error(f"Health check error: {e}")
901
 
902
- sleep(60) # 主循环每分钟运行一次
903
 
904
  # 为了兼容 Flask CLI 和 Gunicorn,修改启动逻辑
905
  if __name__ != "__main__":
 
273
  return ""
274
 
275
  def _get_headers(self, with_content_type: bool = False) -> Dict[str, str]:
276
+ """成请求头。"""
277
  headers = {
278
  'apikey': self._fetch_apikey(),
279
  'user-agent': _USER_AGENT
 
309
  self.current_index = 0
310
  self._last_rotation = time.time()
311
  self._rotation_interval = 300 # 5分钟轮转间隔
312
+ self._last_successful_index = 0 # 记录上一次成功的账号索引
313
+ self._last_success_date = datetime.now().date() # 记录上一次成功的日期
314
 
315
  def _should_rotate(self) -> bool:
316
  """检查是否应该轮转到下一个账号"""
317
+ current_date = datetime.now().date()
318
+ # 如果是新的一天且上一个成功的账号失败了,才进行完整轮询
319
+ return current_date > self._last_success_date
320
 
321
  def get_next_auth_manager(self, model):
322
+ """改进的账号选择逻辑,优先使用上次成功的账号"""
323
+ current_date = datetime.now().date()
324
+
325
+ # 首先尝试上次成功的账号
326
+ auth_manager = self.auth_managers[self._last_successful_index]
327
+ if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
328
+ return auth_manager
329
+
330
+ # 如果是新的一天,或者上次成功的账号不可用,才进行完整轮询
331
  if self._should_rotate():
332
+ start_index = (self._last_successful_index + 1) % len(self.auth_managers)
333
+ current_index = start_index
334
+
335
+ while True:
336
+ auth_manager = self.auth_managers[current_index]
337
+ if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
338
+ self._last_successful_index = current_index
339
+ return auth_manager
340
+
341
+ current_index = (current_index + 1) % len(self.auth_managers)
342
+ if current_index == start_index:
343
+ break
344
+
345
  return None
346
 
347
+ def update_last_success(self, auth_manager):
348
+ """更新最后一次成功使用的账号信息"""
349
+ for i, manager in enumerate(self.auth_managers):
350
+ if manager._email == auth_manager._email:
351
+ self._last_successful_index = i
352
+ self._last_success_date = datetime.now().date()
353
+ break
354
+
355
  def ensure_valid_token(self, model):
356
+ auth_manager = self.get_next_auth_manager(model)
357
+ if auth_manager and auth_manager.ensure_valid_token():
358
+ return auth_manager
 
359
  return None
360
 
361
  def reset_all_model_status(self):
 
800
  return payload
801
 
802
  def make_request(payload, auth_manager, model_id):
803
+ """发送请求并处理可能的认证刷新和模型定错误。"""
804
  global multi_auth_manager
805
  max_retries = 3
806
  retry_delay = 1
 
849
 
850
  if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
851
  logger.info(f"请求成功,使用账号 {auth_manager._email}")
852
+ # 更新最后一次成功使用的账号信息
853
+ multi_auth_manager.update_last_success(auth_manager)
854
  return response
855
 
856
  headers_cache.clear()
 
886
  raise Exception("所有账号均不可用,且注册新账号失败")
887
 
888
  def health_check():
889
+ """改进的健康检查函数,每60秒只检查一个账号"""
890
+ check_index = 0 # 用于跟踪当前要检查的账号索引
891
 
892
  while True:
893
  try:
894
+ if multi_auth_manager and multi_auth_manager.auth_managers:
895
  current_time = time.time()
896
+ current_time_local = time.localtime()
897
 
898
+ # 获取当前要检查的账号
899
+ auth_manager = multi_auth_manager.auth_managers[check_index]
900
+ email = auth_manager._email
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
901
 
902
+ # 检查当前账号
903
+ if auth_manager._should_attempt_auth():
904
+ if not auth_manager.ensure_valid_token():
905
+ logger.warning(f"Auth token validation failed during health check for {email}")
906
+ auth_manager.clear_auth()
907
+ else:
908
+ logger.info(f"Health check passed for {email}")
909
+ else:
910
+ logger.info(f"Skipping health check for {email} due to rate limiting")
911
+
912
+ # 更新索引,准备下一次检查下一个账号
913
+ check_index = (check_index + 1) % len(multi_auth_manager.auth_managers)
914
+
915
+ # 每天凌晨重置所有账号的模型使用状态
916
  if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
917
  multi_auth_manager.reset_all_model_status()
918
  logger.info("Reset model status for all accounts")
 
920
  except Exception as e:
921
  logger.error(f"Health check error: {e}")
922
 
923
+ sleep(60) # 每60秒检查一个账号
924
 
925
  # 为了兼容 Flask CLI 和 Gunicorn,修改启动逻辑
926
  if __name__ != "__main__":