DoraeHope/GrokPythonTo2
0
1import os2import json3import uuid4import time5import base646import sys7import inspect8from loguru import logger9 10import requests11from flask import Flask, request, Response, jsonify, stream_with_context12from curl_cffi import requests as curl_requests13from werkzeug.middleware.proxy_fix import ProxyFix14 15 16class Logger:17 def __init__(self, level="INFO", colorize=True, format=None):18 logger.remove()19 20 if format is None:21 format = (22 "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "23 "<level>{level: <8}</level> | "24 "<cyan>{extra[filename]}</cyan>:<cyan>{extra[function]}</cyan>:<cyan>{extra[lineno]}</cyan> | "25 "<level>{message}</level>"26 )27 28 logger.add(29 sys.stderr,30 level=level,31 format=format,32 colorize=colorize,33 backtrace=True,34 diagnose=True35 )36 37 self.logger = logger38 39 def _get_caller_info(self):40 frame = inspect.currentframe()41 try:42 caller_frame = frame.f_back.f_back43 full_path = caller_frame.f_code.co_filename44 function = caller_frame.f_code.co_name45 lineno = caller_frame.f_lineno46 47 filename = os.path.basename(full_path)48 49 return {50 'filename': filename,51 'function': function,52 'lineno': lineno53 }54 finally:55 del frame56 57 def info(self, message, source="API"):58 caller_info = self._get_caller_info()59 self.logger.bind(**caller_info).info(f"[{source}] {message}")60 61 def error(self, message, source="API"):62 caller_info = self._get_caller_info()63 64 if isinstance(message, Exception):65 self.logger.bind(**caller_info).exception(f"[{source}] {str(message)}")66 else:67 self.logger.bind(**caller_info).error(f"[{source}] {message}")68 69 def warning(self, message, source="API"):70 caller_info = self._get_caller_info()71 self.logger.bind(**caller_info).warning(f"[{source}] {message}")72 73 def debug(self, message, source="API"):74 caller_info = self._get_caller_info()75 self.logger.bind(**caller_info).debug(f"[{source}] {message}")76 77 async def request_logger(self, request):78 caller_info = self._get_caller_info()79 self.logger.bind(**caller_info).info(f"请求: {request.method} {request.path}", "Request")80 81logger = Logger(level="INFO")82 83 84CONFIG = {85 "MODELS": {86 'grok-2': 'grok-latest',87 'grok-2-imageGen': 'grok-latest',88 'grok-2-search': 'grok-latest',89 "grok-3": "grok-3",90 "grok-3-search": "grok-3",91 "grok-3-imageGen": "grok-3",92 "grok-3-deepsearch": "grok-3",93 "grok-3-reasoning": "grok-3"94 },95 "API": {96 "IS_TEMP_CONVERSATION": os.environ.get("IS_TEMP_CONVERSATION", "true").lower() == "true",97 "IS_CUSTOM_SSO": os.environ.get("IS_CUSTOM_SSO", "false").lower() == "true",98 "BASE_URL": "https://grok.com",99 "API_KEY": os.environ.get("API_KEY", "sk-123456"),100 "SIGNATURE_COOKIE": None,101 "PICGO_KEY": os.environ.get("PICGO_KEY") or None,102 "TUMY_KEY": os.environ.get("TUMY_KEY") or None,103 "RETRY_TIME": 1000,104 "PROXY": os.environ.get("PROXY") or None105 },106 "SERVER": {107 "PORT": int(os.environ.get("PORT", 5200))108 },109 "RETRY": {110 "MAX_ATTEMPTS": 2111 },112 "SHOW_THINKING": os.environ.get("SHOW_THINKING") == "true",113 "IS_THINKING": False,114 "IS_IMG_GEN": False,115 "IS_IMG_GEN2": False,116 "ISSHOW_SEARCH_RESULTS": os.environ.get("ISSHOW_SEARCH_RESULTS", "true").lower() == "true"117}118 119 120DEFAULT_HEADERS = {121 'Accept': '*/*',122 'Accept-Language': 'zh-CN,zh;q=0.9',123 'Accept-Encoding': 'gzip, deflate, br, zstd',124 'Content-Type': 'text/plain;charset=UTF-8',125 'Connection': 'keep-alive',126 'Origin': 'https://grok.com',127 'Priority': 'u=1, i',128 'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36',129 'Sec-Ch-Ua': '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"',130 'Sec-Ch-Ua-Mobile': '?0',131 'Sec-Ch-Ua-Platform': '"macOS"',132 'Sec-Fetch-Dest': 'empty',133 'Sec-Fetch-Mode': 'cors',134 'Sec-Fetch-Site': 'same-origin',135 'Baggage': 'sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c'136}137 138class AuthTokenManager:139 def __init__(self):140 self.token_model_map = {}141 self.expired_tokens = set()142 self.token_status_map = {}143 144 self.model_config = {145 "grok-2": {146 "RequestFrequency": 30,147 "ExpirationTime": 1 * 60 * 60 * 1000 # 1小时148 },149 "grok-3": {150 "RequestFrequency": 20,151 "ExpirationTime": 2 * 60 * 60 * 1000 # 2小时152 },153 "grok-3-deepsearch": {154 "RequestFrequency": 10,155 "ExpirationTime": 24 * 60 * 60 * 1000 # 24小时156 },157 "grok-3-reasoning": {158 "RequestFrequency": 10,159 "ExpirationTime": 24 * 60 * 60 * 1000 # 24小时160 }161 }162 self.token_reset_switch = False163 self.token_reset_timer = None164 165 def add_token(self, token):166 sso = token.split("sso=")[1].split(";")[0]167 for model in self.model_config.keys():168 if model not in self.token_model_map:169 self.token_model_map[model] = []170 if sso not in self.token_status_map:171 self.token_status_map[sso] = {}172 173 existing_token_entry = next((entry for entry in self.token_model_map[model] if entry["token"] == token), None)174 175 if not existing_token_entry:176 self.token_model_map[model].append({177 "token": token,178 "RequestCount": 0,179 "AddedTime": int(time.time() * 1000),180 "StartCallTime": None181 })182 183 if model not in self.token_status_map[sso]:184 self.token_status_map[sso][model] = {185 "isValid": True,186 "invalidatedTime": None,187 "totalRequestCount": 0188 }189 190 def set_token(self, token):191 models = list(self.model_config.keys())192 self.token_model_map = {model: [{193 "token": token,194 "RequestCount": 0,195 "AddedTime": int(time.time() * 1000),196 "StartCallTime": None197 }] for model in models}198 199 sso = token.split("sso=")[1].split(";")[0]200 self.token_status_map[sso] = {model: {201 "isValid": True,202 "invalidatedTime": None,203 "totalRequestCount": 0204 } for model in models}205 206 def delete_token(self, token):207 try:208 sso = token.split("sso=")[1].split(";")[0]209 for model in self.token_model_map:210 self.token_model_map[model] = [entry for entry in self.token_model_map[model] if entry["token"] != token]211 212 if sso in self.token_status_map:213 del self.token_status_map[sso]214 215 logger.info(f"令牌已成功移除: {token}", "TokenManager")216 return True217 except Exception as error:218 logger.error(f"令牌删除失败: {str(error)}")219 return False220 221 def get_next_token_for_model(self, model_id):222 normalized_model = self.normalize_model_name(model_id)223 224 if normalized_model not in self.token_model_map or not self.token_model_map[normalized_model]:225 return None226 227 token_entry = self.token_model_map[normalized_model][0]228 229 if token_entry:230 if token_entry["StartCallTime"] is None:231 token_entry["StartCallTime"] = int(time.time() * 1000)232 233 if not self.token_reset_switch:234 self.start_token_reset_process()235 self.token_reset_switch = True236 237 token_entry["RequestCount"] += 1238 239 if token_entry["RequestCount"] > self.model_config[normalized_model]["RequestFrequency"]:240 self.remove_token_from_model(normalized_model, token_entry["token"])241 next_token_entry = self.token_model_map[normalized_model][0] if self.token_model_map[normalized_model] else None242 return next_token_entry["token"] if next_token_entry else None243 244 sso = token_entry["token"].split("sso=")[1].split(";")[0]245 if sso in self.token_status_map and normalized_model in self.token_status_map[sso]:246 if token_entry["RequestCount"] == self.model_config[normalized_model]["RequestFrequency"]:247 self.token_status_map[sso][normalized_model]["isValid"] = False248 self.token_status_map[sso][normalized_model]["invalidatedTime"] = int(time.time() * 1000)249 self.token_status_map[sso][normalized_model]["totalRequestCount"] += 1250 251 return token_entry["token"]252 253 return None254 255 def remove_token_from_model(self, model_id, token):256 normalized_model = self.normalize_model_name(model_id)257 258 if normalized_model not in self.token_model_map:259 logger.error(f"模型 {normalized_model} 不存在", "TokenManager")260 return False261 262 model_tokens = self.token_model_map[normalized_model]263 token_index = next((i for i, entry in enumerate(model_tokens) if entry["token"] == token), -1)264 265 if token_index != -1:266 removed_token_entry = model_tokens.pop(token_index)267 self.expired_tokens.add((268 removed_token_entry["token"],269 normalized_model,270 int(time.time() * 1000)271 ))272 273 if not self.token_reset_switch:274 self.start_token_reset_process()275 self.token_reset_switch = True276 277 logger.info(f"模型{model_id}的令牌已失效，已成功移除令牌: {token}", "TokenManager")278 return True279 280 logger.error(f"在模型 {normalized_model} 中未找到 token: {token}", "TokenManager")281 return False282 283 def get_expired_tokens(self):284 return list(self.expired_tokens)285 286 def normalize_model_name(self, model):287 if model.startswith('grok-') and 'deepsearch' not in model and 'reasoning' not in model:288 return '-'.join(model.split('-')[:2])289 return model290 291 def get_token_count_for_model(self, model_id):292 normalized_model = self.normalize_model_name(model_id)293 return len(self.token_model_map.get(normalized_model, []))294 295 def get_remaining_token_request_capacity(self):296 remaining_capacity_map = {}297 298 for model in self.model_config.keys():299 model_tokens = self.token_model_map.get(model, [])300 model_request_frequency = self.model_config[model]["RequestFrequency"]301 302 total_used_requests = sum(token_entry.get("RequestCount", 0) for token_entry in model_tokens)303 304 remaining_capacity = (len(model_tokens) * model_request_frequency) - total_used_requests305 remaining_capacity_map[model] = max(0, remaining_capacity)306 307 return remaining_capacity_map308 309 def get_token_array_for_model(self, model_id):310 normalized_model = self.normalize_model_name(model_id)311 return self.token_model_map.get(normalized_model, [])312 313 def start_token_reset_process(self):314 def reset_expired_tokens():315 now = int(time.time() * 1000)316 317 tokens_to_remove = set()318 for token_info in self.expired_tokens:319 token, model, expired_time = token_info320 expiration_time = self.model_config[model]["ExpirationTime"]321 322 if now - expired_time >= expiration_time:323 if not any(entry["token"] == token for entry in self.token_model_map.get(model, [])):324 if model not in self.token_model_map:325 self.token_model_map[model] = []326 327 self.token_model_map[model].append({328 "token": token,329 "RequestCount": 0,330 "AddedTime": now,331 "StartCallTime": None332 })333 334 sso = token.split("sso=")[1].split(";")[0]335 if sso in self.token_status_map and model in self.token_status_map[sso]:336 self.token_status_map[sso][model]["isValid"] = True337 self.token_status_map[sso][model]["invalidatedTime"] = None338 self.token_status_map[sso][model]["totalRequestCount"] = 0339 340 tokens_to_remove.add(token_info)341 342 self.expired_tokens -= tokens_to_remove343 344 for model in self.model_config.keys():345 if model not in self.token_model_map:346 continue347 348 for token_entry in self.token_model_map[model]:349 if not token_entry.get("StartCallTime"):350 continue351 352 expiration_time = self.model_config[model]["ExpirationTime"]353 if now - token_entry["StartCallTime"] >= expiration_time:354 sso = token_entry["token"].split("sso=")[1].split(";")[0]355 if sso in self.token_status_map and model in self.token_status_map[sso]:356 self.token_status_map[sso][model]["isValid"] = True357 self.token_status_map[sso][model]["invalidatedTime"] = None358 self.token_status_map[sso][model]["totalRequestCount"] = 0359 360 token_entry["RequestCount"] = 0361 token_entry["StartCallTime"] = None362 363 import threading364 # 启动一个线程执行定时任务，每小时执行一次365 def run_timer():366 while True:367 reset_expired_tokens()368 time.sleep(3600)369 370 timer_thread = threading.Thread(target=run_timer)371 timer_thread.daemon = True372 timer_thread.start()373 374 def get_all_tokens(self):375 all_tokens = set()376 for model_tokens in self.token_model_map.values():377 for entry in model_tokens:378 all_tokens.add(entry["token"])379 return list(all_tokens)380 381 def get_token_status_map(self):382 return self.token_status_map383 384class Utils:385 @staticmethod386 def organize_search_results(search_results):387 if not search_results or 'results' not in search_results:388 return ''389 390 results = search_results['results']391 formatted_results = []392 393 for index, result in enumerate(results):394 title = result.get('title', '未知标题')395 url = result.get('url', '#')396 preview = result.get('preview', '无预览内容')397 398 formatted_result = f"\r\n<details><summary>资料[{index}]: {title}</summary>\r\n{preview}\r\n\n[Link]({url})\r\n</details>"399 formatted_results.append(formatted_result)400 401 return '\n\n'.join(formatted_results)402 403 @staticmethod404 def create_auth_headers(model):405 return token_manager.get_next_token_for_model(model)406 407 @staticmethod408 def get_proxy_options():409 proxy = CONFIG["API"]["PROXY"]410 proxy_options = {}411 412 if proxy:413 logger.info(f"使用代理: {proxy}", "Server")414 proxy_options["proxies"] = {"https": proxy, "http": proxy}415 416 if proxy.startswith("socks5://"):417 proxy_options["proxies"] = {"https": proxy, "http": proxy}418 proxy_options["proxy_type"] = "socks5"419 420 return proxy_options421 422class GrokApiClient:423 def __init__(self, model_id):424 if model_id not in CONFIG["MODELS"]:425 raise ValueError(f"不支持的模型: {model_id}")426 self.model_id = CONFIG["MODELS"][model_id]427 428 def process_message_content(self, content):429 if isinstance(content, str):430 return content431 return None432 433 def get_image_type(self, base64_string):434 mime_type = 'image/jpeg'435 if 'data:image' in base64_string:436 import re437 matches = re.search(r'data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,', base64_string)438 if matches:439 mime_type = matches.group(1)440 441 extension = mime_type.split('/')[1]442 file_name = f"image.{extension}"443 444 return {445 "mimeType": mime_type,446 "fileName": file_name447 }448 449 def upload_base64_image(self, base64_data, url):450 try:451 if 'data:image' in base64_data:452 image_buffer = base64_data.split(',')[1]453 else:454 image_buffer = base64_data455 456 image_info = self.get_image_type(base64_data)457 mime_type = image_info["mimeType"]458 file_name = image_info["fileName"]459 460 upload_data = {461 "rpc": "uploadFile",462 "req": {463 "fileName": file_name,464 "fileMimeType": mime_type,465 "content": image_buffer466 }467 }468 469 logger.info("发送图片请求", "Server")470 471 proxy_options = Utils.get_proxy_options()472 response = curl_requests.post(473 url,474 headers={475 **DEFAULT_HEADERS,476 "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]477 },478 json=upload_data,479 impersonate="chrome133a",480 **proxy_options481 )482 483 if response.status_code != 200:484 logger.error(f"上传图片失败,状态码:{response.status_code}", "Server")485 return ''486 487 result = response.json()488 logger.info(f"上传图片成功: {result}", "Server")489 return result.get("fileMetadataId", "")490 491 except Exception as error:492 logger.error(str(error), "Server")493 return ''494 495 def prepare_chat_request(self, request):496 if ((request["model"] == 'grok-2-imageGen' or request["model"] == 'grok-3-imageGen') and497 not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"] and498 request.get("stream", False)):499 raise ValueError("该模型流式输出需要配置PICGO或者TUMY图床密钥!")500 501 todo_messages = request["messages"]502 if request["model"] in ['grok-2-imageGen', 'grok-3-imageGen', 'grok-3-deepsearch']:503 last_message = todo_messages[-1]504 if last_message["role"] != 'user':505 raise ValueError('此模型最后一条消息必须是用户消息!')506 todo_messages = [last_message]507 508 file_attachments = []509 messages = ''510 last_role = None511 last_content = ''512 search = request["model"] in ['grok-2-search', 'grok-3-search']513 514 # 移除<think>标签及其内容和base64图片515 def remove_think_tags(text):516 import re517 text = re.sub(r'<think>[\s\S]*?<\/think>', '', text).strip()518 text = re.sub(r'!\[image\]\(data:.*?base64,.*?\)', '[图片]', text)519 return text520 521 def process_content(content):522 if isinstance(content, list):523 text_content = ''524 for item in content:525 if item["type"] == 'image_url':526 text_content += ("[图片]" if not text_content else '\n[图片]')527 elif item["type"] == 'text':528 text_content += (remove_think_tags(item["text"]) if not text_content else '\n' + remove_think_tags(item["text"]))529 return text_content530 elif isinstance(content, dict) and content is not None:531 if content["type"] == 'image_url':532 return "[图片]"533 elif content["type"] == 'text':534 return remove_think_tags(content["text"])535 return remove_think_tags(self.process_message_content(content))536 537 for current in todo_messages:538 role = 'assistant' if current["role"] == 'assistant' else 'user'539 is_last_message = current == todo_messages[-1]540 541 if is_last_message and "content" in current:542 if isinstance(current["content"], list):543 for item in current["content"]:544 if item["type"] == 'image_url':545 processed_image = self.upload_base64_image(546 item["image_url"]["url"],547 f"{CONFIG['API']['BASE_URL']}/api/rpc"548 )549 if processed_image:550 file_attachments.append(processed_image)551 elif isinstance(current["content"], dict) and current["content"].get("type") == 'image_url':552 processed_image = self.upload_base64_image(553 current["content"]["image_url"]["url"],554 f"{CONFIG['API']['BASE_URL']}/api/rpc"555 )556 if processed_image:557 file_attachments.append(processed_image)558 559 560 text_content = process_content(current.get("content", ""))561 562 if text_content or (is_last_message and file_attachments):563 if role == last_role and text_content:564 last_content += '\n' + text_content565 messages = messages[:messages.rindex(f"{role.upper()}: ")] + f"{role.upper()}: {last_content}\n"566 else:567 messages += f"{role.upper()}: {text_content or '[图片]'}\n"568 last_content = text_content569 last_role = role570 571 return {572 "temporary": CONFIG["API"]["IS_TEMP_CONVERSATION"],573 "modelName": self.model_id,574 "message": messages.strip(),575 "fileAttachments": file_attachments[:4],576 "imageAttachments": [],577 "disableSearch": False,578 "enableImageGeneration": True,579 "returnImageBytes": False,580 "returnRawGrokInXaiRequest": False,581 "enableImageStreaming": False,582 "imageGenerationCount": 1,583 "forceConcise": False,584 "toolOverrides": {585 "imageGen": request["model"] in ['grok-2-imageGen', 'grok-3-imageGen'],586 "webSearch": search,587 "xSearch": search,588 "xMediaSearch": search,589 "trendsSearch": search,590 "xPostAnalyze": search591 },592 "enableSideBySide": True,593 "isPreset": False,594 "sendFinalMetadata": True,595 "customInstructions": "",596 "deepsearchPreset": "default" if request["model"] == 'grok-3-deepsearch' else "",597 "isReasoning": request["model"] == 'grok-3-reasoning'598 }599 600class MessageProcessor:601 @staticmethod602 def create_chat_response(message, model, is_stream=False):603 base_response = {604 "id": f"chatcmpl-{uuid.uuid4()}",605 "created": int(time.time()),606 "model": model607 }608 609 if is_stream:610 return {611 **base_response,612 "object": "chat.completion.chunk",613 "choices": [{614 "index": 0,615 "delta": {616 "content": message617 }618 }]619 }620 621 return {622 **base_response,623 "object": "chat.completion",624 "choices": [{625 "index": 0,626 "message": {627 "role": "assistant",628 "content": message629 },630 "finish_reason": "stop"631 }],632 "usage": None633 }634 635def process_model_response(response, model):636 result = {"token": None, "imageUrl": None}637 638 if CONFIG["IS_IMG_GEN"]:639 if response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:640 result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]641 return result642 643 if model == 'grok-2':644 result["token"] = response.get("token")645 elif model in ['grok-2-search', 'grok-3-search']:646 if response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:647 result["token"] = f"\r\n<think>{Utils.organize_search_results(response['webSearchResults'])}</think>\r\n"648 else:649 result["token"] = response.get("token")650 elif model == 'grok-3':651 result["token"] = response.get("token")652 elif model == 'grok-3-deepsearch':653 if response.get("messageStepId") and not CONFIG["SHOW_THINKING"]:654 return result655 if response.get("messageStepId") and not CONFIG["IS_THINKING"]:656 result["token"] = "<think>" + response.get("token", "")657 CONFIG["IS_THINKING"] = True658 elif not response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "final":659 result["token"] = "</think>" + response.get("token", "")660 CONFIG["IS_THINKING"] = False661 elif (response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "assistant") or response.get("messageTag") == "final":662 result["token"] = response.get("token")663 elif model == 'grok-3-reasoning':664 if response.get("isThinking") and not CONFIG["SHOW_THINKING"]:665 return result666 667 if response.get("isThinking") and not CONFIG["IS_THINKING"]:668 result["token"] = "<think>" + response.get("token", "")669 CONFIG["IS_THINKING"] = True670 elif not response.get("isThinking") and CONFIG["IS_THINKING"]:671 result["token"] = "</think>" + response.get("token", "")672 CONFIG["IS_THINKING"] = False673 else:674 result["token"] = response.get("token")675 676 return result677 678def handle_image_response(image_url):679 max_retries = 2680 retry_count = 0681 image_base64_response = None682 683 while retry_count < max_retries:684 try:685 proxy_options = Utils.get_proxy_options()686 image_base64_response = curl_requests.get(687 f"https://assets.grok.com/{image_url}",688 headers={689 **DEFAULT_HEADERS,690 "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]691 },692 impersonate="chrome120",693 **proxy_options694 )695 696 if image_base64_response.status_code == 200:697 break698 699 retry_count += 1700 if retry_count == max_retries:701 raise Exception(f"上游服务请求失败! status: {image_base64_response.status_code}")702 703 time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)704 705 except Exception as error:706 logger.error(str(error), "Server")707 retry_count += 1708 if retry_count == max_retries:709 raise710 711 time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)712 713 image_buffer = image_base64_response.content714 715 if not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"]:716 base64_image = base64.b64encode(image_buffer).decode('utf-8')717 image_content_type = image_base64_response.headers.get('content-type', 'image/jpeg')718 return f""719 720 logger.info("开始上传图床", "Server")721 722 if CONFIG["API"]["PICGO_KEY"]:723 files = {'source': ('image.jpg', image_buffer, 'image/jpeg')}724 headers = {725 "X-API-Key": CONFIG["API"]["PICGO_KEY"]726 }727 728 response_url = requests.post(729 "https://www.picgo.net/api/1/upload",730 files=files,731 headers=headers732 )733 734 if response_url.status_code != 200:735 return "生图失败，请查看PICGO图床密钥是否设置正确"736 else:737 logger.info("生图成功", "Server")738 result = response_url.json()739 return f""740 741 742 elif CONFIG["API"]["TUMY_KEY"]:743 files = {'file': ('image.jpg', image_buffer, 'image/jpeg')}744 headers = {745 "Accept": "application/json",746 'Authorization': f"Bearer {CONFIG['API']['TUMY_KEY']}"747 }748 749 response_url = requests.post(750 "https://tu.my/api/v1/upload",751 files=files,752 headers=headers753 )754 755 if response_url.status_code != 200:756 return "生图失败，请查看TUMY图床密钥是否设置正确"757 else:758 try:759 result = response_url.json()760 logger.info("生图成功", "Server")761 return f""762 except Exception as error:763 logger.error(str(error), "Server")764 return "生图失败，请查看TUMY图床密钥是否设置正确"765 766def handle_non_stream_response(response, model):767 try:768 logger.info("开始处理非流式响应", "Server")769 770 stream = response.iter_lines()771 full_response = ""772 773 CONFIG["IS_THINKING"] = False774 CONFIG["IS_IMG_GEN"] = False775 CONFIG["IS_IMG_GEN2"] = False776 777 for chunk in stream:778 if not chunk:779 continue780 try:781 line_json = json.loads(chunk.decode("utf-8").strip())782 if line_json.get("error"):783 logger.error(json.dumps(line_json, indent=2), "Server")784 return json.dumps({"error": "RateLimitError"}) + "\n\n"785 786 response_data = line_json.get("result", {}).get("response")787 if not response_data:788 continue789 790 if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):791 CONFIG["IS_IMG_GEN"] = True792 793 result = process_model_response(response_data, model)794 795 if result["token"]:796 full_response += result["token"]797 798 if result["imageUrl"]:799 CONFIG["IS_IMG_GEN2"] = True800 return handle_image_response(result["imageUrl"])801 802 except json.JSONDecodeError:803 continue804 except Exception as e:805 logger.error(f"处理流式响应行时出错: {str(e)}", "Server")806 continue807 808 return full_response809 except Exception as error:810 logger.error(str(error), "Server")811 raise812def handle_stream_response(response, model):813 def generate():814 logger.info("开始处理流式响应", "Server")815 816 stream = response.iter_lines()817 CONFIG["IS_THINKING"] = False818 CONFIG["IS_IMG_GEN"] = False819 CONFIG["IS_IMG_GEN2"] = False820 821 for chunk in stream:822 if not chunk:823 continue824 try:825 line_json = json.loads(chunk.decode("utf-8").strip())826 if line_json.get("error"):827 logger.error(json.dumps(line_json, indent=2), "Server")828 yield json.dumps({"error": "RateLimitError"}) + "\n\n"829 return830 831 response_data = line_json.get("result", {}).get("response")832 if not response_data:833 continue834 835 if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):836 CONFIG["IS_IMG_GEN"] = True837 838 result = process_model_response(response_data, model)839 840 if result["token"]:841 yield f"data: {json.dumps(MessageProcessor.create_chat_response(result['token'], model, True))}\n\n"842 843 if result["imageUrl"]:844 CONFIG["IS_IMG_GEN2"] = True845 image_data = handle_image_response(result["imageUrl"])846 yield f"data: {json.dumps(MessageProcessor.create_chat_response(image_data, model, True))}\n\n"847 848 except json.JSONDecodeError:849 continue850 except Exception as e:851 logger.error(f"处理流式响应行时出错: {str(e)}", "Server")852 continue853 854 yield "data: [DONE]\n\n"855 return generate()856 857def initialization():858 sso_array = os.environ.get("SSO", "").split(',')859 logger.info("开始加载令牌", "Server")860 for sso in sso_array:861 if sso:862 token_manager.add_token(f"sso-rw={sso};sso={sso}")863 864 logger.info(f"成功加载令牌: {json.dumps(token_manager.get_all_tokens(), indent=2)}", "Server")865 logger.info(f"令牌加载完成，共加载: {len(token_manager.get_all_tokens())}个令牌", "Server")866 867 if CONFIG["API"]["PROXY"]:868 logger.info(f"代理已设置: {CONFIG['API']['PROXY']}", "Server")869 870logger.info("初始化完成", "Server")871 872 873app = Flask(__name__)874app.wsgi_app = ProxyFix(app.wsgi_app)875 876 877@app.before_request878def log_request_info():879 logger.info(f"{request.method} {request.path}", "Request")880 881@app.route('/get/tokens', methods=['GET'])882def get_tokens():883 auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')884 if CONFIG["API"]["IS_CUSTOM_SSO"]:885 return jsonify({"error": '自定义的SSO令牌模式无法获取轮询sso令牌状态'}), 403886 elif auth_token != CONFIG["API"]["API_KEY"]:887 return jsonify({"error": 'Unauthorized'}), 401888 889 return jsonify(token_manager.get_token_status_map())890 891@app.route('/add/token', methods=['POST'])892def add_token():893 auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')894 if CONFIG["API"]["IS_CUSTOM_SSO"]:895 return jsonify({"error": '自定义的SSO令牌模式无法添加sso令牌'}), 403896 elif auth_token != CONFIG["API"]["API_KEY"]:897 return jsonify({"error": 'Unauthorized'}), 401898 899 try:900 sso = request.json.get('sso')901 token_manager.add_token(f"sso-rw={sso};sso={sso}")902 return jsonify(token_manager.get_token_status_map().get(sso, {})), 200903 except Exception as error:904 logger.error(str(error), "Server")905 return jsonify({"error": '添加sso令牌失败'}), 500906 907@app.route('/delete/token', methods=['POST'])908def delete_token():909 auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')910 if CONFIG["API"]["IS_CUSTOM_SSO"]:911 return jsonify({"error": '自定义的SSO令牌模式无法删除sso令牌'}), 403912 elif auth_token != CONFIG["API"]["API_KEY"]:913 return jsonify({"error": 'Unauthorized'}), 401914 915 try:916 sso = request.json.get('sso')917 token_manager.delete_token(f"sso-rw={sso};sso={sso}")918 return jsonify({"message": '删除sso令牌成功'}), 200919 except Exception as error:920 logger.error(str(error), "Server")921 return jsonify({"error": '删除sso令牌失败'}), 500922 923@app.route('/v1/models', methods=['GET'])924def get_models():925 return jsonify({926 "object": "list",927 "data": [928 {929 "id": model,930 "object": "model",931 "created": int(time.time()),932 "owned_by": "grok"933 }934 for model in CONFIG["MODELS"].keys()935 ]936 })937 938@app.route('/v1/chat/completions', methods=['POST'])939def chat_completions():940 try:941 auth_token = request.headers.get('Authorization',942 '').replace('Bearer ', '')943 if auth_token:944 if CONFIG["API"]["IS_CUSTOM_SSO"]:945 result = f"sso={auth_token};sso-rw={auth_token}"946 token_manager.set_token(result)947 elif auth_token != CONFIG["API"]["API_KEY"]:948 return jsonify({"error": 'Unauthorized'}), 401949 else:950 return jsonify({"error": 'API_KEY缺失'}), 401951 952 data = request.json953 model = data.get("model")954 stream = data.get("stream", False)955 956 retry_count = 0957 grok_client = GrokApiClient(model)958 request_payload = grok_client.prepare_chat_request(data)959 960 while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:961 retry_count += 1962 CONFIG["API"]["SIGNATURE_COOKIE"] = Utils.create_auth_headers(963 model)964 965 if not CONFIG["API"]["SIGNATURE_COOKIE"]:966 raise ValueError('该模型无可用令牌')967 968 logger.info(969 f"当前令牌: {json.dumps(CONFIG['API']['SIGNATURE_COOKIE'], indent=2)}",970 "Server")971 logger.info(972 f"当前可用模型的全部可用数量: {json.dumps(token_manager.get_remaining_token_request_capacity(), indent=2)}",973 "Server")974 975 try:976 proxy_options = Utils.get_proxy_options()977 response = curl_requests.post(978 f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",979 headers={980 **DEFAULT_HEADERS, "Cookie":981 CONFIG["API"]["SIGNATURE_COOKIE"]982 },983 data=json.dumps(request_payload),984 impersonate="chrome133a",985 stream=True,986 **proxy_options)987 if response.status_code == 200:988 logger.info("请求成功", "Server")989 logger.info(990 f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}",991 "Server")992 993 try:994 if stream:995 return Response(stream_with_context(996 handle_stream_response(response, model)),997 content_type='text/event-stream')998 else:999 content = handle_non_stream_response(1000 response, model)1001 return jsonify(1002 MessageProcessor.create_chat_response(1003 content, model))1004 1005 except Exception as error:1006 logger.error(str(error), "Server")1007 if CONFIG["API"]["IS_CUSTOM_SSO"]:1008 raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")1009 1010 token_manager.remove_token_from_model(1011 model, CONFIG["API"]["SIGNATURE_COOKIE"])1012 if token_manager.get_token_count_for_model(model) == 0:1013 raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")1014 1015 elif response.status_code == 429:1016 if CONFIG["API"]["IS_CUSTOM_SSO"]:1017 raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")1018 1019 token_manager.remove_token_from_model(1020 model, CONFIG["API"]["SIGNATURE_COOKIE"])1021 if token_manager.get_token_count_for_model(model) == 0:1022 raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")1023 1024 else:1025 if CONFIG["API"]["IS_CUSTOM_SSO"]:1026 raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")1027 1028 logger.error(f"令牌异常错误状态!status: {response.status_code}",1029 "Server")1030 token_manager.remove_token_from_model(1031 model, CONFIG["API"]["SIGNATURE_COOKIE"])1032 logger.info(1033 f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}",1034 "Server")1035 1036 except Exception as e:1037 logger.error(f"请求处理异常: {str(e)}", "Server")1038 if CONFIG["API"]["IS_CUSTOM_SSO"]:1039 raise1040 continue1041 1042 raise ValueError('当前模型所有令牌都已耗尽')1043 1044 except Exception as error:1045 logger.error(str(error), "ChatAPI")1046 return jsonify(1047 {"error": {1048 "message": str(error),1049 "type": "server_error"1050 }}), 5001051 1052@app.route('/', defaults={'path': ''})1053@app.route('/<path:path>')1054def catch_all(path):1055 return 'api运行正常', 2001056 1057if __name__ == '__main__':1058 token_manager = AuthTokenManager()1059 initialization()1060 1061 app.run(1062 host='0.0.0.0',1063 port=CONFIG["SERVER"]["PORT"],1064 debug=False1065 )1066