dan92 commited on
Commit
c6d1804
·
verified ·
1 Parent(s): 3eabaf3

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +153 -180
app.py CHANGED
@@ -7,7 +7,7 @@ import uuid
7
  import re
8
  import socket
9
  from concurrent.futures import ThreadPoolExecutor
10
- from functools import wraps
11
  from typing import Dict, Any, Callable, List, Tuple
12
  import requests
13
  import tiktoken
@@ -46,7 +46,7 @@ if not _PASTE_API_URL:
46
  app = Flask(__name__)
47
  logging.basicConfig(level=logging.INFO)
48
  logger = logging.getLogger(__name__)
49
- CORS(app, resources={r"/": {"origins": "*"}})
50
  executor = ThreadPoolExecutor(max_workers=10)
51
 
52
  proxy_url = os.getenv('PROXY_URL')
@@ -64,7 +64,7 @@ def require_api_key(f):
64
  auth_header = request.headers.get('Authorization')
65
  if not auth_header:
66
  return jsonify({'error': 'No API key provided'}), 401
67
-
68
  try:
69
  # 从 Bearer token 中提取API密钥
70
  provided_key = auth_header.split('Bearer ')[-1].strip()
@@ -72,12 +72,12 @@ def require_api_key(f):
72
  return jsonify({'error': 'Invalid API key'}), 401
73
  except Exception:
74
  return jsonify({'error': 'Invalid Authorization header format'}), 401
75
-
76
  return f(*args, **kwargs)
77
  return decorated_function
78
 
79
  refresh_token_cache = TTLCache(maxsize=1000, ttl=3600)
80
- headers_cache = TTLCache(maxsize=100, ttl=3600) # 增加缓存大小
81
  token_refresh_lock = threading.Lock()
82
 
83
  # 自定义连接函数
@@ -114,58 +114,6 @@ AUTH_CHECK_INTERVAL = 300 # 健康检查间隔(秒)
114
  AUTH_RATE_LIMIT_WINDOW = 3600 # 速率限制窗口(秒)
115
  AUTH_MAX_REQUESTS = 100 # 每个窗口最大请求数
116
 
117
- # 模型信息
118
- MODEL_INFO = {
119
- "gpt-4o-mini": {
120
- "provider": "openai",
121
- "mapping": "gpt-4o-mini"
122
- },
123
- "gpt-4o": {
124
- "provider": "openai",
125
- "mapping": "gpt-4o"
126
- },
127
- "gpt-4-turbo": {
128
- "provider": "openai",
129
- "mapping": "gpt-4-turbo-2024-04-09"
130
- },
131
- "chatgpt-4o-latest": {
132
- "provider": "openai",
133
- "mapping": "chatgpt-4o-latest"
134
- },
135
- "gemini-1.5-pro-latest": {
136
- "provider": "google",
137
- "mapping": "models/gemini-1.5-pro-latest"
138
- },
139
- "gemini-1.5-flash-latest": {
140
- "provider": "google",
141
- "mapping": "models/gemini-1.5-flash-latest"
142
- },
143
- "llama-3.1-70b-instruct": {
144
- "provider": "togetherai",
145
- "mapping": "meta.llama3-1-70b-instruct-v1:0"
146
- },
147
- "llama-3.1-405b-instruct": {
148
- "provider": "togetherai",
149
- "mapping": "meta.llama3-1-405b-instruct-v1:0"
150
- },
151
- "claude-3-5-sonnet-20241022": {
152
- "provider": "anthropic",
153
- "mapping": "anthropic.claude-3-5-sonnet-20241022-v2:0"
154
- },
155
- "claude-3-5-haiku-20241022": {
156
- "provider": "anthropic",
157
- "mapping": "anthropic.claude-3-5-haiku-20241022-v1:0"
158
- },
159
- "perplexity": {
160
- "provider": "perplexity",
161
- "mapping": "llama-3.1-sonar-large-128k-online"
162
- },
163
- "mistral-large-2407": {
164
- "provider": "mistral",
165
- "mapping": "mistral.mistral-large-2407-v1:0"
166
- }
167
- }
168
-
169
  class AuthManager:
170
  def __init__(self, email: str, password: str):
171
  self._email: str = email
@@ -185,44 +133,38 @@ class AuthManager:
185
  self._auth_attempts = 0
186
  self._auth_window_start = time.time()
187
  self._backoff_delay = AUTH_RETRY_DELAY
188
- # 标记账号不可用直到特定时间
189
- self.unavailable_until = 0
190
 
191
  def _should_attempt_auth(self) -> bool:
192
  """检查是否应该尝试认证请求"""
193
  current_time = time.time()
194
-
195
- # 检查是否在不可用期内
196
- if current_time < self.unavailable_until:
197
- return False
198
-
199
  # 检查是否在退避期内
200
  if current_time - self._last_auth_attempt < self._backoff_delay:
201
  return False
202
-
203
  # 检查速率限制窗口
204
  if current_time - self._auth_window_start > AUTH_RATE_LIMIT_WINDOW:
205
  # 重置窗口
206
  self._auth_window_start = current_time
207
  self._auth_attempts = 0
208
  self._backoff_delay = AUTH_RETRY_DELAY
