isididiidid commited on
Commit
22cf637
·
verified ·
1 Parent(s): 9aabd84

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +566 -626
app.py CHANGED
@@ -1,60 +1,24 @@
1
  import os
 
 
 
 
2
  import json
3
  import uuid
 
4
  import base64
5
  import sys
6
  import inspect
7
  from loguru import logger
8
- import os
9
- import asyncio
10
- import time
11
- import aiohttp
12
- import io
13
- from datetime import datetime
14
- from functools import partial
15
 
16
- from quart import Quart, request, jsonify, Response
17
- from quart_cors import cors
18
- import cloudscraper
19
- from dotenv import load_dotenv
20
 
21
- load_dotenv()
22
-
23
- CONFIG = {
24
- "MODELS": {
25
- 'grok-2': 'grok-latest',
26
- 'grok-2-imageGen': 'grok-latest',
27
- 'grok-2-search': 'grok-latest',
28
- "grok-3": "grok-3",
29
- "grok-3-search": "grok-3",
30
- "grok-3-imageGen": "grok-3",
31
- "grok-3-deepsearch": "grok-3",
32
- "grok-3-reasoning": "grok-3"
33
- },
34
- "API": {
35
- "BASE_URL": "https://grok.com",
36
- "API_KEY": os.getenv("API_KEY", "sk-123456"),
37
- "IS_TEMP_CONVERSATION": os.getenv("IS_TEMP_CONVERSATION", "false").lower() == "true",
38
- "PICGO_KEY": os.getenv("PICGO_KEY", None), # 想要流式生图的话需要填入这个PICGO图床的key
39
- "TUMY_KEY": os.getenv("TUMY_KEY", None), # 想要流式生图的话需要填入这个TUMY图床的key
40
- "IS_CUSTOM_SSO": os.getenv("IS_CUSTOM_SSO", "false").lower() == "true"
41
- },
42
- "SERVER": {
43
- "PORT": int(os.getenv("PORT", 3000))
44
- },
45
- "RETRY": {
46
- "MAX_ATTEMPTS": 2
47
- },
48
- "SHOW_THINKING": os.getenv("SHOW_THINKING", "false").lower() == "true",
49
- "IS_THINKING": False,
50
- "IS_IMG_GEN": False,
51
- "IS_IMG_GEN2": False,
52
- "ISSHOW_SEARCH_RESULTS": os.getenv("ISSHOW_SEARCH_RESULTS", "true").lower() == "true"
53
- }
54
 
55
  class Logger:
56
  def __init__(self, level="INFO", colorize=True, format=None):
57
- # 移除默认的日志处理器
58
  logger.remove()
59
 
60
  if format is None:
@@ -120,51 +84,102 @@ class Logger:
120
 
121
  logger = Logger(level="INFO")
122
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
123
  class AuthTokenManager:
124
  def __init__(self):
125
  self.token_model_map = {}
126
  self.expired_tokens = set()
127
  self.token_status_map = {}
128
- self.token_reset_switch = False
129
- self.token_reset_timer = None
130
- self.is_custom_sso = os.getenv("IS_CUSTOM_SSO", "false").lower() == "true"
131
 
132
  self.model_config = {
133
  "grok-2": {
134
- "RequestFrequency": 2,
135
- "ExpirationTime": 1 * 60 * 60
136
  },
137
  "grok-3": {
138
  "RequestFrequency": 20,
139
- "ExpirationTime": 2 * 60 * 60 #
140
  },
141
  "grok-3-deepsearch": {
142
  "RequestFrequency": 10,
143
- "ExpirationTime": 24 * 60 * 60
144
  },
145
  "grok-3-reasoning": {
146
  "RequestFrequency": 10,
147
- "ExpirationTime": 24 * 60 * 60
148
  }
149
  }
150
-
151
- async def add_token(self, token):
 
 
152
  sso = token.split("sso=")[1].split(";")[0]
153
  for model in self.model_config.keys():
154
  if model not in self.token_model_map:
155
  self.token_model_map[model] = []
156
-
157
  if sso not in self.token_status_map:
158
  self.token_status_map[sso] = {}
159
-
160
- existing_token_entry = next((entry for entry in self.token_model_map[model]
161
- if entry.get("token") == token), None)
162
-
163
  if not existing_token_entry:
164
  self.token_model_map[model].append({
165
  "token": token,
166
  "RequestCount": 0,
167
- "AddedTime": time.time(),
168
  "StartCallTime": None
169
  })
170
 
@@ -174,279 +189,258 @@ class AuthTokenManager:
174
  "invalidatedTime": None,
175
  "totalRequestCount": 0
176
  }
177
- logger.info(f"添加令牌成功: {token}", "TokenManager")
178
-
179
- async def set_token(self, token):
180
  models = list(self.model_config.keys())
181
- for model in models:
182
- self.token_model_map[model] = [{
183
- "token": token,
184
- "RequestCount": 0,
185
- "AddedTime": time.time(),
186
- "StartCallTime": None
187
- }]
188
 
189
  sso = token.split("sso=")[1].split(";")[0]
190
- self.token_status_map[sso] = {}
191
- for model in models:
192
- self.token_status_map[sso][model] = {
193
- "isValid": True,
194
- "invalidatedTime": None,
195
- "totalRequestCount": 0
196
- }
197
- logger.info(f"设置令牌成功: {token}", "TokenManager")
198
-
199
- async def delete_token(self, token):
200
  try:
201
  sso = token.split("sso=")[1].split(";")[0]
202
-
203
  for model in self.token_model_map:
204
- self.token_model_map[model] = [
205
- entry for entry in self.token_model_map[model]
206
- if entry.get("token") != token
207
- ]
208
-
209
  if sso in self.token_status_map:
210
  del self.token_status_map[sso]
211
 
212
  logger.info(f"令牌已成功移除: {token}", "TokenManager")
213
  return True
214
  except Exception as error:
215
- logger.error(f"令牌删除失败: {error}", "TokenManager")
216
  return False
217
-
218
  def get_next_token_for_model(self, model_id):
219
  normalized_model = self.normalize_model_name(model_id)
220
 
221
  if normalized_model not in self.token_model_map or not self.token_model_map[normalized_model]:
222
  return None
223
-
224
  token_entry = self.token_model_map[normalized_model][0]
225
 
226
  if token_entry:
227
- if self.is_custom_sso:
228
- return token_entry["token"]
229
-
230
  if token_entry["StartCallTime"] is None:
231
- token_entry["StartCallTime"] = time.time()
232
-
233
  if not self.token_reset_switch:
234
  self.start_token_reset_process()
235
  self.token_reset_switch = True
236
-
237
  token_entry["RequestCount"] += 1
238
 
239
  if token_entry["RequestCount"] > self.model_config[normalized_model]["RequestFrequency"]:
240
  self.remove_token_from_model(normalized_model, token_entry["token"])
241
- if not self.token_model_map[normalized_model]:
242
- return None
243
- next_token_entry = self.token_model_map[normalized_model][0]
244
  return next_token_entry["token"] if next_token_entry else None
245
 
246
  sso = token_entry["token"].split("sso=")[1].split(";")[0]
247
  if sso in self.token_status_map and normalized_model in self.token_status_map[sso]:
248
  if token_entry["RequestCount"] == self.model_config[normalized_model]["RequestFrequency"]:
249
  self.token_status_map[sso][normalized_model]["isValid"] = False
250
- self.token_status_map[sso][normalized_model]["invalidatedTime"] = time.time()
251
-
252
  self.token_status_map[sso][normalized_model]["totalRequestCount"] += 1
253
-
254
  return token_entry["token"]
255
 
256
  return None
257
-
258
  def remove_token_from_model(self, model_id, token):
259
  normalized_model = self.normalize_model_name(model_id)
260
 
261
  if normalized_model not in self.token_model_map:
262
  logger.error(f"模型 {normalized_model} 不存在", "TokenManager")
263
  return False
264
-
265
  model_tokens = self.token_model_map[normalized_model]
266
- token_index = -1
267
 
268
- for i, entry in enumerate(model_tokens):
269
- if entry["token"] == token:
270
- token_index = i
271
- break
272
-
273
  if token_index != -1:
274
  removed_token_entry = model_tokens.pop(token_index)
275
  self.expired_tokens.add((
276
  removed_token_entry["token"],
277
  normalized_model,
278
- time.time()
279
  ))
280
 
281
  if not self.token_reset_switch:
282
  self.start_token_reset_process()
283
  self.token_reset_switch = True
284
-
285
  logger.info(f"模型{model_id}的令牌已失效,已成功移除令牌: {token}", "TokenManager")
286
  return True
287
-
288
  logger.error(f"在模型 {normalized_model} 中未找到 token: {token}", "TokenManager")
289
  return False
290
-
291
  def get_expired_tokens(self):
292
  return list(self.expired_tokens)
293
-
294
  def normalize_model_name(self, model):
295
  if model.startswith('grok-') and 'deepsearch' not in model and 'reasoning' not in model:
296
  return '-'.join(model.split('-')[:2])
297
  return model
298
-
299
  def get_token_count_for_model(self, model_id):
300
  normalized_model = self.normalize_model_name(model_id)
301
  return len(self.token_model_map.get(normalized_model, []))
302
-
303
  def get_remaining_token_request_capacity(self):
304
  remaining_capacity_map = {}
305
 
306
- for model in self.model_config:
307
  model_tokens = self.token_model_map.get(model, [])
308
  model_request_frequency = self.model_config[model]["RequestFrequency"]
309
 
310
- total_used_requests = sum(entry.get("RequestCount", 0) for entry in model_tokens)
 
311
  remaining_capacity = (len(model_tokens) * model_request_frequency) - total_used_requests
