yxmiler commited on
Commit
cd3fccd
·
verified ·
1 Parent(s): 2040788

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +336 -743
app.py CHANGED
@@ -1,25 +1,49 @@
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',
@@ -35,16 +59,35 @@ CONFIG = {
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,
@@ -52,358 +95,27 @@ CONFIG = {
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:
61
- format = (
62
- "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
63
- "<level>{level: <8}</level> | "
64
- "<cyan>{extra[filename]}</cyan>:<cyan>{extra[function]}</cyan>:<cyan>{extra[lineno]}</cyan> | "
65
- "<level>{message}</level>"
66
- )
67
-
68
- logger.add(
69
- sys.stderr,
70
- level=level,
71
- format=format,
72
- colorize=colorize,
73
- backtrace=True,
74
- diagnose=True
75
- )
76
-
77
- self.logger = logger
78
-
79
- def _get_caller_info(self):
80
- frame = inspect.currentframe()
81
- try:
82
- caller_frame = frame.f_back.f_back
83
- full_path = caller_frame.f_code.co_filename
84
- function = caller_frame.f_code.co_name
85
- lineno = caller_frame.f_lineno
86
-
87
- filename = os.path.basename(full_path)
88
-
89
- return {
90
- 'filename': filename,
91
- 'function': function,
92
- 'lineno': lineno
93
- }
94
- finally:
95
- del frame
96
-
97
- def info(self, message, source="API"):
98
- caller_info = self._get_caller_info()
99
- self.logger.bind(**caller_info).info(f"[{source}] {message}")
100
-
101
- def error(self, message, source="API"):
102
- caller_info = self._get_caller_info()
103
-
104
- if isinstance(message, Exception):
105
- self.logger.bind(**caller_info).exception(f"[{source}] {str(message)}")
106
- else:
107
- self.logger.bind(**caller_info).error(f"[{source}] {message}")
108
-
109
- def warning(self, message, source="API"):
110
- caller_info = self._get_caller_info()
111
- self.logger.bind(**caller_info).warning(f"[{source}] {message}")
112
-
113
- def debug(self, message, source="API"):
114
- caller_info = self._get_caller_info()
115
- self.logger.bind(**caller_info).debug(f"[{source}] {message}")
116
-
117
- async def request_logger(self, request):
118
- caller_info = self._get_caller_info()
119
- self.logger.bind(**caller_info).info(f"请求: {request.method} {request.path}", "Request")
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
-
171
- if model not in self.token_status_map[sso]:
172
- self.token_status_map[sso][model] = {
173
- "isValid": True,
174
- "invalidatedTime": None,
175
- "totalRequestCount": 0
176
- }
177
- logger.info(f"添加令牌成功: {token}", "TokenManager")
178
-
179
- 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
 
@@ -411,171 +123,155 @@ class Utils:
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
-
453
  return {
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"]
468
 
 
469
  upload_data = {
470
  "rpc": "uploadFile",
471
  "req": {
472
- "fileName": file_name,
473
- "fileMimeType": mime_type,
474
  "content": image_buffer
475
  }
476
  }
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": "*/*",
489
- "Accept-Encoding": "gzip, deflate, br, zstd",
490
- "Content-Type": "text/plain;charset=UTF-8",
491
- "Cookie": token,
492
- "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c"
493
- }
494
 
495
- response = await Utils.run_in_executor(
496
- self.scraper.post,
497
- url,
498
- headers=headers,
499
- data=json.dumps(upload_data),
 
 
 
500
  )
501
-
502
  if response.status_code != 200:
503
- logger.error(f"上传图片失败,状态码:{response.status_code},原因:{response.text}", "Server")
504
  return ''
505
-
506
  result = response.json()
507
- logger.info(f'上传图片成功: {result}', "Server")
508
  return result["fileMetadataId"]
509
-
510
  except Exception as error:
511
- logger.error(error, "Server")
512
  return ''
513
-
514
  async def prepare_chat_request(self, request_data):
 
515
  todo_messages = request_data["messages"]
516
- if request_data["model"] in ["grok-2-imageGen", "grok-3-imageGen", "grok-3-deepsearch"]:
517
  last_message = todo_messages[-1]
518
- if last_message["role"] != "user":
519
- raise ValueError("画图模型的最后一条消息必须是用户消息!")
520
  todo_messages = [last_message]
521
-
522
  file_attachments = []
523
  messages = ''
524
  last_role = None
525
  last_content = ''
526
- search = request_data["model"] in ["grok-2-search", "grok-3-search"]
527
-
 
528
  def remove_think_tags(text):
529
  import re
530
  text = re.sub(r'<think>[\s\S]*?<\/think>', '', text).strip()
531
  text = re.sub(r'!\[image\]\(data:.*?base64,.*?\)', '[图片]', text)
532
  return text
533
-
534
  async def process_image_url(content):
535
- if content["type"] == "image_url" and "data:image" in content["image_url"]["url"]:
536
  image_response = await self.upload_base64_image(
537
  content["image_url"]["url"],
538
  f"{CONFIG['API']['BASE_URL']}/api/rpc"
539
  )
540
  return image_response
541
  return None
542
-
543
  async def process_content(content):
544
  if isinstance(content, list):
545
  text_content = ''
546
  for item in content:
547
- if item["type"] == "image_url":
548
- text_content += ("[图片]" if text_content else '') + "\n" if text_content else "[图片]"
549
- elif item["type"] == "text":
550
- text_content += ("\n" + remove_think_tags(item["text"]) if text_content else remove_think_tags(item["text"]))
551
  return text_content
552
  elif isinstance(content, dict) and content is not None:
553
- if content["type"] == "image_url":
554
  return "[图片]"
555
- elif content["type"] == "text":
556
  return remove_think_tags(content["text"])
557
  return remove_think_tags(self.process_message_content(content))
558
-
559
  for current in todo_messages:
560
- role = "assistant" if current["role"] == "assistant" else "user"
561
  is_last_message = current == todo_messages[-1]
562
-
563
- logger.info(json.dumps(current, indent=2, ensure_ascii=False), "Server")
564
  if is_last_message and "content" in current:
565
  if isinstance(current["content"], list):
566
  for item in current["content"]:
567
- if item["type"] == "image_url":
568
- logger.info("处理图片附件", "Server")
569
  processed_image = await process_image_url(item)
570
  if processed_image:
571
  file_attachments.append(processed_image)
572
- elif isinstance(current["content"], dict) and current["content"].get("type") == "image_url":
573
  processed_image = await process_image_url(current["content"])
574
  if processed_image:
575
  file_attachments.append(processed_image)
576
-
577
- text_content = await process_content(current["content"])
578
-
 
579
  if text_content or (is_last_message and file_attachments):
580
  if role == last_role and text_content:
581
  last_content += '\n' + text_content
@@ -583,7 +279,8 @@ class GrokApiClient:
583
  else:
584
  messages += f"{role.upper()}: {text_content or '[图片]'}\n"
585
  last_content = text_content
586
- last_role = role
 
587
  return {
588
  "temporary": CONFIG["API"]["IS_TEMP_CONVERSATION"],
589
  "modelName": self.model_id,
@@ -598,7 +295,7 @@ class GrokApiClient:
598
  "imageGenerationCount": 1,
599
  "forceConcise": False,
600
  "toolOverrides": {
601
- "imageGen": request_data["model"] in ["grok-2-imageGen", "grok-3-imageGen"],
602
  "webSearch": search,
603
  "xSearch": search,
604
  "xMediaSearch": search,
@@ -609,19 +306,20 @@ class GrokApiClient:
609
  "isPreset": False,
610
  "sendFinalMetadata": True,
611
  "customInstructions": "",
612
- "deepsearchPreset": "default" if request_data["model"] == "grok-3-deepsearch" else "",
613
- "isReasoning": request_data["model"] == "grok-3-reasoning"
614
  }
615
 
616
  class MessageProcessor:
617
  @staticmethod
618
  def create_chat_response(message, model, is_stream=False):
 
619
  base_response = {
620
  "id": f"chatcmpl-{str(uuid.uuid4())}",
621
- "created": int(datetime.now().timestamp()),
622
  "model": model
623
  }
624
-
625
  if is_stream:
626
  return {
627
  **base_response,
@@ -648,34 +346,36 @@ class MessageProcessor:
648
  "usage": None
649
  }
650
 
651
- async def process_model_response(response, model):
 
652
  result = {"token": None, "imageUrl": None}
653
 
654
  if CONFIG["IS_IMG_GEN"]:
655
  if response and response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:
656
  result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]
657
  return result
658
-
659
- if model == "grok-2":
 
660
  result["token"] = response.get("token")
661
- elif model in ["grok-2-search", "grok-3-search"]:
662
  if response and response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:
663
- result["token"] = f"\r\n<think>{await Utils.organize_search_results(response['webSearchResults'])}</think>\r\n"
664
  else:
665
  result["token"] = response.get("token")
666
- elif model == "grok-3":
667
  result["token"] = response.get("token")
668
- elif model == "grok-3-deepsearch":
669
  if response and response.get("messageTag") == "final":
670
- result["token"] = response.get("token")
671
- elif model == "grok-3-reasoning":
672
- if response and response.get("isThinking", False) and not CONFIG["SHOW_THINKING"]:
673
  return result
674
-
675
- if response and response.get("isThinking", False) and not CONFIG["IS_THINKING"]:
676
  result["token"] = "<think>" + response.get("token", "")
677
  CONFIG["IS_THINKING"] = True
678
- elif response and not response.get("isThinking", True) and CONFIG["IS_THINKING"]:
679
  result["token"] = "</think>" + response.get("token", "")
680
  CONFIG["IS_THINKING"] = False
681
  else:
@@ -683,96 +383,72 @@ async def process_model_response(response, model):
683
 
684
  return result
685
 
686
- async def stream_response_generator(response, model):
 
687
  try:
 
 
688
  CONFIG["IS_THINKING"] = False
689
  CONFIG["IS_IMG_GEN"] = False
690
  CONFIG["IS_IMG_GEN2"] = False
691
- logger.info("开始处理流式响应", "Server")
692
-
693
- async def iter_lines():
694
- line_iter = response.iter_lines()
695
- while True:
 
 
696
  try:
697
- line = await Utils.run_in_executor(lambda: next(line_iter, None))
698
- if line is None:
699
- break
700
- yield line
701
- except StopIteration:
702
- break
703
- except Exception as e:
704
- logger.error(f"迭代行时出错: {str(e)}", "Server")
705
- break
706
-
707
- async for line in iter_lines():
708
- if not line:
709
- continue
710
-
711
- try:
712
- line_str = line.decode('utf-8')
713
- line_json = json.loads(line_str)
714
-
715
- if line_json and line_json.get("error"):
716
- raise ValueError("RateLimitError")
717
 
718
- response_data = line_json.get("result", {}).get("response")
719
- if not response_data:
 
 
 
 
 
 
 
 
 
 
720
  continue
721
-
722
- if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
723
- CONFIG["IS_IMG_GEN"] = True
724
-
725
- result = await process_model_response(response_data, model)
726
-
727
- if result["token"]:
728
- yield f"data: {json.dumps(MessageProcessor.create_chat_response(result['token'], model, True))}\n\n"
729
-
730
- if result["imageUrl"]:
731
- CONFIG["IS_IMG_GEN2"] = True
732
- data_image = await handle_image_response(result["imageUrl"], model)
733
- yield f"data: {json.dumps(MessageProcessor.create_chat_response(data_image, model, True))}\n\n"
734
-
735
- except Exception as error:
736
- logger.error(error, "Server")
737
- continue
738
-
739
- yield "data: [DONE]\n\n"
740
-
741
  except Exception as error:
742
- logger.error(error, "Server")
743
  raise error
744
 
745
- async def handle_normal_response(response, model):
 
746
  try:
747
- full_response = ''
 
 
748
  CONFIG["IS_THINKING"] = False
749
  CONFIG["IS_IMG_GEN"] = False
750
  CONFIG["IS_IMG_GEN2"] = False
751
- logger.info("开始处理非流式响应", "Server")
752
- image_url = None
753
-
754
- async def iter_lines():
755
- line_iter = response.iter_lines()
756
- while True:
757
- try:
758
- line = await Utils.run_in_executor(lambda: next(line_iter, None))
759
- if line is None:
760
- break
761
- yield line
762
- except StopIteration:
763
- break
764
- except Exception as e:
765
- logger.error(f"迭代行时出错: {str(e)}", "Server")
766
- break
767
-
768
- async for line in iter_lines():
769
  if not line:
770
  continue
771
 
772
  try:
773
- line_str = line.decode('utf-8')
774
- line_json = json.loads(line_str)
775
-
776
  if line_json and line_json.get("error"):
777
  raise ValueError("RateLimitError")
778
 
@@ -783,49 +459,43 @@ async def handle_normal_response(response, model):
783
  if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
784
  CONFIG["IS_IMG_GEN"] = True
785
 
786
- result = await process_model_response(response_data, model)
787
 
788
  if result["token"]:
789
  full_response += result["token"]
790
 
791
  if result["imageUrl"]:
792
  CONFIG["IS_IMG_GEN2"] = True
793
- image_url = result["imageUrl"]
 
794
 
795
  except Exception as error:
796
- logger.error(error, "Server")
797
  continue
798
-
799
- if CONFIG["IS_IMG_GEN2"] and image_url:
800
- data_image = await handle_image_response(image_url, model)
801
- return MessageProcessor.create_chat_response(data_image, model)
802
- else:
803
  return MessageProcessor.create_chat_response(full_response, model)
804
-
805
  except Exception as error:
806
- logger.error(error, "Server")
807
  raise error
808
 
809
- async def handle_image_response(image_url,model):
 
810
  MAX_RETRIES = 2
811
  retry_count = 0
812
- scraper = cloudscraper.create_scraper()
813
 
814
  while retry_count < MAX_RETRIES:
815
  try:
816
- token = token_manager.get_next_token_for_model(model)
817
- if not token:
818
- raise ValueError("没有可用的token")
819
-
820
- image_response = await Utils.run_in_executor(
821
- scraper.get,
822
  f"https://assets.grok.com/{image_url}",
823
  headers={
824
  **CONFIG["DEFAULT_HEADERS"],
825
- "cookie": token
826
  }
827
  )
828
-
829
  if image_response.status_code == 200:
830
  break
831
 
@@ -833,244 +503,167 @@ async def handle_image_response(image_url,model):
833
  if retry_count == MAX_RETRIES:
834
  raise ValueError(f"上游服务请求失败! status: {image_response.status_code}")
835
 
836
- await asyncio.sleep(1 * retry_count)
837
-
838
  except Exception as error:
839
- logger.error(error, "Server")
840
  retry_count += 1
841
  if retry_count == MAX_RETRIES:
842
  raise error
843
 
844
- await asyncio.sleep(1 * retry_count)
845
-
846
- image_content = image_response.content
 
847
 
848
  if CONFIG["API"]["PICGO_KEY"]:
849
- form = aiohttp.FormData()
850
- form.add_field('source',
851
- io.BytesIO(image_content),
852
- filename=f'image-{int(datetime.now().timestamp())}.jpg',
853
- content_type='image/jpeg')
854
 
855
- async with aiohttp.ClientSession() as session:
856
- async with session.post(
857
- "https://www.picgo.net/api/1/upload",
858
- data=form,
859
- headers={"X-API-Key": CONFIG["API"]["PICGO_KEY"]}
860
- ) as response_url:
861
- if response_url.status != 200:
862
- return "生图失败,请查看PICGO图床密钥是否设置正确"
863
- else:
864
- logger.info("生图成功", "Server")
865
- result = await response_url.json()
866
- return f"![image]({result['image']['url']})"
867
- elif CONFIG["API"]["TUMY_KEY"]:
868
- form = aiohttp.FormData()
869
- form.add_field('file',
870
- io.BytesIO(image_content),
871
- filename=f'image-{int(datetime.now().timestamp())}.jpg',
872
- content_type='image/jpeg')
873
 
874
- async with aiohttp.ClientSession() as session:
875
- async with session.post(
876
- "https://tu.my/api/v1/upload",
877
- data=form,
878
- headers={
879
- "Accept": "application/json",
880
- "Authorization": f"Bearer {CONFIG['API']['TUMY_KEY']}"
881
- }
882
- ) as response_url:
883
- if response_url.status != 200:
884
- return "生图失败,请查看TUMY图床密钥是否设置正确"
885
- else:
886
- logger.info("生图成功", "Server")
887
- result = await response_url.json()
888
- return f"![image]({result['image']['url']})"
889
- # 如果没有PICGO_KEY或者TUMY_KEY则返回base64图片
890
- image_base64 = base64.b64encode(image_content).decode('utf-8')
891
- return f"![image](data:image/jpeg;base64,{image_base64})"
892
-
 
 
 
 
893
 
894
- app = Quart(__name__)
895
- app = cors(app, allow_origin="*", allow_methods=["GET", "POST", "OPTIONS"], allow_headers=["Content-Type", "Authorization"])
 
896
 
897
  @app.before_request
898
- async def before_request():
899
- await logger.request_logger(request)
 
900
 
901
  @app.route('/v1/models', methods=['GET'])
902
- async def models():
 
903
  return jsonify({
904
  "object": "list",
905
  "data": [
906
  {
907
  "id": model,
908
  "object": "model",
909
- "created": int(datetime.now().timestamp()),
910
  "owned_by": "grok"
911
  } for model in CONFIG["MODELS"].keys()
912
  ]
913
  })
914
 
915
-
916
- @app.route('/get/tokens', methods=['GET'])
917
- async def get_tokens():
918
- auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
919
-
920
- if CONFIG["API"]["IS_CUSTOM_SSO"]:
921
- return jsonify({"error": '自定义的SSO令牌模式无法获取轮询sso令牌状态'}), 403
922
- elif auth_token != CONFIG["API"]["API_KEY"]:
923
- return jsonify({"error": 'Unauthorized'}), 401
924
-
925
- return jsonify(token_manager.get_token_status_map())
926
-
927
- @app.route('/add/token', methods=['POST'])
928
- async def add_token():
929
- auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
930
-
931
- if CONFIG["API"]["IS_CUSTOM_SSO"]:
932
- return jsonify({"error": '自定义的SSO令牌模式无法添加sso令牌'}), 403
933
- elif auth_token != CONFIG["API"]["API_KEY"]:
934
- return jsonify({"error": 'Unauthorized'}), 401
935
-
936
- try:
937
- data = await request.get_json()
938
- sso = data.get('sso')
939
- if not sso:
940
- return jsonify({"error": 'SSO令牌不能为空'}), 400
941
-
942
- await token_manager.add_token(f"sso-rw={sso};sso={sso}")
943
- return jsonify(token_manager.get_token_status_map().get(sso, {}))
944
- except Exception as error:
945
- logger.error(error, "Server")
946
- return jsonify({"error": '添加sso令牌失败'}), 500
947
-
948
- @app.route('/delete/token', methods=['POST'])
949
- async def delete_token():
950
- auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
951
-
952
- if CONFIG["API"]["IS_CUSTOM_SSO"]:
953
- return jsonify({"error": '自定义的SSO令牌模式无法删除sso令牌'}), 403
954
- elif auth_token != CONFIG["API"]["API_KEY"]:
955
- return jsonify({"error": 'Unauthorized'}), 401
956
-
957
- try:
958
- data = await request.get_json()
959
- sso = data.get('sso')
960
- if not sso:
961
- return jsonify({"error": 'SSO令牌不能为空'}), 400
962
-
963
- success = await token_manager.delete_token(f"sso-rw={sso};sso={sso}")
964
- if success:
965
- return jsonify({"message": '删除sso令牌成功'})
966
- else:
967
- return jsonify({"error": '删除sso令牌失败'}), 500
968
- except Exception as error:
969
- logger.error(error, "Server")
970
- return jsonify({"error": '删除sso令牌失败'}), 500
971
-
972
  @app.route('/v1/chat/completions', methods=['POST'])
973
  async def chat_completions():
 
974
  try:
975
- data = await request.get_json()
976
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
977
-
978
- if auth_token:
979
- if CONFIG["API"]["IS_CUSTOM_SSO"]:
980
- await token_manager.set_token(f"sso-rw={auth_token};sso={auth_token}")
981
- elif auth_token != CONFIG["API"]["API_KEY"]:
982
- return jsonify({"error": "Unauthorized"}), 401
983
- else:
984
  return jsonify({"error": "Unauthorized"}), 401
985
-
 
 
986
  model = data.get("model")
987
  stream = data.get("stream", False)
988
- retry_count = 0
989
 
990
- try:
991
- grok_client = GrokApiClient(model)
992
- request_payload = await grok_client.prepare_chat_request(data)
 
 
 
 
 
 
 
993
 
994
- while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:
995
- retry_count += 1
996
- logger.info(f"开始请求(第{retry_count}次尝试)", "Server")
997
-
998
- token = token_manager.get_next_token_for_model(model)
999
- if not token:
1000
- logger.error(f"没有可用的{model}模型令牌", "Server")
1001
- if retry_count == CONFIG["RETRY"]["MAX_ATTEMPTS"]:
1002
- raise ValueError(f"没有可用的{model}模型令牌")
1003
- continue
1004
-
1005
- scraper = cloudscraper.create_scraper()
 
 
 
 
 
 
 
 
1006
 
1007
- try:
1008
- headers = {
1009
- "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",
1010
- "Connection": "keep-alive",
1011
- "Accept": "*/*",
1012
- "Accept-Encoding": "gzip, deflate, br, zstd",
1013
- "Content-Type": "text/plain;charset=UTF-8",
1014
- "Cookie": token,
1015
- "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c"
1016
- }
1017
- logger.info(f"使用令牌: {token}", "Server")
1018
-
1019
- response = await Utils.run_in_executor(
1020
- scraper.post,
1021
- f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",
1022
- headers=headers,
1023
- data=json.dumps(request_payload),
1024
- stream=True
1025
- )
1026
-
1027
- if response.status_code == 200:
1028
- logger.info("请求成功", "Server")
1029
-
1030
- if stream:
1031
- return Response(
1032
- stream_response_generator(response, model),
1033
- content_type='text/event-stream',
1034
- headers={
1035
- 'Cache-Control': 'no-cache',
1036
- 'Connection': 'keep-alive'
1037
- }
1038
- )
1039
- else:
1040
- result = await handle_normal_response(response, model)
1041
- return jsonify(result)
1042
- else:
1043
- logger.error(f"请求失败: 状态码 {response.status_code}", "Server")
1044
- token_manager.remove_token_from_model(model, token)
1045
-
1046
- except Exception as e:
1047
- logger.error(f"请求异常: {str(e)}", "Server")
1048
- token_manager.remove_token_from_model(model, token)
1049
 
1050
- raise ValueError("请求失败,已达到最大重试次数")
 
1051
 
1052
- except Exception as e:
1053
- logger.error(e, "ChatAPI")
1054
- return jsonify({
1055
- "error": {
1056
- "message": str(e),
1057
- "type": "server_error"
1058
- }
1059
- }), 500
1060
 
1061
- except Exception as e:
1062
- logger.error(e, "ChatAPI")
1063
  return jsonify({
1064
  "error": {
1065
- "message": str(e),
1066
  "type": "server_error"
1067
  }
1068
  }), 500
1069
 
1070
- @app.route('/')
1071
- async def index():
1072
- return "api运行正常"
 
 
1073
 
 
1074
  if __name__ == "__main__":
1075
- asyncio.run(initialize_tokens())
1076
- app.run(host="0.0.0.0", port=5200)
 
 
 
 
 
 
 
1
  import os
2
  import json
3
  import uuid
4
+ import logging
 
 
 
 
 
5
  import time
 
 
 
 
6
 
7
+ from flask import Flask, request, Response, jsonify, stream_with_context
8
+ from flask_cors import CORS
 
9
  from dotenv import load_dotenv
10
+ from curl_cffi import requests
11
+
12
+ # 配置日志
13
+ class CustomLogger:
14
+ def __init__(self):
15
+ self.logger = logging.getLogger("grok_api")
16
+ self.logger.setLevel(logging.INFO)
17
+
18
+ # 创建控制台处理器
19
+ console_handler = logging.StreamHandler()
20
+ console_handler.setLevel(logging.INFO)
21
+
22
+ # 设置日志格式
23
+ formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
24
+ console_handler.setFormatter(formatter)
25
+
26
+ # 添加处理器到日志器
27
+ self.logger.addHandler(console_handler)
28
+
29
+ def info(self, message, component="App"):
30
+ self.logger.info(f"[{component}] {message}")
31
+
32
+ def error(self, message, component="App"):
33
+ self.logger.error(f"[{component}] {message}")
34
+
35
+ def request_logger(self):
36
+ def middleware():
37
+ # 在请求处理之前记录
38
+ self.info(f"{request.method} {request.path}", "Request")
39
+ # 继续处理请求
40
+ return None
41
+ return middleware
42
 
43
+ # 加载环境变量
44
  load_dotenv()
45
 
46
+ # 全局配置
47
  CONFIG = {
48
  "MODELS": {
49
  'grok-2': 'grok-latest',
 
59
  "BASE_URL": "https://grok.com",
60
  "API_KEY": os.getenv("API_KEY", "sk-123456"),
61
  "IS_TEMP_CONVERSATION": os.getenv("IS_TEMP_CONVERSATION", "false").lower() == "true",
62
+ "PICGO_KEY": os.getenv("PICGO_KEY", None),
63
+ "SIGNATURE_COOKIE": ""
 
64
  },
65
  "SERVER": {
66
+ "PORT": int(os.getenv("PORT", 3000)),
67
+ "BODY_LIMIT": "5mb"
68
  },
69
  "RETRY": {
70
+ "MAX_ATTEMPTS": 2 # 重试次数
71
+ },
72
+ "DEFAULT_HEADERS": {
73
+ 'Accept': '*/*',
74
+ 'Accept-Language': 'zh-CN,zh;q=0.9',
75
+ 'Accept-Encoding': 'gzip, deflate, br, zstd',
76
+ 'Content-Type': 'text/plain;charset=UTF-8',
77
+ 'Connection': 'keep-alive',
78
+ 'Origin': 'https://grok.com',
79
+ 'Priority': 'u=1, i',
80
+ 'Sec-Ch-Ua': '"Chromium";v="130", "Google Chrome";v="130", "Not?A_Brand";v="99"',
81
+ 'Sec-Ch-Ua-Mobile': '?0',
82
+ 'Sec-Ch-Ua-Platform': '"Windows"',
83
+ 'Sec-Fetch-Dest': 'empty',
84
+ 'Sec-Fetch-Mode': 'cors',
85
+ 'Sec-Fetch-Site': 'same-origin',
86
+ 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36',
87
+ 'Baggage': 'sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c'
88
  },
89
+ "SIGNATUREARRAY": [],
90
+ "SIGNATUREINDEX": 0,
91
  "SHOW_THINKING": os.getenv("SHOW_THINKING", "false").lower() == "true",
92
  "IS_THINKING": False,
93
  "IS_IMG_GEN": False,
 
95
  "ISSHOW_SEARCH_RESULTS": os.getenv("ISSHOW_SEARCH_RESULTS", "true").lower() == "true"
96
  }
97
 
98
+ # 初始化Logger
99
+ Logger = CustomLogger()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
100
 
101
+ async def initialization():
102
+ """初始化函数,加载令牌"""
103
+ sso_array = os.getenv("SSO", "").split(",")
104
+ Logger.info("开始加载令牌", "Server")
 
105
 
106
  for sso in sso_array:
107
+ if sso.strip(): # 确保不添加空值
108
+ CONFIG["SIGNATUREARRAY"].append(f"sso-rw={sso};sso={sso}")
109
 
110
+ Logger.info(f"成功加载令牌: {json.dumps(CONFIG['SIGNATUREARRAY'], indent=2)}", "Server")
111
+ Logger.info(f"令牌加载完成,共加载: {len(CONFIG['SIGNATUREARRAY'])}个令牌", "Server")
112
+ Logger.info("初始化完成", "Server")
113
 
114
  class Utils:
115
  @staticmethod
116
+ def organize_search_results(search_results):
117
+ """格式化搜索结果"""
118
+ # 确保传入的是有效的搜索结果对象
119
  if not search_results or "results" not in search_results:
120
  return ''
121
 
 
123
  formatted_results = []
124
 
125
  for index, result in enumerate(results):
126
+ # 处理可能为空的字段
127
  title = result.get("title", "未知标题")
128
  url = result.get("url", "#")
129
  preview = result.get("preview", "无预览内容")
130
+
131
  formatted_result = f"\r\n<details><summary>资料[{index}]: {title}</summary>\r\n{preview}\r\n\n[Link]({url})\r\n</details>"
132
  formatted_results.append(formatted_result)
133
 
134
  return '\n\n'.join(formatted_results)
 
 
 
 
 
 
135
 
136
  class GrokApiClient:
137
  def __init__(self, model_id):
138
  if model_id not in CONFIG["MODELS"]:
139
  raise ValueError(f"不支持的模型: {model_id}")
 
140
  self.model_id = CONFIG["MODELS"][model_id]
141
+
 
142
  def process_message_content(self, content):
143
+ """处理消息内容"""
144
  if isinstance(content, str):
145
  return content
146
  return None
147
+
148
  def get_image_type(self, base64_string):
149
+ """获取图片类型"""
150
  mime_type = 'image/jpeg'
151
  if 'data:image' in base64_string:
152
  import re
153
+ matches = re.search(r'data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,', base64_string)
154
  if matches:
155
  mime_type = matches.group(1)
156
+
157
  extension = mime_type.split('/')[1]
158
  file_name = f"image.{extension}"
159
+
160
  return {
161
  "mimeType": mime_type,
162
  "fileName": file_name
163
  }
164
+
165
  async def upload_base64_image(self, base64_data, url):
166
+ """上传Base64图片"""
167
  try:
168
+ # 处理 base64 数据
169
  if 'data:image' in base64_data:
170
  image_buffer = base64_data.split(',')[1]
171
  else:
172
  image_buffer = base64_data
 
 
 
 
173
 
174
+ image_info = self.get_image_type(base64_data)
175
  upload_data = {
176
  "rpc": "uploadFile",
177
  "req": {
178
+ "fileName": image_info["fileName"],
179
+ "fileMimeType": image_info["mimeType"],
180
  "content": image_buffer
181
  }
182
  }
183
 
184
+ Logger.info("发送图片请求", "Server")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
185
 
186
+ # 使用curl_cffi发送请求
187
+ response = requests.post(
188
+ url,
189
+ headers={
190
+ **CONFIG["DEFAULT_HEADERS"],
191
+ "cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
192
+ },
193
+ json=upload_data
194
  )
195
+
196
  if response.status_code != 200:
197
+ Logger.error(f"上传图片失败,状态码:{response.status_code}", "Server")
198
  return ''
199
+
200
  result = response.json()
201
+ Logger.info(f"上传图片成功: {result}", "Server")
202
  return result["fileMetadataId"]
203
+
204
  except Exception as error:
205
+ Logger.error(str(error), "Server")
206
  return ''
207
+
208
  async def prepare_chat_request(self, request_data):
209
+ """准备聊天请求"""
210
  todo_messages = request_data["messages"]
211
+ if request_data["model"] in ['grok-2-imageGen', 'grok-3-imageGen']:
212
  last_message = todo_messages[-1]
213
+ if last_message["role"] != 'user':
214
+ raise ValueError('画图模型的最后一条消息必须是用户消息!')
215
  todo_messages = [last_message]
216
+
217
  file_attachments = []
218
  messages = ''
219
  last_role = None
220
  last_content = ''
221
+ search = request_data["model"] in ['grok-2-search', 'grok-3-search']
222
+
223
+ # 移除<think>标签及其内容和base64图片
224
  def remove_think_tags(text):
225
  import re
226
  text = re.sub(r'<think>[\s\S]*?<\/think>', '', text).strip()
227
  text = re.sub(r'!\[image\]\(data:.*?base64,.*?\)', '[图片]', text)
228
  return text
229
+
230
  async def process_image_url(content):
231
+ if content["type"] == 'image_url' and 'data:image' in content["image_url"]["url"]:
232
  image_response = await self.upload_base64_image(
233
  content["image_url"]["url"],
234
  f"{CONFIG['API']['BASE_URL']}/api/rpc"
235
  )
236
  return image_response
237
  return None
238
+
239
  async def process_content(content):
240
  if isinstance(content, list):
241
  text_content = ''
242
  for item in content:
243
+ if item["type"] == 'image_url':
244
+ text_content += (text_content + '\n' if text_content else '') + "[图片]"
245
+ elif item["type"] == 'text':
246
+ text_content += (text_content + '\n' if text_content else '') + remove_think_tags(item["text"])
247
  return text_content
248
  elif isinstance(content, dict) and content is not None:
249
+ if content["type"] == 'image_url':
250
  return "[图片]"
251
+ elif content["type"] == 'text':
252
  return remove_think_tags(content["text"])
253
  return remove_think_tags(self.process_message_content(content))
254
+
255
  for current in todo_messages:
256
+ role = 'assistant' if current["role"] == 'assistant' else 'user'
257
  is_last_message = current == todo_messages[-1]
258
+
259
+ # 处理图片附件
260
  if is_last_message and "content" in current:
261
  if isinstance(current["content"], list):
262
  for item in current["content"]:
263
+ if item.get("type") == 'image_url':
 
264
  processed_image = await process_image_url(item)
265
  if processed_image:
266
  file_attachments.append(processed_image)
267
+ elif isinstance(current["content"], dict) and current["content"].get("type") == 'image_url':
268
  processed_image = await process_image_url(current["content"])
269
  if processed_image:
270
  file_attachments.append(processed_image)
271
+
272
+ # 处理文本内容
273
+ text_content = await process_content(current.get("content", ""))
274
+
275
  if text_content or (is_last_message and file_attachments):
276
  if role == last_role and text_content:
277
  last_content += '\n' + text_content
 
279
  else:
280
  messages += f"{role.upper()}: {text_content or '[图片]'}\n"
281
  last_content = text_content
282
+ last_role = role
283
+
284
  return {
285
  "temporary": CONFIG["API"]["IS_TEMP_CONVERSATION"],
286
  "modelName": self.model_id,
 
295
  "imageGenerationCount": 1,
296
  "forceConcise": False,
297
  "toolOverrides": {
298
+ "imageGen": request_data["model"] in ['grok-2-imageGen', 'grok-3-imageGen'],
299
  "webSearch": search,
300
  "xSearch": search,
301
  "xMediaSearch": search,
 
306
  "isPreset": False,
307
  "sendFinalMetadata": True,
308
  "customInstructions": "",
309
+ "deepsearchPreset": "default" if request_data["model"] == 'grok-3-deepsearch' else "",
310
+ "isReasoning": request_data["model"] == 'grok-3-reasoning'
311
  }
312
 
313
  class MessageProcessor:
314
  @staticmethod
315
  def create_chat_response(message, model, is_stream=False):
316
+ """创建聊天响应"""
317
  base_response = {
318
  "id": f"chatcmpl-{str(uuid.uuid4())}",
319
+ "created": int(time.time()),
320
  "model": model
321
  }
322
+
323
  if is_stream:
324
  return {
325
  **base_response,
 
346
  "usage": None
347
  }
348
 
349
+ def process_model_response(response, model):
350
+ """处理模型响应"""
351
  result = {"token": None, "imageUrl": None}
352
 
353
  if CONFIG["IS_IMG_GEN"]:
354
  if response and response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:
355
  result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]
356
  return result
357
+
358
+ # 非生图模型的处理
359
+ if model == 'grok-2':
360
  result["token"] = response.get("token")
361
+ elif model in ['grok-2-search', 'grok-3-search']:
362
  if response and response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:
363
+ result["token"] = f"\r\n<think>{Utils.organize_search_results(response['webSearchResults'])}</think>\r\n"
364
  else:
365
  result["token"] = response.get("token")
366
+ elif model == 'grok-3':
367
  result["token"] = response.get("token")
368
+ elif model == 'grok-3-deepsearch':
369
  if response and response.get("messageTag") == "final":
370
+ result["token"] = response.get("token")
371
+ elif model == 'grok-3-reasoning':
372
+ if response and response.get("isThinking") and not CONFIG["SHOW_THINKING"]:
373
  return result
374
+
375
+ if response and response.get("isThinking") and not CONFIG["IS_THINKING"]:
376
  result["token"] = "<think>" + response.get("token", "")
377
  CONFIG["IS_THINKING"] = True
378
+ elif response and not response.get("isThinking") and CONFIG["IS_THINKING"]:
379
  result["token"] = "</think>" + response.get("token", "")
380
  CONFIG["IS_THINKING"] = False
381
  else:
 
383
 
384
  return result
385
 
386
+ async def handle_stream_response(response, model, flask_response):
387
+ """处理流式响应"""
388
  try:
389
+ stream = response.iter_lines()
390
+
391
  CONFIG["IS_THINKING"] = False
392
  CONFIG["IS_IMG_GEN"] = False
393
  CONFIG["IS_IMG_GEN2"] = False
394
+ Logger.info("开始处理流式响应", "Server")
395
+
396
+ def generate():
397
+ for line in stream:
398
+ if not line:
399
+ continue
400
+ line_json = json.loads(line.decode("utf-8").strip())
401
  try:
402
+ if line_json and line_json.get("error"):
403
+ raise ValueError("RateLimitError")
404
+
405
+ response_data = line_json.get("result", {}).get("response")
406
+ if not response_data:
407
+ continue
408
+
409
+ if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
410
+ CONFIG["IS_IMG_GEN"] = True
411
+
412
+ result = process_model_response(response_data, model)
 
 
 
 
 
 
 
 
 
413
 
414
+ if result["token"]:
415
+ chat_response = MessageProcessor.create_chat_response(result["token"], model, True)
416
+ yield f"data: {json.dumps(chat_response)}\n\n"
417
+
418
+ if result["imageUrl"]:
419
+ CONFIG["IS_IMG_GEN2"] = True
420
+ data_image = handle_image_response(result["imageUrl"])
421
+ image_response = MessageProcessor.create_chat_response(data_image, model, True)
422
+ yield f"data: {json.dumps(image_response)}\n\n"
423
+
424
+ except Exception as error:
425
+ Logger.error(str(error), "Server")
426
  continue
427
+
428
+ yield "data: [DONE]\n\n"
429
+
430
+ return Response(generate(), mimetype="text/event-stream")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
431
  except Exception as error:
432
+ Logger.error(str(error), "Server")
433
  raise error
434
 
435
+ async def handle_non_stream_response(response, model):
436
+ """处理非流式响应"""
437
  try:
438
+ stream = response.iter_lines()
439
+ full_response = ""
440
+
441
  CONFIG["IS_THINKING"] = False
442
  CONFIG["IS_IMG_GEN"] = False
443
  CONFIG["IS_IMG_GEN2"] = False
444
+ Logger.info("开始处理非流式响应", "Server")
445
+
446
+ for line in stream:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
447
  if not line:
448
  continue
449
 
450
  try:
451
+ line_json = json.loads(line.decode("utf-8").strip())
 
 
452
  if line_json and line_json.get("error"):
453
  raise ValueError("RateLimitError")
454
 
 
459
  if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
460
  CONFIG["IS_IMG_GEN"] = True
461
 
462
+ result = process_model_response(response_data, model)
463
 
464
  if result["token"]:
465
  full_response += result["token"]
466
 
467
  if result["imageUrl"]:
468
  CONFIG["IS_IMG_GEN2"] = True
469
+ data_image = await handle_image_response(result["imageUrl"])
470
+ return MessageProcessor.create_chat_response(data_image, model)
471
 
472
  except Exception as error:
473
+ Logger.error(str(error), "Server")
474
  continue
475
+
476
+ if not CONFIG["IS_IMG_GEN2"]:
 
 
 
477
  return MessageProcessor.create_chat_response(full_response, model)
478
+ return None
479
  except Exception as error:
480
+ Logger.error(str(error), "Server")
481
  raise error
482
 
483
+ async def handle_image_response(image_url):
484
+ """处理图片响应"""
485
  MAX_RETRIES = 2
486
  retry_count = 0
 
487
 
488
  while retry_count < MAX_RETRIES:
489
  try:
490
+ # 使用curl_cffi获取图片
491
+ image_response = requests.get(
 
 
 
 
492
  f"https://assets.grok.com/{image_url}",
493
  headers={
494
  **CONFIG["DEFAULT_HEADERS"],
495
+ "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
496
  }
497
  )
498
+
499
  if image_response.status_code == 200:
500
  break
501
 
 
503
  if retry_count == MAX_RETRIES:
504
  raise ValueError(f"上游服务请求失败! status: {image_response.status_code}")
505
 
506
+ time.sleep(1 * retry_count) # 简单的退避策略
507
+
508
  except Exception as error:
509
+ Logger.error(str(error), "Server")
510
  retry_count += 1
511
  if retry_count == MAX_RETRIES:
512
  raise error
513
 
514
+ time.sleep(1 * retry_count)
515
+
516
+ # 图片处理逻辑
517
+ image_buffer = image_response.content
518
 
519
  if CONFIG["API"]["PICGO_KEY"]:
520
+ # 使用curl_cffi上传到PICGO
521
+ import io
522
+ from curl_cffi.requests import AsyncSession
 
 
523
 
524
+ files = {
525
+ "source": ("image.jpg", image_buffer, "image/jpeg")
526
+ }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
527
 
528
+ headers = {
529
+ "X-API-Key": CONFIG["API"]["PICGO_KEY"]
530
+ }
531
+
532
+ try:
533
+ pic_response = requests.post(
534
+ "https://www.picgo.net/api/1/upload",
535
+ headers=headers,
536
+ files=files
537
+ )
538
+
539
+ if pic_response.status_code != 200:
540
+ return "生图失败,请查看PICGO图床密钥是否设置正确"
541
+
542
+ Logger.info("生图成功", "Server")
543
+ result = pic_response.json()
544
+ return f"![image]({result['image']['url']})"
545
+ except Exception as e:
546
+ Logger.error(f"上传PICGO失败: {str(e)}", "Server")
547
+ return "生图上传失败,请检查网络连接和PICGO配置"
548
+
549
+ # 如果没有PICGO_KEY,可以返回一个适当的消息或其他逻辑
550
+ return "生图成功,但未配置PICGO图床,无法显示图片"
551
 
552
+ # 创建Flask应用
553
+ app = Flask(__name__)
554
+ CORS(app, resources={r"/*": {"origins": "*", "methods": ["GET", "POST", "OPTIONS"], "allow_headers": ["Content-Type", "Authorization"]}})
555
 
556
  @app.before_request
557
+ def log_request_info():
558
+ """请求日志中间件"""
559
+ Logger.info(f"{request.method} {request.path}", "Request")
560
 
561
  @app.route('/v1/models', methods=['GET'])
562
+ async def get_models():
563
+ """获取模型列表"""
564
  return jsonify({
565
  "object": "list",
566
  "data": [
567
  {
568
  "id": model,
569
  "object": "model",
570
+ "created": int(time.time()),
571
  "owned_by": "grok"
572
  } for model in CONFIG["MODELS"].keys()
573
  ]
574
  })
575
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
576
  @app.route('/v1/chat/completions', methods=['POST'])
577
  async def chat_completions():
578
+ """处理聊天完成请求"""
579
  try:
580
+ # 验证API密钥
581
  auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
582
+ if auth_token != CONFIG["API"]["API_KEY"]:
 
 
 
 
 
 
583
  return jsonify({"error": "Unauthorized"}), 401
584
+
585
+ # 解析请求
586
+ data = request.get_json()
587
  model = data.get("model")
588
  stream = data.get("stream", False)
 
589
 
590
+ # 创建Grok客户端
591
+ grok_client = GrokApiClient(model)
592
+ request_payload = await grok_client.prepare_chat_request(data)
593
+ Logger.info(json.dumps(request_payload, indent=2), "Server")
594
+
595
+ # 重试逻辑
596
+ retry_count = 0
597
+ while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:
598
+ retry_count += 1
599
+ Logger.info("开始请求", "Server")
600
 
601
+ # 设置cookie
602
+ CONFIG["API"]["SIGNATURE_COOKIE"] = CONFIG["SIGNATUREARRAY"][CONFIG["SIGNATUREINDEX"]]
603
+
604
+ # 发送请求
605
+ response = requests.post(
606
+ f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",
607
+ headers={
608
+ "Accept": "text/event-stream",
609
+ "Baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
610
+ "Content-Type": "text/plain;charset=UTF-8",
611
+ "Connection": "keep-alive",
612
+ "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
613
+ },
614
+ data=json.dumps(request_payload),
615
+ impersonate="chrome110",
616
+ stream=True
617
+ )
618
+
619
+ if response.status_code == 200:
620
+ Logger.info("请求成功", "Server")
621
 
622
+ if stream:
623
+ # 处理流式响应
624
+ stream_response = await handle_stream_response(response, model, Response())
625
+ CONFIG["SIGNATUREINDEX"] = (CONFIG["SIGNATUREINDEX"] + 1) % len(CONFIG["SIGNATUREARRAY"])
626
+ return stream_response
627
+ else:
628
+ # 处理非流式响应
629
+ non_stream_result = await handle_non_stream_response(response, model)
630
+ CONFIG["SIGNATUREINDEX"] = (CONFIG["SIGNATUREINDEX"] + 1) % len(CONFIG["SIGNATUREARRAY"])
631
+ if non_stream_result:
632
+ return jsonify(non_stream_result)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
633
 
634
+ # 轮换token
635
+ CONFIG["SIGNATUREINDEX"] = (CONFIG["SIGNATUREINDEX"] + 1) % len(CONFIG["SIGNATUREARRAY"])
636
 
637
+ # 如果所有重试都失败
638
+ return jsonify({
639
+ "error": {
640
+ "message": "请求失败,所有令牌均已尝试",
641
+ "type": "server_error"
642
+ }
643
+ }), 500
 
644
 
645
+ except Exception as error:
646
+ Logger.error(str(error), "ChatAPI")
647
  return jsonify({
648
  "error": {
649
+ "message": str(error),
650
  "type": "server_error"
651
  }
652
  }), 500
653
 
654
+ @app.route('/', defaults={'path': ''})
655
+ @app.route('/<path:path>')
656
+ def catch_all(path):
657
+ """处理所有其他路由"""
658
+ return "api运行正常", 200
659
 
660
+ # 主程序入口
661
  if __name__ == "__main__":
662
+ # 初始化应用
663
+ import asyncio
664
+ asyncio.run(initialization())
665
+
666
+ # 启动服务器
667
+ Logger.info(f"服务器已启动,监听端口: {CONFIG['SERVER']['PORT']}", "Server")
668
+ from waitress import serve
669
+ serve(app, host="0.0.0.0", port=CONFIG["SERVER"]["PORT"])