209
-
210
  # 检查请求数量
211
  if self._auth_attempts >= AUTH_MAX_REQUESTS:
212
  return False
213
-
214
  return True
215
 
216
  def login(self) -> bool:
217
  """改进的登录方法,包含速率限制和退避机制"""
218
  if not self._should_attempt_auth():
219
- self._logger.warning(f"Rate limit reached for {self._email}, waiting {self._backoff_delay}s")
220
  return False
221
 
222
  try:
223
  self._last_auth_attempt = time.time()
224
  self._auth_attempts += 1
225
-
226
  url = f"{_API_BASE_URL}/auth/v1/token?grant_type=password"
227
  headers = self._get_headers(with_content_type=True)
228
  data = {
@@ -230,32 +172,31 @@ class AuthManager:
230
  "password": self._password,
231
  "gotrue_meta_security": {}
232
  }
233
-
234
  response = self._make_request('POST', url, headers=headers, json=data)
235
-
236
  if response.status_code == 429:
237
  self._backoff_delay *= AUTH_BACKOFF_FACTOR
238
- self._logger.warning(f"Rate limit hit, increasing backoff to {self._backoff_delay}s")
239
  return False
240
-
241
  response.raise_for_status()
242
  self._user_info = response.json()
243
  self._refresh_token = self._user_info.get('refresh_token', '')
244
  self._access_token = self._user_info.get('access_token', '')
245
  self._token_expiry = time.time() + self._user_info.get('expires_in', 3600)
246
-
247
  # 重置退避延迟
248
  self._backoff_delay = AUTH_RETRY_DELAY
249
  self._log_values()
250
  return True
251
-
252
  except requests.RequestException as e:
253
- self._logger.error(f"\033[91m登录请求错误: {e}\033[0m")
254
  self._backoff_delay *= AUTH_BACKOFF_FACTOR
255
  return False
256
 
257
  def refresh_user_token(self) -> bool:
258
- """刷新用户令牌"""
259
  url = f"{_API_BASE_URL}/auth/v1/token?grant_type=refresh_token"
260
  headers = self._get_headers(with_content_type=True)
261
  data = {"refresh_token": self._refresh_token}
@@ -286,13 +227,13 @@ class AuthManager:
286
  """改进的token验证方法"""
287
  if self.is_token_valid():
288
  return True
289
-
290
  if not self._should_attempt_auth():
291
  return False
292
-
293
  if self._refresh_token and self.refresh_user_token():
294
  return True
295
-
296
  return self.login()
297
 
298
  def clear_auth(self) -> None:
@@ -302,14 +243,6 @@ class AuthManager:
302
  self._access_token = ""
303
  self._token_expiry = 0
304
 
305
- def set_unavailable_until_next_day(self) -> None:
306
- """将账号标记为不可用,直到次日"""
307
- now = datetime.now()
308
- next_day = now + timedelta(days=1)
309
- next_day_start = datetime(year=next_day.year, month=next_day.month, day=next_day.day)
310
- self.unavailable_until = next_day_start.timestamp()
311
- self._logger.info(f"Account {self._email} marked as unavailable until {next_day_start}")
312
-
313
  def _log_values(self) -> None:
314
  """记录刷新令牌到日志中。"""
315
  self._logger.info(f"\033[92mRefresh Token: {self._refresh_token}\033[0m")
@@ -322,17 +255,17 @@ class AuthManager:
322
  try:
323
  login_url = f"{_BASE_URL}/login"
324
  response = self._make_request('GET', login_url)
325
-
326
  match = re.search(r'<script src="(/_next/static/chunks/app/layout-[^"]+\.js)"', response.text)
327
  if not match:
328
  raise ValueError("未找到匹配的脚本标签")
329
  js_url = f"{_BASE_URL}{match.group(1)}"
330
  js_response = self._make_request('GET', js_url)
331
-
332
- api_key_match = re.search(r'\$\$"https://spuckhogycrxcbomznwo\.supabase\.co","([^"]+)"\$\$', js_response.text)
333
  if not api_key_match:
334
  raise ValueError("未能匹配API key")
335
-
336
  self._api_key = api_key_match.group(1)
337
  return self._api_key
338
  except (requests.RequestException, ValueError) as e:
@@ -376,25 +309,21 @@ class MultiAuthManager:
376
  self.current_index = 0
377
  self._last_rotation = time.time()
378
  self._rotation_interval = 300 # 5分钟轮转间隔
379
- self.last_successful_index = -1 # 上一次成功的账号索引
380
 
381
  def _should_rotate(self) -> bool:
382
  """检查是否应该轮转到下一个账号"""
383
  return time.time() - self._last_rotation >= self._rotation_interval
384
 
385
  def get_next_auth_manager(self, model):
386
- """改进的账号选择逻辑,从上一次成功的账号开始"""
387
- if self.last_successful_index == -1:
388
- self.current_index = 0
389
- else:
390
- self.current_index = (self.last_successful_index + 1) % len(self.auth_managers)
391
 
392
  start_index = self.current_index
393
  for _ in range(len(self.auth_managers)):
394
  auth_manager = self.auth_managers[self.current_index]
395
  if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