312
  remaining_capacity_map[model] = max(0, remaining_capacity)
313
-
314
- return remaining_capacity_map
315
 
 
 
316
  def get_token_array_for_model(self, model_id):
317
  normalized_model = self.normalize_model_name(model_id)
318
  return self.token_model_map.get(normalized_model, [])
319
-
320
  def start_token_reset_process(self):
321
- if hasattr(self, '_reset_task') and self._reset_task:
322
- pass
323
- else:
324
- self._reset_task = asyncio.create_task(self._token_reset_worker())
325
-
326
- async def _token_reset_worker(self):
327
- while True:
328
- try:
329
- current_time = time.time()
330
 
331
- expired_tokens_to_remove = set()
332
- for token_info in self.expired_tokens:
333
- token, model, expired_time = token_info
334
- expiration_time = self.model_config[model]["ExpirationTime"]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
335
 
336
- if current_time - expired_time >= expiration_time:
337
- if not any(entry["token"] == token for entry in self.token_model_map[model]):
338
- self.token_model_map[model].append({
339
- "token": token,
340
- "RequestCount": 0,
341
- "AddedTime": current_time,
342
- "StartCallTime": None
343
- })
344
-
345
- sso = token.split("sso=")[1].split(";")[0]
346
  if sso in self.token_status_map and model in self.token_status_map[sso]:
347
  self.token_status_map[sso][model]["isValid"] = True
348
  self.token_status_map[sso][model]["invalidatedTime"] = None
349
  self.token_status_map[sso][model]["totalRequestCount"] = 0
350
-
351
- expired_tokens_to_remove.add(token_info)
352
-
353
- for token_info in expired_tokens_to_remove:
354
- self.expired_tokens.remove(token_info)
355
-
356
- for model in self.model_config:
357
- if model not in self.token_model_map:
358
- continue
359
 
360
- for token_entry in self.token_model_map[model]:
361
- if token_entry["StartCallTime"] is None:
362
- continue
363
-
364
- expiration_time = self.model_config[model]["ExpirationTime"]
365
- if current_time - token_entry["StartCallTime"] >= expiration_time:
366
- sso = token_entry["token"].split("sso=")[1].split(";")[0]
367
- if sso in self.token_status_map and model in self.token_status_map[sso]:
368
- self.token_status_map[sso][model]["isValid"] = True
369
- self.token_status_map[sso][model]["invalidatedTime"] = None
370
- self.token_status_map[sso][model]["totalRequestCount"] = 0
371
-
372
- token_entry["RequestCount"] = 0
373
- token_entry["StartCallTime"] = None
374
-
375
- await asyncio.sleep(3600)
376
- except Exception as e:
377
- logger.error(f"令牌重置过程中出错: {e}", "TokenManager")
378
- await asyncio.sleep(3600)
379
-
380
  def get_all_tokens(self):
381
  all_tokens = set()
382
  for model_tokens in self.token_model_map.values():
383
  for entry in model_tokens:
384
  all_tokens.add(entry["token"])
385
  return list(all_tokens)
386
-
387
  def get_token_status_map(self):
388
  return self.token_status_map
389
-
390
- token_manager = AuthTokenManager()
391
-
392
- async def initialize_tokens():
393
- sso_array = os.getenv("SSO", "").split(',')
394
- logger.info("开始加载令牌", "Server")
395
 
396
- for sso in sso_array:
397
- if sso.strip():
398
- await token_manager.add_token(f"sso-rw={sso};sso={sso}")
399
-
400
- logger.info(f"成功加载令牌: {json.dumps(token_manager.get_all_tokens(), indent=2)}", "Server")
401
- logger.info(f"令牌加载完成,共加载: {len(token_manager.get_all_tokens())}个令牌", "Server")
402
- logger.info("初始化完成", "Server")
403
-
404
  class Utils:
405
  @staticmethod
406
- async def organize_search_results(search_results):
407
- if not search_results or "results" not in search_results:
408
  return ''
409
-
410
- results = search_results["results"]
411
  formatted_results = []
412
 
413
  for index, result in enumerate(results):
414
- title = result.get("title", "未知标题")
415
- url = result.get("url", "#")
416
- preview = result.get("preview", "无预览内容")
417
 
418
  formatted_result = f"\r\n<details><summary>资料[{index}]: {title}</summary>\r\n{preview}\r\n\n[Link]({url})\r\n</details>"
419
  formatted_results.append(formatted_result)
420
-
421
  return '\n\n'.join(formatted_results)
422
 
423
  @staticmethod
424
- async def run_in_executor(func, *args, **kwargs):
425
- return await asyncio.get_event_loop().run_in_executor(
426
- None, partial(func, *args, **kwargs)
427
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
428
 
429
  class GrokApiClient:
430
  def __init__(self, model_id):
431
  if model_id not in CONFIG["MODELS"]:
432
  raise ValueError(f"不支持的模型: {model_id}")
433
- self.model = model_id
434
  self.model_id = CONFIG["MODELS"][model_id]
435
- self.scraper = cloudscraper.create_scraper()
436
-
437
  def process_message_content(self, content):
438
  if isinstance(content, str):
439
  return content
440
  return None
441
-
442
  def get_image_type(self, base64_string):
443
  mime_type = 'image/jpeg'
444
  if 'data:image' in base64_string:
445
  import re
446
- matches = re.match(r'data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,', base64_string)
447
  if matches:
448
  mime_type = matches.group(1)
449
-
450
  extension = mime_type.split('/')[1]
451
  file_name = f"image.{extension}"
452
 
@@ -454,14 +448,14 @@ class GrokApiClient:
454
  "mimeType": mime_type,
455
  "fileName": file_name
456
  }
457
-
458
- async def upload_base64_image(self, base64_data, url):
459
  try:
460
  if 'data:image' in base64_data:
461
  image_buffer = base64_data.split(',')[1]
462
  else:
463
  image_buffer = base64_data
464
-
465
  image_info = self.get_image_type(base64_data)
466
  mime_type = image_info["mimeType"]
467
  file_name = image_info["fileName"]
@@ -477,103 +471,96 @@ class GrokApiClient:
477
 
478
  logger.info("发送图片请求", "Server")
479
 
480
- token = token_manager.get_next_token_for_model(self.model)
481
- if not token:
482
- logger.error("没有可用的token", "Server")
483
- return ''
484
-
485
- headers = {
486
- "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36",
487
- "Connection": "keep-alive",
488
- "Accept": "text/event-stream",
489
- "Content-Type": "text/plain;charset=UTF-8",
490
- "Cookie": token,
491
- "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c"
492
- }
493
-
494
- response = await Utils.run_in_executor(
495
- self.scraper.post,
496
- url,
497
- headers=headers,
498
- data=json.dumps(upload_data),
499
  )
500
 
501
  if response.status_code != 200:
502
- logger.error(f"上传图片失败,状态码:{response.status_code},原因:{response.text}", "Server")
503
  return ''
504
-
505
- result = response.json()
506
- logger.info(f'上传图片成功: {result}', "Server")
507
- return result["fileMetadataId"]
508
 
 
 
 
 
509
  except Exception as error:
510
- logger.error(error, "Server")
511
  return ''
512
-
513
- async def prepare_chat_request(self, request_data):
514
- todo_messages = request_data["messages"]
515
- if request_data["model"] in ["grok-2-imageGen", "grok-3-imageGen", "grok-3-deepsearch"]:
 
 
 
 
 
516
  last_message = todo_messages[-1]
517
- if last_message["role"] != "user":
518
- raise ValueError("画图模型的最后一条消息必须是用户消息!")
519
  todo_messages = [last_message]
520
-
521
  file_attachments = []
522
  messages = ''
523
  last_role = None
524
  last_content = ''
525
- search = request_data["model"] in ["grok-2-search", "grok-3-search"]
526
 
 
527
  def remove_think_tags(text):
528
  import re
529
  text = re.sub(r'<think>[\s\S]*?<\/think>', '', text).strip()
530
  text = re.sub(r'!\[image\]\(data:.*?base64,.*?\)', '[图片]', text)
531
  return text
532
-
533
- async def process_image_url(content):
534
- if content["type"] == "image_url" and "data:image" in content["image_url"]["url"]:
535
- image_response = await self.upload_base64_image(
536
- content["image_url"]["url"],
537
- f"{CONFIG['API']['BASE_URL']}/api/rpc"
538
- )
539
- return image_response
540
- return None
541
-
542
- async def process_content(content):
543
  if isinstance(content, list):
544
  text_content = ''
545
  for item in content:
546
- if item["type"] == "image_url":
547
- text_content += ("[图片]" if text_content else '') + "\n" if text_content else "[图片]"
548
- elif item["type"] == "text":
549
- text_content += ("\n" + remove_think_tags(item["text"]) if text_content else remove_think_tags(item["text"]))
550
  return text_content
551
  elif isinstance(content, dict) and content is not None:
552
- if content["type"] == "image_url":
553
  return "[图片]"
554
- elif content["type"] == "text":
555
  return remove_think_tags(content["text"])
556
  return remove_think_tags(self.process_message_content(content))
557
-
558
  for current in todo_messages:
559
- role = "assistant" if current["role"] == "assistant" else "user"
560
  is_last_message = current == todo_messages[-1]
561
 
562
- logger.info(json.dumps(current, indent=2, ensure_ascii=False), "Server")
563
  if is_last_message and "content" in current:
564
  if isinstance(current["content"], list):
565
  for item in current["content"]:
566
- if item["type"] == "image_url":
567
- logger.info("处理图片附件", "Server")
568
- processed_image = await process_image_url(item)
 
 
569
  if processed_image:
570
  file_attachments.append(processed_image)
571
- elif isinstance(current["content"], dict) and current["content"].get("type") == "image_url":
572
- processed_image = await process_image_url(current["content"])
 
 
 
573
  if processed_image:
574
  file_attachments.append(processed_image)
575
-
576
- text_content = await process_content(current["content"])
 
577
 
578
  if text_content or (is_last_message and file_attachments):
579
  if role == last_role and text_content:
@@ -582,9 +569,10 @@ class GrokApiClient:
582
  else:
583
  messages += f"{role.upper()}: {text_content or '[图片]'}\n"
584
  last_content = text_content
585
- last_role = role
 
586
  return {
587
- "temporary": CONFIG["API"]["IS_TEMP_CONVERSATION"],
588
  "modelName": self.model_id,
589
  "message": messages.strip(),
590
  "fileAttachments": file_attachments[:4],
@@ -597,7 +585,7 @@ class GrokApiClient:
597
  "imageGenerationCount": 1,
598
  "forceConcise": False,
599
  "toolOverrides": {
600
- "imageGen": request_data["model"] in ["grok-2-imageGen", "grok-3-imageGen"],
601
  "webSearch": search,
602
  "xSearch": search,
603
  "xMediaSearch": search,
@@ -608,16 +596,16 @@ class GrokApiClient:
608
  "isPreset": False,
609
  "sendFinalMetadata": True,
610
  "customInstructions": "",
611
- "deepsearchPreset": "default" if request_data["model"] == "grok-3-deepsearch" else "",
612
- "isReasoning": request_data["model"] == "grok-3-reasoning"
613
  }
614
 
615
  class MessageProcessor:
616
  @staticmethod
617
  def create_chat_response(message, model, is_stream=False):
618
  base_response = {
619
- "id": f"chatcmpl-{str(uuid.uuid4())}",
620
- "created": int(datetime.now().timestamp()),
621
  "model": model
622
  }
623
 
@@ -647,473 +635,425 @@ class MessageProcessor:
647
  "usage": None
648
  }
649
 
650
- async def process_model_response(response, model):
651
  result = {"token": None, "imageUrl": None}
652
 
653
  if CONFIG["IS_IMG_GEN"]:
654
- if response and response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:
655
  result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]
656
  return result
657
 
658
- if model == "grok-2":
659
  result["token"] = response.get("token")
660
- elif model in ["grok-2-search", "grok-3-search"]:
661
- if response and response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:
662
- result["token"] = f"\r\n<think>{await Utils.organize_search_results(response['webSearchResults'])}</think>\r\n"
663
  else:
664
  result["token"] = response.get("token")
665
- elif model == "grok-3":
666
  result["token"] = response.get("token")
667
- elif model == "grok-3-deepsearch":
668
- if response and response.get("messageTag") == "final":
669
- result["token"] = response.get("token")
670
- elif model == "grok-3-reasoning":
671
- if response and response.get("isThinking", False) and not CONFIG["SHOW_THINKING"]:
672
  return result
673
-
674
- if response and response.get("isThinking", False) and not CONFIG["IS_THINKING"]:
 
 
 
 
 
 
 
 
 
 
 
675
  result["token"] = "<think>" + response.get("token", "")
676
  CONFIG["IS_THINKING"] = True
677
- elif response and not response.get("isThinking", True) and CONFIG["IS_THINKING"]:
678
  result["token"] = "</think>" + response.get("token", "")
679
  CONFIG["IS_THINKING"] = False
680
  else:
681
  result["token"] = response.get("token")
682
-
683
  return result
684
 
685
- async def stream_response_generator(response, model):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
686
  try:
 
 
 
 
 
687
  CONFIG["IS_THINKING"] = False
688
  CONFIG["IS_IMG_GEN"] = False
689
  CONFIG["IS_IMG_GEN2"] = False
690
- logger.info("开始处理流式响应", "Server")
691
-
692
- async def iter_lines():
693
- line_iter = response.iter_lines()
694
- while True:
695
- try:
696
- line = await Utils.run_in_executor(lambda: next(line_iter, None))
697
- if line is None:
698
- break
699
- yield line
700
- except StopIteration:
701
- break
702
- except Exception as e:
703
- logger.error(f"迭代行时出错: {str(e)}", "Server")
704
- break
705
-
706
- async for line in iter_lines():
707
- if not line:
708
  continue
709
-
710
  try:
711
- try:
712
- line_str = line.decode('utf-8', errors='replace')
713
- # 添加检查,确保解码后的字符串是有效的JSON格式
714
- if not line_str.strip() or not line_str.strip()[0] in ['{', '[']:
715
- logger.warning(f"无效的JSON数据: {line_str}", "Server")
716
- continue
717
- except UnicodeDecodeError:
718
- logger.warning("解码失败,跳过此行数据", "Server")
719
- continue
720
 
721
- try:
722
- line_json = json.loads(line_str)
723
- except json.JSONDecodeError as e:
724
- logger.warning(f"JSON解析失败: {e}, 数据: {line_str[:50]}...", "Server")
725
- continue
726
-
727
- if line_json and line_json.get("error"):
728
- raise ValueError("RateLimitError")
729
-
730
  response_data = line_json.get("result", {}).get("response")
731
  if not response_data:
732
  continue
733
-
734
  if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
735
  CONFIG["IS_IMG_GEN"] = True
736
-
737
- result = await process_model_response(response_data, model)
738
 
739
  if result["token"]:
740
- yield f"data: {json.dumps(MessageProcessor.create_chat_response(result['token'], model, True))}\n\n"
741
-
742
  if result["imageUrl"]:
743
  CONFIG["IS_IMG_GEN2"] = True
744
- data_image = await handle_image_response(result["imageUrl"], model)
745
- yield f"data: {json.dumps(MessageProcessor.create_chat_response(data_image, model, True))}\n\n"
746
-
747
- except Exception as error:
748
- logger.error(f"处理行数据错误: {str(error)}", "Server")
749
  continue
750
-
751
- yield "data: [DONE]\n\n"
752
-
 
 
753
  except Exception as error:
754
- logger.error(f"流式响应总体错误: {str(error)}", "Server")
755
- raise error
 
 
 
756
 
757
- async def handle_normal_response(response, model):
758
- try:
759
- full_response = ''
760
  CONFIG["IS_THINKING"] = False
761
  CONFIG["IS_IMG_GEN"] = False
762
  CONFIG["IS_IMG_GEN2"] = False
763
- logger.info("开始处理非流式响应", "Server")
764
- image_url = None
765
-
766
- async def iter_lines():
767
- line_iter = response.iter_lines()
768
- while True:
769
- try:
770
- line = await Utils.run_in_executor(lambda: next(line_iter, None))
771
- if line is None:
772
- break
773
- yield line
774
- except StopIteration:
775
- break
776
- except Exception as e:
777
- logger.error(f"迭代行时出错: {str(e)}", "Server")
778
- break
779
-
780
- async for line in iter_lines():
781
- if not line:
782
  continue
783
-
784
  try:
785
- try:
786
- line_str = line.decode('utf-8', errors='replace')
787
- # 添加检查,确保解码后的字符串是有效的JSON格式
788
- if not line_str.strip() or not line_str.strip()[0] in ['{', '[']:
789
- logger.warning(f"无效的JSON数据: {line_str}", "Server")
790
- continue
791
- except UnicodeDecodeError:
792
- logger.warning("解码失败,跳过此行数据", "Server")
793
- continue
794
 
795
- try:
796
- line_json = json.loads(line_str)
797
- except json.JSONDecodeError as e:
798
- logger.warning(f"JSON解析失败: {e}, 数据: {line_str[:50]}...", "Server")
799
- continue
800
-
801
- if line_json and line_json.get("error"):
802
- raise ValueError("RateLimitError")
803
-
804
  response_data = line_json.get("result", {}).get("response")
805
  if not response_data:
806
  continue
807
-
808
  if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
809
  CONFIG["IS_IMG_GEN"] = True
810
-
811
- result = await process_model_response(response_data, model)
812
 
813
  if result["token"]:
814
- full_response += result["token"]
815
-
816
  if result["imageUrl"]:
817
  CONFIG["IS_IMG_GEN2"] = True
818
- image_url = result["imageUrl"]
819
-
820
- except Exception as error:
821
- logger.error(f"处理行数据错误: {str(error)}", "Server")
822
- continue
823
-
824
- if CONFIG["IS_IMG_GEN2"] and image_url:
825
- data_image = await handle_image_response(image_url, model)
826
- return MessageProcessor.create_chat_response(data_image, model)
827
- else:
828
- return MessageProcessor.create_chat_response(full_response, model)
829
 
830
- except Exception as error:
831
- logger.error(f"非流式响应总体错误: {str(error)}", "Server")
832
- raise error
 
 
 
 
 
833
 
834
- async def handle_image_response(image_url,model):
835
- MAX_RETRIES = 2
836
- retry_count = 0
837
- scraper = cloudscraper.create_scraper()
 
 
838
 
839
- while retry_count < MAX_RETRIES:
840
- try:
841
- token = token_manager.get_next_token_for_model(model)
842
- if not token:
843
- raise ValueError("没有可用的token")
844
-
845
- image_response = await Utils.run_in_executor(
846
- scraper.get,
847
- f"https://assets.grok.com/{image_url}",
848
- headers={
849
- **CONFIG["DEFAULT_HEADERS"],
850
- "cookie": token
851
- }
852
- )
853
-
854
- if image_response.status_code == 200:
855
- break
856
-
857
- retry_count += 1
858
- if retry_count == MAX_RETRIES:
859
- raise ValueError(f"上游服务请求失败! status: {image_response.status_code}")
860
-
861
- await asyncio.sleep(1 * retry_count)
862
-
863
- except Exception as error:
864
- logger.error(error, "Server")
865
- retry_count += 1
866
- if retry_count == MAX_RETRIES:
867
- raise error
868
-
869
- await asyncio.sleep(1 * retry_count)
870
 
