lidark/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 def reduce_token_request_count(self, model_id, count):221 try:222 normalized_model = self.normalize_model_name(model_id)223 224 if normalized_model not in self.token_model_map:225 logger.error(f"模型 {normalized_model} 不存在", "TokenManager")226 return False227 228 if not self.token_model_map[normalized_model]:229 logger.error(f"模型 {normalized_model} 没有可用的token", "TokenManager")230 return False231 232 token_entry = self.token_model_map[normalized_model][0]233 234 # 确保RequestCount不会小于0235 new_count = max(0, token_entry["RequestCount"] - count)236 reduction = token_entry["RequestCount"] - new_count237 238 token_entry["RequestCount"] = new_count239 240 # 更新token状态241 if token_entry["token"]:242 sso = token_entry["token"].split("sso=")[1].split(";")[0]243 if sso in self.token_status_map and normalized_model in self.token_status_map[sso]:244 self.token_status_map[sso][normalized_model]["totalRequestCount"] = max(245 0, 246 self.token_status_map[sso][normalized_model]["totalRequestCount"] - reduction247 )248 return True249 250 except Exception as error:251 logger.error(f"重置校对token请求次数时发生错误: {str(error)}", "TokenManager")252 return False253 def get_next_token_for_model(self, model_id):254 normalized_model = self.normalize_model_name(model_id)255 256 if normalized_model not in self.token_model_map or not self.token_model_map[normalized_model]:257 return None258 259 token_entry = self.token_model_map[normalized_model][0]260 261 if token_entry:262 if token_entry["StartCallTime"] is None:263 token_entry["StartCallTime"] = int(time.time() * 1000)264 265 if not self.token_reset_switch:266 self.start_token_reset_process()267 self.token_reset_switch = True268 269 token_entry["RequestCount"] += 1270 271 if token_entry["RequestCount"] > self.model_config[normalized_model]["RequestFrequency"]:272 self.remove_token_from_model(normalized_model, token_entry["token"])273 next_token_entry = self.token_model_map[normalized_model][0] if self.token_model_map[normalized_model] else None274 return next_token_entry["token"] if next_token_entry else None275 276 sso = token_entry["token"].split("sso=")[1].split(";")[0]277 if sso in self.token_status_map and normalized_model in self.token_status_map[sso]:278 if token_entry["RequestCount"] == self.model_config[normalized_model]["RequestFrequency"]:279 self.token_status_map[sso][normalized_model]["isValid"] = False280 self.token_status_map[sso][normalized_model]["invalidatedTime"] = int(time.time() * 1000)281 self.token_status_map[sso][normalized_model]["totalRequestCount"] += 1282 283 return token_entry["token"]284 285 return None286 287 def remove_token_from_model(self, model_id, token):288 normalized_model = self.normalize_model_name(model_id)289 290 if normalized_model not in self.token_model_map:291 logger.error(f"模型 {normalized_model} 不存在", "TokenManager")292 return False293 294 model_tokens = self.token_model_map[normalized_model]295 token_index = next((i for i, entry in enumerate(model_tokens) if entry["token"] == token), -1)296 297 if token_index != -1:298 removed_token_entry = model_tokens.pop(token_index)299 self.expired_tokens.add((300 removed_token_entry["token"],301 normalized_model,302 int(time.time() * 1000)303 ))304 305 if not self.token_reset_switch:306 self.start_token_reset_process()307 self.token_reset_switch = True308 309 logger.info(f"模型{model_id}的令牌已失效，已成功移除令牌: {token}", "TokenManager")310 return True311 312 logger.error(f"在模型 {normalized_model} 中未找到 token: {token}", "TokenManager")313 return False314 315 def get_expired_tokens(self):316 return list(self.expired_tokens)317 318 def normalize_model_name(self, model):319 if model.startswith('grok-') and 'deepsearch' not in model and 'reasoning' not in model:320 return '-'.join(model.split('-')[:2])321 return model322 323 def get_token_count_for_model(self, model_id):324 normalized_model = self.normalize_model_name(model_id)325 return len(self.token_model_map.get(normalized_model, []))326 327 def get_remaining_token_request_capacity(self):328 remaining_capacity_map = {}329 330 for model in self.model_config.keys():331 model_tokens = self.token_model_map.get(model, [])332 model_request_frequency = self.model_config[model]["RequestFrequency"]333 334 total_used_requests = sum(token_entry.get("RequestCount", 0) for token_entry in model_tokens)335 336 remaining_capacity = (len(model_tokens) * model_request_frequency) - total_used_requests337 remaining_capacity_map[model] = max(0, remaining_capacity)338 339 return remaining_capacity_map340 341 def get_token_array_for_model(self, model_id):342 normalized_model = self.normalize_model_name(model_id)343 return self.token_model_map.get(normalized_model, [])344 345 def start_token_reset_process(self):346 def reset_expired_tokens():347 now = int(time.time() * 1000)348 349 tokens_to_remove = set()350 for token_info in self.expired_tokens:351 token, model, expired_time = token_info352 expiration_time = self.model_config[model]["ExpirationTime"]353 354 if now - expired_time >= expiration_time:355 if not any(entry["token"] == token for entry in self.token_model_map.get(model, [])):356 if model not in self.token_model_map:357 self.token_model_map[model] = []358 359 self.token_model_map[model].append({360 "token": token,361 "RequestCount": 0,362 "AddedTime": now,363 "StartCallTime": None364 })365 366 sso = 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"] = True369 self.token_status_map[sso][model]["invalidatedTime"] = None370 self.token_status_map[sso][model]["totalRequestCount"] = 0371 372 tokens_to_remove.add(token_info)373 374 self.expired_tokens -= tokens_to_remove375 376 for model in self.model_config.keys():377 if model not in self.token_model_map:378 continue379 380 for token_entry in self.token_model_map[model]:381 if not token_entry.get("StartCallTime"):382 continue383 384 expiration_time = self.model_config[model]["ExpirationTime"]385 if now - token_entry["StartCallTime"] >= expiration_time:386 sso = token_entry["token"].split("sso=")[1].split(";")[0]387 if sso in self.token_status_map and model in self.token_status_map[sso]:388 self.token_status_map[sso][model]["isValid"] = True389 self.token_status_map[sso][model]["invalidatedTime"] = None390 self.token_status_map[sso][model]["totalRequestCount"] = 0391 392 token_entry["RequestCount"] = 0393 token_entry["StartCallTime"] = None394 395 import threading396 # 启动一个线程执行定时任务，每小时执行一次397 def run_timer():398 while True:399 reset_expired_tokens()400 time.sleep(3600)401 402 timer_thread = threading.Thread(target=run_timer)403 timer_thread.daemon = True404 timer_thread.start()405 406 def get_all_tokens(self):407 all_tokens = set()408 for model_tokens in self.token_model_map.values():409 for entry in model_tokens:410 all_tokens.add(entry["token"])411 return list(all_tokens)412 413 def get_token_status_map(self):414 return self.token_status_map415 416class Utils:417 @staticmethod418 def organize_search_results(search_results):419 if not search_results or 'results' not in search_results:420 return ''421 422 results = search_results['results']423 formatted_results = []424 425 for index, result in enumerate(results):426 title = result.get('title', '未知标题')427 url = result.get('url', '#')428 preview = result.get('preview', '无预览内容')429 430 formatted_result = f"\r\n<details><summary>资料[{index}]: {title}</summary>\r\n{preview}\r\n\n[Link]({url})\r\n</details>"431 formatted_results.append(formatted_result)432 433 return '\n\n'.join(formatted_results)434 435 @staticmethod436 def create_auth_headers(model):437 return token_manager.get_next_token_for_model(model)438 439 @staticmethod440 def get_proxy_options():441 proxy = CONFIG["API"]["PROXY"]442 proxy_options = {}443 444 if proxy:445 logger.info(f"使用代理: {proxy}", "Server")446 if proxy.startswith("socks5://"):447 proxy_options["proxy"] = proxy448 449 if '@' in proxy:450 auth_part = proxy.split('@')[0].split('://')[1]451 if ':' in auth_part:452 username, password = auth_part.split(':')453 proxy_options["proxy_auth"] = (username, password)454 else:455 proxy_options["proxies"] = {"https": proxy, "http": proxy}456 print(proxy_options) 457 return proxy_options458 459class GrokApiClient:460 def __init__(self, model_id):461 if model_id not in CONFIG["MODELS"]:462 raise ValueError(f"不支持的模型: {model_id}")463 self.model_id = CONFIG["MODELS"][model_id]464 465 def process_message_content(self, content):466 if isinstance(content, str):467 return content468 return None469 470 def get_image_type(self, base64_string):471 mime_type = 'image/jpeg'472 if 'data:image' in base64_string:473 import re474 matches = re.search(r'data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,', base64_string)475 if matches:476 mime_type = matches.group(1)477 478 extension = mime_type.split('/')[1]479 file_name = f"image.{extension}"480 481 return {482 "mimeType": mime_type,483 "fileName": file_name484 }485 486 def upload_base64_image(self, base64_data, url):487 try:488 if 'data:image' in base64_data:489 image_buffer = base64_data.split(',')[1]490 else:491 image_buffer = base64_data492 493 image_info = self.get_image_type(base64_data)494 mime_type = image_info["mimeType"]495 file_name = image_info["fileName"]496 497 upload_data = {498 "rpc": "uploadFile",499 "req": {500 "fileName": file_name,501 "fileMimeType": mime_type,502 "content": image_buffer503 }504 }505 506 logger.info("发送图片请求", "Server")507 508 proxy_options = Utils.get_proxy_options()509 response = curl_requests.post(510 url,511 headers={512 **DEFAULT_HEADERS,513 "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]514 },515 json=upload_data,516 impersonate="chrome133a",517 **proxy_options518 )519 520 if response.status_code != 200:521 logger.error(f"上传图片失败,状态码:{response.status_code}", "Server")522 return ''523 524 result = response.json()525 logger.info(f"上传图片成功: {result}", "Server")526 return result.get("fileMetadataId", "")527 528 except Exception as error:529 logger.error(str(error), "Server")530 return ''531 532 def prepare_chat_request(self, request):533 if ((request["model"] == 'grok-2-imageGen' or request["model"] == 'grok-3-imageGen') and534 not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"] and535 request.get("stream", False)):536 raise ValueError("该模型流式输出需要配置PICGO或者TUMY图床密钥!")537 538 todo_messages = request["messages"]539 if request["model"] in ['grok-2-imageGen', 'grok-3-imageGen', 'grok-3-deepsearch']:540 last_message = todo_messages[-1]541 if last_message["role"] != 'user':542 raise ValueError('此模型最后一条消息必须是用户消息!')543 todo_messages = [last_message]544 545 file_attachments = []546 messages = ''547 last_role = None548 last_content = ''549 search = request["model"] in ['grok-2-search', 'grok-3-search']550 551 # 移除<think>标签及其内容和base64图片552 def remove_think_tags(text):553 import re554 text = re.sub(r'<think>[\s\S]*?<\/think>', '', text).strip()555 text = re.sub(r'!\[image\]\(data:.*?base64,.*?\)', '[图片]', text)556 return text557 558 def process_content(content):559 if isinstance(content, list):560 text_content = ''561 for item in content:562 if item["type"] == 'image_url':563 text_content += ("[图片]" if not text_content else '\n[图片]')564 elif item["type"] == 'text':565 text_content += (remove_think_tags(item["text"]) if not text_content else '\n' + remove_think_tags(item["text"]))566 return text_content567 elif isinstance(content, dict) and content is not None:568 if content["type"] == 'image_url':569 return "[图片]"570 elif content["type"] == 'text':571 return remove_think_tags(content["text"])572 return remove_think_tags(self.process_message_content(content))573 574 for current in todo_messages:575 role = 'assistant' if current["role"] == 'assistant' else 'user'576 is_last_message = current == todo_messages[-1]577 578 if is_last_message and "content" in current:579 if isinstance(current["content"], list):580 for item in current["content"]:581 if item["type"] == 'image_url':582 processed_image = self.upload_base64_image(583 item["image_url"]["url"],584 f"{CONFIG['API']['BASE_URL']}/api/rpc"585 )586 if processed_image:587 file_attachments.append(processed_image)588 elif isinstance(current["content"], dict) and current["content"].get("type") == 'image_url':589 processed_image = self.upload_base64_image(590 current["content"]["image_url"]["url"],591 f"{CONFIG['API']['BASE_URL']}/api/rpc"592 )593 if processed_image:594 file_attachments.append(processed_image)595 596 597 text_content = process_content(current.get("content", ""))598 599 if text_content or (is_last_message and file_attachments):600 if role == last_role and text_content:601 last_content += '\n' + text_content602 messages = messages[:messages.rindex(f"{role.upper()}: ")] + f"{role.upper()}: {last_content}\n"603 else:604 messages += f"{role.upper()}: {text_content or '[图片]'}\n"605 last_content = text_content606 last_role = role607 608 return {609 "temporary": CONFIG["API"].get("IS_TEMP_CONVERSATION", False),610 "modelName": self.model_id,611 "message": messages.strip(),612 "fileAttachments": file_attachments[:4],613 "imageAttachments": [],614 "disableSearch": False,615 "enableImageGeneration": True,616 "returnImageBytes": False,617 "returnRawGrokInXaiRequest": False,618 "enableImageStreaming": False,619 "imageGenerationCount": 1,620 "forceConcise": False,621 "toolOverrides": {622 "imageGen": request["model"] in ['grok-2-imageGen', 'grok-3-imageGen'],623 "webSearch": search,624 "xSearch": search,625 "xMediaSearch": search,626 "trendsSearch": search,627 "xPostAnalyze": search628 },629 "enableSideBySide": True,630 "isPreset": False,631 "sendFinalMetadata": True,632 "customInstructions": "",633 "deepsearchPreset": "default" if request["model"] == 'grok-3-deepsearch' else "",634 "isReasoning": request["model"] == 'grok-3-reasoning'635 }636 637class MessageProcessor:638 @staticmethod639 def create_chat_response(message, model, is_stream=False):640 base_response = {641 "id": f"chatcmpl-{uuid.uuid4()}",642 "created": int(time.time()),643 "model": model644 }645 646 if is_stream:647 return {648 **base_response,649 "object": "chat.completion.chunk",650 "choices": [{651 "index": 0,652 "delta": {653 "content": message654 }655 }]656 }657 658 return {659 **base_response,660 "object": "chat.completion",661 "choices": [{662 "index": 0,663 "message": {664 "role": "assistant",665 "content": message666 },667 "finish_reason": "stop"668 }],669 "usage": None670 }671 672def process_model_response(response, model):673 result = {"token": None, "imageUrl": None}674 675 if CONFIG["IS_IMG_GEN"]:676 if response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:677 result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]678 return result679 680 if model == 'grok-2':681 result["token"] = response.get("token")682 elif model in ['grok-2-search', 'grok-3-search']:683 if response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:684 result["token"] = f"\r\n<think>{Utils.organize_search_results(response['webSearchResults'])}</think>\r\n"685 else:686 result["token"] = response.get("token")687 elif model == 'grok-3':688 result["token"] = response.get("token")689 elif model == 'grok-3-deepsearch':690 if response.get("messageStepId") and not CONFIG["SHOW_THINKING"]:691 return result692 if response.get("messageStepId") and not CONFIG["IS_THINKING"]:693 result["token"] = "<think>" + response.get("token", "")694 CONFIG["IS_THINKING"] = True695 elif not response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "final":696 result["token"] = "</think>" + response.get("token", "")697 CONFIG["IS_THINKING"] = False698 elif (response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "assistant") or response.get("messageTag") == "final":699 result["token"] = response.get("token")700 elif model == 'grok-3-reasoning':701 if response.get("isThinking") and not CONFIG["SHOW_THINKING"]:702 return result703 704 if response.get("isThinking") and not CONFIG["IS_THINKING"]:705 result["token"] = "<think>" + response.get("token", "")706 CONFIG["IS_THINKING"] = True707 elif not response.get("isThinking") and CONFIG["IS_THINKING"]:708 result["token"] = "</think>" + response.get("token", "")709 CONFIG["IS_THINKING"] = False710 else:711 result["token"] = response.get("token")712 713 return result714 715def handle_image_response(image_url):716 max_retries = 2717 retry_count = 0718 image_base64_response = None719 720 while retry_count < max_retries:721 try:722 proxy_options = Utils.get_proxy_options()723 image_base64_response = curl_requests.get(724 f"https://assets.grok.com/{image_url}",725 headers={726 **DEFAULT_HEADERS,727 "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]728 },729 impersonate="chrome120",730 **proxy_options731 )732 733 if image_base64_response.status_code == 200:734 break735 736 retry_count += 1737 if retry_count == max_retries:738 raise Exception(f"上游服务请求失败! status: {image_base64_response.status_code}")739 740 time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)741 742 except Exception as error:743 logger.error(str(error), "Server")744 retry_count += 1745 if retry_count == max_retries:746 raise747 748 time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)749 750 image_buffer = image_base64_response.content751 752 if not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"]:753 base64_image = base64.b64encode(image_buffer).decode('utf-8')754 image_content_type = image_base64_response.headers.get('content-type', 'image/jpeg')755 return f""756 757 logger.info("开始上传图床", "Server")758 759 if CONFIG["API"]["PICGO_KEY"]:760 files = {'source': ('image.jpg', image_buffer, 'image/jpeg')}761 headers = {762 "X-API-Key": CONFIG["API"]["PICGO_KEY"]763 }764 765 response_url = requests.post(766 "https://www.picgo.net/api/1/upload",767 files=files,768 headers=headers769 )770 771 if response_url.status_code != 200:772 return "生图失败，请查看PICGO图床密钥是否设置正确"773 else:774 logger.info("生图成功", "Server")775 result = response_url.json()776 return f""777 778 779 elif CONFIG["API"]["TUMY_KEY"]:780 files = {'file': ('image.jpg', image_buffer, 'image/jpeg')}781 headers = {782 "Accept": "application/json",783 'Authorization': f"Bearer {CONFIG['API']['TUMY_KEY']}"784 }785 786 response_url = requests.post(787 "https://tu.my/api/v1/upload",788 files=files,789 headers=headers790 )791 792 if response_url.status_code != 200:793 return "生图失败，请查看TUMY图床密钥是否设置正确"794 else:795 try:796 result = response_url.json()797 logger.info("生图成功", "Server")798 return f""799 except Exception as error:800 logger.error(str(error), "Server")801 return "生图失败，请查看TUMY图床密钥是否设置正确"802 803def handle_non_stream_response(response, model):804 try:805 logger.info("开始处理非流式响应", "Server")806 807 stream = response.iter_lines()808 full_response = ""809 810 CONFIG["IS_THINKING"] = False811 CONFIG["IS_IMG_GEN"] = False812 CONFIG["IS_IMG_GEN2"] = False813 814 for chunk in stream:815 if not chunk:816 continue817 try:818 line_json = json.loads(chunk.decode("utf-8").strip())819 if line_json.get("error"):820 logger.error(json.dumps(line_json, indent=2), "Server")821 return json.dumps({"error": "RateLimitError"}) + "\n\n"822 823 response_data = line_json.get("result", {}).get("response")824 if not response_data:825 continue826 827 if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):828 CONFIG["IS_IMG_GEN"] = True829 830 result = process_model_response(response_data, model)831 832 if result["token"]:833 full_response += result["token"]834 835 if result["imageUrl"]:836 CONFIG["IS_IMG_GEN2"] = True837 return handle_image_response(result["imageUrl"])838 839 except json.JSONDecodeError:840 continue841 except Exception as e:842 logger.error(f"处理流式响应行时出错: {str(e)}", "Server")843 continue844 845 return full_response846 except Exception as error:847 logger.error(str(error), "Server")848 raise849def handle_stream_response(response, model):850 def generate():851 logger.info("开始处理流式响应", "Server")852 853 stream = response.iter_lines()854 CONFIG["IS_THINKING"] = False855 CONFIG["IS_IMG_GEN"] = False856 CONFIG["IS_IMG_GEN2"] = False857 858 for chunk in stream:859 if not chunk:860 continue861 try:862 line_json = json.loads(chunk.decode("utf-8").strip())863 if line_json.get("error"):864 logger.error(json.dumps(line_json, indent=2), "Server")865 yield json.dumps({"error": "RateLimitError"}) + "\n\n"866 return867 868 response_data = line_json.get("result", {}).get("response")869 if not response_data:870 continue871 872 if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):873 CONFIG["IS_IMG_GEN"] = True874 875 result = process_model_response(response_data, model)876 877 if result["token"]:878 yield f"data: {json.dumps(MessageProcessor.create_chat_response(result['token'], model, True))}\n\n"879 880 if result["imageUrl"]:881 CONFIG["IS_IMG_GEN2"] = True882 image_data = handle_image_response(result["imageUrl"])883 yield f"data: {json.dumps(MessageProcessor.create_chat_response(image_data, model, True))}\n\n"884 885 except json.JSONDecodeError:886 continue887 except Exception as e:888 logger.error(f"处理流式响应行时出错: {str(e)}", "Server")889 continue890 891 yield "data: [DONE]\n\n"892 return generate()893 894def initialization():895 sso_array = os.environ.get("SSO", "").split(',')896 logger.info("开始加载令牌", "Server")897 for sso in sso_array:898 if sso:899 token_manager.add_token(f"sso-rw={sso};sso={sso}")900 901 logger.info(f"成功加载令牌: {json.dumps(token_manager.get_all_tokens(), indent=2)}", "Server")902 logger.info(f"令牌加载完成，共加载: {len(token_manager.get_all_tokens())}个令牌", "Server")903 904 if CONFIG["API"]["PROXY"]:905 logger.info(f"代理已设置: {CONFIG['API']['PROXY']}", "Server")906 907logger.info("初始化完成", "Server")908 909 910app = Flask(__name__)911app.wsgi_app = ProxyFix(app.wsgi_app)912 913 914@app.before_request915def log_request_info():916 logger.info(f"{request.method} {request.path}", "Request")917 918@app.route('/get/tokens', methods=['GET'])919def get_tokens():920 auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')921 if CONFIG["API"]["IS_CUSTOM_SSO"]:922 return jsonify({"error": '自定义的SSO令牌模式无法获取轮询sso令牌状态'}), 403923 elif auth_token != CONFIG["API"]["API_KEY"]:924 return jsonify({"error": 'Unauthorized'}), 401925 926 return jsonify(token_manager.get_token_status_map())927 928@app.route('/add/token', methods=['POST'])929def add_token():930 auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')931 if CONFIG["API"]["IS_CUSTOM_SSO"]:932 return jsonify({"error": '自定义的SSO令牌模式无法添加sso令牌'}), 403933 elif auth_token != CONFIG["API"]["API_KEY"]:934 return jsonify({"error": 'Unauthorized'}), 401935 936 try:937 sso = request.json.get('sso')938 token_manager.add_token(f"sso-rw={sso};sso={sso}")939 return jsonify(token_manager.get_token_status_map().get(sso, {})), 200940 except Exception as error:941 logger.error(str(error), "Server")942 return jsonify({"error": '添加sso令牌失败'}), 500943 944@app.route('/delete/token', methods=['POST'])945def delete_token():946 auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')947 if CONFIG["API"]["IS_CUSTOM_SSO"]:948 return jsonify({"error": '自定义的SSO令牌模式无法删除sso令牌'}), 403949 elif auth_token != CONFIG["API"]["API_KEY"]:950 return jsonify({"error": 'Unauthorized'}), 401951 952 try:953 sso = request.json.get('sso')954 token_manager.delete_token(f"sso-rw={sso};sso={sso}")955 return jsonify({"message": '删除sso令牌成功'}), 200956 except Exception as error:957 logger.error(str(error), "Server")958 return jsonify({"error": '删除sso令牌失败'}), 500959 960@app.route('/v1/models', methods=['GET'])961def get_models():962 return jsonify({963 "object": "list",964 "data": [965 {966 "id": model,967 "object": "model",968 "created": int(time.time()),969 "owned_by": "grok"970 }971 for model in CONFIG["MODELS"].keys()972 ]973 })974 975@app.route('/v1/chat/completions', methods=['POST'])976def chat_completions():977 try:978 auth_token = request.headers.get('Authorization',979 '').replace('Bearer ', '')980 if auth_token:981 if CONFIG["API"]["IS_CUSTOM_SSO"]:982 result = f"sso={auth_token};sso-rw={auth_token}"983 token_manager.set_token(result)984 elif auth_token != CONFIG["API"]["API_KEY"]:985 return jsonify({"error": 'Unauthorized'}), 401986 else:987 return jsonify({"error": 'API_KEY缺失'}), 401988 989 data = request.json990 model = data.get("model")991 stream = data.get("stream", False)992 993 retry_count = 0994 grok_client = GrokApiClient(model)995 request_payload = grok_client.prepare_chat_request(data)996 response_status_code = 500997 998 while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:999 retry_count += 11000 CONFIG["API"]["SIGNATURE_COOKIE"] = Utils.create_auth_headers(1001 model)1002 1003 if not CONFIG["API"]["SIGNATURE_COOKIE"]:1004 raise ValueError('该模型无可用令牌')1005 1006 logger.info(1007 f"当前令牌: {json.dumps(CONFIG['API']['SIGNATURE_COOKIE'], indent=2)}",1008 "Server")1009 logger.info(1010 f"当前可用模型的全部可用数量: {json.dumps(token_manager.get_remaining_token_request_capacity(), indent=2)}",1011 "Server")1012 1013 try:1014 proxy_options = Utils.get_proxy_options()1015 response = curl_requests.post(1016 f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",1017 headers={1018 **DEFAULT_HEADERS, "Cookie":1019 CONFIG["API"]["SIGNATURE_COOKIE"]1020 },1021 data=json.dumps(request_payload),1022 impersonate="chrome133a",1023 stream=True,1024 **proxy_options)1025 if response.status_code == 200:1026 response_status_code = 2001027 logger.info("请求成功", "Server")1028 logger.info(1029 f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}",1030 "Server")1031 1032 try:1033 if stream:1034 return Response(stream_with_context(1035 handle_stream_response(response, model)),1036 content_type='text/event-stream')1037 else:1038 content = handle_non_stream_response(1039 response, model)1040 return jsonify(1041 MessageProcessor.create_chat_response(1042 content, model))1043 1044 except Exception as error:1045 logger.error(str(error), "Server")1046 if CONFIG["API"]["IS_CUSTOM_SSO"]:1047 raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")1048 1049 token_manager.remove_token_from_model(1050 model, CONFIG["API"]["SIGNATURE_COOKIE"])1051 if token_manager.get_token_count_for_model(model) == 0:1052 raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")1053 elif response.status_code == 403:1054 response_status_code = 4031055 token_manager.reduce_token_request_count(model,1)#重置去除当前因为错误未成功请求的次数，确保不会因为错误未成功请求的次数导致次数上限1056 if token_manager.get_token_count_for_model(model) == 0:1057 raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")1058 raise ValueError(f"IP暂时被封黑无法破盾，请稍后重试或者更换ip")1059 elif response.status_code == 429:1060 response_status_code = 4291061 token_manager.reduce_token_request_count(model,1)1062 if CONFIG["API"]["IS_CUSTOM_SSO"]:1063 raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")1064 1065 token_manager.remove_token_from_model(1066 model, CONFIG["API"]["SIGNATURE_COOKIE"])1067 if token_manager.get_token_count_for_model(model) == 0:1068 raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")1069 1070 else:1071 if CONFIG["API"]["IS_CUSTOM_SSO"]:1072 raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")1073 1074 logger.error(f"令牌异常错误状态!status: {response.status_code}",1075 "Server")1076 token_manager.remove_token_from_model(1077 model, CONFIG["API"]["SIGNATURE_COOKIE"])1078 logger.info(1079 f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}",1080 "Server")1081 1082 except Exception as e:1083 logger.error(f"请求处理异常: {str(e)}", "Server")1084 if CONFIG["API"]["IS_CUSTOM_SSO"]:1085 raise1086 continue1087 if response_status_code == 403:1088 raise ValueError('IP暂时被封黑无法破盾，请稍后重试或者更换ip')1089 elif response_status_code == 500:1090 raise ValueError('当前模型所有令牌暂无可用，请稍后重试') 1091 1092 except Exception as error:1093 logger.error(str(error), "ChatAPI")1094 return jsonify(1095 {"error": {1096 "message": str(error),1097 "type": "server_error"1098 }}), response_status_code1099 1100@app.route('/', defaults={'path': ''})1101@app.route('/<path:path>')1102def catch_all(path):1103 return 'api运行正常', 2001104 1105if __name__ == '__main__':1106 token_manager = AuthTokenManager()1107 initialization()1108 1109 app.run(1110 host='0.0.0.0',1111 port=CONFIG["SERVER"]["PORT"],1112 debug=False1113 )