396
- self.last_successful_index = self.current_index
397
- self.current_index = (self.current_index + 1) % len(self.auth_managers)
398
  return auth_manager
399
  self.current_index = (self.current_index + 1) % len(self.auth_managers)
400
  if self.current_index == start_index:
@@ -402,7 +331,6 @@ class MultiAuthManager:
402
  return None
403
 
404
  def ensure_valid_token(self, model):
405
- """确保有有效的token"""
406
  for _ in range(len(self.auth_managers)):
407
  auth_manager = self.get_next_auth_manager(model)
408
  if auth_manager and auth_manager.ensure_valid_token():
@@ -413,18 +341,13 @@ class MultiAuthManager:
413
  for auth_manager in self.auth_managers:
414
  auth_manager.reset_model_status()
415
 
416
- def mark_account_unavailable(self, auth_manager):
417
- """标记账号为不可用直到次日"""
418
- auth_manager.set_unavailable_until_next_day()
419
-
420
  def require_auth(func: Callable) -> Callable:
421
  """装饰器,确保在调用API之前有有效的token。"""
422
  @wraps(func)
423
- def wrapper(*args, **kwargs):
424
- auth_manager = multi_auth_manager.ensure_valid_token(kwargs.get('model_id', ''))
425
- if not auth_manager:
426
  raise Exception("无法获取有效的授权token")
427
- return func(*args, **kwargs)
428
  return wrapper
429
 
430
  # 全局的 MultiAuthManager 对象
@@ -439,7 +362,7 @@ def get_notdiamond_url():
439
  def get_notdiamond_headers(auth_manager):
440
  """返回用于 notdiamond API 请求的头信息。"""
441
  cache_key = f'notdiamond_headers_{auth_manager.get_jwt_value()}'
442
-
443
  try:
444
  return headers_cache[cache_key]
445
  except KeyError:
@@ -453,6 +376,57 @@ def get_notdiamond_headers(auth_manager):
453
  headers_cache[cache_key] = headers
454
  return headers
455
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
456
  def generate_system_fingerprint():
457
  """生成并返回唯一的系统指纹。"""
458
  return f"fp_{uuid.uuid4().hex[:10]}"
@@ -476,10 +450,10 @@ def create_openai_chunk(content, model, finish_reason=None, usage=None):
476
  }
477
  ]
478
  }
479
-
480
  if usage is not None:
481
  chunk["usage"] = usage
482
-
483
  return chunk
484
 
485
  def count_tokens(text, model="gpt-3.5-turbo-0301"):
@@ -497,28 +471,28 @@ def stream_notdiamond_response(response, model):
497
  """改进的流式响应处理,确保保持上下文完整性。"""
498
  buffer = ""
499
  full_content = ""
500
-
501
  for chunk in response.iter_content(chunk_size=1024):
502
  if chunk:
503
  try:
504
  new_content = chunk.decode('utf-8')
505
  buffer += new_content
506
  full_content += new_content
507
-
508
  # 创建完整的响应块
509
  chunk_data = create_openai_chunk(new_content, model)
510
-
511
  # 确保响应块包含完整的上下文
512
  if 'choices' in chunk_data and chunk_data['choices']:
513
  chunk_data['choices'][0]['delta']['content'] = new_content
514
  chunk_data['choices'][0]['context'] = full_content # 添加完整上下文
515
-
516
  yield chunk_data
517
-
518
  except Exception as e:
519
  logger.error(f"Error processing chunk: {e}")
520
  continue
521
-
522
  # 发送完成标记
523
  final_chunk = create_openai_chunk('', model, 'stop')
524
  if 'choices' in final_chunk and final_chunk['choices']:
@@ -529,17 +503,17 @@ def handle_non_stream_response(response, model, prompt_tokens):
529
  """改进的非流式响应处理,确保保持完整上下文。"""
530
  full_content = ""
531
  context_buffer = []
532
-
533
  try:
534
  for chunk in response.iter_content(chunk_size=1024):
535
  if chunk:
536
  content = chunk.decode('utf-8')
537
  full_content += content
538
  context_buffer.append(content)
539
-
540
  completion_tokens = count_tokens(full_content, model)
541
  total_tokens = prompt_tokens + completion_tokens
542
-
543
  # 创建包含完整上下文的响应
544
  response_data = {
545
  "id": f"chatcmpl-{uuid.uuid4()}",
@@ -564,9 +538,9 @@ def handle_non_stream_response(response, model, prompt_tokens):
564
  "total_tokens": total_tokens
565
  }
566
  }
567
-
568
  return jsonify(response_data)
569
-
570
  except Exception as e:
571
  logger.error(f"Error processing non-stream response: {e}")
572
  raise
@@ -574,19 +548,19 @@ def handle_non_stream_response(response, model, prompt_tokens):
574
  def generate_stream_response(response, model, prompt_tokens):
575
  """生成流式 HTTP 响应。"""
576
  total_completion_tokens = 0
577
-
578
  for chunk in stream_notdiamond_response(response, model):
579
  content = chunk['choices'][0]['delta'].get('content', '')
580
  total_completion_tokens += count_tokens(content, model)