871
- image_content = image_response.content
 
872
 
873
- if CONFIG["API"]["PICGO_KEY"]:
874
- form = aiohttp.FormData()
875
- form.add_field('source',
876
- io.BytesIO(image_content),
877
- filename=f'image-{int(datetime.now().timestamp())}.jpg',
878
- content_type='image/jpeg')
879
-
880
- async with aiohttp.ClientSession() as session:
881
- async with session.post(
882
- "https://www.picgo.net/api/1/upload",
883
- data=form,
884
- headers={"X-API-Key": CONFIG["API"]["PICGO_KEY"]}
885
- ) as response_url:
886
- if response_url.status != 200:
887
- return "生图失败,请查看PICGO图床密钥是否设置正确"
888
- else:
889
- logger.info("生图成功", "Server")
890
- result = await response_url.json()
891
- return f"![image]({result['image']['url']})"
892
- elif CONFIG["API"]["TUMY_KEY"]:
893
- form = aiohttp.FormData()
894
- form.add_field('file',
895
- io.BytesIO(image_content),
896
- filename=f'image-{int(datetime.now().timestamp())}.jpg',
897
- content_type='image/jpeg')
898
-
899
- async with aiohttp.ClientSession() as session:
900
- async with session.post(
901
- "https://tu.my/api/v1/upload",
902
- data=form,
903
- headers={
904
- "Accept": "application/json",
905
- "Authorization": f"Bearer {CONFIG['API']['TUMY_KEY']}"
906
- }
907
- ) as response_url:
908
- if response_url.status != 200:
909
- return "生图失败,请查看TUMY图床密钥是否设置正确"
910
- else:
911
- logger.info("生图成功", "Server")
912
- result = await response_url.json()
913
- return f"![image]({result['image']['url']})"
914
- # 如果没有PICGO_KEY或者TUMY_KEY则返回base64图片
915
- image_base64 = base64.b64encode(image_content).decode('utf-8')
916
- return f"![image](data:image/jpeg;base64,{image_base64})"
917
 
918
 
919
- app = Quart(__name__)
920
- app = cors(app, allow_origin="*", allow_methods=["GET", "POST", "OPTIONS"], allow_headers=["Content-Type", "Authorization"])
921
 
922
- @app.before_request
923
- async def before_request():
924
- await logger.request_logger(request)
925
-
926
- @app.route('/v1/models', methods=['GET'])
927
- async def models():
928
- return jsonify({
929
- "object": "list",
930
- "data": [
931
- {
932
- "id": model,
933
- "object": "model",
934
- "created": int(datetime.now().timestamp()),
935
- "owned_by": "grok"
936
- } for model in CONFIG["MODELS"].keys()
937
- ]
938
- })
939
 
 
 
 
940
 
941
  @app.route('/get/tokens', methods=['GET'])
942
- async def get_tokens():
943
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
944
-
945
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
946
  return jsonify({"error": '自定义的SSO令牌模式无法获取轮询sso令牌状态'}), 403
947
  elif auth_token != CONFIG["API"]["API_KEY"]:
948
  return jsonify({"error": 'Unauthorized'}), 401
949
-
950
  return jsonify(token_manager.get_token_status_map())
951
 
952
  @app.route('/add/token', methods=['POST'])
953
- async def add_token():
954
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
955
-
956
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
957
  return jsonify({"error": '自定义的SSO令牌模式无法添加sso令牌'}), 403
958
  elif auth_token != CONFIG["API"]["API_KEY"]:
959
  return jsonify({"error": 'Unauthorized'}), 401
960
-
961
  try:
962
- data = await request.get_json()
963
- sso = data.get('sso')
964
- if not sso:
965
- return jsonify({"error": 'SSO令牌不能为空'}), 400
966
-
967
- await token_manager.add_token(f"sso-rw={sso};sso={sso}")
968
- return jsonify(token_manager.get_token_status_map().get(sso, {}))
969
  except Exception as error:
970
- logger.error(error, "Server")
971
  return jsonify({"error": '添加sso令牌失败'}), 500
972
 
973
  @app.route('/delete/token', methods=['POST'])
974
- async def delete_token():
975
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
976
-
977
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
978
  return jsonify({"error": '自定义的SSO���牌模式无法删除sso令牌'}), 403
979
  elif auth_token != CONFIG["API"]["API_KEY"]:
980
  return jsonify({"error": 'Unauthorized'}), 401
981
-
982
  try:
983
- data = await request.get_json()
984
- sso = data.get('sso')
985
- if not sso:
986
- return jsonify({"error": 'SSO令牌不能为空'}), 400
987
-
988
- success = await token_manager.delete_token(f"sso-rw={sso};sso={sso}")
989
- if success:
990
- return jsonify({"message": '删除sso令牌成功'})
991
- else:
992
- return jsonify({"error": '删除sso令牌失败'}), 500
993
  except Exception as error:
994
- logger.error(error, "Server")
995
  return jsonify({"error": '删除sso令牌失败'}), 500
996
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
997
  @app.route('/v1/chat/completions', methods=['POST'])
998
- async def chat_completions():
999
  try:
1000
- data = await request.get_json()
1001
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
1002
-
1003
  if auth_token:
1004
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
1005
- await token_manager.set_token(f"sso-rw={auth_token};sso={auth_token}")
 
1006
  elif auth_token != CONFIG["API"]["API_KEY"]:
1007
- return jsonify({"error": "Unauthorized"}), 401
1008
  else:
1009
- return jsonify({"error": "Unauthorized"}), 401
1010
-
 
1011
  model = data.get("model")
1012
  stream = data.get("stream", False)
 
1013
  retry_count = 0
 
 
1014
 
1015
- try:
1016
- grok_client = GrokApiClient(model)
1017
- request_payload = await grok_client.prepare_chat_request(data)
1018
 
1019
- while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:
1020
- retry_count += 1
1021
- logger.info(f"开始请求(第{retry_count}次尝试)", "Server")
1022
-
1023
- token = token_manager.get_next_token_for_model(model)
1024
- if not token:
1025
- logger.error(f"没有可用的{model}模型令牌", "Server")
1026
- if retry_count == CONFIG["RETRY"]["MAX_ATTEMPTS"]:
1027
- raise ValueError(f"没有可用的{model}模型令牌")
1028
- continue
1029
-
1030
- scraper = cloudscraper.create_scraper()
1031
-
1032
- try:
1033
- headers = {
1034
- "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36",
1035
- "Connection": "keep-alive",
1036
  "Accept": "text/event-stream",
 
1037
  "Content-Type": "text/plain;charset=UTF-8",
1038
- "Cookie": token,
1039
- "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c"
1040
- }
1041
- logger.info(f"使用令牌: {token}", "Server")
1042
-
1043
- response = await Utils.run_in_executor(
1044
- scraper.post,
1045
- f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",
1046
- headers=headers,
1047
- data=json.dumps(request_payload),
1048
- stream=True
1049
- )
1050
 
1051
- if response.status_code == 200:
1052
- logger.info("请求成功", "Server")
1053
-
1054
  if stream:
1055
  return Response(
1056
- stream_response_generator(response, model),
1057
- content_type='text/event-stream',
1058
- headers={
1059
- 'Cache-Control': 'no-cache',
1060
- 'Connection': 'keep-alive'
1061
- }
1062
  )
1063
  else:
1064
- result = await handle_normal_response(response, model)
1065
- return jsonify(result)
1066
- else:
1067
- logger.error(f"请求失败: 状态码 {response.status_code}", "Server")
1068
- token_manager.remove_token_from_model(model, token)
 
 
1069
 
1070
- except Exception as e:
1071
- logger.error(f"请求异常: {str(e)}", "Server")
1072
- token_manager.remove_token_from_model(model, token)
 
 
 
 
1073
 
1074
- raise ValueError("请求失败,已达到最大重试次数")
1075
-
1076
- except Exception as e:
1077
- logger.error(e, "ChatAPI")
1078
- return jsonify({
1079
- "error": {
1080
- "message": str(e),
1081
- "type": "server_error"
1082
- }
1083
- }), 500
 
1084
 
1085
- except Exception as e:
1086
- logger.error(e, "ChatAPI")
 
 
 
 
 
 
 
 
1087
  return jsonify({
1088
  "error": {
1089
- "message": str(e),
1090
  "type": "server_error"
1091
  }
1092
  }), 500
1093
 
1094
- @app.route('/', methods=['GET'])
1095
- async def index():
1096
- return "api运行正常"
 
1097
 
1098
- if __name__ == "__main__":
1099
- asyncio.run(initialize_tokens())
1100
- app.run(host="0.0.0.0", port=CONFIG["SERVER"]["PORT"])
1101
- # 在现有app.py文件的最底部添加或修改这段代码
1102
- if __name__ == "__main__":
1103
- # 初始化令牌
1104
- import asyncio
1105
- asyncio.run(initialize_tokens())
1106
-
1107
- # 获取端口,优先使用环境变量
1108
- port = int(os.getenv("PORT", CONFIG["SERVER"]["PORT"]))
1109
 
1110
- # 启动服务器
1111
- from hypercorn.asyncio import serve
1112
- from hypercorn.config import Config as HyperConfig
1113
-
1114
- config = HyperConfig()
1115
- config.bind = [f"0.0.0.0:{port}"]
1116
- config.use_reloader = True
1117
-
1118
- logger.info(f"服务器正在启动,端口:{port}", "Server")
1119
- asyncio.run(serve(app, config))
 