581
-
582
  chunk['usage'] = {
583
  "prompt_tokens": prompt_tokens,
584
  "completion_tokens": total_completion_tokens,
585
  "total_tokens": prompt_tokens + total_completion_tokens
586
  }
587
-
588
  yield f"data: {json.dumps(chunk)}\n\n"
589
-
590
  yield "data: [DONE]\n\n"
591
 
592
  def get_auth_credentials():
@@ -594,7 +568,7 @@ def get_auth_credentials():
594
  try:
595
  session = create_custom_session()
596
  headers = {
597
- 'accept': '/',
598
  'accept-language': 'zh-CN,zh;q=0.9',
599
  'user-agent': _USER_AGENT,
600
  'x-password': _PASTE_API_PASSWORD
@@ -624,13 +598,13 @@ def get_auth_credentials():
624
  def before_request():
625
  global multi_auth_manager
626
  credentials = get_auth_credentials()
627
-
628
  # 如果没有凭据,尝试自动注册
629
  if not credentials:
630
  try:
631
  # 使用 register_bot 注册新账号
632
  successful_accounts = register_bot.register_and_verify(5) # 注册5个账号
633
-
634
  if successful_accounts:
635
  # 更新凭据
636
  credentials = [(account['email'], account['password']) for account in successful_accounts]
@@ -643,7 +617,7 @@ def before_request():
643
  logger.error(f"自动注册过程发生错误: {e}")
644
  multi_auth_manager = None
645
  return
646
-
647
  if credentials:
648
  multi_auth_manager = MultiAuthManager(credentials)
649
  else:
@@ -698,11 +672,11 @@ def handle_request():
698
  global multi_auth_manager
699
  if not multi_auth_manager:
700
  return jsonify({'error': 'Unauthorized'}), 401
701
-
702
  try:
703
  request_data = request.get_json()
704
  model_id = request_data.get('model', '')
705
-
706
  auth_manager = multi_auth_manager.ensure_valid_token(model_id)
707
  if not auth_manager:
708
  return jsonify({'error': 'No available accounts for this model'}), 403
@@ -721,7 +695,7 @@ def handle_request():
721
  )
722
  else:
723
  return handle_non_stream_response(response, model_id, prompt_tokens)
724
-
725
  except requests.RequestException as e:
726
  logger.error("Request error: %s", str(e), exc_info=True)
727
  return jsonify({
@@ -759,10 +733,10 @@ def handle_request():
759
  def build_payload(request_data, model_id):
760
  """构建请求有效负载,确保保持完整的上下文。"""
761
  messages = request_data.get('messages', [])
762
-
763
  # 检查是否已经存在系统消息
764
  has_system_message = any(message.get('role') == 'system' for message in messages)
765
-
766
  # 如果没有系统消息,添加默认的系统消息
767
  if not has_system_message:
768
  system_message = {
@@ -780,11 +754,11 @@ def build_payload(request_data, model_id):
780
  )
781
  }
782
  messages.insert(0, system_message)
783
-
784
  # 获取模型映射
785
  model_info = MODEL_INFO.get(model_id, {})
786
  mapping = model_info.get('mapping', model_id)
787
-
788
  # 构建完整的payload
789
  payload = {
790
  'model': mapping,
@@ -796,12 +770,12 @@ def build_payload(request_data, model_id):
796
  'frequency_penalty': request_data.get('frequency_penalty'),
797
  'top_p': request_data.get('top_p', 1),
798
  }
799
-
800
  # 添加其他自定义参数
801
  for key, value in request_data.items():
802
  if key not in ['messages', 'model', 'stream', 'temperature'] and value is not None:
803
  payload[key] = value
804
-
805
  return payload
806
 
807
  def make_request(payload, auth_manager, model_id):
@@ -809,9 +783,9 @@ def make_request(payload, auth_manager, model_id):
809
  global multi_auth_manager
810
  max_retries = 3
811
  retry_delay = 1
812
-
813
  logger.info(f"尝试发送请求,模型:{model_id}")
814
-
815
  # 确保 multi_auth_manager 存在
816
  if not multi_auth_manager:
817
  logger.error("MultiAuthManager 不存在,尝试重新初始化")
@@ -827,16 +801,16 @@ def make_request(payload, auth_manager, model_id):
827
 
828
  # 记录已尝试的账号
829
  tried_accounts = set()
830
-
831
  while len(tried_accounts) < len(multi_auth_manager.auth_managers):
832
  auth_manager = multi_auth_manager.get_next_auth_manager(model_id)
833
  if not auth_manager:
834
  break
835
-
836
  # 如果这个账号已经尝试过,继续下一个
837
  if auth_manager._email in tried_accounts:
838
  continue
839
-
840
  tried_accounts.add(auth_manager._email)
841
  logger.info(f"尝试使用账号 {auth_manager._email}")
842
 
@@ -851,29 +825,25 @@ def make_request(payload, auth_manager, model_id):
851
  json=payload,
852
  stream=True
853
  ).result()
854
-
855
  if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
856
  logger.info(f"请求成功,使用账号 {auth_manager._email}")
857
- # 记录最后成功的账号索引
858
- multi_auth_manager.last_successful_index = multi_auth_manager.auth_managers.index(auth_manager)
859
  return response
860
-
861
  headers_cache.clear()
862
-
863
  if response.status_code == 401: # Unauthorized
864
  logger.info(f"Token expired for account {auth_manager._email}, attempting refresh")
865
  if auth_manager.ensure_valid_token():
866
  continue
867
-
868
  if response.status_code == 403: # Forbidden, 模型使用限制
869
  logger.warning(f"Model {model_id} usage limit reached for account {auth_manager._email}")
870
  auth_manager.set_model_unavailable(model_id)
871
- # 标记账号为不可用直到次日
872
- multi_auth_manager.mark_account_unavailable(auth_manager)
873
  break # 跳出重试循环,尝试下一个账号
874
-
875
  logger.error(f"Request failed with status {response.status_code} for account {auth_manager._email}")
876
-
877
  except Exception as e:
878
  logger.error(f"Request attempt {attempt + 1} failed for account {auth_manager._email}: {e}")
879
  if attempt < max_retries - 1:
@@ -889,45 +859,47 @@ def make_request(payload, auth_manager, model_id):
889
  multi_auth_manager = MultiAuthManager(credentials)
890
  # 使用新注册的账号重试请求
891
  return make_request(payload, None, model_id)
892
-
893
  raise Exception("所有账号均不可用,且注册新账号失败")
894
 
895
  def health_check():
896
- """改进的健康检查函数,每60秒只检测一个账号"""
897
- last_check_index = 0 # 用于跟踪下一个要检查的账号索引
898
-
899
  while True:
900
  try:
901
- if multi_auth_manager and multi_auth_manager.auth_managers:
902
- auth_manager = multi_auth_manager.auth_managers[last_check_index % len(multi_auth_manager.auth_managers)]
903
  current_time = time.time()
904
-
905
- # 如果账号被标记为不可用,检查是否可以恢复
906
- if current_time >= auth_manager.unavailable_until:
907
- auth_manager.unavailable_until = 0 # 重置不可用状态
908
-
909
- # 进行健康检查
910
- if auth_manager._should_attempt_auth():
 
 
 
 
 
911
  if not auth_manager.ensure_valid_token():
912
- logger.warning(f"Auth token validation failed during health check for {auth_manager._email}")
913
  auth_manager.clear_auth()
914
  else:
915
- logger.info(f"Health check passed for {auth_manager._email}")
916
- else:
917
- logger.info(f"Account {auth_manager._email} is still unavailable until {datetime.fromtimestamp(auth_manager.unavailable_until)}")
918
-
919
- last_check_index += 1 # 更新下一个要检查的账��索引
920
-
921
  # 每天重置所有账号的模型使用状态
922
  current_time_local = time.localtime()
923
  if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
924
  multi_auth_manager.reset_all_model_status()
925
  logger.info("Reset model status for all accounts")
926
-
927
  except Exception as e:
928
  logger.error(f"Health check error: {e}")
929
-
930
- sleep(60) # 主循环每分钟运行一次,检测一个账号
931
 
932
  # 为了兼容 Flask CLI 和 Gunicorn,修改启动逻辑
933
  if __name__ != "__main__":
@@ -937,6 +909,7 @@ if __name__ != "__main__":
937
  if __name__ == "__main__":
938
  health_check_thread = threading.Thread(target=health_check, daemon=True)
939
  health_check_thread.start()
940
-
941
  port = int(os.environ.get("PORT", 3000))
942
- app.run(debug=False, host='0.0.0.0', port=port, threaded=True)
 
 
7
  import re
8
  import socket
9
  from concurrent.futures import ThreadPoolExecutor
10
+ from functools import lru_cache, wraps
11
  from typing import Dict, Any, Callable, List, Tuple
12
  import requests
13
  import tiktoken
 
46
  app = Flask(__name__)
47
  logging.basicConfig(level=logging.INFO)
48
  logger = logging.getLogger(__name__)
49
+ CORS(app, resources={r"/*": {"origins": "*"}})
50
  executor = ThreadPoolExecutor(max_workers=10)
51
 
52
  proxy_url = os.getenv('PROXY_URL')
 
64
  auth_header = request.headers.get('Authorization')
65
  if not auth_header:
66
  return jsonify({'error': 'No API key provided'}), 401
67
+
68
  try:
69
  # 从 Bearer token 中提取API密钥
70
  provided_key = auth_header.split('Bearer ')[-1].strip()
 
72
  return jsonify({'error': 'Invalid API key'}), 401
73
  except Exception:
74
  return jsonify({'error': 'Invalid Authorization header format'}), 401
75
+
76
  return f(*args, **kwargs)
77
  return decorated_function
78
 
79
  refresh_token_cache = TTLCache(maxsize=1000, ttl=3600)
80
+ headers_cache = TTLCache(maxsize=1, ttl=3600) # 1小时过期
81
  token_refresh_lock = threading.Lock()
82
 
83
  # 自定义连接函数
 
114
  AUTH_RATE_LIMIT_WINDOW = 3600 # 速率限制窗口(秒)
115
  AUTH_MAX_REQUESTS = 100 # 每个窗口最大请求数
116
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
117
  class AuthManager:
118
  def __init__(self, email: str, password: str):
119
  self._email: str = email
 
133
  self._auth_attempts = 0