1
  import os
2
+ from dotenv import load_dotenv
3
+
4
+ load_dotenv() # 加载.env文件中的环境变量
5
+
6
  import json
7
  import uuid
8
+ import time
9
  import base64
10
  import sys
11
  import inspect
12
  from loguru import logger
 
 
 
 
 
 
 
13
 
14
+ import requests
15
+ from flask import Flask, request, Response, jsonify, stream_with_context
16
+ from curl_cffi import requests as curl_requests
17
+ from werkzeug.middleware.proxy_fix import ProxyFix
18
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
 
20
  class Logger:
21
  def __init__(self, level="INFO", colorize=True, format=None):
 
22
  logger.remove()
23
 
24
  if format is None:
 
84
 
85
  logger = Logger(level="INFO")
86
 
87
+
88
+ CONFIG = {
89
+ "MODELS": {
90
+ 'grok-2': 'grok-latest',
91
+ 'grok-2-imageGen': 'grok-latest',
92
+ 'grok-2-search': 'grok-latest',
93
+ "grok-3": "grok-3",
94
+ "grok-3-search": "grok-3",
95
+ "grok-3-imageGen": "grok-3",
96
+ "grok-3-deepsearch": "grok-3",
97
+ "grok-3-reasoning": "grok-3"
98
+ },
99
+ "API": {
100
+ "IS_CUSTOM_SSO": os.environ.get("IS_CUSTOM_SSO", "false").lower() == "true",
101
+ "BASE_URL": "https://grok.com",
102
+ "API_KEY": os.environ.get("API_KEY", "sk-123456"),
103
+ "SIGNATURE_COOKIE": None,
104
+ "PICGO_KEY": os.environ.get("PICGO_KEY") or None,
105
+ "TUMY_KEY": os.environ.get("TUMY_KEY") or None,
106
+ "RETRY_TIME": 1000,
107
+ "PROXY": os.environ.get("PROXY") or None,
108
+ "IS_TEMP_CONVERSATION": os.environ.get("IS_TEMP_CONVERSATION", "false").lower() == "true"
109
+ },
110
+ "SERVER": {
111
+ "PORT": int(os.environ.get("PORT", 7860)) # 修改默认端口为7860
112
+ },
113
+ "RETRY": {
114
+ "MAX_ATTEMPTS": 2
115
+ },
116
+ "SHOW_THINKING": os.environ.get("SHOW_THINKING", "true").lower() == "true",
117
+ "IS_THINKING": False,
118
+ "IS_IMG_GEN": False,
119
+ "IS_IMG_GEN2": False,
120
+ "ISSHOW_SEARCH_RESULTS": os.environ.get("ISSHOW_SEARCH_RESULTS", "false").lower() == "true"
121
+ }
122
+
123
+
124
+ DEFAULT_HEADERS = {
125
+ 'Accept': '*/*',
126
+ 'Accept-Language': 'zh-CN,zh;q=0.9',
127
+ 'Accept-Encoding': 'gzip, deflate, br, zstd',
128
+ 'Content-Type': 'text/plain;charset=UTF-8',
129
+ 'Connection': 'keep-alive',
130
+ 'Origin': 'https://grok.com',
131
+ 'Priority': 'u=1, i',
132
+ 'Sec-Ch-Ua': '"Chromium";v="130", "Google Chrome";v="130", "Not?A_Brand";v="99"',
133
+ 'Sec-Ch-Ua-Mobile': '?0',
134
+ 'Sec-Ch-Ua-Platform': '"Windows"',
135
+ 'Sec-Fetch-Dest': 'empty',
136
+ 'Sec-Fetch-Mode': 'cors',
137
+ 'Sec-Fetch-Site': 'same-origin',
138
+ 'Baggage': 'sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c'
139
+ }
140
+
141
  class AuthTokenManager:
142
  def __init__(self):
143
  self.token_model_map = {}
144
  self.expired_tokens = set()
145
  self.token_status_map = {}
 
 
 
146
 
147
  self.model_config = {
148
  "grok-2": {
149
+ "RequestFrequency": 30,
150
+ "ExpirationTime": 1 * 60 * 60 * 1000 # 1小时
151
  },
152
  "grok-3": {
153
  "RequestFrequency": 20,
154
+ "ExpirationTime": 2 * 60 * 60 * 1000 # 2小时
155
  },
156
  "grok-3-deepsearch": {
157
  "RequestFrequency": 10,
158
+ "ExpirationTime": 24 * 60 * 60 * 1000 # 24小时
159
  },
160
  "grok-3-reasoning": {
161
  "RequestFrequency": 10,
162
+ "ExpirationTime": 24 * 60 * 60 * 1000 # 24小时
163
  }
164
  }
165
+ self.token_reset_switch = False
166
+ self.token_reset_timer = None
167
+
168
+ def add_token(self, token):
169
  sso = token.split("sso=")[1].split(";")[0]
170
  for model in self.model_config.keys():
171
  if model not in self.token_model_map:
172
  self.token_model_map[model] = []
 
173
  if sso not in self.token_status_map:
174
  self.token_status_map[sso] = {}
175
+
176
+ existing_token_entry = next((entry for entry in self.token_model_map[model] if entry["token"] == token), None)
177
+
 
178
  if not existing_token_entry:
179
  self.token_model_map[model].append({
180
  "token": token,
181
  "RequestCount": 0,
182
+ "AddedTime": int(time.time() * 1000),
183
  "StartCallTime": None
184
  })
185
 
 
189
  "invalidatedTime": None,
190
  "totalRequestCount": 0
191
  }
192
+
193
+ def set_token(self, token):
 
194
  models = list(self.model_config.keys())
195
+ self.token_model_map = {model: [{
196
+ "token": token,
197
+ "RequestCount": 0,
198
+ "AddedTime": int(time.time() * 1000),
199
+ "StartCallTime": None
200
+ }] for model in models}
 
201
 
202
  sso = token.split("sso=")[1].split(";")[0]
203
+ self.token_status_map[sso] = {model: {
204
+ "isValid": True,
205
+ "invalidatedTime": None,
206
+ "totalRequestCount": 0
207
+ } for model in models}
208
+
209
+ def delete_token(self, token):
 
 
 
210
  try:
211
  sso = token.split("sso=")[1].split(";")[0]
 
212
  for model in self.token_model_map:
213
+ self.token_model_map[model] = [entry for entry in self.token_model_map[model] if entry["token"] != token]
214
+
 
 
 
215
  if sso in self.token_status_map:
216
  del self.token_status_map[sso]
217
 
218
  logger.info(f"令牌已成功移除: {token}", "TokenManager")
219
  return True
220
  except Exception as error:
221
+ logger.error(f"令牌删除失败: {str(error)}")
222
  return False
223
+
224
  def get_next_token_for_model(self, model_id):
225
  normalized_model = self.normalize_model_name(model_id)
226
 
227
  if normalized_model not in self.token_model_map or not self.token_model_map[normalized_model]:
228
  return None
229
+
230
  token_entry = self.token_model_map[normalized_model][0]
231
 
232
  if token_entry:
 
 
 
233
  if token_entry["StartCallTime"] is None:
234
+ token_entry["StartCallTime"] = int(time.time() * 1000)
235
+
236
  if not self.token_reset_switch:
237
  self.start_token_reset_process()
238
  self.token_reset_switch = True
239
+
240
  token_entry["RequestCount"] += 1
241
 
242
  if token_entry["RequestCount"] > self.model_config[normalized_model]["RequestFrequency"]:
243
  self.remove_token_from_model(normalized_model, token_entry["token"])
244
+ next_token_entry = self.token_model_map[normalized_model][0] if self.token_model_map[normalized_model] else None
 
 
245
  return next_token_entry["token"] if next_token_entry else None
246
 
247
  sso = token_entry["token"].split("sso=")[1].split(";")[0]
248
  if sso in self.token_status_map and normalized_model in self.token_status_map[sso]:
249
  if token_entry["RequestCount"] == self.model_config[normalized_model]["RequestFrequency"]:
250
  self.token_status_map[sso][normalized_model]["isValid"] = False
251
+ self.token_status_map[sso][normalized_model]["invalidatedTime"] = int(time.time() * 1000)
 
252
  self.token_status_map[sso][normalized_model]["totalRequestCount"] += 1
253
+
254
  return token_entry["token"]
255
 
256
  return None
257
+
258
  def remove_token_from_model(self, model_id, token):
259
  normalized_model = self.normalize_model_name(model_id)
260
 
261
  if normalized_model not in self.token_model_map:
262
  logger.error(f"模型 {normalized_model} 不存在", "TokenManager")
263
  return False
264
+
265
  model_tokens = self.token_model_map[normalized_model]
266
+ token_index = next((i for i, entry in enumerate(model_tokens) if entry["token"] == token), -1)
267
 
 
 
 
 
 
268
  if token_index != -1:
269
  removed_token_entry = model_tokens.pop(token_index)
270
  self.expired_tokens.add((
271
  removed_token_entry["token"],
272
  normalized_model,
273
+ int(time.time() * 1000)
274
  ))
275
 
276
  if not self.token_reset_switch:
277
  self.start_token_reset_process()
278
  self.token_reset_switch = True
279
+
280
  logger.info(f"模型{model_id}的令牌已失效,已成功移除令牌: {token}", "TokenManager")
281
  return True
282
+
283
  logger.error(f"在模型 {normalized_model} 中未找到 token: {token}", "TokenManager")
284
  return False
285
+
286
  def get_expired_tokens(self):
287
  return list(self.expired_tokens)
288
+
289
  def normalize_model_name(self, model):