134
  self._auth_window_start = time.time()
135
  self._backoff_delay = AUTH_RETRY_DELAY
 
 
136
 
137
  def _should_attempt_auth(self) -> bool:
138
  """检查是否应该尝试认证请求"""
139
  current_time = time.time()
140
+
 
 
 
 
141
  # 检查是否在退避期内
142
  if current_time - self._last_auth_attempt < self._backoff_delay:
143
  return False
144
+
145
  # 检查速率限制窗口
146
  if current_time - self._auth_window_start > AUTH_RATE_LIMIT_WINDOW:
147
  # 重置窗口
148
  self._auth_window_start = current_time
149
  self._auth_attempts = 0
150
  self._backoff_delay = AUTH_RETRY_DELAY
151
+
152
  # 检查请求数量
153
  if self._auth_attempts >= AUTH_MAX_REQUESTS:
154
  return False
155
+
156
  return True
157
 
158
  def login(self) -> bool:
159
  """改进的登录方法,包含速率限制和退避机制"""
160
  if not self._should_attempt_auth():
161
+ logger.warning(f"Rate limit reached for {self._email}, waiting {self._backoff_delay}s")
162
  return False
163
 
164
  try:
165
  self._last_auth_attempt = time.time()
166
  self._auth_attempts += 1
167
+
168
  url = f"{_API_BASE_URL}/auth/v1/token?grant_type=password"
169
  headers = self._get_headers(with_content_type=True)
170
  data = {
 
172
  "password": self._password,
173
  "gotrue_meta_security": {}
174
  }
175
+
176
  response = self._make_request('POST', url, headers=headers, json=data)
177
+
178
  if response.status_code == 429:
179
  self._backoff_delay *= AUTH_BACKOFF_FACTOR
180
+ logger.warning(f"Rate limit hit, increasing backoff to {self._backoff_delay}s")
181
  return False
182
+
183
  response.raise_for_status()
184
  self._user_info = response.json()
185
  self._refresh_token = self._user_info.get('refresh_token', '')
186
  self._access_token = self._user_info.get('access_token', '')
187
  self._token_expiry = time.time() + self._user_info.get('expires_in', 3600)
188
+
189
  # 重置退避延迟
190
  self._backoff_delay = AUTH_RETRY_DELAY
191
  self._log_values()
192
  return True
193
+
194
  except requests.RequestException as e:
195
+ logger.error(f"\033[91m登录请求错误: {e}\033[0m")
196
  self._backoff_delay *= AUTH_BACKOFF_FACTOR
197
  return False
198
 
199
  def refresh_user_token(self) -> bool:
 
200
  url = f"{_API_BASE_URL}/auth/v1/token?grant_type=refresh_token"
201
  headers = self._get_headers(with_content_type=True)
202
  data = {"refresh_token": self._refresh_token}
 
227
  """改进的token验证方法"""
228
  if self.is_token_valid():
229
  return True
230
+
231
  if not self._should_attempt_auth():
232
  return False
233
+
234
  if self._refresh_token and self.refresh_user_token():
235
  return True
236
+
237
  return self.login()
238
 
239
  def clear_auth(self) -> None:
 
243
  self._access_token = ""
244
  self._token_expiry = 0
245
 
 
 
 
 
 
 
 
 
246
  def _log_values(self) -> None:
247
  """记录刷新令牌到日志中。"""
248
  self._logger.info(f"\033[92mRefresh Token: {self._refresh_token}\033[0m")
 
255
  try:
256
  login_url = f"{_BASE_URL}/login"
257
  response = self._make_request('GET', login_url)
258
+
259
  match = re.search(r'<script src="(/_next/static/chunks/app/layout-[^"]+\.js)"', response.text)
260
  if not match:
261
  raise ValueError("未找到匹配的脚本标签")
262
  js_url = f"{_BASE_URL}{match.group(1)}"
263
  js_response = self._make_request('GET', js_url)
264
+
265
+ api_key_match = re.search(r'\("https://spuckhogycrxcbomznwo\.supabase\.co","([^"]+)"\)', js_response.text)
266
  if not api_key_match:
267
  raise ValueError("未能匹配API key")
268
+
269
  self._api_key = api_key_match.group(1)
270
  return self._api_key
271
  except (requests.RequestException, ValueError) as e:
 
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:
 
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():
 
341
  for auth_manager in self.auth_managers:
342
  auth_manager.reset_model_status()
343
 
 
 
 
 
344
  def require_auth(func: Callable) -> Callable:
345
  """装饰器,确保在调用API之前有有效的token。"""
346
  @wraps(func)
347
+ def wrapper(self, *args, **kwargs):
348
+ if not self.ensure_valid_token():
 
349
  raise Exception("无法获取有效的授权token")
350
+ return func(self, *args, **kwargs)
351
  return wrapper
352
 
353
  # 全局的 MultiAuthManager 对象
 
362
  def get_notdiamond_headers(auth_manager):
363
  """返回用于 notdiamond API 请求的头信息。"""
364
  cache_key = f'notdiamond_headers_{auth_manager.get_jwt_value()}'
365
+
366
  try:
367
  return headers_cache[cache_key]
368
  except KeyError:
 
376
  headers_cache[cache_key] = headers
377
  return headers
378
 
379
+ MODEL_INFO = {
380
+ "gpt-4o-mini": {
381
+ "provider": "openai",
382
+ "mapping": "gpt-4o-mini"
383
+ },
384
+ "gpt-4o": {
385
+ "provider": "openai",
386
+ "mapping": "gpt-4o"
387
+ },
388
+ "gpt-4-turbo": {
389
+ "provider": "openai",
390
+ "mapping": "gpt-4-turbo-2024-04-09"
391
+ },
392
+ "chatgpt-4o-latest": {
393
+ "provider": "openai",
394
+ "mapping": "chatgpt-4o-latest"
395
+ },
396
+ "gemini-1.5-pro-latest": {
397
+ "provider": "google",
398
+ "mapping": "models/gemini-1.5-pro-latest"
399
+ },
400
+ "gemini-1.5-flash-latest": {
401
+ "provider": "google",
402
+ "mapping": "models/gemini-1.5-flash-latest"
403
+ },
404
+ "llama-3.1-70b-instruct": {
405
+ "provider": "togetherai",
406
+ "mapping": "meta.llama3-1-70b-instruct-v1:0"
407
+ },
408
+ "llama-3.1-405b-instruct": {
409
+ "provider": "togetherai",
410
+ "mapping": "meta.llama3-1-405b-instruct-v1:0"
411
+ },
412
+ "claude-3-5-sonnet-20241022": {
413
+ "provider": "anthropic",
414
+ "mapping": "anthropic.claude-3-5-sonnet-20241022-v2:0"
415
+ },
416
+ "claude-3-5-haiku-20241022": {
417
+ "provider": "anthropic",
418
+ "mapping": "anthropic.claude-3-5-haiku-20241022-v1:0"
419
+ },
420
+ "perplexity": {
421
+ "provider": "perplexity",
422
+ "mapping": "llama-3.1-sonar-large-128k-online"
423
+ },
424
+ "mistral-large-2407": {
425
+ "provider": "mistral",
426
+ "mapping": "mistral.mistral-large-2407-v1:0"
427
+ }
428
+ }
429
+
430
  def generate_system_fingerprint():
431
  """生成并返回唯一的系统指纹。"""
432
  return f"fp_{uuid.uuid4().hex[:10]}"
 
450
  }
451
  ]
452
  }
453
+
454
  if usage is not None:
455
  chunk["usage"] = usage
456
+
457
  return chunk
458
 
459
  def count_tokens(text, model="gpt-3.5-turbo-0301"):
 
471
  """改进的流式响应处理,确保保持上下文完整性。"""
472
  buffer = ""
473
  full_content = ""
474
+
475
  for chunk in response.iter_content(chunk_size=1024):
476
  if chunk:
477
  try:
478
  new_content = chunk.decode('utf-8')
479
  buffer += new_content
480
  full_content += new_content
481
+
482
  # 创建完整的响应块
483
  chunk_data = create_openai_chunk(new_content, model)
484
+
485
  # 确保响应块包含完整的上下文
486
  if 'choices' in chunk_data and chunk_data['choices']:
487
  chunk_data['choices'][0]['delta']['content'] = new_content
488
  chunk_data['choices'][0]['context'] = full_content # 添加完整上下文
489
+
490
  yield chunk_data
491
+
492
  except Exception as e:
493
  logger.error(f"Error processing chunk: {e}")
494
  continue
495
+
496
  # 发送完成标记
497
  final_chunk = create_openai_chunk('', model, 'stop')
498
  if 'choices' in final_chunk and final_chunk['choices']:
 
503
  """改进的非流式响应处理,确保保持完整上下文。"""
504
  full_content = ""
505
  context_buffer = []
506
+
507
  try:
508
  for chunk in response.iter_content(chunk_size=1024):
509
  if chunk:
510
  content = chunk.decode('utf-8')
511
  full_content += content
512
  context_buffer.append(content)
513
+
514
  completion_tokens = count_tokens(full_content, model)
515
  total_tokens = prompt_tokens + completion_tokens
516
+
517
  # 创建包含完整上下文的响应
518
  response_data = {
519
  "id": f"chatcmpl-{uuid.uuid4()}",
 
538
  "total_tokens": total_tokens
539
  }
540
  }
541
+
542
  return jsonify(response_data)
543
+
544
  except Exception as e:
545
  logger.error(f"Error processing non-stream response: {e}")
546
  raise
 
548
  def generate_stream_response(response, model, prompt_tokens):
549
  """生成流式 HTTP 响应。"""
550
  total_completion_tokens = 0
551
+
552
  for chunk in stream_notdiamond_response(response, model):
553
  content = chunk['choices'][0]['delta'].get('content', '')
554
  total_completion_tokens += count_tokens(content, model)
555
+
556
  chunk['usage'] = {
557
  "prompt_tokens": prompt_tokens,
558
  "completion_tokens": total_completion_tokens,
559
  "total_tokens": prompt_tokens + total_completion_tokens
560
  }
561
+
562
  yield f"data: {json.dumps(chunk)}\n\n"
563
+
564
  yield "data: [DONE]\n\n"
565
 