290
  if model.startswith('grok-') and 'deepsearch' not in model and 'reasoning' not in model:
291
  return '-'.join(model.split('-')[:2])
292
  return model
293
+
294
  def get_token_count_for_model(self, model_id):
295
  normalized_model = self.normalize_model_name(model_id)
296
  return len(self.token_model_map.get(normalized_model, []))
297
+
298
  def get_remaining_token_request_capacity(self):
299
  remaining_capacity_map = {}
300
 
301
+ for model in self.model_config.keys():
302
  model_tokens = self.token_model_map.get(model, [])
303
  model_request_frequency = self.model_config[model]["RequestFrequency"]
304
 
305
+ total_used_requests = sum(token_entry.get("RequestCount", 0) for token_entry in model_tokens)
306
+
307
  remaining_capacity = (len(model_tokens) * model_request_frequency) - total_used_requests
308
  remaining_capacity_map[model] = max(0, remaining_capacity)
 
 
309
 
310
+ return remaining_capacity_map
311
+
312
  def get_token_array_for_model(self, model_id):
313
  normalized_model = self.normalize_model_name(model_id)
314
  return self.token_model_map.get(normalized_model, [])
315
+
316
  def start_token_reset_process(self):
317
+ def reset_expired_tokens():
318
+ now = int(time.time() * 1000)
319
+
320
+ tokens_to_remove = set()
321
+ for token_info in self.expired_tokens:
322
+ token, model, expired_time = token_info
323
+ expiration_time = self.model_config[model]["ExpirationTime"]
 
 
324
 
325
+ if now - expired_time >= expiration_time:
326
+ if not any(entry["token"] == token for entry in self.token_model_map.get(model, [])):
327
+ if model not in self.token_model_map:
328
+ self.token_model_map[model] = []
329
+
330
+ self.token_model_map[model].append({
331
+ "token": token,
332
+ "RequestCount": 0,
333
+ "AddedTime": now,
334
+ "StartCallTime": None
335
+ })
336
+
337
+ sso = token.split("sso=")[1].split(";")[0]
338
+ if sso in self.token_status_map and model in self.token_status_map[sso]:
339
+ self.token_status_map[sso][model]["isValid"] = True
340
+ self.token_status_map[sso][model]["invalidatedTime"] = None
341
+ self.token_status_map[sso][model]["totalRequestCount"] = 0
342
+
343
+ tokens_to_remove.add(token_info)
344
+
345
+ self.expired_tokens -= tokens_to_remove
346
+
347
+ for model in self.model_config.keys():
348
+ if model not in self.token_model_map:
349
+ continue
350
+
351
+ for token_entry in self.token_model_map[model]:
352
+ if not token_entry.get("StartCallTime"):
353
+ continue
354
 
355
+ expiration_time = self.model_config[model]["ExpirationTime"]
356
+ if now - token_entry["StartCallTime"] >= expiration_time:
357
+ sso = token_entry["token"].split("sso=")[1].split(";")[0]
 
 
 
 
 
 
 
358
  if sso in self.token_status_map and model in self.token_status_map[sso]:
359
  self.token_status_map[sso][model]["isValid"] = True
360
  self.token_status_map[sso][model]["invalidatedTime"] = None
361
  self.token_status_map[sso][model]["totalRequestCount"] = 0
 
 
 
 
 
 
 
 
 
362
 
363
+ token_entry["RequestCount"] = 0
364
+ token_entry["StartCallTime"] = None
365
+
366
+ import threading
367
+ # 启动一个线程执行定时任务,每小时执行一次
368
+ def run_timer():
369
+ while True:
370
+ reset_expired_tokens()
371
+ time.sleep(3600)
372
+
373
+ timer_thread = threading.Thread(target=run_timer)
374
+ timer_thread.daemon = True
375
+ timer_thread.start()
376
+
 
 
 
 
 
 
377
  def get_all_tokens(self):
378
  all_tokens = set()
379
  for model_tokens in self.token_model_map.values():
380
  for entry in model_tokens:
381
  all_tokens.add(entry["token"])
382
  return list(all_tokens)
383
+
384
  def get_token_status_map(self):
385
  return self.token_status_map
 
 
 
 
 
 
386
 
 
 
 
 
 
 
 
 
387
  class Utils:
388
  @staticmethod
389
+ def organize_search_results(search_results):
390
+ if not search_results or 'results' not in search_results:
391
  return ''
392
+
393
+ results = search_results['results']
394
  formatted_results = []
395
 
396
  for index, result in enumerate(results):
397
+ title = result.get('title', '未知标题')
398
+ url = result.get('url', '#')
399
+ preview = result.get('preview', '无预览内容')
400
 
401
  formatted_result = f"\r\n<details><summary>资料[{index}]: {title}</summary>\r\n{preview}\r\n\n[Link]({url})\r\n</details>"
402
  formatted_results.append(formatted_result)
403
+
404
  return '\n\n'.join(formatted_results)
405
 
406
  @staticmethod
407
+ def create_auth_headers(model):
408
+ return token_manager.get_next_token_for_model(model)
409
+
410
+ @staticmethod
411
+ def get_proxy_options():
412
+ proxy = CONFIG["API"]["PROXY"]
413
+ proxy_options = {}
414
+
415
+ if proxy:
416
+ logger.info(f"使用代理: {proxy}", "Server")
417
+ proxy_options["proxies"] = {"https": proxy, "http": proxy}
418
+
419
+ if proxy.startswith("socks5://"):
420
+ proxy_options["proxies"] = {"https": proxy, "http": proxy}
421
+ proxy_options["proxy_type"] = "socks5"
422
+
423
+ return proxy_options
424
 
425
  class GrokApiClient:
426
  def __init__(self, model_id):
427
  if model_id not in CONFIG["MODELS"]:
428
  raise ValueError(f"不支持的模型: {model_id}")
 
429
  self.model_id = CONFIG["MODELS"][model_id]
430
+
 
431
  def process_message_content(self, content):
432
  if isinstance(content, str):
433
  return content
434
  return None
435
+
436
  def get_image_type(self, base64_string):
437
  mime_type = 'image/jpeg'
438
  if 'data:image' in base64_string:
439
  import re
440
+ matches = re.search(r'data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,', base64_string)
441
  if matches:
442
  mime_type = matches.group(1)
443
+
444
  extension = mime_type.split('/')[1]
445
  file_name = f"image.{extension}"
446
 
 
448
  "mimeType": mime_type,
449
  "fileName": file_name
450
  }
451
+
452
+ def upload_base64_image(self, base64_data, url):
453
  try:
454
  if 'data:image' in base64_data:
455
  image_buffer = base64_data.split(',')[1]
456
  else:
457
  image_buffer = base64_data
458
+
459
  image_info = self.get_image_type(base64_data)
460
  mime_type = image_info["mimeType"]
461
  file_name = image_info["fileName"]
 
471
 
472
  logger.info("发送图片请求", "Server")
473
 
474
+ proxy_options = Utils.get_proxy_options()
475
+ response = curl_requests.post(
476
+ url,
477
+ headers={
478
+ **DEFAULT_HEADERS,
479
+ "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
480
+ },
481
+ json=upload_data,
482
+ impersonate="chrome120",
483
+ **proxy_options
 
 
 
 
 
 
 
 
 
484
  )
485
 
486
  if response.status_code != 200:
487
+ logger.error(f"上传图片失败,状态码:{response.status_code}", "Server")
488
  return ''
 
 
 
 
489
 
490
+ result = response.json()
491
+ logger.info(f"上传图片成功: {result}", "Server")
492
+ return result.get("fileMetadataId", "")
493
+
494
  except Exception as error:
495
+ logger.error(str(error), "Server")
496
  return ''
497
+
498
+ def prepare_chat_request(self, request):
499
+ if ((request["model"] == 'grok-2-imageGen' or request["model"] == 'grok-3-imageGen') and
500
+ not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"] and
501
+ request.get("stream", False)):
502
+ raise ValueError("该模型流式输出需要配置PICGO或者TUMY图床密钥!")
503
+
504
+ todo_messages = request["messages"]
505
+ if request["model"] in ['grok-2-imageGen', 'grok-3-imageGen', 'grok-3-deepsearch']:
506
  last_message = todo_messages[-1]
507
+ if last_message["role"] != 'user':
508
+ raise ValueError('此模型最后一条消息必须是用户消息!')
509
  todo_messages = [last_message]
510
+
511
  file_attachments = []
512
  messages = ''
513
  last_role = None
514
  last_content = ''
515
+ search = request["model"] in ['grok-2-search', 'grok-3-search']
516
 
517
+ # 移除<think>标签及其内容和base64图片
518
  def remove_think_tags(text):
519
  import re
520
  text = re.sub(r'<think>[\s\S]*?<\/think>', '', text).strip()
521
  text = re.sub(r'!\[image\]\(data:.*?base64,.*?\)', '[图片]', text)
522
  return text
523
+
524
+ def process_content(content):
 
 
 
 
 
 
 
 
 
525
  if isinstance(content, list):
526
  text_content = ''
527
  for item in content:
528
+ if item["type"] == 'image_url':
529
+ text_content += ("[图片]" if not text_content else '\n[图片]')
530
+ elif item["type"] == 'text':
531
+ text_content += (remove_think_tags(item["text"]) if not text_content else '\n' + remove_think_tags(item["text"]))
532
  return text_content
533
  elif isinstance(content, dict) and content is not None:
534
+ if content["type"] == 'image_url':
535
  return "[图片]"
536
+ elif content["type"] == 'text':
537
  return remove_think_tags(content["text"])
538
  return remove_think_tags(self.process_message_content(content))
539
+
540
  for current in todo_messages:
541
+ role = 'assistant' if current["role"] == 'assistant' else 'user'
542
  is_last_message = current == todo_messages[-1]
543
 
 
544
  if is_last_message and "content" in current:
545
  if isinstance(current["content"], list):
546
  for item in current["content"]:
547
+ if item["type"] == 'image_url':
548
+ processed_image = self.upload_base64_image(
549
+ item["image_url"]["url"],
550
+ f"{CONFIG['API']['BASE_URL']}/api/rpc"
551
+ )
552
  if processed_image:
553
  file_attachments.append(processed_image)
554
+ elif isinstance(current["content"], dict) and current["content"].get("type") == 'image_url':
555
+ processed_image = self.upload_base64_image(
556
+ current["content"]["image_url"]["url"],
557
+ f"{CONFIG['API']['BASE_URL']}/api/rpc"
558
+ )
559
  if processed_image:
560
  file_attachments.append(processed_image)
561
+
562
+
563
+ text_content = process_content(current.get("content", ""))
564
 
565
  if text_content or (is_last_message and file_attachments):
566
  if role == last_role and text_content:
 
569
  else:
570
  messages += f"{role.upper()}: {text_content or '[图片]'}\n"
571
  last_content = text_content
572
+ last_role = role
573
+
574
  return {
575
+ "temporary": CONFIG["API"].get("IS_TEMP_CONVERSATION", False),
576
  "modelName": self.model_id,
577
  "message": messages.strip(),
578
  "fileAttachments": file_attachments[:4],
 
585
  "imageGenerationCount": 1,
586
  "forceConcise": False,
587
  "toolOverrides": {
588
+ "imageGen": request["model"] in ['grok-2-imageGen', 'grok-3-imageGen'],
589
  "webSearch": search,
590
  "xSearch": search,
591
  "xMediaSearch": search,
 
596
  "isPreset": False,
597
  "sendFinalMetadata": True,
598
  "customInstructions": "",
599
+ "deepsearchPreset": "default" if request["model"] == 'grok-3-deepsearch' else "",
600
+ "isReasoning": request["model"] == 'grok-3-reasoning'
601
  }
602
 
603
  class MessageProcessor:
604
  @staticmethod
605
  def create_chat_response(message, model, is_stream=False):
606
  base_response = {
607
+ "id": f"chatcmpl-{uuid.uuid4()}",
608
+ "created": int(time.time()),
609
  "model": model
610
  }
611
 
 
635
  "usage": None
636
  }
637
 
638
+ def process_model_response(response, model):
639
  result = {"token": None, "imageUrl": None}
640
 
641
  if CONFIG["IS_IMG_GEN"]:
642
+ if response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:
643
  result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]
644
  return result
645
 
646
+ if model == 'grok-2':
647
  result["token"] = response.get("token")
648
+ elif model in ['grok-2-search', 'grok-3-search']:
649
+ if response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:
650
+ result["token"] = f"\r\n<think>{Utils.organize_search_results(response['webSearchResults'])}</think>\r\n"
651
  else:
652
  result["token"] = response.get("token")
653
+ elif model == 'grok-3':
654
  result["token"] = response.get("token")
655
+ elif model == 'grok-3-deepsearch':
656
+ if response.get("messageStepId") and not CONFIG["SHOW_THINKING"]:
 
 
 
657
  return result
658
+ if response.get("messageStepId") and not CONFIG["IS_THINKING"]:
659
+ result["token"] = "<think>" + response.get("token", "")
660
+ CONFIG["IS_THINKING"] = True
661
+ elif not response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "final":
662
+ result["token"] = "</think>" + response.get("token", "")
663
+ CONFIG["IS_THINKING"] = False
664
+ elif (response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "assistant") or response.get("messageTag") == "final":
665
+ result["token"] = response.get("token")
666
+ elif model == 'grok-3-reasoning':
667
+ if response.get("isThinking") and not CONFIG["SHOW_THINKING"]:
668
+ return result
669
+
670
+ if response.get("isThinking") and not CONFIG["IS_THINKING"]:
671
  result["token"] = "<think>" + response.get("token", "")
672
  CONFIG["IS_THINKING"] = True
673
+ elif not response.get("isThinking") and CONFIG["IS_THINKING"]:
674
  result["token"] = "</think>" + response.get("token", "")
675
  CONFIG["IS_THINKING"] = False
676
  else:
677
  result["token"] = response.get("token")
678
+
679
  return result
680
 
681
+ def handle_image_response(image_url):
682
+ max_retries = 2
683
+ retry_count = 0
684
+ image_base64_response = None
685
+
686
+ while retry_count < max_retries:
687
+ try:
688
+ proxy_options = Utils.get_proxy_options()
689
+ image_base64_response = curl_requests.get(
690
+ f"https://assets.grok.com/{image_url}",
691
+ headers={
692
+ **DEFAULT_HEADERS,
693
+ "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
694
+ },
695
+ impersonate="chrome120",
696
+ **proxy_options
697
+ )
698
+
699
+ if image_base64_response.status_code == 200:
700
+ break
701
+
702
+ retry_count += 1
703
+ if retry_count == max_retries:
704
+ raise Exception(f"上游服务请求失败! status: {image_base64_response.status_code}")
705
+
706
+ time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)
707
+
708
+ except Exception as error:
709
+ logger.error(str(error), "Server")
710
+ retry_count += 1
711
+ if retry_count == max_retries:
712
+ raise
713
+
714
+ time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)
715
+
716
+ image_buffer = image_base64_response.content
717
+
718
+ if not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"]:
719
+ base64_image = base64.b64encode(image_buffer).decode('utf-8')
720
+ image_content_type = image_base64_response.headers.get('content-type', 'image/jpeg')
721
+ return f"![image](data:{image_content_type};base64,{base64_image})"
722
+
723
+ logger.info("开始上传图床", "Server")
724
+
725
+ if CONFIG["API"]["PICGO_KEY"]:
726
+ files = {'source': ('image.jpg', image_buffer, 'image/jpeg')}
727
+ headers = {
728
+ "X-API-Key": CONFIG["API"]["PICGO_KEY"]
729
+ }
730
+
731
+ response_url = requests.post(
732
+ "https://www.picgo.net/api/1/upload",
733
+ files=files,
734
+ headers=headers
735
+ )
736
+
737
+ if response_url.status_code != 200:
738
+ return "生图失败,请查看PICGO图床密钥是否设置正确"
739
+ else:
740
+ logger.info("生图成功", "Server")
741
+ result = response_url.json()
742
+ return f"![image]({result['image']['url']})"
743
+
744
+
745
+ elif CONFIG["API"]["TUMY_KEY"]:
746
+ files = {'file': ('image.jpg', image_buffer, 'image/jpeg')}
747
+ headers = {
748
+ "Accept": "application/json",
749
+ 'Authorization': f"Bearer {CONFIG['API']['TUMY_KEY']}"
750
+ }
751
+
752
+ response_url = requests.post(
753
+ "https://tu.my/api/v1/upload",
754
+ files=files,
755
+ headers=headers
756
+ )
757
+
758
+ if response_url.status_code != 200:
759
+ return "生图失败,请查看TUMY图床密钥是否设置正确"
760
+ else:
761
+ try:
762
+ result = response_url.json()
763
+ logger.info("生图成功", "Server")
764
+ return f"![image]({result['data']['links']['url']})"
765
+ except Exception as error:
766
+ logger.error(str(error), "Server")
767
+ return "生图失败,请查看TUMY图床密钥是否设置正确"
768
+
769
+ def handle_non_stream_response(response, model):
770
  try:
771
+ logger.info("开始处理非流式响应", "Server")
772
+
773
+ stream = response.iter_lines()
774
+ full_response = ""
775
+
776
  CONFIG["IS_THINKING"] = False
777
  CONFIG["IS_IMG_GEN"] = False
778
  CONFIG["IS_IMG_GEN2"] = False
779
+
780
+ for chunk in stream:
781
+ if not chunk:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
782
  continue
 
783
  try:
784
+ line_json = json.loads(chunk.decode("utf-8").strip())
785
+ if line_json.get("error"):
786
+ logger.error(json.dumps(line_json, indent=2), "Server")
787
+ return json.dumps({"error": "RateLimitError"}) + "\n\n"
 
 
 
 
 
788
 
 
 
 
 
 
 
 
 
 
789
  response_data = line_json.get("result", {}).get("response")
790
  if not response_data:
791
  continue
792
+
793
  if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
794
  CONFIG["IS_IMG_GEN"] = True
795
+
796
+ result = process_model_response(response_data, model)
797
 
798
  if result["token"]:
799
+ full_response += result["token"]
800
+
801
  if result["imageUrl"]:
802
  CONFIG["IS_IMG_GEN2"] = True
803
+ return handle_image_response(result["imageUrl"])
804
+
805
+ except json.JSONDecodeError:
 
 
806
  continue
807
+ except Exception as e:
808
+ logger.error(f"处理流式响应行时出错: {str(e)}", "Server")
809
+ continue
810
+
811
+ return full_response
812
  except Exception as error:
813
+ logger.error(str(error), "Server")
814
+ raise
815
+ def handle_stream_response(response, model):
816
+ def generate():
817
+ logger.info("开始处理流式响应", "Server")
818
 
819
+ stream = response.iter_lines()
 
 
820
  CONFIG["IS_THINKING"] = False
821
  CONFIG["IS_IMG_GEN"] = False
822
  CONFIG["IS_IMG_GEN2"] = False
823
+
824
+ for chunk in stream:
825
+ if not chunk:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
826
  continue
 
827
  try:
828
+ line_json = json.loads(chunk.decode("utf-8").strip())
829
+ if line_json.get("error"):
830
+ logger.error(json.dumps(line_json, indent=2), "Server")
831
+ yield json.dumps({"error": "RateLimitError"}) + "\n\n"
832
+ return
 
 
 
 
833
 
 
 
 
 
 
 
 
 
 
834
  response_data = line_json.get("result", {}).get("response")
835
  if not response_data:
836
  continue
837
+
838
  if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
839
  CONFIG["IS_IMG_GEN"] = True
840
+
841
+ result = process_model_response(response_data, model)
842
 
843
  if result["token"]:
844
+ yield f"data: {json.dumps(MessageProcessor.create_chat_response(result["token"], model, True))}\n\n"
845
+
846
  if result["imageUrl"]:
847
  CONFIG["IS_IMG_GEN2"] = True
848
+ image_data = handle_image_response(result["imageUrl"])
849
+ yield f"data: {json.dumps(MessageProcessor.create_chat_response(image_data, model, True))}\n\n"
 
 
 
 
 
 
 
 
 
850
 
851
+ except json.JSONDecodeError:
852
+ continue
853
+ except Exception as e:
854
+ logger.error(f"处理流式响应行时出错: {str(e)}", "Server")
855
+ continue
856
+
857
+ yield "data: [DONE]\n\n"
858
+ return generate()
859
 
860
+ def initialization():
861
+ sso_array = os.environ.get("SSO", "").split(',')
862
+ logger.info("开始加载令牌", "Server")
863
+ for sso in sso_array:
864
+ if sso:
865
+ token_manager.add_token(f"sso-rw={sso};sso={sso}")
866
 
867
+ logger.info(f"成功加载令牌: {json.dumps(token_manager.get_all_tokens(), indent=2)}", "Server")
868
+ logger.info(f"令牌加载完成,共加载: {len(token_manager.get_all_tokens())}个令牌", "Server")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
869
 
870
+ if CONFIG["API"]["PROXY"]:
871
+ logger.info(f"代理已设置: {CONFIG['API']['PROXY']}", "Server")
872
 
873
+ logger.info("初始化完成", "Server")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
874
 
875
 
876
+ app = Flask(__name__)
877
+ app.wsgi_app = ProxyFix(app.wsgi_app)
878
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
879
 
880
+ @app.before_request
881
+ def log_request_info():
882
+ logger.info(f"{request.method} {request.path}", "Request")
883
 
884
  @app.route('/get/tokens', methods=['GET'])
885
+ def get_tokens():
886
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
 
887
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
888
  return jsonify({"error": '自定义的SSO令牌模式无法获取轮询sso令牌状态'}), 403
889
  elif auth_token != CONFIG["API"]["API_KEY"]:
890
  return jsonify({"error": 'Unauthorized'}), 401
891
+
892
  return jsonify(token_manager.get_token_status_map())
893
 
894
  @app.route('/add/token', methods=['POST'])
895
+ def add_token():
896
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
 
897
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
898
  return jsonify({"error": '自定义的SSO令牌模式无法添加sso令牌'}), 403
899
  elif auth_token != CONFIG["API"]["API_KEY"]:
900
  return jsonify({"error": 'Unauthorized'}), 401
901
+
902
  try:
903
+ sso = request.json.get('sso')
904
+ token_manager.add_token(f"sso-rw={sso};sso={sso}")
905
+ return jsonify(token_manager.get_token_status_map().get(sso, {})), 200
 
 
 
 
906
  except Exception as error:
907
+ logger.error(str(error), "Server")
908
  return jsonify({"error": '添加sso令牌失败'}), 500
909
 
910
  @app.route('/delete/token', methods=['POST'])
911
+ def delete_token():
912
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
 
913
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
914
  return jsonify({"error": '自定义的SSO���牌模式无法删除sso令牌'}), 403
915
  elif auth_token != CONFIG["API"]["API_KEY"]:
916
  return jsonify({"error": 'Unauthorized'}), 401
917
+
918
  try:
919
+ sso = request.json.get('sso')
920
+ token_manager.delete_token(f"sso-rw={sso};sso={sso}")
921
+ return jsonify({"message": '删除sso令牌成功'}), 200
 
 
 
 
 
 
 
922
  except Exception as error:
923
+ logger.error(str(error), "Server")
924
  return jsonify({"error": '删除sso令牌失败'}), 500
925
 
926
+ @app.route('/v1/models', methods=['GET'])
927
+ def get_models():
928
+ return jsonify({
929
+ "object": "list",
930
+ "data": [
931
+ {
932
+ "id": model,
933
+ "object": "model",
934
+ "created": int(time.time()),
935
+ "owned_by": "grok"
936
+ }
937
+ for model in CONFIG["MODELS"].keys()
938
+ ]
939
+ })
940
+
941
  @app.route('/v1/chat/completions', methods=['POST'])
942
+ def chat_completions():
943
  try:
 
944
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
 
945
  if auth_token:
946
  if CONFIG["API"]["IS_CUSTOM_SSO"]:
947
+ result = f"sso={auth_token};sso-rw={auth_token}"
948
+ token_manager.set_token(result)
949
  elif auth_token != CONFIG["API"]["API_KEY"]:
950
+ return jsonify({"error": 'Unauthorized'}), 401
951
  else:
952
+ return jsonify({"error": 'API_KEY缺失'}), 401
953
+
954
+ data = request.json
955
  model = data.get("model")
956
  stream = data.get("stream", False)
957
+
958
  retry_count = 0
959
+ grok_client = GrokApiClient(model)
960
+ request_payload = grok_client.prepare_chat_request(data)
961
 
962
+ while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:
963
+ retry_count += 1
964
+ CONFIG["API"]["SIGNATURE_COOKIE"] = Utils.create_auth_headers(model)
965
 
966
+ if not CONFIG["API"]["SIGNATURE_COOKIE"]:
967
+ raise ValueError('该模型无可用令牌')
968
+
969
+ logger.info(f"当前令牌: {json.dumps(CONFIG['API']['SIGNATURE_COOKIE'], indent=2)}", "Server")
970
+ logger.info(f"当前可用模型的全部可用数量: {json.dumps(token_manager.get_remaining_token_request_capacity(), indent=2)}", "Server")
971
+
972
+ try:
973
+ proxy_options = Utils.get_proxy_options()
974
+ response = curl_requests.post(
975
+ f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",
976
+ headers={
 
 
 
 
 
 
977
  "Accept": "text/event-stream",
978
+ "Baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
979
  "Content-Type": "text/plain;charset=UTF-8",
980
+ "Connection": "keep-alive",
981
+ "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
982
+ },
983
+ data=json.dumps(request_payload),
984
+ impersonate="chrome120",
985
+ stream=True,
986
+ **proxy_options
987
+ )
988
+
989
+ if response.status_code == 200:
990
+ logger.info("请求成功", "Server")
991
+ logger.info(f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}", "Server")
992
 
993
+ try:
 
 
994
  if stream:
995
  return Response(
996
+ stream_with_context(handle_stream_response(response, model)),
997
+ content_type='text/event-stream'
 
 
 
 
998
  )
999
  else:
1000
+ content = handle_non_stream_response(response, model)
1001
+ return jsonify(MessageProcessor.create_chat_response(content, model))
1002
+
1003
+ except Exception as error:
1004
+ logger.error(str(error), "Server")
1005
+ if CONFIG["API"]["IS_CUSTOM_SSO"]:
1006
+ raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")
1007
 
1008
+ token_manager.remove_token_from_model(model, CONFIG["API"]["SIGNATURE_COOKIE"])
1009
+ if token_manager.get_token_count_for_model(model) == 0:
1010
+ raise ValueError(f"{model} 次数已达上限,请切换其他模型或者重新对话")
1011
+
1012
+ elif response.status_code == 429:
1013
+ if CONFIG["API"]["IS_CUSTOM_SSO"]:
1014
+ raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")
1015
 
1016
+ token_manager.remove_token_from_model(model, CONFIG["API"]["SIGNATURE_COOKIE"])
1017
+ if token_manager.get_token_count_for_model(model) == 0:
1018
+ raise ValueError(f"{model} 次数已达上限,请切换其他模型或者重新对话")
1019
+
1020
+ else:
1021
+ if CONFIG["API"]["IS_CUSTOM_SSO"]:
1022
+ raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")
1023
+
1024
+ logger.error(f"令牌异常错误状态!status: {response.status_code}", "Server")
1025
+ token_manager.remove_token_from_model(model, CONFIG["API"]["SIGNATURE_COOKIE"])
1026
+ logger.info(f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}", "Server")
1027
 
1028
+ except Exception as e:
1029
+ logger.error(f"请求处理异常: {str(e)}", "Server")
1030
+ if CONFIG["API"]["IS_CUSTOM_SSO"]:
1031
+ raise
1032
+ continue
1033
+
1034
+ raise ValueError('当前模型所有令牌都已耗尽')
1035
+
1036
+ except Exception as error:
1037
+ logger.error(str(error), "ChatAPI")
1038
  return jsonify({
1039
  "error": {
1040
+ "message": str(error),
1041
  "type": "server_error"
1042
  }
1043
  }), 500
1044
 
1045
+ @app.route('/', defaults={'path': ''})
1046
+ @app.route('/<path:path>')
1047
+ def catch_all(path):
1048
+ return 'api运行正常', 200
1049
 
1050
+ if __name__ == '__main__':
1051
+ token_manager = AuthTokenManager()
1052
+ initialization()
 
 
 
 
 
 
 
 
1053
 
1054
+ app.run(
1055
+ host='0.0.0.0',
1056
+ port=CONFIG["SERVER"]["PORT"],
1057
+ debug=False
1058
+ )
1059
+