566
  def get_auth_credentials():
 
568
  try:
569
  session = create_custom_session()
570
  headers = {
571
+ 'accept': '*/*',
572
  'accept-language': 'zh-CN,zh;q=0.9',
573
  'user-agent': _USER_AGENT,
574
  'x-password': _PASTE_API_PASSWORD
 
598
  def before_request():
599
  global multi_auth_manager
600
  credentials = get_auth_credentials()
601
+
602
  # 如果没有凭据,尝试自动注册
603
  if not credentials:
604
  try:
605
  # 使用 register_bot 注册新账号
606
  successful_accounts = register_bot.register_and_verify(5) # 注册5个账号
607
+
608
  if successful_accounts:
609
  # 更新凭据
610
  credentials = [(account['email'], account['password']) for account in successful_accounts]
 
617
  logger.error(f"自动注册过程发生错误: {e}")
618
  multi_auth_manager = None
619
  return
620
+
621
  if credentials:
622
  multi_auth_manager = MultiAuthManager(credentials)
623
  else:
 
672
  global multi_auth_manager
673
  if not multi_auth_manager:
674
  return jsonify({'error': 'Unauthorized'}), 401
675
+
676
  try:
677
  request_data = request.get_json()
678
  model_id = request_data.get('model', '')
679
+
680
  auth_manager = multi_auth_manager.ensure_valid_token(model_id)
681
  if not auth_manager:
682
  return jsonify({'error': 'No available accounts for this model'}), 403
 
695
  )
696
  else:
697
  return handle_non_stream_response(response, model_id, prompt_tokens)
698
+
699
  except requests.RequestException as e:
700
  logger.error("Request error: %s", str(e), exc_info=True)
701
  return jsonify({
 
733
  def build_payload(request_data, model_id):
734
  """构建请求有效负载,确保保持完整的上下文。"""
735
  messages = request_data.get('messages', [])
736
+
737
  # 检查是否已经存在系统消息
738
  has_system_message = any(message.get('role') == 'system' for message in messages)
739
+
740
  # 如果没有系统消息,添加默认的系统消息
741
  if not has_system_message:
742
  system_message = {
 
754
  )
755
  }
756
  messages.insert(0, system_message)
757
+
758
  # 获取模型映射
759
  model_info = MODEL_INFO.get(model_id, {})
760
  mapping = model_info.get('mapping', model_id)
761
+
762
  # 构建完整的payload
763
  payload = {
764
  'model': mapping,
 
770
  'frequency_penalty': request_data.get('frequency_penalty'),
771
  'top_p': request_data.get('top_p', 1),
772
  }
773
+
774
  # 添加其他自定义参数
775
  for key, value in request_data.items():
776
  if key not in ['messages', 'model', 'stream', 'temperature'] and value is not None:
777
  payload[key] = value
778
+
779
  return payload
780
 
781
  def make_request(payload, auth_manager, model_id):
 
783
  global multi_auth_manager
784
  max_retries = 3
785
  retry_delay = 1
786
+
787
  logger.info(f"尝试发送请求,模型:{model_id}")
788
+
789
  # 确保 multi_auth_manager 存在
790
  if not multi_auth_manager:
791
  logger.error("MultiAuthManager 不存在,尝试重新初始化")
 
801
 
802
  # 记录已尝试的账号
803
  tried_accounts = set()
804
+
805
  while len(tried_accounts) < len(multi_auth_manager.auth_managers):
806
  auth_manager = multi_auth_manager.get_next_auth_manager(model_id)
807
  if not auth_manager:
808
  break
809
+
810
  # 如果这个账号已经尝试过,继续下一个
811
  if auth_manager._email in tried_accounts:
812
  continue
813
+
814
  tried_accounts.add(auth_manager._email)
815
  logger.info(f"尝试使用账号 {auth_manager._email}")
816
 
 
825
  json=payload,
826
  stream=True
827
  ).result()
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()
834
+
835
  if response.status_code == 401: # Unauthorized
836
  logger.info(f"Token expired for account {auth_manager._email}, attempting refresh")
837
  if auth_manager.ensure_valid_token():
838
  continue
839
+
840
  if response.status_code == 403: # Forbidden, 模型使用限制
841
  logger.warning(f"Model {model_id} usage limit reached for account {auth_manager._email}")
842
  auth_manager.set_model_unavailable(model_id)
 
 
843
  break # 跳出重试循环,尝试下一个账号
844
+
845
  logger.error(f"Request failed with status {response.status_code} for account {auth_manager._email}")
846
+
847
  except Exception as e:
848
  logger.error(f"Request attempt {attempt + 1} failed for account {auth_manager._email}: {e}")
849
  if attempt < max_retries - 1:
 
859
  multi_auth_manager = MultiAuthManager(credentials)
860
  # 使用新注册的账号重试请求
861
  return make_request(payload, None, model_id)
862
+
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")
898
+
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__":
 
909
  if __name__ == "__main__":
910
  health_check_thread = threading.Thread(target=health_check, daemon=True)
911
  health_check_thread.start()
912
+
913
  port = int(os.environ.get("PORT", 3000))
914
+ app.run(debug=False, host='0.0.0.0', port=port, threaded=True)
915
+