From b132590045ef20aec463bc7fbef3ac052938a084 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sat, 7 Feb 2026 11:32:25 +0800 Subject: [PATCH 1/6] feat: enhance configuration options across services for improved flexibility --- app/api/v1/admin.py | 93 +++++++++++++++------- app/api/v1/image.py | 21 +++-- app/core/auth.py | 7 +- app/services/grok/processors/processor.py | 43 +++++----- app/services/grok/services/assets.py | 95 +++++++++++++---------- app/services/grok/services/chat.py | 27 ++++--- app/services/grok/services/image.py | 39 +++++----- app/services/grok/services/media.py | 57 ++++++++------ app/services/grok/services/nsfw.py | 11 ++- app/services/grok/services/usage.py | 40 +++++----- app/services/grok/services/voice.py | 8 +- app/services/grok/utils/headers.py | 5 +- app/services/grok/utils/retry.py | 60 +++++--------- app/services/grok/utils/statsig.py | 5 +- app/services/token/manager.py | 45 ++++++----- data/config.toml | 3 + 16 files changed, 324 insertions(+), 235 deletions(-) diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 09346ee..7436409 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -15,12 +15,21 @@ from app.core.logger import logger from app.services.grok.services.voice import VoiceService from app.services.token import get_token_manager - -router = APIRouter() - +DEFAULT_STORAGE_TYPE = "" +DEFAULT_USAGE_MAX_TOKENS = 1000 +DEFAULT_USAGE_MAX_CONCURRENT = 25 +DEFAULT_USAGE_BATCH_SIZE = 50 +DEFAULT_NSFW_MAX_TOKENS = 1000 +DEFAULT_NSFW_MAX_CONCURRENT = 10 +DEFAULT_NSFW_BATCH_SIZE = 50 +DEFAULT_ASSETS_MAX_CONCURRENT = 25 +DEFAULT_ASSETS_BATCH_SIZE = 10 +DEFAULT_ASSETS_MAX_TOKENS = 1000 TEMPLATE_DIR = Path(__file__).parent.parent.parent / "static" +router = APIRouter() + async def render_template(filename: str): """渲染指定模板""" template_path = TEMPLATE_DIR / filename @@ -211,7 +220,7 @@ async def get_storage_info(): storage_type = os.getenv("SERVER_STORAGE_TYPE", "local").lower() logger.info(f"Storage type: {storage_type}") if not storage_type: - storage_type = str(get_config("storage.type", "")).lower() + storage_type = str(get_config("storage.type", DEFAULT_STORAGE_TYPE)).lower() if not storage_type: storage = get_storage() if isinstance(storage, LocalStorage): @@ -327,7 +336,7 @@ async def refresh_tokens_api(data: dict): unique_tokens = list(dict.fromkeys(tokens)) # 最大数量限制 - max_tokens = get_config("performance.usage_max_tokens", 1000) + max_tokens = get_config("performance.usage_max_tokens", DEFAULT_USAGE_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -343,8 +352,12 @@ async def refresh_tokens_api(data: dict): ) # 批量执行配置 - max_concurrent = get_config("performance.usage_max_concurrent", 25) - batch_size = get_config("performance.usage_batch_size", 50) + max_concurrent = get_config( + "performance.usage_max_concurrent", DEFAULT_USAGE_MAX_CONCURRENT + ) + batch_size = get_config( + "performance.usage_batch_size", DEFAULT_USAGE_BATCH_SIZE + ) async def _refresh_one(t): return await mgr.sync_usage( @@ -395,7 +408,7 @@ async def refresh_tokens_api_async(data: dict): unique_tokens = list(dict.fromkeys(tokens)) - max_tokens = get_config("performance.usage_max_tokens", 1000) + max_tokens = get_config("performance.usage_max_tokens", DEFAULT_USAGE_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -410,8 +423,10 @@ async def refresh_tokens_api_async(data: dict): f"Usage refresh: truncated from {original_count} to {max_tokens} tokens" ) - max_concurrent = get_config("performance.usage_max_concurrent", 25) - batch_size = get_config("performance.usage_batch_size", 50) + max_concurrent = get_config( + "performance.usage_max_concurrent", DEFAULT_USAGE_MAX_CONCURRENT + ) + batch_size = get_config("performance.usage_batch_size", DEFAULT_USAGE_BATCH_SIZE) task = create_task(len(unique_tokens)) @@ -515,7 +530,7 @@ async def enable_nsfw_api(data: dict): unique_tokens = list(dict.fromkeys(tokens)) # 限制最大数量（超出时截取前 N 个） - max_tokens = get_config("performance.nsfw_max_tokens", 1000) + max_tokens = get_config("performance.nsfw_max_tokens", DEFAULT_NSFW_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -531,8 +546,10 @@ async def enable_nsfw_api(data: dict): ) # 批量执行配置 - max_concurrent = get_config("performance.nsfw_max_concurrent", 10) - batch_size = get_config("performance.nsfw_batch_size", 50) + max_concurrent = get_config( + "performance.nsfw_max_concurrent", DEFAULT_NSFW_MAX_CONCURRENT + ) + batch_size = get_config("performance.nsfw_batch_size", DEFAULT_NSFW_BATCH_SIZE) # 定义 worker async def _enable(token: str): @@ -621,7 +638,7 @@ async def enable_nsfw_api_async(data: dict): unique_tokens = list(dict.fromkeys(tokens)) - max_tokens = get_config("performance.nsfw_max_tokens", 1000) + max_tokens = get_config("performance.nsfw_max_tokens", DEFAULT_NSFW_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -636,8 +653,10 @@ async def enable_nsfw_api_async(data: dict): f"NSFW enable: truncated from {original_count} to {max_tokens} tokens" ) - max_concurrent = get_config("performance.nsfw_max_concurrent", 10) - batch_size = get_config("performance.nsfw_batch_size", 50) + max_concurrent = get_config( + "performance.nsfw_max_concurrent", DEFAULT_NSFW_MAX_CONCURRENT + ) + batch_size = get_config("performance.nsfw_batch_size", DEFAULT_NSFW_BATCH_SIZE) task = create_task(len(unique_tokens)) @@ -772,8 +791,12 @@ async def get_cache_stats_api(request: Request): } online_details = [] account_map = {a["token"]: a for a in accounts} - max_concurrent = get_config("performance.assets_max_concurrent", 25) - batch_size = get_config("performance.assets_batch_size", 10) + max_concurrent = get_config( + "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT + ) + batch_size = get_config( + "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE + ) try: max_concurrent = int(max_concurrent) except Exception: @@ -785,7 +808,7 @@ async def get_cache_stats_api(request: Request): max_concurrent = max(1, max_concurrent) batch_size = max(1, batch_size) - max_tokens = get_config("performance.assets_max_tokens", 1000) + max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -1006,7 +1029,7 @@ async def load_online_cache_api_async(data: dict): selected_tokens = list(dict.fromkeys(selected_tokens)) - max_tokens = get_config("performance.assets_max_tokens", 1000) + max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -1018,8 +1041,12 @@ async def load_online_cache_api_async(data: dict): selected_tokens = selected_tokens[:max_tokens] truncated = True - max_concurrent = get_config("performance.assets_max_concurrent", 25) - batch_size = get_config("performance.assets_batch_size", 10) + max_concurrent = get_config( + "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT + ) + batch_size = get_config( + "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE + ) task = create_task(len(selected_tokens)) @@ -1193,7 +1220,9 @@ async def clear_online_cache_api(data: dict): token_list = list(dict.fromkeys(token_list)) # 最大数量限制 - max_tokens = get_config("performance.assets_max_tokens", 1000) + max_tokens = get_config( + "performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS + ) try: max_tokens = int(max_tokens) except Exception: @@ -1205,8 +1234,12 @@ async def clear_online_cache_api(data: dict): truncated = True results = {} - max_concurrent = get_config("performance.assets_max_concurrent", 25) - batch_size = get_config("performance.assets_batch_size", 10) + max_concurrent = get_config( + "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT + ) + batch_size = get_config( + "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE + ) try: max_concurrent = int(max_concurrent) except Exception: @@ -1281,7 +1314,7 @@ async def clear_online_cache_api_async(data: dict): token_list = list(dict.fromkeys(token_list)) - max_tokens = get_config("performance.assets_max_tokens", 1000) + max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS) try: max_tokens = int(max_tokens) except Exception: @@ -1292,8 +1325,12 @@ async def clear_online_cache_api_async(data: dict): token_list = token_list[:max_tokens] truncated = True - max_concurrent = get_config("performance.assets_max_concurrent", 25) - batch_size = get_config("performance.assets_batch_size", 10) + max_concurrent = get_config( + "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT + ) + batch_size = get_config( + "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE + ) task = create_task(len(token_list)) diff --git a/app/api/v1/image.py b/app/api/v1/image.py index 929bc3f..9d61d01 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -31,9 +31,14 @@ from app.core.config import get_config from app.core.logger import logger -router = APIRouter(tags=["Images"]) +DEFAULT_IMAGE_WS = False +DEFAULT_IMAGE_FORMAT = "url" +DEFAULT_IMAGE_WS_NSFW = True +DEFAULT_GROK_TEMPORARY = True +router = APIRouter(tags=["Images"]) + class ImageGenerationRequest(BaseModel): """图片生成请求 - OpenAI 兼容""" @@ -92,8 +97,8 @@ def _validate_common_request( if allow_ws_stream: # WS 流式仅支持 b64_json (base64 视为同义) if ( - request.stream - and get_config("grok.image_ws", False) + request.stream + and get_config("grok.image_ws", DEFAULT_IMAGE_WS) and request.response_format and request.response_format not in {"b64_json", "base64"} ): @@ -140,7 +145,7 @@ def validate_generation_request(request: ImageGenerationRequest): def resolve_response_format(response_format: Optional[str]) -> str: - fmt = response_format or get_config("app.image_format", "url") + fmt = response_format or get_config("app.image_format", DEFAULT_IMAGE_FORMAT) if isinstance(fmt, str): fmt = fmt.lower() if fmt in ("b64_json", "base64", "url"): @@ -313,13 +318,13 @@ async def create_image(request: ImageGenerationRequest): # 获取模型信息 model_info = ModelService.get(request.model) - use_ws = bool(get_config("grok.image_ws", False)) + use_ws = bool(get_config("grok.image_ws", DEFAULT_IMAGE_WS)) # 流式模式 if request.stream: if use_ws: aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("grok.image_ws_nsfw", True)) + enable_nsfw = bool(get_config("grok.image_ws_nsfw", DEFAULT_IMAGE_WS_NSFW)) upstream = image_service.stream( token=token, prompt=request.prompt, @@ -404,7 +409,7 @@ async def create_image(request: ImageGenerationRequest): usage_override = None if use_ws: aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("grok.image_ws_nsfw", True)) + enable_nsfw = bool(get_config("grok.image_ws_nsfw", DEFAULT_IMAGE_WS_NSFW)) all_images = [] seen = set() expected_per_call = 6 @@ -709,7 +714,7 @@ async def edit_image( ] = parent_post_id raw_payload = { - "temporary": bool(get_config("grok.temporary", True)), + "temporary": bool(get_config("grok.temporary", DEFAULT_GROK_TEMPORARY)), "modelName": model_info.grok_model, "message": edit_request.prompt, "enableImageGeneration": True, diff --git a/app/core/auth.py b/app/core/auth.py index e17804c..e6bb3c3 100644 --- a/app/core/auth.py +++ b/app/core/auth.py @@ -8,6 +8,9 @@ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from app.core.config import get_config +DEFAULT_API_KEY = "" +DEFAULT_APP_KEY = "grok2api" + # 定义 Bearer Scheme security = HTTPBearer( auto_error=False, @@ -22,7 +25,7 @@ def get_admin_api_key() -> str: 为空时表示不启用后台接口认证。 """ - api_key = get_config("app.api_key", "") + api_key = get_config("app.api_key", DEFAULT_API_KEY) return api_key or "" @@ -63,7 +66,7 @@ async def verify_app_key( app_key 必须配置，否则拒绝登录。 """ - app_key = get_config("app.app_key", "") + app_key = get_config("app.app_key", DEFAULT_APP_KEY) if not app_key: raise HTTPException( diff --git a/app/services/grok/processors/processor.py b/app/services/grok/processors/processor.py index ae3dc41..49ac9d1 100644 --- a/app/services/grok/processors/processor.py +++ b/app/services/grok/processors/processor.py @@ -11,7 +11,6 @@ import re from pathlib import Path import orjson from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar, Dict - from curl_cffi.requests.errors import RequestsError from app.core.config import get_config @@ -20,6 +19,15 @@ from app.core.exceptions import UpstreamException from app.services.grok.services.assets import DownloadService +DEFAULT_APP_URL = "" +DEFAULT_FILTER_TAGS = ["grok:render", "xaiartifact", "xai:tool_usage_card"] +DEFAULT_IMAGE_FORMAT = "url" +DEFAULT_THINKING = False +DEFAULT_STREAM_IDLE_TIMEOUT = 45.0 +DEFAULT_VIDEO_FORMAT = "html" +DEFAULT_VIDEO_IDLE_TIMEOUT = 90.0 + + def _is_http2_stream_error(e: Exception) -> bool: """检查是否为 HTTP/2 流错误""" err_str = str(e).lower() @@ -133,7 +141,7 @@ class BaseProcessor: self.model = model self.token = token self.created = int(time.time()) - self.app_url = get_config("app.app_url", "") + self.app_url = get_config("app.app_url", DEFAULT_APP_URL) self._dl_service: Optional[DownloadService] = None def _get_dl(self) -> DownloadService: @@ -204,14 +212,14 @@ class StreamProcessor(BaseProcessor): self.fingerprint: str = "" self.think_opened: bool = False self.role_sent: bool = False - self.filter_tags = get_config("grok.filter_tags", []) - self.image_format = get_config("app.image_format", "url") + self.filter_tags = get_config("grok.filter_tags", DEFAULT_FILTER_TAGS) + self.image_format = get_config("app.image_format", DEFAULT_IMAGE_FORMAT) # 用于过滤跨 token 的标签 self._tag_buffer: str = "" self._in_filter_tag: bool = False if think is None: - self.show_think = get_config("grok.thinking", False) + self.show_think = get_config("grok.thinking", DEFAULT_THINKING) else: self.show_think = think @@ -281,7 +289,7 @@ class StreamProcessor(BaseProcessor): ) -> AsyncGenerator[str, None]: """处理流式响应""" # 获取空闲超时配置 - idle_timeout = get_config("grok.stream_idle_timeout", 45.0) + idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -407,13 +415,10 @@ class StreamProcessor(BaseProcessor): class CollectProcessor(BaseProcessor): """非流式响应处理器""" - # 需要过滤的标签 - FILTER_TAGS = ["grok:render", "xaiartifact", "xai:tool_usage_card"] - def __init__(self, model: str, token: str = ""): super().__init__(model, token) - self.image_format = get_config("app.image_format", "url") - self.filter_tags = get_config("grok.filter_tags", self.FILTER_TAGS) + self.image_format = get_config("app.image_format", DEFAULT_IMAGE_FORMAT) + self.filter_tags = get_config("grok.filter_tags", DEFAULT_FILTER_TAGS) def _filter_content(self, content: str) -> str: """过滤内容中的特殊标签""" @@ -436,7 +441,7 @@ class CollectProcessor(BaseProcessor): fingerprint = "" content = "" # 获取空闲超时配置 - idle_timeout = get_config("grok.stream_idle_timeout", 45.0) + idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -552,10 +557,10 @@ class VideoStreamProcessor(BaseProcessor): self.response_id: Optional[str] = None self.think_opened: bool = False self.role_sent: bool = False - self.video_format = str(get_config("app.video_format", "html")).lower() + self.video_format = str(get_config("app.video_format", DEFAULT_VIDEO_FORMAT)).lower() if think is None: - self.show_think = get_config("grok.thinking", False) + self.show_think = get_config("grok.thinking", DEFAULT_THINKING) else: self.show_think = think @@ -575,7 +580,7 @@ class VideoStreamProcessor(BaseProcessor): ) -> AsyncGenerator[str, None]: """处理视频流式响应""" # 视频生成使用更长的空闲超时 - idle_timeout = get_config("grok.video_idle_timeout", 90.0) + idle_timeout = get_config("grok.video_idle_timeout", DEFAULT_VIDEO_IDLE_TIMEOUT) try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -684,7 +689,7 @@ class VideoCollectProcessor(BaseProcessor): def __init__(self, model: str, token: str = ""): super().__init__(model, token) - self.video_format = str(get_config("app.video_format", "html")).lower() + self.video_format = str(get_config("app.video_format", DEFAULT_VIDEO_FORMAT)).lower() def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" @@ -697,7 +702,7 @@ class VideoCollectProcessor(BaseProcessor): response_id = "" content = "" # 视频生成使用更长的空闲超时 - idle_timeout = get_config("grok.video_idle_timeout", 90.0) + idle_timeout = get_config("grok.video_idle_timeout", DEFAULT_VIDEO_IDLE_TIMEOUT) try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -807,7 +812,7 @@ class ImageStreamProcessor(BaseProcessor): """处理流式响应""" final_images = [] # 图片生成使用标准空闲超时 - idle_timeout = get_config("grok.stream_idle_timeout", 45.0) + idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -937,7 +942,7 @@ class ImageCollectProcessor(BaseProcessor): """处理并收集图片""" images = [] # 图片生成使用标准空闲超时 - idle_timeout = get_config("grok.stream_idle_timeout", 45.0) + idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) try: async for line in _with_idle_timeout(response, idle_timeout, self.model): diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index 15791f4..6e53634 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -28,51 +28,27 @@ from app.services.grok.utils.headers import apply_statsig, build_sso_cookie from app.services.token.service import TokenService -# ==================== 常量 ==================== - UPLOAD_API = "https://grok.com/rest/app-chat/upload-file" LIST_API = "https://grok.com/rest/assets" DELETE_API = "https://grok.com/rest/assets-metadata" DOWNLOAD_API = "https://assets.grok.com" LOCK_DIR = Path(__file__).parent.parent.parent.parent.parent / "data" / ".locks" -TIMEOUT = 120 +DEFAULT_TIMEOUT = 120 DEFAULT_BROWSER = "chrome136" DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" DEFAULT_MIME = "application/octet-stream" - -# 并发控制 +DEFAULT_BASE_PROXY_URL = "" +DEFAULT_ASSET_PROXY_URL = "" +DEFAULT_CACHE_ENABLE_AUTO_CLEAN = True +DEFAULT_CACHE_LIMIT_MB = 1024 +DEFAULT_APP_URL = "" DEFAULT_MAX_CONCURRENT = 25 DEFAULT_DELETE_BATCH_SIZE = 10 _ASSETS_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) _ASSETS_SEM_VALUE = DEFAULT_MAX_CONCURRENT -def _get_assets_semaphore() -> asyncio.Semaphore: - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - return _ASSETS_SEMAPHORE - - -def _get_delete_batch_size() -> int: - value = get_config( - "performance.assets_delete_batch_size", DEFAULT_DELETE_BATCH_SIZE - ) - try: - value = int(value) - except Exception: - value = DEFAULT_DELETE_BATCH_SIZE - return max(1, value) - - @asynccontextmanager async def _file_lock(name: str, timeout: int = 10): if fcntl is None: @@ -172,10 +148,10 @@ class BaseService: def __init__(self, proxy: str = None): self.proxy = ( proxy - or get_config("grok.asset_proxy_url") - or get_config("grok.base_proxy_url", "") + or get_config("grok.asset_proxy_url", DEFAULT_ASSET_PROXY_URL) + or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) ) - self.timeout = get_config("grok.timeout", TIMEOUT) + self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) self._session: Optional[AsyncSession] = None def _headers(self, token: str, referer: str = "https://grok.com/") -> dict: @@ -331,7 +307,17 @@ class UploadService(BaseService): ValidationException: 输入无效 UpstreamException: 上传失败 """ - async with _get_assets_semaphore(): + value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) + try: + value = int(value) + except Exception: + value = DEFAULT_MAX_CONCURRENT + value = max(1, value) + global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE + if value != _ASSETS_SEM_VALUE: + _ASSETS_SEM_VALUE = value + _ASSETS_SEMAPHORE = asyncio.Semaphore(value) + async with _ASSETS_SEMAPHORE: try: # 处理输入 if self.is_url(file_input): @@ -520,7 +506,17 @@ class DeleteService(BaseService): Raises: UpstreamException: 删除失败 """ - async with _get_assets_semaphore(): + value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) + try: + value = int(value) + except Exception: + value = DEFAULT_MAX_CONCURRENT + value = max(1, value) + global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE + if value != _ASSETS_SEM_VALUE: + _ASSETS_SEM_VALUE = value + _ASSETS_SEMAPHORE = asyncio.Semaphore(value) + async with _ASSETS_SEMAPHORE: try: headers = self._headers(token, referer="https://grok.com/files") url = f"{DELETE_API}/{asset_id}" @@ -578,7 +574,14 @@ class DeleteService(BaseService): return False return False - batch_size = _get_delete_batch_size() + batch_size = get_config( + "performance.assets_delete_batch_size", DEFAULT_DELETE_BATCH_SIZE + ) + try: + batch_size = int(batch_size) + except Exception: + batch_size = DEFAULT_DELETE_BATCH_SIZE + batch_size = max(1, batch_size) for i in range(0, len(assets), batch_size): batch = assets[i : i + batch_size] results = await asyncio.gather( @@ -631,7 +634,17 @@ class DownloadService(BaseService): Raises: UpstreamException: 下载失败 """ - async with _get_assets_semaphore(): + value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) + try: + value = int(value) + except Exception: + value = DEFAULT_MAX_CONCURRENT + value = max(1, value) + global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE + if value != _ASSETS_SEM_VALUE: + _ASSETS_SEM_VALUE = value + _ASSETS_SEMAPHORE = asyncio.Semaphore(value) + async with _ASSETS_SEMAPHORE: try: cache_path = self._cache_path(file_path, media_type) @@ -847,10 +860,12 @@ class DownloadService(BaseService): self._cleanup_running = True try: async with _file_lock("cache_cleanup", timeout=5): - if not get_config("cache.enable_auto_clean", True): + if not get_config( + "cache.enable_auto_clean", DEFAULT_CACHE_ENABLE_AUTO_CLEAN + ): return - limit_mb = get_config("cache.limit_mb", 1024) + limit_mb = get_config("cache.limit_mb", DEFAULT_CACHE_LIMIT_MB) # 统计总大小 total_size = 0 @@ -906,7 +921,7 @@ class DownloadService(BaseService): 如果配置了 app_url，则返回自托管 URL，否则返回 Grok 原始 URL """ - app_url = get_config("app.app_url", "") + app_url = get_config("app.app_url", DEFAULT_APP_URL) if not app_url: return f"{DOWNLOAD_API}{file_path if file_path.startswith('/') else '/' + file_path}" diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index 237cd38..f0c1b3b 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -25,9 +25,18 @@ from app.services.token import get_token_manager, EffortType CHAT_API = "https://grok.com/rest/app-chat/conversations/new" -TIMEOUT = 120 + +DEFAULT_TIMEOUT = 120 DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" +DEFAULT_USER_AGENT = ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/136.0.0.0 Safari/537.36" +) +DEFAULT_TEMPORARY = True +DEFAULT_DISABLE_MEMORY = True +DEFAULT_STREAM = True +DEFAULT_BASE_PROXY_URL = "" @dataclass @@ -210,7 +219,7 @@ class ChatRequestBuilder: file_attachments: 文件附件 ID 列表 image_attachments: 图片附件 ID 列表（合并到 fileAttachments） """ - temporary = get_config("grok.temporary", True) + temporary = get_config("grok.temporary", DEFAULT_TEMPORARY) merged_attachments: List[str] = [] if file_attachments: @@ -240,7 +249,7 @@ class ChatRequestBuilder: "modelConfigOverride": {"modelMap": {}}, "requestModelDetails": {"modelId": model}, }, - "disableMemory": get_config("grok.disable_memory", True), + "disableMemory": get_config("grok.disable_memory", DEFAULT_DISABLE_MEMORY), "forceSideBySide": False, "isAsyncChat": False, "disableSelfHarmShortCircuit": False, @@ -267,7 +276,7 @@ class GrokChatService: """Grok API 调用服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", "") + self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) async def chat( self, @@ -296,7 +305,7 @@ class GrokChatService: UpstreamException: 当 Grok API 返回错误且重试耗尽时 """ if stream is None: - stream = get_config("grok.stream", True) + stream = get_config("grok.stream", DEFAULT_STREAM) headers = ChatRequestBuilder.build_headers(token) if raw_payload is not None: @@ -310,7 +319,7 @@ class GrokChatService: image_attachments, ) proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None - timeout = get_config("grok.timeout", TIMEOUT) + timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) # 状态码提取器 def extract_status(e: Exception) -> int | None: @@ -455,7 +464,7 @@ class GrokChatService: stream = ( request.stream if request.stream is not None - else get_config("grok.stream", True) + else get_config("grok.stream", DEFAULT_STREAM) ) response = await self.chat( @@ -561,7 +570,7 @@ class ChatService: elif thinking == "disabled": think = False - is_stream = stream if stream is not None else get_config("grok.stream", True) + is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM) # 构造请求 chat_request = ChatRequest( diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 8d0ebbb..7ca9ade 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -19,18 +19,18 @@ from app.core.logger import logger from app.services.grok.utils.headers import build_sso_cookie -TIMEOUT = 120 -BLOCKED_SECONDS = 15 -FINAL_MIN_BYTES = 100000 -MEDIUM_MIN_BYTES = 30000 - WS_URL = "wss://grok.com/ws/imagine/listen" -USER_AGENT = ( - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " +DEFAULT_TIMEOUT = 120 +DEFAULT_BLOCKED_SECONDS = 15 +DEFAULT_FINAL_MIN_BYTES = 100000 +DEFAULT_MEDIUM_MIN_BYTES = 30000 +DEFAULT_USER_AGENT = ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/131.0.0.0 Safari/537.36" + "Chrome/136.0.0.0 Safari/537.36" ) +DEFAULT_BASE_PROXY_URL = "" class _BlockedError(Exception): @@ -47,14 +47,8 @@ class ImageService: def _get_ws_url(self) -> str: return WS_URL - def _get_timeout(self) -> float: - return float(get_config("grok.timeout", TIMEOUT)) - - def _get_blocked_seconds(self) -> float: - return float(get_config("grok.image_ws_blocked_seconds", BLOCKED_SECONDS)) - def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]: - proxy_url = get_config("grok.base_proxy_url", "") + proxy_url = get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) if not proxy_url: return aiohttp.TCPConnector(ssl=self._ssl_context), None @@ -68,10 +62,11 @@ class ImageService: def _get_ws_headers(self, token: str) -> Dict[str, str]: cookie = build_sso_cookie(token, include_rw=True) + user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) return { "Cookie": cookie, "Origin": "https://grok.com", - "User-Agent": USER_AGENT, + "User-Agent": user_agent, "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8", "Cache-Control": "no-cache", "Pragma": "no-cache", @@ -84,7 +79,9 @@ class ImageService: return None def _is_final_image(self, url: str, blob_size: int) -> bool: - min_bytes = int(get_config("grok.image_ws_final_min_bytes", FINAL_MIN_BYTES)) + min_bytes = int( + get_config("grok.image_ws_final_min_bytes", DEFAULT_FINAL_MIN_BYTES) + ) return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > min_bytes def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]: @@ -96,7 +93,7 @@ class ImageService: is_final = self._is_final_image(url, blob_size) if is_final: stage = "final" - elif blob_size > MEDIUM_MIN_BYTES: + elif blob_size > DEFAULT_MEDIUM_MIN_BYTES: stage = "medium" else: stage = "preview" @@ -165,8 +162,10 @@ class ImageService: request_id = str(uuid.uuid4()) headers = self._get_ws_headers(token) - timeout = self._get_timeout() - blocked_seconds = self._get_blocked_seconds() + timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT)) + blocked_seconds = float( + get_config("grok.image_ws_blocked_seconds", DEFAULT_BLOCKED_SECONDS) + ) try: connector, proxy = self._resolve_proxy() diff --git a/app/services/grok/services/media.py b/app/services/grok/services/media.py index 7b0eab7..0780957 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/media.py @@ -21,39 +21,30 @@ from app.services.token import get_token_manager, EffortType from app.services.grok.processors.processor import VideoStreamProcessor, VideoCollectProcessor from app.services.grok.utils.headers import apply_statsig, build_sso_cookie -# API 端点 + CREATE_POST_API = "https://grok.com/rest/media/post/create" CHAT_API = "https://grok.com/rest/app-chat/conversations/new" -# 常量 DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" -TIMEOUT = 300 +DEFAULT_USER_AGENT = ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/136.0.0.0 Safari/537.36" +) +DEFAULT_TIMEOUT = 300 DEFAULT_MAX_CONCURRENT = 50 +DEFAULT_BASE_PROXY_URL = "" +DEFAULT_STREAM = True _MEDIA_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) _MEDIA_SEM_VALUE = DEFAULT_MAX_CONCURRENT -def _get_media_semaphore() -> asyncio.Semaphore: - global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE - value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - if value != _MEDIA_SEM_VALUE: - _MEDIA_SEM_VALUE = value - _MEDIA_SEMAPHORE = asyncio.Semaphore(value) - return _MEDIA_SEMAPHORE - - class VideoService: """视频生成服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", "") - self.timeout = get_config("grok.timeout", TIMEOUT) + self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) def _build_headers( self, token: str, referer: str = "https://grok.com/imagine" @@ -234,7 +225,17 @@ class VideoService: Raises: UpstreamException: 连接失败时 """ - async with _get_media_semaphore(): + value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) + try: + value = int(value) + except Exception: + value = DEFAULT_MAX_CONCURRENT + value = max(1, value) + global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE + if value != _MEDIA_SEM_VALUE: + _MEDIA_SEM_VALUE = value + _MEDIA_SEMAPHORE = asyncio.Semaphore(value) + async with _MEDIA_SEMAPHORE: session = None try: # Step 1: 创建帖子 @@ -315,7 +316,17 @@ class VideoService: Returns: AsyncGenerator，流式传输 """ - async with _get_media_semaphore(): + value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) + try: + value = int(value) + except Exception: + value = DEFAULT_MAX_CONCURRENT + value = max(1, value) + global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE + if value != _MEDIA_SEM_VALUE: + _MEDIA_SEM_VALUE = value + _MEDIA_SEMAPHORE = asyncio.Semaphore(value) + async with _MEDIA_SEMAPHORE: session = None try: # Step 1: 创建帖子 @@ -463,7 +474,7 @@ class VideoService: elif thinking == "disabled": think = False - is_stream = stream if stream is not None else get_config("grok.stream", True) + is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM) # 提取内容 from app.services.grok.services.chat import MessageExtractor diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py index be797a4..88411be 100644 --- a/app/services/grok/services/nsfw.py +++ b/app/services/grok/services/nsfw.py @@ -23,9 +23,11 @@ from app.services.grok.utils.headers import build_sso_cookie NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" BIRTH_DATE_API = "https://grok.com/rest/auth/set-birth-date" + DEFAULT_BROWSER = "chrome136" DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36" -TIMEOUT = 30 +DEFAULT_TIMEOUT = 30 +DEFAULT_BASE_PROXY_URL = "" @dataclass @@ -43,7 +45,8 @@ class NSFWService: """NSFW 模式服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", "") + self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT)) @staticmethod def _random_birth_date() -> str: @@ -109,7 +112,7 @@ class NSFWService: BIRTH_DATE_API, json=payload, headers=headers, - timeout=TIMEOUT, + timeout=self.timeout, proxies=proxies, ) if response.status_code in (200, 204): @@ -145,7 +148,7 @@ class NSFWService: NSFW_API, data=payload, headers=headers, - timeout=TIMEOUT, + timeout=self.timeout, proxies=proxies, ) diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index a339962..dc2ff52 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -15,34 +15,26 @@ from app.services.grok.utils.retry import retry_on_status LIMITS_API = "https://grok.com/rest/rate-limits" + DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" -TIMEOUT = 10 +DEFAULT_USER_AGENT = ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/136.0.0.0 Safari/537.36" +) +DEFAULT_TIMEOUT = 10 DEFAULT_MAX_CONCURRENT = 25 +DEFAULT_BASE_PROXY_URL = "" _USAGE_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) _USAGE_SEM_VALUE = DEFAULT_MAX_CONCURRENT -def _get_usage_semaphore() -> asyncio.Semaphore: - global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE - value = get_config("performance.usage_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - if value != _USAGE_SEM_VALUE: - _USAGE_SEM_VALUE = value - _USAGE_SEMAPHORE = asyncio.Semaphore(value) - return _USAGE_SEMAPHORE - - class UsageService: """用量查询服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", "") - self.timeout = get_config("grok.timeout", TIMEOUT) + self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) def _build_headers(self, token: str) -> dict: """构建请求头""" @@ -93,7 +85,17 @@ class UsageService: Raises: UpstreamException: 当获取失败且重试耗尽时 """ - async with _get_usage_semaphore(): + value = get_config("performance.usage_max_concurrent", DEFAULT_MAX_CONCURRENT) + try: + value = int(value) + except Exception: + value = DEFAULT_MAX_CONCURRENT + value = max(1, value) + global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE + if value != _USAGE_SEM_VALUE: + _USAGE_SEM_VALUE = value + _USAGE_SEMAPHORE = asyncio.Semaphore(value) + async with _USAGE_SEMAPHORE: # 定义状态码提取器 def extract_status(e: Exception) -> int | None: if isinstance(e, UpstreamException) and e.details: diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 8b87991..17a4de7 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -14,15 +14,16 @@ from app.services.grok.utils.headers import apply_statsig, build_sso_cookie LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens" -TIMEOUT = 30 +DEFAULT_TIMEOUT = 30 DEFAULT_BROWSER = "chrome136" +DEFAULT_BASE_PROXY_URL = "" class VoiceService: """Voice Mode Service (LiveKit)""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", "") + self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) async def get_token( self, @@ -46,12 +47,13 @@ class VoiceService: try: browser = get_config("grok.browser", DEFAULT_BROWSER) + timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) async with AsyncSession(impersonate=browser) as session: response = await session.post( LIVEKIT_TOKEN_API, headers=headers, data=orjson.dumps(payload), - timeout=TIMEOUT, + timeout=timeout, proxies=proxies, ) diff --git a/app/services/grok/utils/headers.py b/app/services/grok/utils/headers.py index 71dc0e5..44c94ae 100644 --- a/app/services/grok/utils/headers.py +++ b/app/services/grok/utils/headers.py @@ -11,13 +11,16 @@ from app.core.config import get_config from app.services.grok.utils.statsig import StatsigService +DEFAULT_CF_CLEARANCE = "" + + def _normalize_token(token: str) -> str: return token[4:] if token.startswith("sso=") else token def build_sso_cookie(token: str, include_rw: bool = False) -> str: token = _normalize_token(token) - cf = get_config("grok.cf_clearance", "") + cf = get_config("grok.cf_clearance", DEFAULT_CF_CLEARANCE) cookie = f"sso={token}" if include_rw: cookie = f"{cookie}; sso-rw={token}" diff --git a/app/services/grok/utils/retry.py b/app/services/grok/utils/retry.py index 69299da..9a69e87 100644 --- a/app/services/grok/utils/retry.py +++ b/app/services/grok/utils/retry.py @@ -10,7 +10,7 @@ Grok API 重试工具 import asyncio import random -from typing import Callable, Any, Optional, List +from typing import Callable, Any, Optional from functools import wraps from app.core.logger import logger @@ -18,38 +18,12 @@ from app.core.config import get_config from app.core.exceptions import UpstreamException -class RetryConfig: - """重试配置""" - - @staticmethod - def get_max_retry() -> int: - """获取最大重试次数""" - return get_config("grok.max_retry", 3) - - @staticmethod - def get_retry_codes() -> List[int]: - """获取可重试的状态码""" - return get_config("grok.retry_status_codes", [401, 429, 403]) - - @staticmethod - def get_backoff_base() -> float: - """获取退避基础时间(秒)""" - return get_config("grok.retry_backoff_base", 0.5) - - @staticmethod - def get_backoff_factor() -> float: - """获取退避因子""" - return get_config("grok.retry_backoff_factor", 2.0) - - @staticmethod - def get_backoff_max() -> float: - """获取最大退避时间(秒)""" - return get_config("grok.retry_backoff_max", 30.0) - - @staticmethod - def get_retry_budget() -> float: - """获取重试总预算时间(秒)""" - return get_config("grok.retry_budget", 90.0) +DEFAULT_MAX_RETRY = 3 +DEFAULT_RETRY_STATUS_CODES = [401, 429, 403] +DEFAULT_RETRY_BACKOFF_BASE = 0.5 +DEFAULT_RETRY_BACKOFF_FACTOR = 2.0 +DEFAULT_RETRY_BACKOFF_MAX = 30.0 +DEFAULT_RETRY_BUDGET = 90.0 class RetryContext: @@ -57,17 +31,25 @@ class RetryContext: def __init__(self): self.attempt = 0 - self.max_retry = RetryConfig.get_max_retry() - self.retry_codes = RetryConfig.get_retry_codes() + self.max_retry = int(get_config("grok.max_retry", DEFAULT_MAX_RETRY)) + self.retry_codes = get_config( + "grok.retry_status_codes", DEFAULT_RETRY_STATUS_CODES + ) self.last_error = None self.last_status = None self.total_delay = 0.0 - self.retry_budget = RetryConfig.get_retry_budget() + self.retry_budget = float(get_config("grok.retry_budget", DEFAULT_RETRY_BUDGET)) # 退避参数 - self.backoff_base = RetryConfig.get_backoff_base() - self.backoff_factor = RetryConfig.get_backoff_factor() - self.backoff_max = RetryConfig.get_backoff_max() + self.backoff_base = float( + get_config("grok.retry_backoff_base", DEFAULT_RETRY_BACKOFF_BASE) + ) + self.backoff_factor = float( + get_config("grok.retry_backoff_factor", DEFAULT_RETRY_BACKOFF_FACTOR) + ) + self.backoff_max = float( + get_config("grok.retry_backoff_max", DEFAULT_RETRY_BACKOFF_MAX) + ) # decorrelated jitter 状态 self._last_delay = self.backoff_base diff --git a/app/services/grok/utils/statsig.py b/app/services/grok/utils/statsig.py index 7dd1b80..d5813b6 100644 --- a/app/services/grok/utils/statsig.py +++ b/app/services/grok/utils/statsig.py @@ -9,6 +9,9 @@ import string from app.core.config import get_config +DEFAULT_DYNAMIC_STATSIG = True + + class StatsigService: """Statsig ID 生成服务""" @@ -31,7 +34,7 @@ class StatsigService: Base64 编码的 ID """ # 读取配置 - dynamic = get_config("grok.dynamic_statsig", True) + dynamic = get_config("grok.dynamic_statsig", DEFAULT_DYNAMIC_STATSIG) if not dynamic: return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" diff --git a/app/services/token/manager.py b/app/services/token/manager.py index c7f3da7..8ff5e74 100644 --- a/app/services/token/manager.py +++ b/app/services/token/manager.py @@ -18,9 +18,13 @@ from app.core.storage import get_storage from app.core.config import get_config from app.services.token.pool import TokenPool -# 批量刷新配置 -REFRESH_BATCH_SIZE = 10 -REFRESH_CONCURRENCY = 5 + +DEFAULT_REFRESH_BATCH_SIZE = 10 +DEFAULT_REFRESH_CONCURRENCY = 5 +DEFAULT_SUPER_REFRESH_INTERVAL_HOURS = 2 +DEFAULT_REFRESH_INTERVAL_HOURS = 8 +DEFAULT_RELOAD_INTERVAL_SEC = 30 +DEFAULT_SAVE_DELAY_MS = 500 SUPER_POOL_NAME = "ssoSuper" BASIC_POOL_NAME = "ssoBasic" @@ -32,12 +36,6 @@ def _default_quota_for_pool(pool_name: str) -> int: return BASIC__DEFAULT_QUOTA -def _refresh_interval_hours_for_pool(pool_name: str) -> float: - if pool_name == SUPER_POOL_NAME: - return get_config("token.super_refresh_interval_hours", 2) - return get_config("token.refresh_interval_hours", 8) - - class TokenManager: """管理 Token 的增删改查和配额同步""" @@ -50,7 +48,7 @@ class TokenManager: self._save_lock = asyncio.Lock() self._dirty = False self._save_task: Optional[asyncio.Task] = None - self._save_delay = 0.5 + self._save_delay = DEFAULT_SAVE_DELAY_MS / 1000.0 self._last_reload_at = 0.0 @classmethod @@ -131,11 +129,11 @@ class TokenManager: async def reload_if_stale(self): """在多 worker 场景下保持短周期一致性""" - interval = get_config("token.reload_interval_sec", 30) + interval = get_config("token.reload_interval_sec", DEFAULT_RELOAD_INTERVAL_SEC) try: interval = float(interval) except Exception: - interval = 30.0 + interval = float(DEFAULT_RELOAD_INTERVAL_SEC) if interval <= 0: return if time.monotonic() - self._last_reload_at < interval: @@ -158,11 +156,11 @@ class TokenManager: def _schedule_save(self): """合并高频保存请求，减少写入开销""" - delay_ms = get_config("token.save_delay_ms", 500) + delay_ms = get_config("token.save_delay_ms", DEFAULT_SAVE_DELAY_MS) try: delay_ms = float(delay_ms) except Exception: - delay_ms = 500 + delay_ms = float(DEFAULT_SAVE_DELAY_MS) self._save_delay = max(0.0, delay_ms / 1000.0) self._dirty = True if self._save_delay == 0: @@ -523,7 +521,16 @@ class TokenManager: # 收集需要刷新的 token to_refresh: List[TokenInfo] = [] for pool in self.pools.values(): - interval_hours = _refresh_interval_hours_for_pool(pool.name) + if pool.name == SUPER_POOL_NAME: + interval_hours = get_config( + "token.super_refresh_interval_hours", + DEFAULT_SUPER_REFRESH_INTERVAL_HOURS, + ) + else: + interval_hours = get_config( + "token.refresh_interval_hours", + DEFAULT_REFRESH_INTERVAL_HOURS, + ) for token in pool: if token.need_refresh(interval_hours): to_refresh.append(token) @@ -535,7 +542,7 @@ class TokenManager: logger.info(f"Refresh check: found {len(to_refresh)} cooling tokens to refresh") # 批量并发刷新 - semaphore = asyncio.Semaphore(REFRESH_CONCURRENCY) + semaphore = asyncio.Semaphore(DEFAULT_REFRESH_CONCURRENCY) usage_service = UsageService() refreshed = 0 recovered = 0 @@ -606,15 +613,15 @@ class TokenManager: return {"recovered": False, "expired": False} # 批量处理 - for i in range(0, len(to_refresh), REFRESH_BATCH_SIZE): - batch = to_refresh[i : i + REFRESH_BATCH_SIZE] + for i in range(0, len(to_refresh), DEFAULT_REFRESH_BATCH_SIZE): + batch = to_refresh[i : i + DEFAULT_REFRESH_BATCH_SIZE] results = await asyncio.gather(*[_refresh_one(t) for t in batch]) refreshed += len(batch) recovered += sum(r["recovered"] for r in results) expired += sum(r["expired"] for r in results) # 批次间延迟 - if i + REFRESH_BATCH_SIZE < len(to_refresh): + if i + DEFAULT_REFRESH_BATCH_SIZE < len(to_refresh): await asyncio.sleep(1) await self._save() diff --git a/data/config.toml b/data/config.toml index 97cd4f5..3db391c 100644 --- a/data/config.toml +++ b/data/config.toml @@ -8,6 +8,9 @@ timeout = 120 base_proxy_url = "" asset_proxy_url = "" cf_clearance = "" +browser = "chrome136" +user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" +disable_memory = true max_retry = 3 retry_status_codes = [401,429,403] image_ws = true From 2e201b5ebfb545b665744bfef6140b5918412ad7 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sat, 7 Feb 2026 12:40:42 +0800 Subject: [PATCH 2/6] feat: enhance configuration management and service initialization --- app/api/v1/image.py | 2 +- app/core/config.py | 16 +- app/services/grok/defaults.py | 71 + app/services/grok/processors/__init__.py | 22 + app/services/grok/processors/base.py | 159 +++ .../grok/processors/chat_processors.py | 326 +++++ .../grok/processors/image_processors.py | 211 +++ .../grok/processors/image_ws_processors.py | 231 +++ app/services/grok/processors/processor.py | 1262 ----------------- .../grok/processors/video_processors.py | 236 +++ app/services/grok/services/assets.py | 996 +++++-------- app/services/grok/services/chat.py | 424 ++---- app/services/grok/services/image.py | 155 +- app/services/grok/services/media.py | 482 ++----- app/services/grok/services/nsfw.py | 63 +- app/services/grok/services/usage.py | 27 +- app/services/grok/services/voice.py | 11 +- app/services/grok/utils/headers.py | 5 +- app/services/grok/utils/retry.py | 28 +- app/services/grok/utils/statsig.py | 6 +- app/services/grok/utils/stream.py | 43 + data/config.toml | 51 +- main.py | 14 +- tests/test_model.py | 2 +- 24 files changed, 2014 insertions(+), 2829 deletions(-) create mode 100644 app/services/grok/defaults.py create mode 100644 app/services/grok/processors/base.py create mode 100644 app/services/grok/processors/chat_processors.py create mode 100644 app/services/grok/processors/image_processors.py create mode 100644 app/services/grok/processors/image_ws_processors.py delete mode 100644 app/services/grok/processors/processor.py create mode 100644 app/services/grok/processors/video_processors.py create mode 100644 app/services/grok/utils/stream.py diff --git a/app/api/v1/image.py b/app/api/v1/image.py index 9d61d01..20c6545 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -19,7 +19,7 @@ from app.services.grok.services.image import image_service from app.services.grok.services.assets import UploadService from app.services.grok.services.media import VideoService from app.services.grok.models.model import ModelService -from app.services.grok.processors.processor import ( +from app.services.grok.processors import ( ImageStreamProcessor, ImageCollectProcessor, ImageWSStreamProcessor, diff --git a/app/core/config.py b/app/core/config.py index 641b1cb..d9d3e8c 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -53,12 +53,19 @@ class Config: def __init__(self): self._config = {} self._defaults = {} + self._code_defaults = {} self._defaults_loaded = False + def register_defaults(self, defaults: Dict[str, Any]): + """注册代码中定义的默认值""" + self._code_defaults = _deep_merge(self._code_defaults, defaults) + def _ensure_defaults(self): if self._defaults_loaded: return - self._defaults = _load_defaults() + file_defaults = _load_defaults() + # 合并文件默认值和代码默认值（代码默认值优先级更低） + self._defaults = _deep_merge(self._code_defaults, file_defaults) self._defaults_loaded = True async def load(self): @@ -140,4 +147,9 @@ def get_config(key: str, default: Any = None) -> Any: return config.get(key, default) -__all__ = ["Config", "config", "get_config"] +def register_defaults(defaults: Dict[str, Any]): + """注册默认配置""" + config.register_defaults(defaults) + + +__all__ = ["Config", "config", "get_config", "register_defaults"] diff --git a/app/services/grok/defaults.py b/app/services/grok/defaults.py new file mode 100644 index 0000000..78ade9a --- /dev/null +++ b/app/services/grok/defaults.py @@ -0,0 +1,71 @@ +""" +Grok 服务默认配置 + +此文件定义所有 Grok 相关服务的默认值，会在应用启动时注册到配置系统中。 +""" + +# Grok 服务默认配置 +GROK_DEFAULTS = { + "grok": { + # 网络配置 + "browser": "chrome136", + "timeout": 120, + "user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36", + "base_proxy_url": "", + "asset_proxy_url": "", + "cf_clearance": "", + + # 对话配置 + "temporary": True, + "disable_memory": True, + "stream": True, + "thinking": False, + "filter_tags": ["grok:render", "xaiartifact", "xai:tool_usage_card"], + + # 重试配置 + "max_retry": 3, + "retry_status_codes": [401, 429, 403], + "retry_backoff_base": 0.5, + "retry_backoff_factor": 2.0, + "retry_backoff_max": 30.0, + "retry_budget": 90.0, + + # 超时配置 + "stream_idle_timeout": 45.0, + "video_idle_timeout": 90.0, + + # 图片配置 + "image_ws_blocked_seconds": 15, + "image_ws_final_min_bytes": 100000, + "image_ws_medium_min_bytes": 30000, + + # Statsig + "dynamic_statsig": True, + }, + + "app": { + "app_url": "", + "image_format": "url", + "video_format": "html", + }, + + "performance": { + "assets_max_concurrent": 25, + "assets_delete_batch_size": 10, + "media_max_concurrent": 50, + "usage_max_concurrent": 25, + }, + + "cache": { + "enable_auto_clean": True, + "limit_mb": 1024, + }, +} + + +def get_grok_defaults(): + """获取 Grok 默认配置""" + return GROK_DEFAULTS + + +__all__ = ["GROK_DEFAULTS", "get_grok_defaults"] diff --git a/app/services/grok/processors/__init__.py b/app/services/grok/processors/__init__.py index e69de29..04773f6 100644 --- a/app/services/grok/processors/__init__.py +++ b/app/services/grok/processors/__init__.py @@ -0,0 +1,22 @@ +""" +OpenAI 响应格式处理器 +""" + +from .base import BaseProcessor, StreamIdleTimeoutError +from .chat_processors import StreamProcessor, CollectProcessor +from .video_processors import VideoStreamProcessor, VideoCollectProcessor +from .image_processors import ImageStreamProcessor, ImageCollectProcessor +from .image_ws_processors import ImageWSStreamProcessor, ImageWSCollectProcessor + +__all__ = [ + "BaseProcessor", + "StreamIdleTimeoutError", + "StreamProcessor", + "CollectProcessor", + "VideoStreamProcessor", + "VideoCollectProcessor", + "ImageStreamProcessor", + "ImageCollectProcessor", + "ImageWSStreamProcessor", + "ImageWSCollectProcessor", +] diff --git a/app/services/grok/processors/base.py b/app/services/grok/processors/base.py new file mode 100644 index 0000000..bfda0d2 --- /dev/null +++ b/app/services/grok/processors/base.py @@ -0,0 +1,159 @@ +""" +响应处理器基类和通用工具 +""" + +import asyncio +import time +from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar + +from app.core.config import get_config +from app.core.logger import logger +from app.services.grok.services.assets import DownloadService + + +ASSET_URL = "https://assets.grok.com/" + +T = TypeVar("T") + + +def _is_http2_stream_error(e: Exception) -> bool: + """检查是否为 HTTP/2 流错误""" + err_str = str(e).lower() + return "http/2" in err_str or "curl: (92)" in err_str or "stream" in err_str + + +def _normalize_stream_line(line: Any) -> Optional[str]: + """规范化流式响应行，兼容 SSE data 前缀与空行""" + if line is None: + return None + if isinstance(line, (bytes, bytearray)): + text = line.decode("utf-8", errors="ignore") + else: + text = str(line) + text = text.strip() + if not text: + return None + if text.startswith("data:"): + text = text[5:].strip() + if text == "[DONE]": + return None + return text + + +def _collect_image_urls(obj: Any) -> List[str]: + """递归收集响应中的图片 URL""" + urls: List[str] = [] + seen = set() + + def add(url: str): + if not url or url in seen: + return + seen.add(url) + urls.append(url) + + def walk(value: Any): + if isinstance(value, dict): + for key, item in value.items(): + if key in {"generatedImageUrls", "imageUrls", "imageURLs"}: + if isinstance(item, list): + for url in item: + if isinstance(url, str): + add(url) + elif isinstance(item, str): + add(item) + continue + walk(item) + elif isinstance(value, list): + for item in value: + walk(item) + + walk(obj) + return urls + + +class StreamIdleTimeoutError(Exception): + """流空闲超时错误""" + + def __init__(self, idle_seconds: float): + self.idle_seconds = idle_seconds + super().__init__(f"Stream idle timeout after {idle_seconds}s") + + +async def _with_idle_timeout( + iterable: AsyncIterable[T], idle_timeout: float, model: str = "" +) -> AsyncGenerator[T, None]: + """ + 包装异步迭代器，添加空闲超时检测 + + Args: + iterable: 原始异步迭代器 + idle_timeout: 空闲超时时间(秒)，0 表示禁用 + model: 模型名称(用于日志) + """ + if idle_timeout <= 0: + async for item in iterable: + yield item + return + + iterator = iterable.__aiter__() + while True: + try: + item = await asyncio.wait_for(iterator.__anext__(), timeout=idle_timeout) + yield item + except asyncio.TimeoutError: + logger.warning( + f"Stream idle timeout after {idle_timeout}s", + extra={"model": model, "idle_timeout": idle_timeout}, + ) + raise StreamIdleTimeoutError(idle_timeout) + except StopAsyncIteration: + break + + +class BaseProcessor: + """基础处理器""" + + def __init__(self, model: str, token: str = ""): + self.model = model + self.token = token + self.created = int(time.time()) + self.app_url = get_config("app.app_url") + self._dl_service: Optional[DownloadService] = None + + def _get_dl(self) -> DownloadService: + """获取下载服务实例（复用）""" + if self._dl_service is None: + self._dl_service = DownloadService() + return self._dl_service + + async def close(self): + """释放下载服务资源""" + if self._dl_service: + await self._dl_service.close() + self._dl_service = None + + async def process_url(self, path: str, media_type: str = "image") -> str: + """处理资产 URL""" + if path.startswith("http"): + from urllib.parse import urlparse + path = urlparse(path).path + + if not path.startswith("/"): + path = f"/{path}" + + if self.app_url: + dl_service = self._get_dl() + await dl_service.download(path, self.token, media_type) + return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}" + else: + return f"{ASSET_URL.rstrip('/')}{path}" + + +__all__ = [ + "BaseProcessor", + "StreamIdleTimeoutError", + "_with_idle_timeout", + "_normalize_stream_line", + "_collect_image_urls", + "_is_http2_stream_error", +] diff --git a/app/services/grok/processors/chat_processors.py b/app/services/grok/processors/chat_processors.py new file mode 100644 index 0000000..a7bf5d4 --- /dev/null +++ b/app/services/grok/processors/chat_processors.py @@ -0,0 +1,326 @@ +""" +聊天响应处理器 +""" + +import asyncio +import uuid +import re +from typing import Any, AsyncGenerator, AsyncIterable + +import orjson +from curl_cffi.requests.errors import RequestsError + +from app.core.config import get_config +from app.core.logger import logger +from app.core.exceptions import UpstreamException +from .base import ( + BaseProcessor, + StreamIdleTimeoutError, + _with_idle_timeout, + _normalize_stream_line, + _collect_image_urls, + _is_http2_stream_error, +) + + +class StreamProcessor(BaseProcessor): + """流式响应处理器""" + + def __init__(self, model: str, token: str = "", think: bool = None): + super().__init__(model, token) + self.response_id: str = None + self.fingerprint: str = "" + self.think_opened: bool = False + self.role_sent: bool = False + self.filter_tags = get_config("grok.filter_tags") + self.image_format = get_config("app.image_format") + self._tag_buffer: str = "" + self._in_filter_tag: bool = False + + if think is None: + self.show_think = get_config("grok.thinking") + else: + self.show_think = think + + def _filter_token(self, token: str) -> str: + """过滤 token 中的特殊标签（如 ...），支持跨 token 的标签过滤""" + if not self.filter_tags: + return token + + result = [] + i = 0 + while i < len(token): + char = token[i] + + if self._in_filter_tag: + self._tag_buffer += char + if char == ">": + if "/>" in self._tag_buffer: + self._in_filter_tag = False + self._tag_buffer = "" + else: + for tag in self.filter_tags: + if f"" in self._tag_buffer: + self._in_filter_tag = False + self._tag_buffer = "" + break + i += 1 + continue + + if char == "<": + remaining = token[i:] + tag_started = False + for tag in self.filter_tags: + if remaining.startswith(f"<{tag}"): + tag_started = True + break + if len(remaining) < len(tag) + 1: + for j in range(1, len(remaining) + 1): + if f"<{tag}".startswith(remaining[:j]): + tag_started = True + break + + if tag_started: + self._in_filter_tag = True + self._tag_buffer = char + i += 1 + continue + + result.append(char) + i += 1 + + return "".join(result) + + def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: + """构建 SSE 响应""" + delta = {} + if role: + delta["role"] = role + delta["content"] = "" + elif content: + delta["content"] = content + + chunk = { + "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion.chunk", + "created": self.created, + "model": self.model, + "system_fingerprint": self.fingerprint, + "choices": [ + {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} + ], + } + return f"data: {orjson.dumps(chunk).decode()}\n\n" + + async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + """处理流式响应""" + idle_timeout = get_config("grok.stream_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + if (llm := resp.get("llmInfo")) and not self.fingerprint: + self.fingerprint = llm.get("modelHash", "") + if rid := resp.get("responseId"): + self.response_id = rid + + if not self.role_sent: + yield self._sse(role="assistant") + self.role_sent = True + + # 图像生成进度 + if img := resp.get("streamingImageGenerationResponse"): + if self.show_think: + if not self.think_opened: + yield self._sse("\n") + self.think_opened = True + idx = img.get("imageIndex", 0) + 1 + progress = img.get("progress", 0) + yield self._sse(f"正在生成第{idx}张图片中，当前进度{progress}%\n") + continue + + # modelResponse + if mr := resp.get("modelResponse"): + if self.think_opened and self.show_think: + if msg := mr.get("message"): + yield self._sse(msg + "\n") + yield self._sse("\n") + self.think_opened = False + + # 处理生成的图片 + for url in _collect_image_urls(mr): + parts = url.split("/") + img_id = parts[-2] if len(parts) >= 2 else "image" + + if self.image_format == "base64": + dl_service = self._get_dl() + base64_data = await dl_service.to_base64(url, self.token, "image") + if base64_data: + yield self._sse(f"![{img_id}]({base64_data})\n") + else: + final_url = await self.process_url(url, "image") + yield self._sse(f"![{img_id}]({final_url})\n") + else: + final_url = await self.process_url(url, "image") + yield self._sse(f"![{img_id}]({final_url})\n") + + if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + self.fingerprint = meta["llm_info"]["modelHash"] + continue + + # 普通 token + if (token := resp.get("token")) is not None: + if token: + filtered = self._filter_token(token) + if filtered: + yield self._sse(filtered) + + if self.think_opened: + yield self._sse("\n") + yield self._sse(finish="stop") + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + logger.debug("Stream cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + raise UpstreamException( + message=f"Stream idle timeout after {e.idle_seconds}s", + status_code=504, + details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}, + ) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error: {e}", extra={"model": self.model}) + raise UpstreamException( + message="Upstream connection closed unexpectedly", + status_code=502, + details={"error": str(e), "type": "http2_stream_error"}, + ) + logger.error(f"Stream request error: {e}", extra={"model": self.model}) + raise UpstreamException( + message=f"Upstream request failed: {e}", + status_code=502, + details={"error": str(e)}, + ) + except Exception as e: + logger.error(f"Stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + raise + finally: + await self.close() + + +class CollectProcessor(BaseProcessor): + """非流式响应处理器""" + + def __init__(self, model: str, token: str = ""): + super().__init__(model, token) + self.image_format = get_config("app.image_format") + self.filter_tags = get_config("grok.filter_tags") + + def _filter_content(self, content: str) -> str: + """过滤内容中的特殊标签""" + if not content or not self.filter_tags: + return content + + result = content + for tag in self.filter_tags: + pattern = rf"<{re.escape(tag)}[^>]*>.*?|<{re.escape(tag)}[^>]*/>" + result = re.sub(pattern, "", result, flags=re.DOTALL) + + return result + + async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: + """处理并收集完整响应""" + response_id = "" + fingerprint = "" + content = "" + idle_timeout = get_config("grok.stream_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + if (llm := resp.get("llmInfo")) and not fingerprint: + fingerprint = llm.get("modelHash", "") + + if mr := resp.get("modelResponse"): + response_id = mr.get("responseId", "") + content = mr.get("message", "") + + if urls := _collect_image_urls(mr): + content += "\n" + for url in urls: + parts = url.split("/") + img_id = parts[-2] if len(parts) >= 2 else "image" + + if self.image_format == "base64": + dl_service = self._get_dl() + base64_data = await dl_service.to_base64(url, self.token, "image") + if base64_data: + content += f"![{img_id}]({base64_data})\n" + else: + final_url = await self.process_url(url, "image") + content += f"![{img_id}]({final_url})\n" + else: + final_url = await self.process_url(url, "image") + content += f"![{img_id}]({final_url})\n" + + if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + fingerprint = meta["llm_info"]["modelHash"] + + except asyncio.CancelledError: + logger.debug("Collect cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model}) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in collect: {e}", extra={"model": self.model}) + else: + logger.error(f"Collect request error: {e}", extra={"model": self.model}) + except Exception as e: + logger.error(f"Collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + finally: + await self.close() + + content = self._filter_content(content) + + return { + "id": response_id, + "object": "chat.completion", + "created": self.created, + "model": self.model, + "system_fingerprint": fingerprint, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content, "refusal": None, "annotations": []}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "prompt_tokens_details": {"cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, + "completion_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0}, + }, + } + + +__all__ = ["StreamProcessor", "CollectProcessor"] diff --git a/app/services/grok/processors/image_processors.py b/app/services/grok/processors/image_processors.py new file mode 100644 index 0000000..8550108 --- /dev/null +++ b/app/services/grok/processors/image_processors.py @@ -0,0 +1,211 @@ +""" +图片生成响应处理器（HTTP） +""" + +import asyncio +import random +from typing import Any, AsyncGenerator, AsyncIterable, List + +import orjson +from curl_cffi.requests.errors import RequestsError + +from app.core.config import get_config +from app.core.logger import logger +from app.core.exceptions import UpstreamException +from .base import ( + BaseProcessor, + StreamIdleTimeoutError, + _with_idle_timeout, + _normalize_stream_line, + _collect_image_urls, + _is_http2_stream_error, +) + + +class ImageStreamProcessor(BaseProcessor): + """图片生成流式响应处理器""" + + def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json"): + super().__init__(model, token) + self.partial_index = 0 + self.n = n + self.target_index = random.randint(0, 1) if n == 1 else None + self.response_format = response_format + if response_format == "url": + self.response_field = "url" + elif response_format == "base64": + self.response_field = "base64" + else: + self.response_field = "b64_json" + + def _sse(self, event: str, data: dict) -> str: + """构建 SSE 响应""" + return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" + + async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + """处理流式响应""" + final_images = [] + idle_timeout = get_config("grok.stream_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + # 图片生成进度 + if img := resp.get("streamingImageGenerationResponse"): + image_index = img.get("imageIndex", 0) + progress = img.get("progress", 0) + + if self.n == 1 and image_index != self.target_index: + continue + + out_index = 0 if self.n == 1 else image_index + + yield self._sse( + "image_generation.partial_image", + { + "type": "image_generation.partial_image", + self.response_field: "", + "index": out_index, + "progress": progress, + }, + ) + continue + + # modelResponse + if mr := resp.get("modelResponse"): + if urls := _collect_image_urls(mr): + for url in urls: + if self.response_format == "url": + processed = await self.process_url(url, "image") + if processed: + final_images.append(processed) + continue + dl_service = self._get_dl() + base64_data = await dl_service.to_base64(url, self.token, "image") + if base64_data: + if "," in base64_data: + b64 = base64_data.split(",", 1)[1] + else: + b64 = base64_data + final_images.append(b64) + continue + + for index, b64 in enumerate(final_images): + if self.n == 1: + if index != self.target_index: + continue + out_index = 0 + else: + out_index = index + + yield self._sse( + "image_generation.completed", + { + "type": "image_generation.completed", + self.response_field: b64, + "index": out_index, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "output_tokens": 0, + "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + }, + }, + ) + except asyncio.CancelledError: + logger.debug("Image stream cancelled by client") + except StreamIdleTimeoutError as e: + raise UpstreamException( + message=f"Image stream idle timeout after {e.idle_seconds}s", + status_code=504, + details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}, + ) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in image: {e}") + raise UpstreamException( + message="Upstream connection closed unexpectedly", + status_code=502, + details={"error": str(e), "type": "http2_stream_error"}, + ) + logger.error(f"Image stream request error: {e}") + raise UpstreamException( + message=f"Upstream request failed: {e}", + status_code=502, + details={"error": str(e)}, + ) + except Exception as e: + logger.error(f"Image stream processing error: {e}", extra={"error_type": type(e).__name__}) + raise + finally: + await self.close() + + +class ImageCollectProcessor(BaseProcessor): + """图片生成非流式响应处理器""" + + def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): + super().__init__(model, token) + self.response_format = response_format + + async def process(self, response: AsyncIterable[bytes]) -> List[str]: + """处理并收集图片""" + images = [] + idle_timeout = get_config("grok.stream_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + if mr := resp.get("modelResponse"): + if urls := _collect_image_urls(mr): + for url in urls: + if self.response_format == "url": + processed = await self.process_url(url, "image") + if processed: + images.append(processed) + continue + dl_service = self._get_dl() + base64_data = await dl_service.to_base64(url, self.token, "image") + if base64_data: + if "," in base64_data: + b64 = base64_data.split(",", 1)[1] + else: + b64 = base64_data + images.append(b64) + + except asyncio.CancelledError: + logger.debug("Image collect cancelled by client") + except StreamIdleTimeoutError as e: + logger.warning(f"Image collect idle timeout: {e}") + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in image collect: {e}") + else: + logger.error(f"Image collect request error: {e}") + except Exception as e: + logger.error(f"Image collect processing error: {e}", extra={"error_type": type(e).__name__}) + finally: + await self.close() + + return images + + +__all__ = ["ImageStreamProcessor", "ImageCollectProcessor"] diff --git a/app/services/grok/processors/image_ws_processors.py b/app/services/grok/processors/image_ws_processors.py new file mode 100644 index 0000000..86a6aa9 --- /dev/null +++ b/app/services/grok/processors/image_ws_processors.py @@ -0,0 +1,231 @@ +""" +图片生成响应处理器（WebSocket） +""" + +import base64 +import time +from pathlib import Path +from typing import AsyncGenerator, AsyncIterable, List, Dict, Optional + +import orjson + +from app.core.config import get_config +from app.core.logger import logger +from app.core.exceptions import UpstreamException +from .base import BaseProcessor + + +class ImageWSBaseProcessor(BaseProcessor): + """WebSocket 图片处理基类""" + + def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): + super().__init__(model, token) + self.response_format = response_format + if response_format == "url": + self.response_field = "url" + elif response_format == "base64": + self.response_field = "base64" + else: + self.response_field = "b64_json" + self._image_dir: Optional[Path] = None + + def _ensure_image_dir(self) -> Path: + if self._image_dir is None: + base_dir = Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" / "image" + base_dir.mkdir(parents=True, exist_ok=True) + self._image_dir = base_dir + return self._image_dir + + def _strip_base64(self, blob: str) -> str: + if not blob: + return "" + if "," in blob and "base64" in blob.split(",", 1)[0]: + return blob.split(",", 1)[1] + return blob + + def _filename(self, image_id: str, is_final: bool) -> str: + ext = "jpg" if is_final else "png" + return f"{image_id}.{ext}" + + def _build_file_url(self, filename: str) -> str: + app_url = get_config("app.app_url") + if app_url: + return f"{app_url.rstrip('/')}/v1/files/image/{filename}" + return f"/v1/files/image/{filename}" + + def _save_blob(self, image_id: str, blob: str, is_final: bool) -> str: + data = self._strip_base64(blob) + if not data: + return "" + image_dir = self._ensure_image_dir() + filename = self._filename(image_id, is_final) + filepath = image_dir / filename + with open(filepath, "wb") as f: + f.write(base64.b64decode(data)) + return self._build_file_url(filename) + + def _pick_best(self, existing: Optional[Dict], incoming: Dict) -> Dict: + if not existing: + return incoming + if incoming.get("is_final") and not existing.get("is_final"): + return incoming + if existing.get("is_final") and not incoming.get("is_final"): + return existing + if incoming.get("blob_size", 0) > existing.get("blob_size", 0): + return incoming + return existing + + def _to_output(self, image_id: str, item: Dict) -> str: + try: + if self.response_format == "url": + return self._save_blob(image_id, item.get("blob", ""), item.get("is_final", False)) + return self._strip_base64(item.get("blob", "")) + except Exception as e: + logger.warning(f"Image output failed: {e}") + return "" + + +class ImageWSStreamProcessor(ImageWSBaseProcessor): + """WebSocket 图片流式响应处理器""" + + def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json", size: str = "1024x1024"): + super().__init__(model, token, "b64_json") + self.n = n + self.size = size + self._target_id: Optional[str] = None + self._index_map: Dict[str, int] = {} + self._partial_map: Dict[str, int] = {} + + def _assign_index(self, image_id: str) -> Optional[int]: + if image_id in self._index_map: + return self._index_map[image_id] + if len(self._index_map) >= self.n: + return None + self._index_map[image_id] = len(self._index_map) + return self._index_map[image_id] + + def _sse(self, event: str, data: dict) -> str: + return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" + + async def process(self, response: AsyncIterable[dict]) -> AsyncGenerator[str, None]: + images: Dict[str, Dict] = {} + + async for item in response: + if item.get("type") == "error": + message = item.get("error") or "Upstream error" + code = item.get("error_code") or "upstream_error" + yield self._sse("error", {"error": {"message": message, "type": "server_error", "code": code}}) + return + if item.get("type") != "image": + continue + + image_id = item.get("image_id") + if not image_id: + continue + + if self.n == 1: + if self._target_id is None: + self._target_id = image_id + index = 0 if image_id == self._target_id else None + else: + index = self._assign_index(image_id) + + images[image_id] = self._pick_best(images.get(image_id), item) + + if index is None: + continue + + if item.get("stage") != "final": + partial_b64 = self._strip_base64(item.get("blob", "")) + if not partial_b64: + continue + partial_index = self._partial_map.get(image_id, 0) + if item.get("stage") == "medium": + partial_index = max(partial_index, 1) + self._partial_map[image_id] = partial_index + yield self._sse( + "image_generation.partial_image", + { + "type": "image_generation.partial_image", + "b64_json": partial_b64, + "created_at": int(time.time()), + "size": self.size, + "index": index, + "partial_image_index": partial_index, + }, + ) + + if self.n == 1: + if self._target_id and self._target_id in images: + selected = [(self._target_id, images[self._target_id])] + else: + selected = [max(images.items(), key=lambda x: (x[1].get("is_final", False), x[1].get("blob_size", 0)))] if images else [] + else: + selected = [(image_id, images[image_id]) for image_id in self._index_map if image_id in images] + + for image_id, item in selected: + output = self._strip_base64(item.get("blob", "")) + if not output: + continue + + if self.n == 1: + index = 0 + else: + index = self._index_map.get(image_id, 0) + yield self._sse( + "image_generation.completed", + { + "type": "image_generation.completed", + "b64_json": output, + "created_at": int(time.time()), + "size": self.size, + "index": index, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "output_tokens": 0, + "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + }, + }, + ) + + +class ImageWSCollectProcessor(ImageWSBaseProcessor): + """WebSocket 图片非流式响应处理器""" + + def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json"): + super().__init__(model, token, response_format) + self.n = n + + async def process(self, response: AsyncIterable[dict]) -> List[str]: + images: Dict[str, Dict] = {} + + async for item in response: + if item.get("type") == "error": + message = item.get("error") or "Upstream error" + raise UpstreamException(message, details=item) + if item.get("type") != "image": + continue + image_id = item.get("image_id") + if not image_id: + continue + images[image_id] = self._pick_best(images.get(image_id), item) + + selected = sorted( + images.values(), + key=lambda x: (x.get("is_final", False), x.get("blob_size", 0)), + reverse=True, + ) + if self.n: + selected = selected[: self.n] + + results: List[str] = [] + for item in selected: + output = self._to_output(item.get("image_id", ""), item) + if output: + results.append(output) + + return results + + +__all__ = ["ImageWSStreamProcessor", "ImageWSCollectProcessor"] diff --git a/app/services/grok/processors/processor.py b/app/services/grok/processors/processor.py deleted file mode 100644 index 49ac9d1..0000000 --- a/app/services/grok/processors/processor.py +++ /dev/null @@ -1,1262 +0,0 @@ -""" -OpenAI 响应格式处理器 -""" - -import asyncio -import base64 -import time -import uuid -import random -import re -from pathlib import Path -import orjson -from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar, Dict -from curl_cffi.requests.errors import RequestsError - -from app.core.config import get_config -from app.core.logger import logger -from app.core.exceptions import UpstreamException -from app.services.grok.services.assets import DownloadService - - -DEFAULT_APP_URL = "" -DEFAULT_FILTER_TAGS = ["grok:render", "xaiartifact", "xai:tool_usage_card"] -DEFAULT_IMAGE_FORMAT = "url" -DEFAULT_THINKING = False -DEFAULT_STREAM_IDLE_TIMEOUT = 45.0 -DEFAULT_VIDEO_FORMAT = "html" -DEFAULT_VIDEO_IDLE_TIMEOUT = 90.0 - - -def _is_http2_stream_error(e: Exception) -> bool: - """检查是否为 HTTP/2 流错误""" - err_str = str(e).lower() - return "http/2" in err_str or "curl: (92)" in err_str or "stream" in err_str - - -def _normalize_stream_line(line: Any) -> Optional[str]: - """规范化流式响应行，兼容 SSE data 前缀与空行""" - if line is None: - return None - if isinstance(line, (bytes, bytearray)): - text = line.decode("utf-8", errors="ignore") - else: - text = str(line) - text = text.strip() - if not text: - return None - if text.startswith("data:"): - text = text[5:].strip() - if text == "[DONE]": - return None - return text - - -def _collect_image_urls(obj: Any) -> List[str]: - """递归收集响应中的图片 URL""" - urls: List[str] = [] - seen = set() - - def add(url: str): - if not url or url in seen: - return - seen.add(url) - urls.append(url) - - def walk(value: Any): - if isinstance(value, dict): - for key, item in value.items(): - if key in {"generatedImageUrls", "imageUrls", "imageURLs"}: - if isinstance(item, list): - for url in item: - if isinstance(url, str): - add(url) - elif isinstance(item, str): - add(item) - continue - walk(item) - elif isinstance(value, list): - for item in value: - walk(item) - - walk(obj) - return urls - - -T = TypeVar("T") - - -class StreamIdleTimeoutError(Exception): - """流空闲超时错误""" - - def __init__(self, idle_seconds: float): - self.idle_seconds = idle_seconds - super().__init__(f"Stream idle timeout after {idle_seconds}s") - - -async def _with_idle_timeout( - iterable: AsyncIterable[T], idle_timeout: float, model: str = "" -) -> AsyncGenerator[T, None]: - """ - 包装异步迭代器，添加空闲超时检测 - - Args: - iterable: 原始异步迭代器 - idle_timeout: 空闲超时时间(秒)，0 表示禁用 - model: 模型名称(用于日志) - - Yields: - 原始迭代器的元素 - - Raises: - StreamIdleTimeoutError: 当空闲超时时 - """ - if idle_timeout <= 0: - async for item in iterable: - yield item - return - - iterator = iterable.__aiter__() - while True: - try: - item = await asyncio.wait_for(iterator.__anext__(), timeout=idle_timeout) - yield item - except asyncio.TimeoutError: - logger.warning( - f"Stream idle timeout after {idle_timeout}s", - extra={"model": model, "idle_timeout": idle_timeout}, - ) - raise StreamIdleTimeoutError(idle_timeout) - except StopAsyncIteration: - break - - -ASSET_URL = "https://assets.grok.com/" - - -class BaseProcessor: - """基础处理器""" - - def __init__(self, model: str, token: str = ""): - self.model = model - self.token = token - self.created = int(time.time()) - self.app_url = get_config("app.app_url", DEFAULT_APP_URL) - self._dl_service: Optional[DownloadService] = None - - def _get_dl(self) -> DownloadService: - """获取下载服务实例（复用）""" - if self._dl_service is None: - self._dl_service = DownloadService() - return self._dl_service - - async def close(self): - """释放下载服务资源""" - if self._dl_service: - await self._dl_service.close() - self._dl_service = None - - async def process_url(self, path: str, media_type: str = "image") -> str: - """处理资产 URL""" - # 处理可能的绝对路径 - if path.startswith("http"): - from urllib.parse import urlparse - - path = urlparse(path).path - - if not path.startswith("/"): - path = f"/{path}" - - if self.app_url: - dl_service = self._get_dl() - await dl_service.download(path, self.token, media_type) - return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}" - else: - return f"{ASSET_URL.rstrip('/')}{path}" - - def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: - """构建 SSE 响应 (StreamProcessor 通用)""" - if not hasattr(self, "response_id"): - self.response_id = None - if not hasattr(self, "fingerprint"): - self.fingerprint = "" - - delta = {} - if role: - delta["role"] = role - delta["content"] = "" - elif content: - delta["content"] = content - - chunk = { - "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", - "object": "chat.completion.chunk", - "created": self.created, - "model": self.model, - "system_fingerprint": self.fingerprint - if hasattr(self, "fingerprint") - else "", - "choices": [ - {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} - ], - } - return f"data: {orjson.dumps(chunk).decode()}\n\n" - - -class StreamProcessor(BaseProcessor): - """流式响应处理器""" - - def __init__(self, model: str, token: str = "", think: bool = None): - super().__init__(model, token) - self.response_id: Optional[str] = None - self.fingerprint: str = "" - self.think_opened: bool = False - self.role_sent: bool = False - self.filter_tags = get_config("grok.filter_tags", DEFAULT_FILTER_TAGS) - self.image_format = get_config("app.image_format", DEFAULT_IMAGE_FORMAT) - # 用于过滤跨 token 的标签 - self._tag_buffer: str = "" - self._in_filter_tag: bool = False - - if think is None: - self.show_think = get_config("grok.thinking", DEFAULT_THINKING) - else: - self.show_think = think - - def _filter_token(self, token: str) -> str: - """ - 过滤 token 中的特殊标签（如 ...） - 支持跨 token 的标签过滤 - """ - if not self.filter_tags: - return token - - result = [] - i = 0 - while i < len(token): - char = token[i] - - # 如果在过滤标签内 - if self._in_filter_tag: - self._tag_buffer += char - # 检查是否到达结束标签 - if char == ">": - # 检查是否是自闭合标签 - if "/>" in self._tag_buffer: - self._in_filter_tag = False - self._tag_buffer = "" - else: - # 检查是否是结束标签 - for tag in self.filter_tags: - if f"" in self._tag_buffer: - self._in_filter_tag = False - self._tag_buffer = "" - break - # 如果不是结束标签，检查是否是开始标签结束（非自闭合） - # 继续等待结束标签 - i += 1 - continue - - # 检查是否开始一个过滤标签 - if char == "<": - # 查看后续字符 - remaining = token[i:] - tag_started = False - for tag in self.filter_tags: - if remaining.startswith(f"<{tag}"): - tag_started = True - break - # 部分匹配（可能跨 token） - if len(remaining) < len(tag) + 1: - for j in range(1, len(remaining) + 1): - if f"<{tag}".startswith(remaining[:j]): - tag_started = True - break - - if tag_started: - self._in_filter_tag = True - self._tag_buffer = char - i += 1 - continue - - result.append(char) - i += 1 - - return "".join(result) - - async def process( - self, response: AsyncIterable[bytes] - ) -> AsyncGenerator[str, None]: - """处理流式响应""" - # 获取空闲超时配置 - idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - # 元数据 - if (llm := resp.get("llmInfo")) and not self.fingerprint: - self.fingerprint = llm.get("modelHash", "") - if rid := resp.get("responseId"): - self.response_id = rid - - # 首次发送 role - if not self.role_sent: - yield self._sse(role="assistant") - self.role_sent = True - - # 图像生成进度 - if img := resp.get("streamingImageGenerationResponse"): - if self.show_think: - if not self.think_opened: - yield self._sse("\n") - self.think_opened = True - idx = img.get("imageIndex", 0) + 1 - progress = img.get("progress", 0) - yield self._sse( - f"正在生成第{idx}张图片中，当前进度{progress}%\n" - ) - continue - - # modelResponse - if mr := resp.get("modelResponse"): - if self.think_opened and self.show_think: - if msg := mr.get("message"): - yield self._sse(msg + "\n") - yield self._sse("\n") - self.think_opened = False - - # 处理生成的图片 - for url in _collect_image_urls(mr): - parts = url.split("/") - img_id = parts[-2] if len(parts) >= 2 else "image" - - if self.image_format == "base64": - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - yield self._sse(f"![{img_id}]({base64_data})\n") - else: - final_url = await self.process_url(url, "image") - yield self._sse(f"![{img_id}]({final_url})\n") - else: - final_url = await self.process_url(url, "image") - yield self._sse(f"![{img_id}]({final_url})\n") - - if ( - (meta := mr.get("metadata", {})) - .get("llm_info", {}) - .get("modelHash") - ): - self.fingerprint = meta["llm_info"]["modelHash"] - continue - - # 普通 token - if (token := resp.get("token")) is not None: - if token: - filtered = self._filter_token(token) - if filtered: - yield self._sse(filtered) - - if self.think_opened: - yield self._sse("\n") - yield self._sse(finish="stop") - yield "data: [DONE]\n\n" - except asyncio.CancelledError: - # 客户端断开连接，静默处理 - logger.debug("Stream cancelled by client", extra={"model": self.model}) - except StreamIdleTimeoutError as e: - # 流空闲超时 - raise UpstreamException( - message=f"Stream idle timeout after {e.idle_seconds}s", - status_code=504, - details={ - "error": str(e), - "type": "stream_idle_timeout", - "idle_seconds": e.idle_seconds, - }, - ) - except RequestsError as e: - # HTTP/2 流错误转换为 UpstreamException - if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error: {e}", extra={"model": self.model}) - raise UpstreamException( - message="Upstream connection closed unexpectedly", - status_code=502, - details={"error": str(e), "type": "http2_stream_error"}, - ) - logger.error(f"Stream request error: {e}", extra={"model": self.model}) - raise UpstreamException( - message=f"Upstream request failed: {e}", - status_code=502, - details={"error": str(e)}, - ) - except Exception as e: - logger.error( - f"Stream processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - raise - finally: - await self.close() - - -class CollectProcessor(BaseProcessor): - """非流式响应处理器""" - - def __init__(self, model: str, token: str = ""): - super().__init__(model, token) - self.image_format = get_config("app.image_format", DEFAULT_IMAGE_FORMAT) - self.filter_tags = get_config("grok.filter_tags", DEFAULT_FILTER_TAGS) - - def _filter_content(self, content: str) -> str: - """过滤内容中的特殊标签""" - import re - - if not content or not self.filter_tags: - return content - - result = content - for tag in self.filter_tags: - # 匹配 ... 或 ，re.DOTALL 使 . 匹配换行符 - pattern = rf"<{re.escape(tag)}[^>]*>.*?|<{re.escape(tag)}[^>]*/>" - result = re.sub(pattern, "", result, flags=re.DOTALL) - - return result - - async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: - """处理并收集完整响应""" - response_id = "" - fingerprint = "" - content = "" - # 获取空闲超时配置 - idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - if (llm := resp.get("llmInfo")) and not fingerprint: - fingerprint = llm.get("modelHash", "") - - if mr := resp.get("modelResponse"): - response_id = mr.get("responseId", "") - content = mr.get("message", "") - - if urls := _collect_image_urls(mr): - content += "\n" - for url in urls: - parts = url.split("/") - img_id = parts[-2] if len(parts) >= 2 else "image" - - if self.image_format == "base64": - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - content += f"![{img_id}]({base64_data})\n" - else: - final_url = await self.process_url(url, "image") - content += f"![{img_id}]({final_url})\n" - else: - final_url = await self.process_url(url, "image") - content += f"![{img_id}]({final_url})\n" - - if ( - (meta := mr.get("metadata", {})) - .get("llm_info", {}) - .get("modelHash") - ): - fingerprint = meta["llm_info"]["modelHash"] - - except asyncio.CancelledError: - logger.debug("Collect cancelled by client", extra={"model": self.model}) - except StreamIdleTimeoutError as e: - logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model}) - # 非流式模式下，超时后返回已收集的内容 - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning( - f"HTTP/2 stream error in collect: {e}", extra={"model": self.model} - ) - else: - logger.error(f"Collect request error: {e}", extra={"model": self.model}) - except Exception as e: - logger.error( - f"Collect processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - finally: - await self.close() - - # 过滤特殊标签 - content = self._filter_content(content) - - return { - "id": response_id, - "object": "chat.completion", - "created": self.created, - "model": self.model, - "system_fingerprint": fingerprint, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": content, - "refusal": None, - "annotations": [], - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0, - "prompt_tokens_details": { - "cached_tokens": 0, - "text_tokens": 0, - "audio_tokens": 0, - "image_tokens": 0, - }, - "completion_tokens_details": { - "text_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - }, - }, - } - - -class VideoStreamProcessor(BaseProcessor): - """视频流式响应处理器""" - - def __init__(self, model: str, token: str = "", think: bool = None): - super().__init__(model, token) - self.response_id: Optional[str] = None - self.think_opened: bool = False - self.role_sent: bool = False - self.video_format = str(get_config("app.video_format", DEFAULT_VIDEO_FORMAT)).lower() - - if think is None: - self.show_think = get_config("grok.thinking", DEFAULT_THINKING) - else: - self.show_think = think - - def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: - """构建视频 HTML 标签""" - import html - - safe_video_url = html.escape(video_url) - safe_thumbnail_url = html.escape(thumbnail_url) - poster_attr = f' poster="{safe_thumbnail_url}"' if safe_thumbnail_url else "" - return f'''''' - - async def process( - self, response: AsyncIterable[bytes] - ) -> AsyncGenerator[str, None]: - """处理视频流式响应""" - # 视频生成使用更长的空闲超时 - idle_timeout = get_config("grok.video_idle_timeout", DEFAULT_VIDEO_IDLE_TIMEOUT) - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - if rid := resp.get("responseId"): - self.response_id = rid - - # 首次发送 role - if not self.role_sent: - yield self._sse(role="assistant") - self.role_sent = True - - # 视频生成进度 - if video_resp := resp.get("streamingVideoGenerationResponse"): - progress = video_resp.get("progress", 0) - - if self.show_think: - if not self.think_opened: - yield self._sse("\n") - self.think_opened = True - yield self._sse(f"正在生成视频中，当前进度{progress}%\n") - - if progress == 100: - video_url = video_resp.get("videoUrl", "") - thumbnail_url = video_resp.get("thumbnailImageUrl", "") - - if self.think_opened and self.show_think: - yield self._sse("\n") - self.think_opened = False - - if video_url: - final_video_url = await self.process_url(video_url, "video") - final_thumbnail_url = "" - if thumbnail_url: - final_thumbnail_url = await self.process_url( - thumbnail_url, "image" - ) - - if self.video_format == "url": - yield self._sse(final_video_url) - else: - video_html = self._build_video_html( - final_video_url, final_thumbnail_url - ) - yield self._sse(video_html) - - logger.info(f"Video generated: {video_url}") - continue - - if self.think_opened: - yield self._sse("\n") - yield self._sse(finish="stop") - yield "data: [DONE]\n\n" - except asyncio.CancelledError: - logger.debug( - "Video stream cancelled by client", extra={"model": self.model} - ) - except StreamIdleTimeoutError as e: - raise UpstreamException( - message=f"Video stream idle timeout after {e.idle_seconds}s", - status_code=504, - details={ - "error": str(e), - "type": "stream_idle_timeout", - "idle_seconds": e.idle_seconds, - }, - ) - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning( - f"HTTP/2 stream error in video: {e}", extra={"model": self.model} - ) - raise UpstreamException( - message="Upstream connection closed unexpectedly", - status_code=502, - details={"error": str(e), "type": "http2_stream_error"}, - ) - logger.error( - f"Video stream request error: {e}", extra={"model": self.model} - ) - raise UpstreamException( - message=f"Upstream request failed: {e}", - status_code=502, - details={"error": str(e)}, - ) - except Exception as e: - logger.error( - f"Video stream processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - finally: - await self.close() - - -class VideoCollectProcessor(BaseProcessor): - """视频非流式响应处理器""" - - def __init__(self, model: str, token: str = ""): - super().__init__(model, token) - self.video_format = str(get_config("app.video_format", DEFAULT_VIDEO_FORMAT)).lower() - - def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: - poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" - return f'''''' - - async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: - """处理并收集视频响应""" - response_id = "" - content = "" - # 视频生成使用更长的空闲超时 - idle_timeout = get_config("grok.video_idle_timeout", DEFAULT_VIDEO_IDLE_TIMEOUT) - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - if video_resp := resp.get("streamingVideoGenerationResponse"): - if video_resp.get("progress") == 100: - response_id = resp.get("responseId", "") - video_url = video_resp.get("videoUrl", "") - thumbnail_url = video_resp.get("thumbnailImageUrl", "") - - if video_url: - final_video_url = await self.process_url(video_url, "video") - final_thumbnail_url = "" - if thumbnail_url: - final_thumbnail_url = await self.process_url( - thumbnail_url, "image" - ) - - if self.video_format == "url": - content = final_video_url - else: - content = self._build_video_html( - final_video_url, final_thumbnail_url - ) - logger.info(f"Video generated: {video_url}") - - except asyncio.CancelledError: - logger.debug( - "Video collect cancelled by client", extra={"model": self.model} - ) - except StreamIdleTimeoutError as e: - logger.warning( - f"Video collect idle timeout: {e}", extra={"model": self.model} - ) - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning( - f"HTTP/2 stream error in video collect: {e}", - extra={"model": self.model}, - ) - else: - logger.error( - f"Video collect request error: {e}", extra={"model": self.model} - ) - except Exception as e: - logger.error( - f"Video collect processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - finally: - await self.close() - - return { - "id": response_id, - "object": "chat.completion", - "created": self.created, - "model": self.model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": content, - "refusal": None, - }, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, - } - - -class ImageStreamProcessor(BaseProcessor): - """图片生成流式响应处理器""" - - def __init__( - self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" - ): - super().__init__(model, token) - self.partial_index = 0 - self.n = n - self.target_index = random.randint(0, 1) if n == 1 else None - self.response_format = response_format - if response_format == "url": - self.response_field = "url" - elif response_format == "base64": - self.response_field = "base64" - else: - self.response_field = "b64_json" - - def _sse(self, event: str, data: dict) -> str: - """构建 SSE 响应 (覆盖基类)""" - return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" - - async def process( - self, response: AsyncIterable[bytes] - ) -> AsyncGenerator[str, None]: - """处理流式响应""" - final_images = [] - # 图片生成使用标准空闲超时 - idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - # 图片生成进度 - if img := resp.get("streamingImageGenerationResponse"): - image_index = img.get("imageIndex", 0) - progress = img.get("progress", 0) - - if self.n == 1 and image_index != self.target_index: - continue - - out_index = 0 if self.n == 1 else image_index - - yield self._sse( - "image_generation.partial_image", - { - "type": "image_generation.partial_image", - self.response_field: "", - "index": out_index, - "progress": progress, - }, - ) - continue - - # modelResponse - if mr := resp.get("modelResponse"): - if urls := _collect_image_urls(mr): - for url in urls: - if self.response_format == "url": - processed = await self.process_url(url, "image") - if processed: - final_images.append(processed) - continue - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - if "," in base64_data: - b64 = base64_data.split(",", 1)[1] - else: - b64 = base64_data - final_images.append(b64) - continue - - for index, b64 in enumerate(final_images): - if self.n == 1: - if index != self.target_index: - continue - out_index = 0 - else: - out_index = index - - yield self._sse( - "image_generation.completed", - { - "type": "image_generation.completed", - self.response_field: b64, - "index": out_index, - "usage": { - "total_tokens": 0, - "input_tokens": 0, - "output_tokens": 0, - "input_tokens_details": { - "text_tokens": 0, - "image_tokens": 0, - }, - }, - }, - ) - except asyncio.CancelledError: - logger.debug("Image stream cancelled by client") - except StreamIdleTimeoutError as e: - raise UpstreamException( - message=f"Image stream idle timeout after {e.idle_seconds}s", - status_code=504, - details={ - "error": str(e), - "type": "stream_idle_timeout", - "idle_seconds": e.idle_seconds, - }, - ) - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in image: {e}") - raise UpstreamException( - message="Upstream connection closed unexpectedly", - status_code=502, - details={"error": str(e), "type": "http2_stream_error"}, - ) - logger.error(f"Image stream request error: {e}") - raise UpstreamException( - message=f"Upstream request failed: {e}", - status_code=502, - details={"error": str(e)}, - ) - except Exception as e: - logger.error( - f"Image stream processing error: {e}", - extra={"error_type": type(e).__name__}, - ) - raise - finally: - await self.close() - - -class ImageCollectProcessor(BaseProcessor): - """图片生成非流式响应处理器""" - - def __init__( - self, model: str, token: str = "", response_format: str = "b64_json" - ): - super().__init__(model, token) - self.response_format = response_format - - async def process(self, response: AsyncIterable[bytes]) -> List[str]: - """处理并收集图片""" - images = [] - # 图片生成使用标准空闲超时 - idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_STREAM_IDLE_TIMEOUT) - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - if mr := resp.get("modelResponse"): - if urls := _collect_image_urls(mr): - for url in urls: - if self.response_format == "url": - processed = await self.process_url(url, "image") - if processed: - images.append(processed) - continue - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - if "," in base64_data: - b64 = base64_data.split(",", 1)[1] - else: - b64 = base64_data - images.append(b64) - - except asyncio.CancelledError: - logger.debug("Image collect cancelled by client") - except StreamIdleTimeoutError as e: - logger.warning(f"Image collect idle timeout: {e}") - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in image collect: {e}") - else: - logger.error(f"Image collect request error: {e}") - except Exception as e: - logger.error( - f"Image collect processing error: {e}", - extra={"error_type": type(e).__name__}, - ) - finally: - await self.close() - - return images - - -class ImageWSBaseProcessor(BaseProcessor): - """WebSocket 图片处理基类""" - - def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): - super().__init__(model, token) - self.response_format = response_format - if response_format == "url": - self.response_field = "url" - elif response_format == "base64": - self.response_field = "base64" - else: - self.response_field = "b64_json" - self._image_dir: Optional[Path] = None - - def _ensure_image_dir(self) -> Path: - if self._image_dir is None: - base_dir = ( - Path(__file__).parent.parent.parent.parent.parent - / "data" - / "tmp" - / "image" - ) - base_dir.mkdir(parents=True, exist_ok=True) - self._image_dir = base_dir - return self._image_dir - - def _strip_base64(self, blob: str) -> str: - if not blob: - return "" - if "," in blob and "base64" in blob.split(",", 1)[0]: - return blob.split(",", 1)[1] - return blob - - def _filename(self, image_id: str, is_final: bool) -> str: - ext = "jpg" if is_final else "png" - return f"{image_id}.{ext}" - - def _build_file_url(self, filename: str) -> str: - if self.app_url: - return f"{self.app_url.rstrip('/')}/v1/files/image/{filename}" - return f"/v1/files/image/{filename}" - - def _save_blob(self, image_id: str, blob: str, is_final: bool) -> str: - data = self._strip_base64(blob) - if not data: - return "" - image_dir = self._ensure_image_dir() - filename = self._filename(image_id, is_final) - filepath = image_dir / filename - with open(filepath, "wb") as f: - f.write(base64.b64decode(data)) - return self._build_file_url(filename) - - def _pick_best(self, existing: Optional[Dict], incoming: Dict) -> Dict: - if not existing: - return incoming - if incoming.get("is_final") and not existing.get("is_final"): - return incoming - if existing.get("is_final") and not incoming.get("is_final"): - return existing - if incoming.get("blob_size", 0) > existing.get("blob_size", 0): - return incoming - return existing - - def _to_output(self, image_id: str, item: Dict) -> str: - try: - if self.response_format == "url": - return self._save_blob( - image_id, item.get("blob", ""), item.get("is_final", False) - ) - return self._strip_base64(item.get("blob", "")) - except Exception as e: - logger.warning(f"Image output failed: {e}") - return "" - - -class ImageWSStreamProcessor(ImageWSBaseProcessor): - """WebSocket 图片流式响应处理器""" - - def __init__( - self, - model: str, - token: str = "", - n: int = 1, - response_format: str = "b64_json", - size: str = "1024x1024", - ): - super().__init__(model, token, "b64_json") - self.n = n - self.size = size - self._target_id: Optional[str] = None - self._index_map: Dict[str, int] = {} - self._partial_map: Dict[str, int] = {} - - def _assign_index(self, image_id: str) -> Optional[int]: - if image_id in self._index_map: - return self._index_map[image_id] - if len(self._index_map) >= self.n: - return None - self._index_map[image_id] = len(self._index_map) - return self._index_map[image_id] - - def _sse(self, event: str, data: dict) -> str: - return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" - - async def process( - self, response: AsyncIterable[dict] - ) -> AsyncGenerator[str, None]: - images: Dict[str, Dict] = {} - - async for item in response: - if item.get("type") == "error": - message = item.get("error") or "Upstream error" - code = item.get("error_code") or "upstream_error" - yield self._sse( - "error", - { - "error": { - "message": message, - "type": "server_error", - "code": code, - } - }, - ) - return - if item.get("type") != "image": - continue - - image_id = item.get("image_id") - if not image_id: - continue - - if self.n == 1: - if self._target_id is None: - self._target_id = image_id - index = 0 if image_id == self._target_id else None - else: - index = self._assign_index(image_id) - - images[image_id] = self._pick_best(images.get(image_id), item) - - if index is None: - continue - - if item.get("stage") != "final": - partial_b64 = self._strip_base64(item.get("blob", "")) - if not partial_b64: - continue - partial_index = self._partial_map.get(image_id, 0) - if item.get("stage") == "medium": - partial_index = max(partial_index, 1) - self._partial_map[image_id] = partial_index - yield self._sse( - "image_generation.partial_image", - { - "type": "image_generation.partial_image", - "b64_json": partial_b64, - "created_at": int(time.time()), - "size": self.size, - "index": index, - "partial_image_index": partial_index, - }, - ) - - if self.n == 1: - if self._target_id and self._target_id in images: - selected = [(self._target_id, images[self._target_id])] - else: - selected = [ - max( - images.items(), - key=lambda x: ( - x[1].get("is_final", False), - x[1].get("blob_size", 0), - ), - ) - ] if images else [] - else: - selected = [ - (image_id, images[image_id]) - for image_id in self._index_map - if image_id in images - ] - - for image_id, item in selected: - output = self._strip_base64(item.get("blob", "")) - - if not output: - continue - - if self.n == 1: - index = 0 - else: - index = self._index_map.get(image_id, 0) - yield self._sse( - "image_generation.completed", - { - "type": "image_generation.completed", - "b64_json": output, - "created_at": int(time.time()), - "size": self.size, - "index": index, - "usage": { - "total_tokens": 0, - "input_tokens": 0, - "output_tokens": 0, - "input_tokens_details": { - "text_tokens": 0, - "image_tokens": 0, - }, - }, - }, - ) - - -class ImageWSCollectProcessor(ImageWSBaseProcessor): - """WebSocket 图片非流式响应处理器""" - - def __init__( - self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" - ): - super().__init__(model, token, response_format) - self.n = n - - async def process(self, response: AsyncIterable[dict]) -> List[str]: - images: Dict[str, Dict] = {} - - async for item in response: - if item.get("type") == "error": - message = item.get("error") or "Upstream error" - raise UpstreamException(message, details=item) - if item.get("type") != "image": - continue - image_id = item.get("image_id") - if not image_id: - continue - images[image_id] = self._pick_best(images.get(image_id), item) - - selected = sorted( - images.values(), - key=lambda x: (x.get("is_final", False), x.get("blob_size", 0)), - reverse=True, - ) - if self.n: - selected = selected[: self.n] - - results: List[str] = [] - for item in selected: - output = self._to_output(item.get("image_id", ""), item) - if output: - results.append(output) - - return results - - -__all__ = [ - "StreamProcessor", - "CollectProcessor", - "VideoStreamProcessor", - "VideoCollectProcessor", - "ImageStreamProcessor", - "ImageCollectProcessor", - "ImageWSStreamProcessor", - "ImageWSCollectProcessor", -] diff --git a/app/services/grok/processors/video_processors.py b/app/services/grok/processors/video_processors.py new file mode 100644 index 0000000..df8750c --- /dev/null +++ b/app/services/grok/processors/video_processors.py @@ -0,0 +1,236 @@ +""" +视频响应处理器 +""" + +import asyncio +import uuid +from typing import Any, AsyncGenerator, AsyncIterable, Optional + +import orjson +from curl_cffi.requests.errors import RequestsError + +from app.core.config import get_config +from app.core.logger import logger +from app.core.exceptions import UpstreamException +from .base import ( + BaseProcessor, + StreamIdleTimeoutError, + _with_idle_timeout, + _normalize_stream_line, + _is_http2_stream_error, +) + + +class VideoStreamProcessor(BaseProcessor): + """视频流式响应处理器""" + + def __init__(self, model: str, token: str = "", think: bool = None): + super().__init__(model, token) + self.response_id: Optional[str] = None + self.think_opened: bool = False + self.role_sent: bool = False + self.video_format = str(get_config("app.video_format")).lower() + + if think is None: + self.show_think = get_config("grok.thinking") + else: + self.show_think = think + + def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: + """构建 SSE 响应""" + delta = {} + if role: + delta["role"] = role + delta["content"] = "" + elif content: + delta["content"] = content + + chunk = { + "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion.chunk", + "created": self.created, + "model": self.model, + "choices": [ + {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} + ], + } + return f"data: {orjson.dumps(chunk).decode()}\n\n" + + def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: + """构建视频 HTML 标签""" + import html + safe_video_url = html.escape(video_url) + safe_thumbnail_url = html.escape(thumbnail_url) + poster_attr = f' poster="{safe_thumbnail_url}"' if safe_thumbnail_url else "" + return f'''''' + + async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + """处理视频流式响应""" + idle_timeout = get_config("grok.video_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + if rid := resp.get("responseId"): + self.response_id = rid + + if not self.role_sent: + yield self._sse(role="assistant") + self.role_sent = True + + # 视频生成进度 + if video_resp := resp.get("streamingVideoGenerationResponse"): + progress = video_resp.get("progress", 0) + + if self.show_think: + if not self.think_opened: + yield self._sse("\n") + self.think_opened = True + yield self._sse(f"正在生成视频中，当前进度{progress}%\n") + + if progress == 100: + video_url = video_resp.get("videoUrl", "") + thumbnail_url = video_resp.get("thumbnailImageUrl", "") + + if self.think_opened and self.show_think: + yield self._sse("\n") + self.think_opened = False + + if video_url: + final_video_url = await self.process_url(video_url, "video") + final_thumbnail_url = "" + if thumbnail_url: + final_thumbnail_url = await self.process_url(thumbnail_url, "image") + + if self.video_format == "url": + yield self._sse(final_video_url) + else: + video_html = self._build_video_html(final_video_url, final_thumbnail_url) + yield self._sse(video_html) + + logger.info(f"Video generated: {video_url}") + continue + + if self.think_opened: + yield self._sse("\n") + yield self._sse(finish="stop") + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + logger.debug("Video stream cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + raise UpstreamException( + message=f"Video stream idle timeout after {e.idle_seconds}s", + status_code=504, + details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}, + ) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in video: {e}", extra={"model": self.model}) + raise UpstreamException( + message="Upstream connection closed unexpectedly", + status_code=502, + details={"error": str(e), "type": "http2_stream_error"}, + ) + logger.error(f"Video stream request error: {e}", extra={"model": self.model}) + raise UpstreamException( + message=f"Upstream request failed: {e}", + status_code=502, + details={"error": str(e)}, + ) + except Exception as e: + logger.error(f"Video stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + finally: + await self.close() + + +class VideoCollectProcessor(BaseProcessor): + """视频非流式响应处理器""" + + def __init__(self, model: str, token: str = ""): + super().__init__(model, token) + self.video_format = str(get_config("app.video_format")).lower() + + def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: + poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" + return f'''''' + + async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: + """处理并收集视频响应""" + response_id = "" + content = "" + idle_timeout = get_config("grok.video_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + if video_resp := resp.get("streamingVideoGenerationResponse"): + if video_resp.get("progress") == 100: + response_id = resp.get("responseId", "") + video_url = video_resp.get("videoUrl", "") + thumbnail_url = video_resp.get("thumbnailImageUrl", "") + + if video_url: + final_video_url = await self.process_url(video_url, "video") + final_thumbnail_url = "" + if thumbnail_url: + final_thumbnail_url = await self.process_url(thumbnail_url, "image") + + if self.video_format == "url": + content = final_video_url + else: + content = self._build_video_html(final_video_url, final_thumbnail_url) + logger.info(f"Video generated: {video_url}") + + except asyncio.CancelledError: + logger.debug("Video collect cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + logger.warning(f"Video collect idle timeout: {e}", extra={"model": self.model}) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in video collect: {e}", extra={"model": self.model}) + else: + logger.error(f"Video collect request error: {e}", extra={"model": self.model}) + except Exception as e: + logger.error(f"Video collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + finally: + await self.close() + + return { + "id": response_id, + "object": "chat.completion", + "created": self.created, + "model": self.model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content, "refusal": None}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + } + + +__all__ = ["VideoStreamProcessor", "VideoCollectProcessor"] diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index 6e53634..9ac3eb7 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -4,29 +4,31 @@ Grok 文件资产服务 import asyncio import base64 -import os -import time import hashlib +import os import re -from pathlib import Path +import time from contextlib import asynccontextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple +from urllib.parse import urlparse try: import fcntl -except ImportError: # pragma: no cover - non-posix platforms +except ImportError: fcntl = None -from typing import Tuple, List, Dict, Optional, Any -from urllib.parse import urlparse import aiofiles from curl_cffi.requests import AsyncSession -from app.core.logger import logger from app.core.config import get_config from app.core.exceptions import AppException, UpstreamException, ValidationException +from app.core.logger import logger from app.services.grok.utils.headers import apply_statsig, build_sso_cookie from app.services.token.service import TokenService +# ==================== 常量 ==================== UPLOAD_API = "https://grok.com/rest/app-chat/upload-file" LIST_API = "https://grok.com/rest/assets" @@ -34,31 +36,48 @@ DELETE_API = "https://grok.com/rest/assets-metadata" DOWNLOAD_API = "https://assets.grok.com" LOCK_DIR = Path(__file__).parent.parent.parent.parent.parent / "data" / ".locks" -DEFAULT_TIMEOUT = 120 -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" -DEFAULT_MIME = "application/octet-stream" -DEFAULT_BASE_PROXY_URL = "" -DEFAULT_ASSET_PROXY_URL = "" -DEFAULT_CACHE_ENABLE_AUTO_CLEAN = True -DEFAULT_CACHE_LIMIT_MB = 1024 -DEFAULT_APP_URL = "" -DEFAULT_MAX_CONCURRENT = 25 -DEFAULT_DELETE_BATCH_SIZE = 10 -_ASSETS_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) -_ASSETS_SEM_VALUE = DEFAULT_MAX_CONCURRENT +# 全局信号量（运行时动态初始化） +_ASSETS_SEMAPHORE = None +_ASSETS_SEM_VALUE = None + +# 常用 MIME 类型（业务数据，非配置） +MIME_TYPES = { + ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".gif": "image/gif", + ".webp": "image/webp", ".bmp": "image/bmp", ".pdf": "application/pdf", ".txt": "text/plain", + ".md": "text/markdown", ".csv": "text/csv", ".json": "application/json", ".xml": "application/xml", + ".py": "text/x-python-script", ".js": "application/javascript", ".html": "text/html", + ".css": "text/css", ".mp4": "video/mp4", ".webm": "video/webm", +} + +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} +VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} + +# ==================== 工具函数 ==================== + +def _get_assets_semaphore() -> asyncio.Semaphore: + """获取全局并发控制信号量""" + value = max(1, int(get_config("performance.assets_max_concurrent"))) + + global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE + if _ASSETS_SEMAPHORE is None or value != _ASSETS_SEM_VALUE: + _ASSETS_SEM_VALUE = value + _ASSETS_SEMAPHORE = asyncio.Semaphore(value) + return _ASSETS_SEMAPHORE @asynccontextmanager async def _file_lock(name: str, timeout: int = 10): + """文件锁""" if fcntl is None: yield return + LOCK_DIR.mkdir(parents=True, exist_ok=True) lock_path = LOCK_DIR / f"{name}.lock" fd = None locked = False start = time.monotonic() + try: fd = open(lock_path, "a+") while True: @@ -78,131 +97,59 @@ async def _file_lock(name: str, timeout: int = 10): fcntl.flock(fd, fcntl.LOCK_UN) except Exception: pass - try: - fd.close() - except Exception: - pass + fd.close() -MIME_TYPES = { - # 图片 - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".png": "image/png", - ".gif": "image/gif", - ".webp": "image/webp", - ".bmp": "image/bmp", - # 文档 - ".pdf": "application/pdf", - ".txt": "text/plain", - ".md": "text/markdown", - ".doc": "application/msword", - ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - ".rtf": "application/rtf", - # 表格 - ".csv": "text/csv", - ".xls": "application/vnd.ms-excel", - ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - # 代码 - ".py": "text/x-python-script", - ".js": "application/javascript", - ".ts": "application/typescript", - ".java": "text/x-java", - ".cpp": "text/x-c++", - ".c": "text/x-c", - ".go": "text/x-go", - ".rs": "text/x-rust", - ".rb": "text/x-ruby", - ".php": "text/x-php", - ".sh": "application/x-sh", - ".html": "text/html", - ".css": "text/css", - ".sql": "application/sql", - # 数据 - ".json": "application/json", - ".xml": "application/xml", - ".yaml": "application/x-yaml", - ".yml": "application/x-yaml", - ".toml": "application/toml", - ".ini": "text/plain", - ".log": "text/plain", - ".tmp": "application/octet-stream", - # 其他 - ".graphql": "application/graphql", - ".proto": "application/x-protobuf", - ".latex": "application/x-latex", - ".wiki": "text/plain", - ".rst": "text/x-rst", -} - -IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} -VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} +@dataclass +class ServiceConfig: + """服务配置""" + proxy: str + timeout: int + browser: str + user_agent: str + + @classmethod + def from_settings(cls, proxy: Optional[str] = None): + return cls( + proxy=proxy or get_config("grok.asset_proxy_url") or get_config("grok.base_proxy_url"), + timeout=get_config("grok.timeout"), + browser=get_config("grok.browser"), + user_agent=get_config("grok.user_agent"), + ) + + def get_proxies(self) -> Optional[dict]: + return {"http": self.proxy, "https": self.proxy} if self.proxy else None # ==================== 基础服务 ==================== - class BaseService: """基础服务类""" - def __init__(self, proxy: str = None): - self.proxy = ( - proxy - or get_config("grok.asset_proxy_url", DEFAULT_ASSET_PROXY_URL) - or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - ) - self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + def __init__(self, proxy: Optional[str] = None): + self.config = ServiceConfig.from_settings(proxy) self._session: Optional[AsyncSession] = None - def _headers(self, token: str, referer: str = "https://grok.com/") -> dict: + def _build_headers(self, token: str, referer: str = "https://grok.com/", download: bool = False) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) - headers = { - "Accept": "*/*", - "Accept-Encoding": "gzip, deflate, br, zstd", - "Accept-Language": "zh-CN,zh;q=0.9", - "Baggage": "sentry-environment=production,sentry-release=d6add6fb0460641fd482d767a335ef72b9b6abb8,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c", - "Cache-Control": "no-cache", - "Content-Type": "application/json", - "Origin": "https://grok.com", - "Pragma": "no-cache", - "Priority": "u=1, i", - "Referer": referer, - "Sec-Ch-Ua": '"Google Chrome";v="136", "Chromium";v="136", "Not(A:Brand";v="24"', - "Sec-Ch-Ua-Arch": "arm", - "Sec-Ch-Ua-Bitness": "64", - "Sec-Ch-Ua-Mobile": "?0", - "Sec-Ch-Ua-Model": "", - "Sec-Ch-Ua-Platform": '"macOS"', - "Sec-Fetch-Dest": "empty", - "Sec-Fetch-Mode": "cors", - "Sec-Fetch-Site": "same-origin", - "User-Agent": user_agent, - } - - apply_statsig(headers) + if download: + headers = { + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + "Sec-Fetch-Dest": "document", + "Sec-Fetch-Mode": "navigate", + "Referer": referer, + } + else: + headers = { + "Accept": "*/*", + "Content-Type": "application/json", + "Origin": "https://grok.com", + "Referer": referer, + "User-Agent": self.config.user_agent, + } + apply_statsig(headers) + headers["Cookie"] = build_sso_cookie(token) - - return headers - - def _proxies(self) -> Optional[dict]: - """构建代理配置""" - return {"http": self.proxy, "https": self.proxy} if self.proxy else None - - def _dl_headers(self, token: str, file_path: str) -> dict: - """构建下载请求头""" - headers = { - "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8", - "Sec-Fetch-Dest": "document", - "Sec-Fetch-Mode": "navigate", - "Sec-Fetch-Site": "same-site", - "Sec-Fetch-User": "?1", - "Upgrade-Insecure-Requests": "1", - "Referer": "https://grok.com/", - } - - headers["Cookie"] = build_sso_cookie(token) - return headers async def _get_session(self) -> AsyncSession: @@ -218,200 +165,139 @@ class BaseService: self._session = None @staticmethod - def is_url(input_str: str) -> bool: + def is_url(s: str) -> bool: """检查是否为 URL""" try: - result = urlparse(input_str) - return all([result.scheme, result.netloc]) and result.scheme in [ - "http", - "https", - ] + r = urlparse(s) + return bool(r.scheme and r.netloc and r.scheme in ["http", "https"]) except Exception: return False @staticmethod async def fetch(url: str) -> Tuple[str, str, str]: - """ - 获取远程资源并转 Base64 - - Raises: - UpstreamException: 当获取失败时 - """ + """获取远程资源并转 Base64""" try: async with AsyncSession() as session: response = await session.get(url, timeout=10) if response.status_code >= 400: raise UpstreamException( - message=f"Failed to fetch resource: {response.status_code}", + message=f"Failed to fetch: {response.status_code}", details={"url": url, "status": response.status_code}, ) - + filename = url.split("/")[-1].split("?")[0] or "download" - content_type = response.headers.get("content-type", DEFAULT_MIME).split( - ";" - )[0] + content_type = response.headers.get("content-type", "application/octet-stream").split(";")[0] b64 = base64.b64encode(response.content).decode() - - logger.debug(f"Fetched: {url} -> {filename}") + + logger.debug(f"Fetched: {url}") return filename, b64, content_type except Exception as e: - logger.error(f"Fetch failed: {url} - {e}") if isinstance(e, AppException): - raise e - raise UpstreamException( - f"Resource fetch failed: {str(e)}", details={"url": url} - ) + raise + logger.error(f"Fetch failed: {url} - {e}") + raise UpstreamException(f"Fetch failed: {str(e)}", details={"url": url}) @staticmethod def parse_b64(data_uri: str) -> Tuple[str, str, str]: """解析 Base64 数据""" - if data_uri.startswith("data:"): - try: - header, b64 = data_uri.split(",", 1) - except ValueError: - return "file.bin", data_uri, DEFAULT_MIME - if ";base64" in header: - mime = header[5:].split(";", 1)[0] or DEFAULT_MIME - b64 = re.sub(r"\\s+", "", b64) - ext = mime.split("/")[-1] if "/" in mime else "bin" - return f"file.{ext}", b64, mime - return "file.bin", data_uri, DEFAULT_MIME + if not data_uri.startswith("data:"): + return "file.bin", data_uri, "application/octet-stream" + + try: + header, b64 = data_uri.split(",", 1) + except ValueError: + return "file.bin", data_uri, "application/octet-stream" + + if ";base64" not in header: + return "file.bin", data_uri, "application/octet-stream" + + mime = header[5:].split(";", 1)[0] or "application/octet-stream" + b64 = re.sub(r"\s+", "", b64) + ext = mime.split("/")[-1] if "/" in mime else "bin" + return f"file.{ext}", b64, mime @staticmethod def to_b64(file_path: Path, mime_type: str) -> str: - """将本地文件转为 base64 data URI""" + """文件转 base64 data URI""" try: b64_data = base64.b64encode(file_path.read_bytes()).decode() return f"data:{mime_type};base64,{b64_data}" except Exception as e: logger.error(f"File to base64 failed: {file_path} - {e}") - raise AppException( - f"Failed to read file: {file_path}", code="file_read_error" - ) + raise AppException(f"Failed to read file: {file_path}", code="file_read_error") # ==================== 上传服务 ==================== - class UploadService(BaseService): """文件上传服务""" async def upload(self, file_input: str, token: str) -> Tuple[str, str]: """ 上传文件到 Grok - + Returns: (file_id, file_uri) - - Raises: - ValidationException: 输入无效 - UpstreamException: 上传失败 """ - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - async with _ASSETS_SEMAPHORE: - try: - # 处理输入 - if self.is_url(file_input): - filename, b64, mime = await self.fetch(file_input) - else: - filename, b64, mime = self.parse_b64(file_input) + async with _get_assets_semaphore(): + # 处理输入 + if self.is_url(file_input): + filename, b64, mime = await self.fetch(file_input) + else: + filename, b64, mime = self.parse_b64(file_input) + + if not b64: + raise ValidationException("Invalid file input: empty content") - if not b64: - raise ValidationException("Invalid file input: empty content") + # 执行上传 + session = await self._get_session() + response = await session.post( + UPLOAD_API, + headers=self._build_headers(token), + json={"fileName": filename, "fileMimeType": mime, "content": b64}, + impersonate=self.config.browser, + timeout=self.config.timeout, + proxies=self.config.get_proxies(), + ) - # 构建请求 - headers = self._headers(token) - payload = { - "fileName": filename, - "fileMimeType": mime, - "content": b64, - } + # 处理响应 + if response.status_code == 200: + result = response.json() + file_id = result.get("fileMetadataId", "") + file_uri = result.get("fileUri", "") + logger.info(f"Upload success: {filename} -> {file_id}") + return file_id, file_uri - # 执行上传 - session = await self._get_session() - response = await session.post( - UPLOAD_API, - headers=headers, - json=payload, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - timeout=self.timeout, - proxies=self._proxies(), - ) - - if response.status_code == 200: - result = response.json() - file_id = result.get("fileMetadataId", "") - file_uri = result.get("fileUri", "") - logger.info( - f"Upload success: {filename} -> {file_id}", - extra={"file_id": file_id}, - ) - return file_id, file_uri - - # 401/403 表示 token 失效，记录失败并标记冷却 - if response.status_code in (401, 403): - logger.warning( - f"Upload token invalid: {response.status_code}, recording failure", - extra={"response": response.text[:200]}, - ) - try: - await TokenService.record_fail( - token, - status_code=response.status_code, - reason="upload_auth_failed", - ) - except Exception as mark_err: - logger.error(f"Failed to record token failure: {mark_err}") - - raise UpstreamException( - message=f"Upload authentication failed: {response.status_code}", - details={ - "status": response.status_code, - "response": response.text[:200], - "token_invalidated": True, - }, - ) - - logger.error( - f"Upload failed: {filename} - {response.status_code}", - extra={"response": response.text[:200]}, - ) + # 认证失败 + if response.status_code in (401, 403): + logger.warning(f"Upload auth failed: {response.status_code}") + try: + await TokenService.record_fail(token, response.status_code, "upload_auth_failed") + except Exception as e: + logger.error(f"Failed to record token failure: {e}") + raise UpstreamException( - message=f"Upload failed with status {response.status_code}", - details={ - "status": response.status_code, - "response": response.text[:200], - }, + message=f"Upload authentication failed: {response.status_code}", + details={"status": response.status_code, "token_invalidated": True}, ) - except Exception as e: - logger.error(f"Upload error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Upload process error: {str(e)}") + # 其他错误 + logger.error(f"Upload failed: {filename} - {response.status_code}") + raise UpstreamException( + message=f"Upload failed: {response.status_code}", + details={"status": response.status_code}, + ) # ==================== 列表服务 ==================== - class ListService(BaseService): """文件列表查询服务""" async def iter_assets(self, token: str): - """ - 分页迭代资产列表 - """ - headers = self._headers(token, referer="https://grok.com/files") - base_params = { + """分页迭代资产列表""" + headers = self._build_headers(token, referer="https://grok.com/files") + params = { "pageSize": 50, "orderBy": "ORDER_BY_LAST_USE_TIME", "source": "SOURCE_ANY", @@ -422,29 +308,27 @@ class ListService(BaseService): async with AsyncSession() as session: while True: - params = dict(base_params) if page_token: if page_token in seen_tokens: - logger.warning( - "List pagination stopped due to repeated page token" - ) + logger.warning("Pagination stopped: repeated page token") break seen_tokens.add(page_token) params["pageToken"] = page_token + else: + params.pop("pageToken", None) response = await session.get( LIST_API, headers=headers, params=params, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - timeout=self.timeout, - proxies=self._proxies(), + impersonate=self.config.browser, + timeout=self.config.timeout, + proxies=self.config.get_proxies(), ) if response.status_code != 200: - logger.error(f"List failed: {response.status_code}") raise UpstreamException( - message=f"List assets failed: {response.status_code}", + message=f"List failed: {response.status_code}", details={"status": response.status_code}, ) @@ -457,161 +341,107 @@ class ListService(BaseService): break async def list(self, token: str) -> List[Dict]: - """ - 查询文件列表 - - Raises: - UpstreamException: 查询失败 - """ - try: - assets: List[Dict] = [] - async for page_assets in self.iter_assets(token): - assets.extend(page_assets) - - logger.info(f"List success: {len(assets)} files") - return assets - - except Exception as e: - logger.error(f"List error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"List assets error: {str(e)}") + """查询文件列表""" + assets = [] + async for page_assets in self.iter_assets(token): + assets.extend(page_assets) + logger.info(f"List success: {len(assets)} files") + return assets async def count(self, token: str) -> int: - """ - 统计资产数量（不保留明细） - """ - try: - total = 0 - async for page_assets in self.iter_assets(token): - total += len(page_assets) - return total - except Exception as e: - logger.error(f"List count error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"List assets error: {str(e)}") + """统计资产数量""" + total = 0 + async for page_assets in self.iter_assets(token): + total += len(page_assets) + return total # ==================== 删除服务 ==================== - class DeleteService(BaseService): """文件删除服务""" async def delete(self, token: str, asset_id: str) -> bool: - """ - 删除单个文件 + """删除单个文件""" + async with _get_assets_semaphore(): + session = await self._get_session() + response = await session.delete( + f"{DELETE_API}/{asset_id}", + headers=self._build_headers(token, referer="https://grok.com/files"), + impersonate=self.config.browser, + timeout=self.config.timeout, + proxies=self.config.get_proxies(), + ) - Raises: - UpstreamException: 删除失败 - """ - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - async with _ASSETS_SEMAPHORE: - try: - headers = self._headers(token, referer="https://grok.com/files") - url = f"{DELETE_API}/{asset_id}" + if response.status_code == 200: + logger.debug(f"Deleted: {asset_id}") + return True - session = await self._get_session() - response = await session.delete( - url, - headers=headers, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - timeout=self.timeout, - proxies=self._proxies(), - ) - - if response.status_code == 200: - logger.debug(f"Delete success: {asset_id}") - return True - - logger.error(f"Delete failed: {asset_id} - {response.status_code}") - #: Note: Returning False or raising Exception? - #: Assuming caller handles Exception for stricter control, or False for loose control. - #: Given "optimization" and "standardization", raising exceptions is better for API feedback. - raise UpstreamException( - message=f"Delete failed: {asset_id}", - details={"status": response.status_code}, - ) - - except Exception as e: - logger.error(f"Delete error: {asset_id} - {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Delete error: {str(e)}") + logger.error(f"Delete failed: {asset_id} - {response.status_code}") + raise UpstreamException( + message=f"Delete failed: {asset_id}", + details={"status": response.status_code}, + ) async def delete_all(self, token: str) -> Dict[str, int]: - """ - 删除所有文件 - """ - total = 0 - success = 0 - failed = 0 - list_service = ListService(self.proxy) + """删除所有文件""" + total = success = failed = 0 + list_service = ListService(self.config.proxy) + try: async for assets in list_service.iter_assets(token): if not assets: continue + total += len(assets) - - # 批量并发删除 - async def _delete_one(asset: Dict, index: int) -> bool: - await asyncio.sleep(0.01 * index) - asset_id = asset.get("assetId", "") - if asset_id: - try: - return await self.delete(token, asset_id) - except Exception: - return False - return False - - batch_size = get_config( - "performance.assets_delete_batch_size", DEFAULT_DELETE_BATCH_SIZE - ) - try: - batch_size = int(batch_size) - except Exception: - batch_size = DEFAULT_DELETE_BATCH_SIZE - batch_size = max(1, batch_size) - for i in range(0, len(assets), batch_size): - batch = assets[i : i + batch_size] - results = await asyncio.gather( - *[_delete_one(asset, idx) for idx, asset in enumerate(batch)] - ) - success += sum(results) - failed += len(batch) - sum(results) + batch_result = await self._delete_batch(token, assets) + success += batch_result["success"] + failed += batch_result["failed"] if total == 0: logger.info("No assets to delete") return {"total": 0, "success": 0, "failed": 0, "skipped": True} - except Exception as e: - logger.error(f"Delete all failed during list: {e}") - return {"total": total, "success": success, "failed": failed} finally: await list_service.close() logger.info(f"Delete all: total={total}, success={success}, failed={failed}") return {"total": total, "success": success, "failed": failed} + async def _delete_batch(self, token: str, assets: List[Dict]) -> Dict[str, int]: + """批量删除""" + batch_size = max(1, int(get_config("performance.assets_delete_batch_size"))) + success = failed = 0 + + for i in range(0, len(assets), batch_size): + batch = assets[i:i + batch_size] + results = await asyncio.gather( + *[self._delete_one(token, asset, idx) for idx, asset in enumerate(batch)], + return_exceptions=True + ) + success += sum(1 for r in results if r is True) + failed += sum(1 for r in results if r is not True) + + return {"success": success, "failed": failed} + + async def _delete_one(self, token: str, asset: Dict, index: int) -> bool: + """删除单个资产（带延迟）""" + await asyncio.sleep(0.01 * index) + asset_id = asset.get("assetId", "") + if not asset_id: + return False + try: + return await self.delete(token, asset_id) + except Exception: + return False + # ==================== 下载服务 ==================== - class DownloadService(BaseService): """文件下载服务""" - def __init__(self, proxy: str = None): + def __init__(self, proxy: Optional[str] = None): super().__init__(proxy) - # 创建缓存目录 self.base_dir = Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" self.image_dir = self.base_dir / "image" self.video_dir = self.base_dir / "video" @@ -625,214 +455,159 @@ class DownloadService(BaseService): filename = file_path.lstrip("/").replace("/", "-") return cache_dir / filename - async def download( - self, file_path: str, token: str, media_type: str = "image" - ) -> Tuple[Optional[Path], str]: - """ - 下载文件到本地 + def _get_mime(self, cache_path: Path, response=None) -> str: + """获取 MIME 类型""" + if response: + return response.headers.get("content-type", "application/octet-stream").split(";")[0] + return MIME_TYPES.get(cache_path.suffix.lower(), "application/octet-stream") - Raises: - UpstreamException: 下载失败 - """ - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - async with _ASSETS_SEMAPHORE: - try: - cache_path = self._cache_path(file_path, media_type) + async def download(self, file_path: str, token: str, media_type: str = "image") -> Tuple[Optional[Path], str]: + """下载文件到本地""" + async with _get_assets_semaphore(): + cache_path = self._cache_path(file_path, media_type) - # 如果已缓存 + # 检查缓存 + if cache_path.exists(): + logger.debug(f"Cache hit: {cache_path}") + return cache_path, self._get_mime(cache_path) + + # 文件锁防止并发下载 + lock_name = f"dl_{media_type}_{hashlib.sha1(str(cache_path).encode()).hexdigest()[:16]}" + async with _file_lock(lock_name, timeout=10): + # 双重检查 if cache_path.exists(): - logger.debug(f"Cache hit: {cache_path}") - mime_type = MIME_TYPES.get(cache_path.suffix.lower(), DEFAULT_MIME) - return cache_path, mime_type + return cache_path, self._get_mime(cache_path) - lock_name = f"download_{media_type}_{hashlib.sha1(str(cache_path).encode('utf-8')).hexdigest()[:16]}" - async with _file_lock(lock_name, timeout=10): - # Double-check after lock - if cache_path.exists(): - logger.debug(f"Cache hit after lock: {cache_path}") - mime_type = MIME_TYPES.get( - cache_path.suffix.lower(), DEFAULT_MIME - ) - return cache_path, mime_type + # 执行下载 + mime = await self._download_file(file_path, token, cache_path) + logger.info(f"Downloaded: {file_path}") - # 下载文件 - if not file_path.startswith("/"): - file_path = f"/{file_path}" + # 异步检查缓存限制 + asyncio.create_task(self.check_limit()) - url = f"{DOWNLOAD_API}{file_path}" - headers = self._dl_headers(token, file_path) + return cache_path, mime - session = await self._get_session() - response = await session.get( - url, - headers=headers, - proxies=self._proxies(), - timeout=self.timeout, - allow_redirects=True, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - stream=True, - ) + async def _download_file(self, file_path: str, token: str, cache_path: Path) -> str: + """执行下载""" + if not file_path.startswith("/"): + file_path = f"/{file_path}" - if response.status_code != 200: - raise UpstreamException( - message=f"Download failed: {response.status_code}", - details={"path": file_path, "status": response.status_code}, - ) + url = f"{DOWNLOAD_API}{file_path}" + headers = self._build_headers(token, download=True) - # 保存文件（分块写入，避免大文件占用内存） - tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp") - try: - async with aiofiles.open(tmp_path, "wb") as f: - if hasattr(response, "aiter_content"): - async for chunk in response.aiter_content(): - if chunk: - await f.write(chunk) - elif hasattr(response, "aiter_bytes"): - async for chunk in response.aiter_bytes(): - if chunk: - await f.write(chunk) - elif hasattr(response, "aiter_raw"): - async for chunk in response.aiter_raw(): - if chunk: - await f.write(chunk) - else: - await f.write(response.content) - os.replace(tmp_path, cache_path) - finally: - if tmp_path.exists() and not cache_path.exists(): - try: - tmp_path.unlink() - except Exception: - pass - mime_type = response.headers.get( - "content-type", DEFAULT_MIME - ).split(";")[0] + session = await self._get_session() + response = await session.get( + url, + headers=headers, + proxies=self.config.get_proxies(), + timeout=self.config.timeout, + allow_redirects=True, + impersonate=self.config.browser, + stream=True, + ) - logger.info(f"Download success: {file_path}") + if response.status_code != 200: + raise UpstreamException( + message=f"Download failed: {response.status_code}", + details={"path": file_path, "status": response.status_code}, + ) - # 检查缓存限制 - asyncio.create_task(self.check_limit()) - - return cache_path, mime_type - - except Exception as e: - logger.error(f"Download failed: {file_path} - {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Download error: {str(e)}") - - async def to_base64( - self, file_path: str, token: str, media_type: str = "image" - ) -> str: - """ - 下载文件并转为 base64 - """ + # 保存文件 + tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp") try: - cache_path, mime_type = await self.download(file_path, token, media_type) - - if not cache_path or not cache_path.exists(): - raise AppException("File download returned invalid path") - - # 使用基础服务的工具方法转换 - data_uri = self.to_b64(cache_path, mime_type) - - # 删除临时文件 - if data_uri: + async with aiofiles.open(tmp_path, "wb") as f: + # 尝试流式写入 + if hasattr(response, "aiter_content"): + async for chunk in response.aiter_content(): + if chunk: + await f.write(chunk) + else: + await f.write(response.content) + os.replace(tmp_path, cache_path) + finally: + if tmp_path.exists() and not cache_path.exists(): try: - cache_path.unlink() - except Exception as e: - logger.warning(f"Delete temp file failed: {e}") + tmp_path.unlink() + except Exception: + pass - return data_uri + return self._get_mime(cache_path, response) - except Exception as e: - logger.error(f"To base64 failed: {file_path} - {e}") - if isinstance(e, AppException): - raise e - raise AppException(f"Base64 conversion failed: {str(e)}") + async def to_base64(self, file_path: str, token: str, media_type: str = "image") -> str: + """下载并转 base64""" + cache_path, mime = await self.download(file_path, token, media_type) + if not cache_path or not cache_path.exists(): + raise AppException("Download failed: invalid path") + + data_uri = self.to_b64(cache_path, mime) + + # 删除临时文件 + if data_uri: + try: + cache_path.unlink() + except Exception: + pass + + return data_uri def get_stats(self, media_type: str = "image") -> Dict[str, Any]: """获取缓存统计""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"count": 0, "size_mb": 0.0} - - files = [f for f in cache_dir.glob("*") if f.is_file()] - allowed_exts = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS - files = [f for f in files if f.suffix.lower() in allowed_exts] + + allowed = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS + files = [f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed] total_size = sum(f.stat().st_size for f in files) - return {"count": len(files), "size_mb": round(total_size / 1024 / 1024, 2)} - def list_files( - self, media_type: str = "image", page: int = 1, page_size: int = 1000 - ) -> Dict[str, Any]: - """列出本地缓存文件""" + def list_files(self, media_type: str = "image", page: int = 1, page_size: int = 1000) -> Dict[str, Any]: + """列出缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"total": 0, "page": page, "page_size": page_size, "items": []} + + allowed = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS + files = [f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed] - files = [f for f in cache_dir.glob("*") if f.is_file()] - allowed_exts = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS - files = [f for f in files if f.suffix.lower() in allowed_exts] + # 构建文件列表 items = [] for f in files: try: stat = f.stat() - items.append( - { - "name": f.name, - "size_bytes": stat.st_size, - "mtime_ms": int(stat.st_mtime * 1000), - } - ) + items.append({ + "name": f.name, + "size_bytes": stat.st_size, + "mtime_ms": int(stat.st_mtime * 1000), + }) except Exception: continue items.sort(key=lambda x: x["mtime_ms"], reverse=True) + + # 分页 total = len(items) start = max(0, (page - 1) * page_size) - end = start + page_size - paged = items[start:end] + paged = items[start:start + page_size] - if media_type == "image": - for item in paged: - item["view_url"] = f"/v1/files/image/{item['name']}" - else: - preview_map = {} - if self.image_dir.exists(): - for img in self.image_dir.glob("*"): - if img.is_file() and img.suffix.lower() in IMAGE_EXTS: - preview_map.setdefault(img.stem, img.name) - for item in paged: - item["view_url"] = f"/v1/files/video/{item['name']}" - preview_name = preview_map.get(Path(item["name"]).stem) - if preview_name: - item["preview_url"] = f"/v1/files/image/{preview_name}" + # 添加 URL + for item in paged: + item["view_url"] = f"/v1/files/{media_type}/{item['name']}" return {"total": total, "page": page, "page_size": page_size, "items": paged} def delete_file(self, media_type: str, name: str) -> Dict[str, Any]: - """删除单个缓存文件""" + """删除缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir - safe_name = name.replace("/", "-") - file_path = cache_dir / safe_name - if not file_path.exists(): - return {"deleted": False} - try: - file_path.unlink() - return {"deleted": True} - except Exception: - return {"deleted": False} + file_path = cache_dir / name.replace("/", "-") + + if file_path.exists(): + try: + file_path.unlink() + return {"deleted": True} + except Exception: + pass + return {"deleted": False} def clear(self, media_type: str = "image") -> Dict[str, Any]: """清空缓存""" @@ -841,61 +616,41 @@ class DownloadService(BaseService): return {"count": 0, "size_mb": 0.0} files = list(cache_dir.glob("*")) - total_size = sum(f.stat().st_size for f in files) + total_size = sum(f.stat().st_size for f in files if f.is_file()) count = 0 for f in files: - try: - f.unlink() - count += 1 - except Exception as e: - logger.error(f"Failed to delete {f}: {e}") + if f.is_file(): + try: + f.unlink() + count += 1 + except Exception: + pass return {"count": count, "size_mb": round(total_size / 1024 / 1024, 2)} async def check_limit(self): - """检查并清理缓存限制""" - if self._cleanup_running: + """检查并清理缓存""" + if self._cleanup_running or not get_config("cache.enable_auto_clean"): return + self._cleanup_running = True try: async with _file_lock("cache_cleanup", timeout=5): - if not get_config( - "cache.enable_auto_clean", DEFAULT_CACHE_ENABLE_AUTO_CLEAN - ): - return - - limit_mb = get_config("cache.limit_mb", DEFAULT_CACHE_LIMIT_MB) - - # 统计总大小 - total_size = 0 - all_files = [] - - for d in [self.image_dir, self.video_dir]: - if d.exists(): - for f in d.glob("*"): - try: - stat = f.stat() - total_size += stat.st_size - all_files.append((f, stat.st_mtime, stat.st_size)) - except Exception: - pass - + limit_mb = get_config("cache.limit_mb") + all_files, total_size = self._collect_files() current_mb = total_size / 1024 / 1024 + if current_mb <= limit_mb: return - # 需要清理 - logger.info( - f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning up..." - ) - - # 按时间排序 - all_files.sort(key=lambda x: x[1]) - + # 清理到 80% + logger.info(f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning...") + all_files.sort(key=lambda x: x[1]) # 按时间排序 + deleted_count = 0 deleted_size = 0 - target_mb = limit_mb * 0.8 # 清理到 80% + target_mb = limit_mb * 0.8 for f, _, size in all_files: try: @@ -903,39 +658,32 @@ class DownloadService(BaseService): deleted_count += 1 deleted_size += size total_size -= size - if (total_size / 1024 / 1024) <= target_mb: break - except Exception as e: - logger.error(f"Cleanup failed for {f}: {e}") + except Exception: + pass - logger.info( - f"Cache cleanup: deleted {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)" - ) + logger.info(f"Cache cleanup: {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)") finally: self._cleanup_running = False - def get_public_url(self, file_path: str) -> str: - """ - 获取文件的公共访问 URL + def _collect_files(self) -> Tuple[List[Tuple[Path, float, int]], int]: + """收集所有缓存文件""" + total_size = 0 + all_files = [] - 如果配置了 app_url，则返回自托管 URL，否则返回 Grok 原始 URL - """ - app_url = get_config("app.app_url", DEFAULT_APP_URL) - if not app_url: - return f"{DOWNLOAD_API}{file_path if file_path.startswith('/') else '/' + file_path}" + for d in [self.image_dir, self.video_dir]: + if d.exists(): + for f in d.glob("*"): + if f.is_file(): + try: + stat = f.stat() + total_size += stat.st_size + all_files.append((f, stat.st_mtime, stat.st_size)) + except Exception: + pass - if not file_path.startswith("/"): - file_path = f"/{file_path}" - - # 自动添加 /v1/files 前缀 - return f"{app_url.rstrip('/')}/v1/files{file_path}" + return all_files, total_size -__all__ = [ - "BaseService", - "UploadService", - "ListService", - "DeleteService", - "DownloadService", -] +__all__ = ["BaseService", "UploadService", "ListService", "DeleteService", "DownloadService"] diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index f0c1b3b..ee7de10 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -10,39 +10,22 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config -from app.core.exceptions import ( - AppException, - UpstreamException, - ValidationException, - ErrorType, -) +from app.core.exceptions import AppException, UpstreamException, ValidationException, ErrorType from app.services.grok.models.model import ModelService from app.services.grok.services.assets import UploadService -from app.services.grok.processors.processor import StreamProcessor, CollectProcessor +from app.services.grok.processors import StreamProcessor, CollectProcessor from app.services.grok.utils.retry import retry_on_status from app.services.grok.utils.headers import apply_statsig, build_sso_cookie +from app.services.grok.utils.stream import wrap_stream_with_usage from app.services.token import get_token_manager, EffortType CHAT_API = "https://grok.com/rest/app-chat/conversations/new" -DEFAULT_TIMEOUT = 120 -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_TEMPORARY = True -DEFAULT_DISABLE_MEMORY = True -DEFAULT_STREAM = True -DEFAULT_BASE_PROXY_URL = "" - @dataclass class ChatRequest: """聊天请求数据""" - model: str messages: List[Dict[str, Any]] stream: bool = None @@ -52,85 +35,47 @@ class ChatRequest: class MessageExtractor: """消息内容提取器""" - # 需要上传的类型 - UPLOAD_TYPES = {"image_url", "input_audio", "file"} - # 视频模式不支持的类型 - VIDEO_UNSUPPORTED = {"input_audio", "file"} - @staticmethod - def extract( - messages: List[Dict[str, Any]], is_video: bool = False - ) -> tuple[str, List[tuple[str, str]]]: - """ - 从 OpenAI 消息格式提取内容 - - Args: - messages: OpenAI 格式消息列表 - is_video: 是否为视频模型 - - Returns: - (text, attachments): 拼接后的文本和需要上传的附件列表 - - Raises: - ValueError: 视频模型遇到不支持的内容类型 - """ + def extract(messages: List[Dict[str, Any]], is_video: bool = False) -> tuple[str, List[tuple[str, str]]]: + """从 OpenAI 消息格式提取内容，返回 (text, attachments)""" texts = [] - attachments = [] # 需要上传的附件 (URL 或 base64) - - # 先抽取每条消息的文本，保留角色信息用于合并 - extracted: List[Dict[str, str]] = [] + attachments = [] + extracted = [] for msg in messages: role = msg.get("role", "") content = msg.get("content", "") parts = [] - # 简单字符串内容 if isinstance(content, str): if content.strip(): parts.append(content) - - # 列表格式内容 elif isinstance(content, list): for item in content: item_type = item.get("type", "") - # 文本类型 if item_type == "text": - text = item.get("text", "") - if text.strip(): + if text := item.get("text", "").strip(): parts.append(text) - # 图片类型 elif item_type == "image_url": image_data = item.get("image_url", {}) - url = ( - image_data.get("url", "") - if isinstance(image_data, dict) - else str(image_data) - ) + url = image_data.get("url", "") if isinstance(image_data, dict) else str(image_data) if url: attachments.append(("image", url)) - # 音频类型 elif item_type == "input_audio": if is_video: raise ValueError("视频模型不支持 input_audio 类型") audio_data = item.get("input_audio", {}) - data = ( - audio_data.get("data", "") - if isinstance(audio_data, dict) - else str(audio_data) - ) + data = audio_data.get("data", "") if isinstance(audio_data, dict) else str(audio_data) if data: attachments.append(("audio", data)) - # 文件类型 elif item_type == "file": if is_video: raise ValueError("视频模型不支持 file 类型") file_data = item.get("file", {}) - # file 可能是 URL 或 base64 url = file_data.get("url", "") or file_data.get("data", "") if isinstance(file_data, str): url = file_data @@ -140,30 +85,15 @@ class MessageExtractor: if parts: extracted.append({"role": role, "text": "\n".join(parts)}) - # 合并文本 - last_user_index = None - for i in range(len(extracted) - 1, -1, -1): - if extracted[i]["role"] == "user": - last_user_index = i - break + # 找到最后一条 user 消息 + last_user_index = next((i for i in range(len(extracted) - 1, -1, -1) if extracted[i]["role"] == "user"), None) for i, item in enumerate(extracted): role = item["role"] or "user" text = item["text"] - if i == last_user_index: - texts.append(text) - else: - texts.append(f"{role}: {text}") + texts.append(text if i == last_user_index else f"{role}: {text}") - # 换行拼接文本 - message = "\n\n".join(texts) - return message, attachments - - @staticmethod - def extract_text_only(messages: List[Dict[str, Any]]) -> str: - """仅提取文本内容""" - text, _ = MessageExtractor.extract(messages, is_video=True) - return text + return "\n\n".join(texts), attachments class ChatRequestBuilder: @@ -172,7 +102,7 @@ class ChatRequestBuilder: @staticmethod def build_headers(token: str) -> Dict[str, str]: """构造请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -202,33 +132,17 @@ class ChatRequestBuilder: return headers @staticmethod - def build_payload( - message: str, - model: str, - mode: str = None, - file_attachments: List[str] = None, - image_attachments: List[str] = None, - ) -> Dict[str, Any]: - """ - 构造请求体 - - Args: - message: 消息文本 - model: 模型名称 - mode: 模型模式 - file_attachments: 文件附件 ID 列表 - image_attachments: 图片附件 ID 列表（合并到 fileAttachments） - """ - temporary = get_config("grok.temporary", DEFAULT_TEMPORARY) - - merged_attachments: List[str] = [] + def build_payload(message: str, model: str, mode: str = None, + file_attachments: List[str] = None, image_attachments: List[str] = None) -> Dict[str, Any]: + """构造请求体""" + merged_attachments = [] if file_attachments: merged_attachments.extend(file_attachments) if image_attachments: merged_attachments.extend(image_attachments) - payload: Dict[str, Any] = { - "temporary": temporary, + payload = { + "temporary": get_config("grok.temporary"), "modelName": model, "message": message, "fileAttachments": merged_attachments, @@ -236,23 +150,17 @@ class ChatRequestBuilder: "disableSearch": False, "enableImageGeneration": True, "returnImageBytes": False, - "returnRawGrokInXaiRequest": False, "enableImageStreaming": True, "imageGenerationCount": 2, "forceConcise": False, "toolOverrides": {}, "enableSideBySide": True, "sendFinalMetadata": True, - "isReasoning": False, - "disableTextFollowUps": False, "responseMetadata": { "modelConfigOverride": {"modelMap": {}}, "requestModelDetails": {"modelId": model}, }, - "disableMemory": get_config("grok.disable_memory", DEFAULT_DISABLE_MEMORY), - "forceSideBySide": False, - "isAsyncChat": False, - "disableSelfHarmShortCircuit": False, + "disableMemory": get_config("grok.disable_memory"), "deviceEnvInfo": { "darkModeEnabled": False, "devicePixelRatio": 2, @@ -269,70 +177,31 @@ class ChatRequestBuilder: return payload -# ==================== Grok 服务 ==================== - - class GrokChatService: """Grok API 调用服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.proxy = proxy or get_config("grok.base_proxy_url") - async def chat( - self, - token: str, - message: str, - model: str = "grok-3", - mode: str = None, - stream: bool = None, - file_attachments: List[str] = None, - image_attachments: List[str] = None, - raw_payload: Dict[str, Any] = None, - ): - """ - 发送聊天请求 - - Args: - token: 认证 Token - message: 消息文本 - model: Grok 模型名称 - mode: 模型模式 - stream: 是否流式 - file_attachments: 文件附件 ID 列表 - image_attachments: 图片附件 URL 列表 - - Raises: - UpstreamException: 当 Grok API 返回错误且重试耗尽时 - """ + async def chat(self, token: str, message: str, model: str = "grok-3", mode: str = None, + stream: bool = None, file_attachments: List[str] = None, + image_attachments: List[str] = None, raw_payload: Dict[str, Any] = None): + """发送聊天请求""" if stream is None: - stream = get_config("grok.stream", DEFAULT_STREAM) + stream = get_config("grok.stream") headers = ChatRequestBuilder.build_headers(token) - if raw_payload is not None: - payload = raw_payload - else: - payload = ChatRequestBuilder.build_payload( - message, - model, - mode, - file_attachments, - image_attachments, - ) + payload = raw_payload if raw_payload is not None else ChatRequestBuilder.build_payload( + message, model, mode, file_attachments, image_attachments + ) proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None - timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) - - # 状态码提取器 - def extract_status(e: Exception) -> int | None: - if isinstance(e, UpstreamException) and e.details: - return e.details.get("status") - return None + timeout = get_config("grok.timeout") logger.debug(f"Payload: {payload}") - # 建立连接函数 + # 建立连接 async def establish_connection(): - """建立连接并返回 response 对象""" - browser = get_config("grok.browser", DEFAULT_BROWSER) + browser = get_config("grok.browser") session = AsyncSession(impersonate=browser) try: response = await session.post( @@ -345,63 +214,44 @@ class GrokChatService: ) if response.status_code != 200: + content = "" try: content = await response.text() - except Exception: - content = "" - - logger.error( - f"Chat failed: {response.status_code}", - extra={ - "status": response.status_code, - "token": token[:10] + "...", - }, - ) - body_for_log = content if content else "" - logger.debug( - "Grok API error response " - f"(status={response.status_code}): {body_for_log}" - ) - # 关闭 session 并抛出异常 - try: - await session.close() except Exception: pass + + logger.error(f"Chat failed: {response.status_code}", + extra={"status": response.status_code, "token": token[:10] + "..."}) + + await session.close() raise UpstreamException( message=f"Grok API request failed: {response.status_code}", - details={ - "status": response.status_code, - "body": content, - }, + details={"status": response.status_code, "body": content}, ) - # 返回 session 和 response return session, response except UpstreamException: - # 已经处理过的异常，直接抛出 raise except Exception as e: - # 其他异常，关闭 session 并包装 logger.error(f"Chat request error: {e}") - try: - await session.close() - except Exception: - pass + await session.close() raise UpstreamException( message=f"Chat connection failed: {str(e)}", details={"error": str(e)}, ) - # 建立连接 + # 重试机制 + def extract_status(e: Exception) -> int | None: + if isinstance(e, UpstreamException) and e.details: + return e.details.get("status") + return None + session = None response = None try: - session, response = await retry_on_status( - establish_connection, extract_status=extract_status - ) + session, response = await retry_on_status(establish_connection, extract_status=extract_status) except Exception as e: - # 记录失败 status_code = extract_status(e) if status_code: token_mgr = await get_token_manager() @@ -436,124 +286,47 @@ class GrokChatService: # 提取消息和附件 try: - message, attachments = MessageExtractor.extract( - request.messages, is_video=is_video - ) + message, attachments = MessageExtractor.extract(request.messages, is_video=is_video) except ValueError as e: raise ValidationException(str(e)) - # 处理附件上传 + # 上传附件 file_ids = [] - if attachments: upload_service = UploadService() try: - for attach_type, attach_data in attachments: - # 获取 ID + for _, attach_data in attachments: file_id, _ = await upload_service.upload(attach_data, token) - - if attach_type == "image": - file_ids.append(file_id) - logger.debug(f"Image uploaded: {file_id}") - else: - file_ids.append(file_id) - logger.debug(f"File uploaded: {file_id}") + file_ids.append(file_id) finally: await upload_service.close() - stream = ( - request.stream - if request.stream is not None - else get_config("grok.stream", DEFAULT_STREAM) - ) + stream = request.stream if request.stream is not None else get_config("grok.stream") response = await self.chat( - token, - message, - grok_model, - mode, - stream, - file_attachments=file_ids, - image_attachments=[], + token, message, grok_model, mode, stream, + file_attachments=file_ids, image_attachments=[] ) return response, stream, request.model -# ==================== Chat 业务服务 ==================== - - class ChatService: """Chat 业务服务""" @staticmethod - async def _wrap_stream(stream: AsyncGenerator, token_mgr, token: str, model: str): - """ - 包装流式响应，在完成时记录使用 - - Args: - stream: 原始 AsyncGenerator - token_mgr: TokenManager 实例 - token: Token 字符串 - model: 模型名称 - """ - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - # 只在成功完成时记录使用，失败/异常时不扣费 - if success: - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record stream usage: {e}") - - @staticmethod - async def completions( - model: str, - messages: List[Dict[str, Any]], - stream: bool = None, - thinking: str = None, - ): - """ - Chat Completions 入口 - - Args: - model: 模型名称 - messages: 消息列表 - stream: 是否流式 - thinking: 思考模式 - - Returns: - AsyncGenerator 或 dict - """ + async def completions(model: str, messages: List[Dict[str, Any]], + stream: bool = None, thinking: str = None): + """Chat Completions 入口""" # 获取 token - try: - token_mgr = await get_token_manager() - await token_mgr.reload_if_stale() - token = None - for pool_name in ModelService.pool_candidates_for_model(model): - token = token_mgr.get_token(pool_name) - if token: - break - except Exception as e: - logger.error(f"Failed to get token: {e}") - raise AppException( - message="Internal service error obtaining token", - error_type=ErrorType.SERVER.value, - code="internal_error", - ) + token_mgr = await get_token_manager() + await token_mgr.reload_if_stale() + + token = None + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name) + if token: + break if not token: raise AppException( @@ -564,61 +337,30 @@ class ChatService: ) # 解析参数 - think = None - if thinking == "enabled": - think = True - elif thinking == "disabled": - think = False - - is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM) + think = {"enabled": True, "disabled": False}.get(thinking) + is_stream = stream if stream is not None else get_config("grok.stream") # 构造请求 - chat_request = ChatRequest( - model=model, messages=messages, stream=is_stream, think=think - ) + chat_request = ChatRequest(model=model, messages=messages, stream=is_stream, think=think) # 请求 Grok service = GrokChatService() - try: - response, _, model_name = await service.chat_openai(token, chat_request) - except AppException: - raise - except Exception as e: - logger.error(f"Chat service error: {e}") - raise UpstreamException( - message=f"Service processing failed: {str(e)}", - details={"error": str(e)}, - ) + response, _, model_name = await service.chat_openai(token, chat_request) # 处理响应 if is_stream: processor = StreamProcessor(model_name, token, think) - return ChatService._wrap_stream( - processor.process(response), token_mgr, token, model - ) - else: - result = await CollectProcessor(model_name, token).process(response) - # 非流式：处理完成后立即记录使用 - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Collect completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record collect usage: {e}") - return result + return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) + + # 非流式 + result = await CollectProcessor(model_name, token).process(response) + try: + model_info = ModelService.get(model) + effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + await token_mgr.consume(token, effort) + except Exception as e: + logger.warning(f"Failed to record usage: {e}") + return result -__all__ = [ - "GrokChatService", - "ChatRequest", - "ChatRequestBuilder", - "MessageExtractor", - "ChatService", -] +__all__ = ["GrokChatService", "ChatRequest", "ChatRequestBuilder", "MessageExtractor", "ChatService"] diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 7ca9ade..06b1dd7 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -18,25 +18,12 @@ from app.core.config import get_config from app.core.logger import logger from app.services.grok.utils.headers import build_sso_cookie - WS_URL = "wss://grok.com/ws/imagine/listen" -DEFAULT_TIMEOUT = 120 -DEFAULT_BLOCKED_SECONDS = 15 -DEFAULT_FINAL_MIN_BYTES = 100000 -DEFAULT_MEDIUM_MIN_BYTES = 30000 -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_BASE_PROXY_URL = "" - class _BlockedError(Exception): pass - class ImageService: """Grok Imagine WebSocket image service.""" @@ -44,25 +31,22 @@ class ImageService: self._ssl_context = ssl.create_default_context() self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)") - def _get_ws_url(self) -> str: - return WS_URL - def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]: - proxy_url = get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + proxy_url = get_config("grok.base_proxy_url") if not proxy_url: return aiohttp.TCPConnector(ssl=self._ssl_context), None scheme = urlparse(proxy_url).scheme.lower() if scheme.startswith("socks"): - logger.info(f"Grok Imagine WebSocket using SOCKS proxy: {proxy_url}") + logger.info(f"Using SOCKS proxy: {proxy_url}") return ProxyConnector.from_url(proxy_url, ssl=self._ssl_context), None - logger.info(f"Grok Imagine WebSocket using HTTP proxy: {proxy_url}") + logger.info(f"Using HTTP proxy: {proxy_url}") return aiohttp.TCPConnector(ssl=self._ssl_context), proxy_url def _get_ws_headers(self, token: str) -> Dict[str, str]: cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") return { "Cookie": cookie, "Origin": "https://grok.com", @@ -74,15 +58,10 @@ class ImageService: def _extract_image_id(self, url: str) -> Optional[str]: match = self._url_pattern.search(url or "") - if match: - return match.group(1) - return None + return match.group(1) if match else None def _is_final_image(self, url: str, blob_size: int) -> bool: - min_bytes = int( - get_config("grok.image_ws_final_min_bytes", DEFAULT_FINAL_MIN_BYTES) - ) - return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > min_bytes + return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > get_config("grok.image_ws_final_min_bytes") def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]: if not url or not blob: @@ -91,12 +70,8 @@ class ImageService: image_id = self._extract_image_id(url) or uuid.uuid4().hex blob_size = len(blob) is_final = self._is_final_image(url, blob_size) - if is_final: - stage = "final" - elif blob_size > DEFAULT_MEDIUM_MIN_BYTES: - stage = "medium" - else: - stage = "preview" + + stage = "final" if is_final else ("medium" if blob_size > get_config("grok.image_ws_medium_min_bytes") else "preview") return { "type": "image", @@ -117,37 +92,23 @@ class ImageService: enable_nsfw: bool = True, max_retries: int = None, ) -> AsyncGenerator[Dict[str, object], None]: - if max_retries is None: - retries = 1 - else: - retries = max_retries - retries = max(1, retries) + retries = max(1, max_retries if max_retries is not None else 1) for attempt in range(retries): try: yielded_any = False - async for item in self._stream_once( - token, prompt, aspect_ratio, n, enable_nsfw - ): + async for item in self._stream_once(token, prompt, aspect_ratio, n, enable_nsfw): yielded_any = True yield item return except _BlockedError: - if yielded_any: + if yielded_any or attempt + 1 >= retries: + if not yielded_any: + yield {"type": "error", "error_code": "blocked", "error": "blocked_no_final_image"} return - if attempt + 1 < retries: - logger.warning( - f"Grok Imagine WebSocket blocked, retry {attempt + 1}/{retries}" - ) - continue - yield { - "type": "error", - "error_code": "blocked", - "error": "blocked_no_final_image", - } - return + logger.warning(f"WebSocket blocked, retry {attempt + 1}/{retries}") except Exception as e: - logger.error(f"Grok Imagine WebSocket stream failed: {e}") + logger.error(f"WebSocket stream failed: {e}") return async def _stream_once( @@ -158,25 +119,21 @@ class ImageService: n: int, enable_nsfw: bool, ) -> AsyncGenerator[Dict[str, object], None]: - ws_url = self._get_ws_url() - request_id = str(uuid.uuid4()) headers = self._get_ws_headers(token) - timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT)) - blocked_seconds = float( - get_config("grok.image_ws_blocked_seconds", DEFAULT_BLOCKED_SECONDS) - ) + timeout = float(get_config("grok.timeout")) + blocked_seconds = float(get_config("grok.image_ws_blocked_seconds")) try: connector, proxy = self._resolve_proxy() except Exception as e: - logger.error(f"Grok Imagine WebSocket proxy setup failed: {e}") + logger.error(f"WebSocket proxy setup failed: {e}") return try: async with aiohttp.ClientSession(connector=connector) as session: async with session.ws_connect( - ws_url, + WS_URL, headers=headers, heartbeat=20, receive_timeout=timeout, @@ -206,27 +163,21 @@ class ImageService: } await ws.send_json(message) - logger.info(f"Grok Imagine WebSocket request sent: {prompt[:80]}...") + logger.info(f"WebSocket request sent: {prompt[:80]}...") - images: Dict[str, Dict[str, object]] = {} + images = {} completed = 0 - start_time = time.time() - last_activity = time.time() + start_time = last_activity = time.time() medium_received_time = None while time.time() - start_time < timeout: try: ws_msg = await asyncio.wait_for(ws.receive(), timeout=5.0) except asyncio.TimeoutError: - if medium_received_time and completed == 0: - if time.time() - medium_received_time > min( - 10, blocked_seconds - ): - raise _BlockedError() + if medium_received_time and completed == 0 and time.time() - medium_received_time > min(10, blocked_seconds): + raise _BlockedError() if completed > 0 and time.time() - last_activity > 10: - logger.info( - f"Grok Imagine WebSocket idle timeout, collected {completed} images" - ) + logger.info(f"WebSocket idle timeout, collected {completed} images") break continue @@ -236,73 +187,43 @@ class ImageService: msg_type = msg.get("type") if msg_type == "image": - info = self._classify_image( - msg.get("url", ""), msg.get("blob", "") - ) + info = self._classify_image(msg.get("url", ""), msg.get("blob", "")) if not info: continue image_id = info["image_id"] existing = images.get(image_id, {}) + if info["stage"] == "medium" and medium_received_time is None: medium_received_time = time.time() if info["is_final"] and not existing.get("is_final"): completed += 1 - images[image_id] = { - "is_final": info["is_final"] or existing.get("is_final") - } - + images[image_id] = {"is_final": info["is_final"] or existing.get("is_final")} yield info elif msg_type == "error": - error_code = msg.get("err_code", "") - error_msg = msg.get("err_msg", "") - logger.warning( - f"Grok Imagine WebSocket error: {error_code} - {error_msg}" - ) - yield { - "type": "error", - "error_code": error_code, - "error": error_msg, - } + logger.warning(f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}") + yield {"type": "error", "error_code": msg.get("err_code", ""), "error": msg.get("err_msg", "")} return if completed >= n: - logger.info( - f"Grok Imagine WebSocket collected {completed} final images" - ) + logger.info(f"WebSocket collected {completed} final images") break - if medium_received_time and completed == 0: - if time.time() - medium_received_time > blocked_seconds: - raise _BlockedError() + if medium_received_time and completed == 0 and time.time() - medium_received_time > blocked_seconds: + raise _BlockedError() - elif ws_msg.type in ( - aiohttp.WSMsgType.CLOSED, - aiohttp.WSMsgType.ERROR, - ): - logger.warning( - f"Grok Imagine WebSocket closed/error: {ws_msg.type}" - ) - yield { - "type": "error", - "error_code": "ws_closed", - "error": f"websocket closed: {ws_msg.type}", - } + elif ws_msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR): + logger.warning(f"WebSocket closed/error: {ws_msg.type}") + yield {"type": "error", "error_code": "ws_closed", "error": f"websocket closed: {ws_msg.type}"} break except aiohttp.ClientError as e: - logger.error(f"Grok Imagine WebSocket connection error: {e}") - yield { - "type": "error", - "error_code": "connection_failed", - "error": str(e), - } - + logger.error(f"WebSocket connection error: {e}") + yield {"type": "error", "error_code": "connection_failed", "error": str(e)} image_service = ImageService() - __all__ = ["image_service", "ImageService"] diff --git a/app/services/grok/services/media.py b/app/services/grok/services/media.py index 0780957..8004028 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/media.py @@ -18,39 +18,39 @@ from app.core.exceptions import ( ) from app.services.grok.models.model import ModelService from app.services.token import get_token_manager, EffortType -from app.services.grok.processors.processor import VideoStreamProcessor, VideoCollectProcessor +from app.services.grok.processors import VideoStreamProcessor, VideoCollectProcessor from app.services.grok.utils.headers import apply_statsig, build_sso_cookie - +from app.services.grok.utils.stream import wrap_stream_with_usage CREATE_POST_API = "https://grok.com/rest/media/post/create" CHAT_API = "https://grok.com/rest/app-chat/conversations/new" -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_TIMEOUT = 300 -DEFAULT_MAX_CONCURRENT = 50 -DEFAULT_BASE_PROXY_URL = "" -DEFAULT_STREAM = True -_MEDIA_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) -_MEDIA_SEM_VALUE = DEFAULT_MAX_CONCURRENT +_MEDIA_SEMAPHORE = None +_MEDIA_SEM_VALUE = 0 + + +def _get_semaphore() -> asyncio.Semaphore: + """获取或更新信号量""" + global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE + value = max(1, int(get_config("performance.media_max_concurrent"))) + if value != _MEDIA_SEM_VALUE: + _MEDIA_SEM_VALUE = value + _MEDIA_SEMAPHORE = asyncio.Semaphore(value) + return _MEDIA_SEMAPHORE class VideoService: """视频生成服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + self.proxy = proxy or get_config("grok.base_proxy_url") + self.timeout = get_config("grok.timeout") def _build_headers( self, token: str, referer: str = "https://grok.com/imagine" ) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -84,27 +84,14 @@ class VideoService: return {"http": self.proxy, "https": self.proxy} if self.proxy else None async def create_post( - self, - token: str, - prompt: str, + self, token: str, prompt: str, media_type: str = "MEDIA_POST_TYPE_VIDEO", - media_url: str = None, + media_url: str = None ) -> str: - """ - 创建媒体帖子 - - Args: - token: 认证 Token - prompt: 提示词（视频生成用） - media_type: 媒体类型 (MEDIA_POST_TYPE_VIDEO 或 MEDIA_POST_TYPE_IMAGE) - media_url: 媒体 URL（图片模式用） - - Returns: - post ID - """ + """创建媒体帖子，返回 post ID""" try: headers = self._build_headers(token) - + # 根据类型构建不同的载荷 if media_type == "MEDIA_POST_TYPE_IMAGE" and media_url: payload = {"mediaType": media_type, "mediaUrl": media_url} @@ -112,352 +99,160 @@ class VideoService: payload = {"mediaType": media_type, "prompt": prompt} async with AsyncSession() as session: - browser = get_config("grok.browser", DEFAULT_BROWSER) response = await session.post( CREATE_POST_API, headers=headers, json=payload, - impersonate=browser, + impersonate=get_config("grok.browser"), timeout=30, proxies=self._build_proxies(), ) if response.status_code != 200: logger.error(f"Create post failed: {response.status_code}") - raise UpstreamException( - f"Failed to create post: {response.status_code}" - ) - - data = response.json() - post_id = data.get("post", {}).get("id", "") + raise UpstreamException(f"Failed to create post: {response.status_code}") + post_id = response.json().get("post", {}).get("id", "") if not post_id: raise UpstreamException("No post ID in response") logger.info(f"Media post created: {post_id} (type={media_type})") return post_id + except AppException: + raise except Exception as e: logger.error(f"Create post error: {e}") - if isinstance(e, AppException): - raise e raise UpstreamException(f"Create post error: {str(e)}") async def create_image_post(self, token: str, image_url: str) -> str: - """ - 创建图片帖子 - - Args: - token: 认证 Token - image_url: 完整的图片 URL (https://assets.grok.com/...) - - Returns: - post ID - """ - return await self.create_post( - token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url - ) + """创建图片帖子，返回 post ID""" + return await self.create_post(token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url) def _build_payload( - self, - prompt: str, - post_id: str, + self, prompt: str, post_id: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution_name: str = "480p", - preset: str = "normal", + preset: str = "normal" ) -> dict: """构建视频生成载荷""" - mode_flag = "--mode=custom" - if preset == "fun": - mode_flag = "--mode=extremely-crazy" - elif preset == "normal": - mode_flag = "--mode=normal" - elif preset == "spicy": - mode_flag = "--mode=extremely-spicy-or-crazy" - - full_prompt = f"{prompt} {mode_flag}" - video_gen_config = { - "parentPostId": post_id, - "aspectRatio": aspect_ratio, - "videoLength": video_length, - "resolutionName": resolution_name, - } + mode_map = {"fun": "--mode=extremely-crazy", "normal": "--mode=normal", "spicy": "--mode=extremely-spicy-or-crazy"} + mode_flag = mode_map.get(preset, "--mode=custom") + return { "temporary": True, "modelName": "grok-3", - "message": full_prompt, + "message": f"{prompt} {mode_flag}", "toolOverrides": {"videoGen": True}, "enableSideBySide": True, "responseMetadata": { "experiments": [], "modelConfigOverride": { "modelMap": { - "videoGenModelConfig": video_gen_config + "videoGenModelConfig": { + "parentPostId": post_id, + "aspectRatio": aspect_ratio, + "videoLength": video_length, + "resolutionName": resolution_name, + } } }, }, } + async def _generate_internal( + self, token: str, post_id: str, prompt: str, + aspect_ratio: str, video_length: int, + resolution_name: str, preset: str + ) -> AsyncGenerator[bytes, None]: + """内部生成逻辑""" + session = None + try: + headers = self._build_headers(token) + payload = self._build_payload(prompt, post_id, aspect_ratio, video_length, resolution_name, preset) + + session = AsyncSession(impersonate=get_config("grok.browser")) + response = await session.post( + CHAT_API, + headers=headers, + data=orjson.dumps(payload), + timeout=self.timeout, + stream=True, + proxies=self._build_proxies(), + ) + + if response.status_code != 200: + logger.error(f"Video generation failed: {response.status_code}") + raise UpstreamException( + message=f"Video generation failed: {response.status_code}", + details={"status": response.status_code}, + ) + + async def stream_response(): + try: + async for line in response.aiter_lines(): + yield line + finally: + await session.close() + + return stream_response() + + except Exception as e: + if session: + try: + await session.close() + except Exception: + pass + logger.error(f"Video generation error: {e}") + if isinstance(e, AppException): + raise + raise UpstreamException(f"Video generation error: {str(e)}") + async def generate( - self, - token: str, - prompt: str, + self, token: str, prompt: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution_name: str = "480p", - preset: str = "normal", + preset: str = "normal" ) -> AsyncGenerator[bytes, None]: - """ - 生成视频 - - Args: - token: 认证 Token - prompt: 视频描述 - aspect_ratio: 宽高比 - video_length: 视频时长 - resolution_name: 分辨率 - preset: 预设 - - Returns: - AsyncGenerator，流式传输 - - Raises: - UpstreamException: 连接失败时 - """ - value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE - if value != _MEDIA_SEM_VALUE: - _MEDIA_SEM_VALUE = value - _MEDIA_SEMAPHORE = asyncio.Semaphore(value) - async with _MEDIA_SEMAPHORE: - session = None - try: - # Step 1: 创建帖子 - post_id = await self.create_post(token, prompt) - - # Step 2: 建立连接 - headers = self._build_headers(token) - payload = self._build_payload( - prompt, post_id, aspect_ratio, video_length, resolution_name, preset - ) - - browser = get_config("grok.browser", DEFAULT_BROWSER) - session = AsyncSession(impersonate=browser) - response = await session.post( - CHAT_API, - headers=headers, - data=orjson.dumps(payload), - timeout=self.timeout, - stream=True, - proxies=self._build_proxies(), - ) - - if response.status_code != 200: - logger.error(f"Video generation failed: {response.status_code}") - try: - await session.close() - except Exception: - pass - raise UpstreamException( - message=f"Video generation failed: {response.status_code}", - details={"status": response.status_code}, - ) - - # Step 3: 流式传输 - async def stream_response(): - try: - async for line in response.aiter_lines(): - yield line - finally: - if session: - await session.close() - - return stream_response() - - except Exception as e: - if session: - try: - await session.close() - except Exception: - pass - logger.error(f"Video generation error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Video generation error: {str(e)}") + """生成视频""" + async with _get_semaphore(): + post_id = await self.create_post(token, prompt) + return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution_name, preset) async def generate_from_image( - self, - token: str, - prompt: str, - image_url: str, + self, token: str, prompt: str, image_url: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "480p", - preset: str = "normal", + preset: str = "normal" ) -> AsyncGenerator[bytes, None]: - """ - 从图片生成视频 - - Args: - token: 认证 Token - prompt: 视频描述 - image_url: 图片 URL - aspect_ratio: 宽高比 - video_length: 视频时长 - resolution: 分辨率 - preset: 预设 - - Returns: - AsyncGenerator，流式传输 - """ - value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE - if value != _MEDIA_SEM_VALUE: - _MEDIA_SEM_VALUE = value - _MEDIA_SEMAPHORE = asyncio.Semaphore(value) - async with _MEDIA_SEMAPHORE: - session = None - try: - # Step 1: 创建帖子 - post_id = await self.create_image_post(token, image_url) - - # Step 2: 建立连接 - headers = self._build_headers(token) - payload = self._build_payload( - prompt, post_id, aspect_ratio, video_length, resolution, preset - ) - - browser = get_config("grok.browser", DEFAULT_BROWSER) - session = AsyncSession(impersonate=browser) - response = await session.post( - CHAT_API, - headers=headers, - data=orjson.dumps(payload), - timeout=self.timeout, - stream=True, - proxies=self._build_proxies(), - ) - - if response.status_code != 200: - logger.error(f"Video from image failed: {response.status_code}") - try: - await session.close() - except Exception: - pass - raise UpstreamException( - message=f"Video from image failed: {response.status_code}", - details={"status": response.status_code}, - ) - - # Step 3: 流式传输 - async def stream_response(): - try: - async for line in response.aiter_lines(): - yield line - finally: - if session: - await session.close() - - return stream_response() - - except Exception as e: - if session: - try: - await session.close() - except Exception: - pass - logger.error(f"Video from image error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Video from image error: {str(e)}") - - @staticmethod - async def _wrap_stream(stream: AsyncGenerator, token_mgr, token: str, model: str): - """ - 包装流式响应，在完成时记录使用 - - Args: - stream: 原始 AsyncGenerator - token_mgr: TokenManager 实例 - token: Token 字符串 - model: 模型名称 - """ - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - # 只在成功完成时记录使用，失败/异常时不扣费 - if success: - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Video stream completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record video stream usage: {e}") + """从图片生成视频""" + async with _get_semaphore(): + post_id = await self.create_image_post(token, image_url) + return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution, preset) @staticmethod async def completions( - model: str, - messages: list, + model: str, messages: list, stream: bool = None, thinking: str = None, aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "480p", - preset: str = "normal", + preset: str = "normal" ): - """ - 视频生成入口 - - Args: - model: 模型名称 - messages: 消息列表 - stream: 是否流式 - thinking: 思考模式 - aspect_ratio: 宽高比 - video_length: 视频时长 - resolution: 分辨率 - preset: 预设模式 - - Returns: - AsyncGenerator (流式) 或 dict (非流式) - """ + """视频生成入口""" # 获取 token - try: - token_mgr = await get_token_manager() - await token_mgr.reload_if_stale() - token = None - for pool_name in ModelService.pool_candidates_for_model(model): - token = token_mgr.get_token(pool_name) - if token: - break - except Exception as e: - logger.error(f"Failed to get token: {e}") - raise AppException( - message="Internal service error obtaining token", - error_type=ErrorType.SERVER.value, - code="internal_error", - ) + token_mgr = await get_token_manager() + await token_mgr.reload_if_stale() + + token = None + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name) + if token: + break if not token: raise AppException( @@ -467,14 +262,8 @@ class VideoService: status_code=429, ) - # 解析参数 - think = None - if thinking == "enabled": - think = True - elif thinking == "disabled": - think = False - - is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM) + think = {"enabled": True, "disabled": False}.get(thinking) + is_stream = stream if stream is not None else get_config("grok.stream") # 提取内容 from app.services.grok.services.chat import MessageExtractor @@ -492,27 +281,18 @@ class VideoService: try: for attach_type, attach_data in attachments: if attach_type == "image": - # 上传图片 _, file_uri = await upload_service.upload(attach_data, token) image_url = f"https://assets.grok.com/{file_uri}" logger.info(f"Image uploaded for video: {image_url}") - break # 视频模型只使用第一张图片 + break finally: await upload_service.close() # 生成视频 service = VideoService() - - # 图片转视频 if image_url: response = await service.generate_from_image( - token, - prompt, - image_url, - aspect_ratio, - video_length, - resolution, - preset, + token, prompt, image_url, aspect_ratio, video_length, resolution, preset ) else: response = await service.generate( @@ -522,26 +302,16 @@ class VideoService: # 处理响应 if is_stream: processor = VideoStreamProcessor(model, token, think) - return VideoService._wrap_stream( - processor.process(response), token_mgr, token, model - ) - else: - result = await VideoCollectProcessor(model, token).process(response) - # 非流式：处理完成后立即记录使用 - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Video completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record video usage: {e}") - return result - + return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) + + result = await VideoCollectProcessor(model, token).process(response) + try: + model_info = ModelService.get(model) + effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + await token_mgr.consume(token, effort) + logger.debug(f"Video completed, recorded usage (effort={effort.value})") + except Exception as e: + logger.warning(f"Failed to record video usage: {e}") + return result __all__ = ["VideoService"] diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py index 88411be..169ed4d 100644 --- a/app/services/grok/services/nsfw.py +++ b/app/services/grok/services/nsfw.py @@ -20,16 +20,9 @@ from app.services.grok.protocols.grpc_web import ( ) from app.services.grok.utils.headers import build_sso_cookie - NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" BIRTH_DATE_API = "https://grok.com/rest/auth/set-birth-date" -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36" -DEFAULT_TIMEOUT = 30 -DEFAULT_BASE_PROXY_URL = "" - - @dataclass class NSFWResult: """NSFW 操作结果""" @@ -40,20 +33,22 @@ class NSFWResult: grpc_message: Optional[str] = None error: Optional[str] = None - class NSFWService: """NSFW 模式服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - self.timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT)) + self.proxy = proxy or get_config("grok.base_proxy_url") + self.timeout = float(get_config("grok.timeout")) + + def _build_proxies(self) -> Optional[dict]: + """构建代理配置""" + return {"http": self.proxy, "https": self.proxy} if self.proxy else None @staticmethod def _random_birth_date() -> str: - """生成随机出生日期""" + """生成随机出生日期（20-40岁）""" today = datetime.date.today() - age = random.randint(20, 40) - birth_year = today.year - age + birth_year = today.year - random.randint(20, 40) birth_month = random.randint(1, 12) birth_day = random.randint(1, 28) hour = random.randint(0, 23) @@ -65,7 +60,7 @@ class NSFWService: def _build_headers(self, token: str) -> dict: """构造 gRPC-Web 请求头""" cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") return { "accept": "*/*", "content-type": "application/grpc-web+proto", @@ -80,7 +75,7 @@ class NSFWService: def _build_birth_headers(self, token: str) -> dict: """构造设置出生日期请求头""" cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") return { "accept": "*/*", "content-type": "application/json", @@ -102,18 +97,20 @@ class NSFWService: protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner return encode_grpc_web_payload(protobuf) - async def _set_birth_date(self, session: AsyncSession, token: str) -> tuple[bool, int, Optional[str]]: + async def _set_birth_date( + self, session: AsyncSession, token: str + ) -> tuple[bool, int, Optional[str]]: """设置出生日期""" headers = self._build_birth_headers(token) payload = {"birthDate": self._random_birth_date()} - proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None + try: response = await session.post( BIRTH_DATE_API, json=payload, headers=headers, timeout=self.timeout, - proxies=proxies, + proxies=self._build_proxies(), ) if response.status_code in (200, 204): return True, response.status_code, None @@ -125,31 +122,27 @@ class NSFWService: """为单个 token 开启 NSFW 模式""" headers = self._build_headers(token) payload = self._build_payload() - logger.debug( - "NSFW payload: len={} hex={}", - len(payload), - payload.hex(), - ) - proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None + logger.debug(f"NSFW payload: len={len(payload)} hex={payload.hex()}") try: - browser = get_config("grok.browser", DEFAULT_BROWSER) + browser = get_config("grok.browser") async with AsyncSession(impersonate=browser) as session: - ok, birth_status, birth_err = await self._set_birth_date( - session, token - ) + # 先设置出生日期 + ok, birth_status, birth_err = await self._set_birth_date(session, token) if not ok: return NSFWResult( success=False, http_status=birth_status, error=f"Set birth date failed: {birth_err}", ) + + # 开启 NSFW response = await session.post( NSFW_API, data=payload, headers=headers, timeout=self.timeout, - proxies=proxies, + proxies=self._build_proxies(), ) if response.status_code != 200: @@ -160,18 +153,15 @@ class NSFWService: ) # 解析 gRPC-Web 响应 - content_type = response.headers.get("content-type") _, trailers = parse_grpc_web_response( - response.content, content_type=content_type + response.content, + content_type=response.headers.get("content-type") ) grpc_status = get_grpc_status(trailers) logger.debug( - "NSFW response: http={} grpc={} msg={} trailers={}", - response.status_code, - grpc_status.code, - grpc_status.message, - trailers, + f"NSFW response: http={response.status_code} grpc={grpc_status.code} " + f"msg={grpc_status.message} trailers={trailers}" ) # HTTP 200 且无 grpc-status（空响应）或 grpc-status=0 都算成功 @@ -188,5 +178,4 @@ class NSFWService: logger.error(f"NSFW enable failed: {e}") return NSFWResult(success=False, http_status=0, error=str(e)[:100]) - __all__ = ["NSFWService", "NSFWResult"] diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index dc2ff52..8c386f4 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -13,32 +13,22 @@ from app.core.exceptions import UpstreamException from app.services.grok.utils.headers import apply_statsig, build_sso_cookie from app.services.grok.utils.retry import retry_on_status - LIMITS_API = "https://grok.com/rest/rate-limits" -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_TIMEOUT = 10 -DEFAULT_MAX_CONCURRENT = 25 -DEFAULT_BASE_PROXY_URL = "" -_USAGE_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) -_USAGE_SEM_VALUE = DEFAULT_MAX_CONCURRENT +_USAGE_SEMAPHORE = asyncio.Semaphore(25) +_USAGE_SEM_VALUE = 25 class UsageService: """用量查询服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + self.proxy = proxy or get_config("grok.base_proxy_url") + self.timeout = get_config("grok.timeout") def _build_headers(self, token: str) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -85,11 +75,11 @@ class UsageService: Raises: UpstreamException: 当获取失败且重试耗尽时 """ - value = get_config("performance.usage_max_concurrent", DEFAULT_MAX_CONCURRENT) + value = get_config("performance.usage_max_concurrent") try: value = int(value) except Exception: - value = DEFAULT_MAX_CONCURRENT + value = 25 value = max(1, value) global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE if value != _USAGE_SEM_VALUE: @@ -107,7 +97,7 @@ class UsageService: try: headers = self._build_headers(token) payload = {"requestKind": "DEFAULT", "modelName": model_name} - browser = get_config("grok.browser", DEFAULT_BROWSER) + browser = get_config("grok.browser") async with AsyncSession() as session: response = await session.post( @@ -152,5 +142,4 @@ class UsageService: # 最后一次失败已经被记录 raise - __all__ = ["UsageService"] diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 17a4de7..b887d31 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -12,18 +12,13 @@ from app.core.config import get_config from app.core.exceptions import UpstreamException from app.services.grok.utils.headers import apply_statsig, build_sso_cookie - LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens" -DEFAULT_TIMEOUT = 30 -DEFAULT_BROWSER = "chrome136" -DEFAULT_BASE_PROXY_URL = "" - class VoiceService: """Voice Mode Service (LiveKit)""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.proxy = proxy or get_config("grok.base_proxy_url") async def get_token( self, @@ -46,8 +41,8 @@ class VoiceService: proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None try: - browser = get_config("grok.browser", DEFAULT_BROWSER) - timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + browser = get_config("grok.browser") + timeout = get_config("grok.timeout") async with AsyncSession(impersonate=browser) as session: response = await session.post( LIVEKIT_TOKEN_API, diff --git a/app/services/grok/utils/headers.py b/app/services/grok/utils/headers.py index 44c94ae..146e704 100644 --- a/app/services/grok/utils/headers.py +++ b/app/services/grok/utils/headers.py @@ -11,16 +11,13 @@ from app.core.config import get_config from app.services.grok.utils.statsig import StatsigService -DEFAULT_CF_CLEARANCE = "" - - def _normalize_token(token: str) -> str: return token[4:] if token.startswith("sso=") else token def build_sso_cookie(token: str, include_rw: bool = False) -> str: token = _normalize_token(token) - cf = get_config("grok.cf_clearance", DEFAULT_CF_CLEARANCE) + cf = get_config("grok.cf_clearance") cookie = f"sso={token}" if include_rw: cookie = f"{cookie}; sso-rw={token}" diff --git a/app/services/grok/utils/retry.py b/app/services/grok/utils/retry.py index 9a69e87..2cd52c5 100644 --- a/app/services/grok/utils/retry.py +++ b/app/services/grok/utils/retry.py @@ -18,38 +18,22 @@ from app.core.config import get_config from app.core.exceptions import UpstreamException -DEFAULT_MAX_RETRY = 3 -DEFAULT_RETRY_STATUS_CODES = [401, 429, 403] -DEFAULT_RETRY_BACKOFF_BASE = 0.5 -DEFAULT_RETRY_BACKOFF_FACTOR = 2.0 -DEFAULT_RETRY_BACKOFF_MAX = 30.0 -DEFAULT_RETRY_BUDGET = 90.0 - - class RetryContext: """重试上下文""" def __init__(self): self.attempt = 0 - self.max_retry = int(get_config("grok.max_retry", DEFAULT_MAX_RETRY)) - self.retry_codes = get_config( - "grok.retry_status_codes", DEFAULT_RETRY_STATUS_CODES - ) + self.max_retry = int(get_config("grok.max_retry")) + self.retry_codes = get_config("grok.retry_status_codes") self.last_error = None self.last_status = None self.total_delay = 0.0 - self.retry_budget = float(get_config("grok.retry_budget", DEFAULT_RETRY_BUDGET)) + self.retry_budget = float(get_config("grok.retry_budget")) # 退避参数 - self.backoff_base = float( - get_config("grok.retry_backoff_base", DEFAULT_RETRY_BACKOFF_BASE) - ) - self.backoff_factor = float( - get_config("grok.retry_backoff_factor", DEFAULT_RETRY_BACKOFF_FACTOR) - ) - self.backoff_max = float( - get_config("grok.retry_backoff_max", DEFAULT_RETRY_BACKOFF_MAX) - ) + self.backoff_base = float(get_config("grok.retry_backoff_base")) + self.backoff_factor = float(get_config("grok.retry_backoff_factor")) + self.backoff_max = float(get_config("grok.retry_backoff_max")) # decorrelated jitter 状态 self._last_delay = self.backoff_base diff --git a/app/services/grok/utils/statsig.py b/app/services/grok/utils/statsig.py index d5813b6..bf78741 100644 --- a/app/services/grok/utils/statsig.py +++ b/app/services/grok/utils/statsig.py @@ -9,9 +9,6 @@ import string from app.core.config import get_config -DEFAULT_DYNAMIC_STATSIG = True - - class StatsigService: """Statsig ID 生成服务""" @@ -33,8 +30,7 @@ class StatsigService: Returns: Base64 编码的 ID """ - # 读取配置 - dynamic = get_config("grok.dynamic_statsig", DEFAULT_DYNAMIC_STATSIG) + dynamic = get_config("grok.dynamic_statsig") if not dynamic: return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" diff --git a/app/services/grok/utils/stream.py b/app/services/grok/utils/stream.py new file mode 100644 index 0000000..fb107b7 --- /dev/null +++ b/app/services/grok/utils/stream.py @@ -0,0 +1,43 @@ +""" +流式响应通用工具 +""" + +from typing import AsyncGenerator + +from app.core.logger import logger +from app.services.grok.models.model import ModelService +from app.services.token import EffortType + + +async def wrap_stream_with_usage( + stream: AsyncGenerator, + token_mgr, + token: str, + model: str +) -> AsyncGenerator: + """ + 包装流式响应，在完成时记录使用 + + Args: + stream: 原始 AsyncGenerator + token_mgr: TokenManager 实例 + token: Token 字符串 + model: 模型名称 + """ + success = False + try: + async for chunk in stream: + yield chunk + success = True + finally: + if success: + try: + model_info = ModelService.get(model) + effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + await token_mgr.consume(token, effort) + logger.debug(f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})") + except Exception as e: + logger.warning(f"Failed to record stream usage: {e}") + + +__all__ = ["wrap_stream_with_usage"] diff --git a/data/config.toml b/data/config.toml index 3db391c..22296cf 100644 --- a/data/config.toml +++ b/data/config.toml @@ -1,47 +1,36 @@ [grok] -temporary = true -stream = true -thinking = true -dynamic_statsig = true -filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] +browser = "chrome136" timeout = 120 +user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" base_proxy_url = "" asset_proxy_url = "" cf_clearance = "" -browser = "chrome136" -user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" +temporary = true disable_memory = true +stream = true +thinking = true +filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] max_retry = 3 retry_status_codes = [401,429,403] -image_ws = true -image_ws_blocked_seconds = 15 -image_ws_final_min_bytes = 100000 -image_ws_nsfw = true retry_backoff_base = 0.5 retry_backoff_factor = 2.0 retry_backoff_max = 30.0 retry_budget = 90.0 stream_idle_timeout = 120.0 video_idle_timeout = 90.0 +image_ws_blocked_seconds = 15 +image_ws_final_min_bytes = 100000 +image_ws_medium_min_bytes = 30000 +dynamic_statsig = true +image_ws = true +image_ws_nsfw = true [app] app_url = "http://127.0.0.1:8000" -app_key = "grok2api" -api_key = "" image_format = "url" video_format = "html" - -[token] -auto_refresh = true -refresh_interval_hours = 8 -super_refresh_interval_hours = 2 -fail_threshold = 5 -save_delay_ms = 200 -reload_interval_sec = 30 - -[cache] -enable_auto_clean = true -limit_mb = 1024 +app_key = "grok2api" +api_key = "" [performance] assets_max_concurrent = 25 @@ -55,3 +44,15 @@ usage_max_tokens = 1000 nsfw_max_concurrent = 10 nsfw_batch_size = 50 nsfw_max_tokens = 1000 + +[cache] +enable_auto_clean = true +limit_mb = 1024 + +[token] +auto_refresh = true +refresh_interval_hours = 8 +super_refresh_interval_hours = 2 +fail_threshold = 5 +save_delay_ms = 200 +reload_interval_sec = 30 diff --git a/main.py b/main.py index bd81ff2..97da6ac 100644 --- a/main.py +++ b/main.py @@ -41,17 +41,21 @@ setup_logging( @asynccontextmanager async def lifespan(app: FastAPI): """应用生命周期管理""" - # 1. 加载配置 - from app.core.config import config - + # 1. 注册服务默认配置 + from app.core.config import config, register_defaults + from app.services.grok.defaults import get_grok_defaults + + register_defaults(get_grok_defaults()) + + # 2. 加载配置 await config.load() - # 2. 启动服务显示 + # 3. 启动服务显示 logger.info("Starting Grok2API...") logger.info(f"Platform: {platform.system()} {platform.release()}") logger.info(f"Python: {sys.version.split()[0]}") - # 3. 启动 Token 刷新调度器 + # 4. 启动 Token 刷新调度器 refresh_enabled = get_config("token.auto_refresh", True) if refresh_enabled: basic_interval = get_config("token.refresh_interval_hours", 8) diff --git a/tests/test_model.py b/tests/test_model.py index e58f263..b4e96e7 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -18,7 +18,7 @@ load_dotenv(ROOT / ".env") from app.core.config import config from app.core.logger import setup_logging from app.services.grok.services.chat import GrokChatService -from app.services.grok.processors.processor import CollectProcessor +from app.services.grok.processors import CollectProcessor from app.services.grok.services.usage import UsageService From 5d369596c513ce6b5cd44851ed20de81069e864b Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sat, 7 Feb 2026 13:11:23 +0800 Subject: [PATCH 3/6] feat: improve token management and logging across services --- app/api/v1/admin.py | 305 +++++++-------------------- app/api/v1/image.py | 235 ++++++--------------- app/services/grok/services/assets.py | 3 + app/services/grok/services/chat.py | 13 +- app/services/grok/services/image.py | 2 + app/services/grok/services/media.py | 6 +- app/services/grok/services/usage.py | 4 +- app/services/grok/services/voice.py | 11 +- 8 files changed, 176 insertions(+), 403 deletions(-) diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 7436409..1c8cbcd 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -6,6 +6,8 @@ from app.core.config import config, get_config from app.core.batch_tasks import create_task, get_task, expire_task from app.core.storage import get_storage, LocalStorage, RedisStorage, SQLStorage from app.core.exceptions import AppException +from app.services.token.manager import get_token_manager +from app.services.grok.utils.batch import run_in_batches import os from pathlib import Path import aiofiles @@ -13,23 +15,41 @@ import asyncio import orjson from app.core.logger import logger from app.services.grok.services.voice import VoiceService -from app.services.token import get_token_manager -DEFAULT_STORAGE_TYPE = "" -DEFAULT_USAGE_MAX_TOKENS = 1000 -DEFAULT_USAGE_MAX_CONCURRENT = 25 -DEFAULT_USAGE_BATCH_SIZE = 50 -DEFAULT_NSFW_MAX_TOKENS = 1000 -DEFAULT_NSFW_MAX_CONCURRENT = 10 -DEFAULT_NSFW_BATCH_SIZE = 50 -DEFAULT_ASSETS_MAX_CONCURRENT = 25 -DEFAULT_ASSETS_BATCH_SIZE = 10 -DEFAULT_ASSETS_MAX_TOKENS = 1000 TEMPLATE_DIR = Path(__file__).parent.parent.parent / "static" router = APIRouter() + +def _collect_tokens(data: dict) -> list[str]: + """从请求数据中收集 token 列表""" + tokens = [] + if isinstance(data.get("token"), str) and data["token"].strip(): + tokens.append(data["token"].strip()) + if isinstance(data.get("tokens"), list): + tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()]) + return tokens + + +def _truncate_tokens(tokens: list[str], max_tokens: int, operation: str = "operation") -> tuple[list[str], bool, int]: + """去重并截断 token 列表，返回 (unique_tokens, truncated, original_count)""" + unique_tokens = list(dict.fromkeys(tokens)) + original_count = len(unique_tokens) + truncated = False + + if len(unique_tokens) > max_tokens: + unique_tokens = unique_tokens[:max_tokens] + truncated = True + logger.warning(f"{operation}: truncated from {original_count} to {max_tokens} tokens") + + return unique_tokens, truncated, original_count + + +def _mask_token(token: str) -> str: + """掩码 token 显示""" + return f"{token[:8]}...{token[-8:]}" if len(token) > 20 else token + async def render_template(filename: str): """渲染指定模板""" template_path = TEMPLATE_DIR / filename @@ -217,10 +237,9 @@ async def update_config_api(data: dict): @router.get("/api/v1/admin/storage", dependencies=[Depends(verify_api_key)]) async def get_storage_info(): """获取当前存储模式""" - storage_type = os.getenv("SERVER_STORAGE_TYPE", "local").lower() - logger.info(f"Storage type: {storage_type}") + storage_type = os.getenv("SERVER_STORAGE_TYPE", "").lower() if not storage_type: - storage_type = str(get_config("storage.type", DEFAULT_STORAGE_TYPE)).lower() + storage_type = str(get_config("storage.type")).lower() if not storage_type: storage = get_storage() if isinstance(storage, LocalStorage): @@ -228,12 +247,10 @@ async def get_storage_info(): elif isinstance(storage, RedisStorage): storage_type = "redis" elif isinstance(storage, SQLStorage): - if storage.dialect in ("mysql", "mariadb"): - storage_type = "mysql" - elif storage.dialect in ("postgres", "postgresql", "pgsql"): - storage_type = "pgsql" - else: - storage_type = storage.dialect + storage_type = { + "mysql": "mysql", "mariadb": "mysql", + "postgres": "pgsql", "postgresql": "pgsql", "pgsql": "pgsql" + }.get(storage.dialect, storage.dialect) return {"type": storage_type or "local"} @@ -318,46 +335,20 @@ async def update_tokens_api(data: dict): @router.post("/api/v1/admin/tokens/refresh", dependencies=[Depends(verify_api_key)]) async def refresh_tokens_api(data: dict): """刷新 Token 状态""" - from app.services.token.manager import get_token_manager - from app.services.grok.utils.batch import run_in_batches - try: mgr = await get_token_manager() - tokens = [] - if "token" in data: - tokens.append(data["token"]) - if "tokens" in data and isinstance(data["tokens"], list): - tokens.extend(data["tokens"]) - + tokens = _collect_tokens(data) + if not tokens: raise HTTPException(status_code=400, detail="No tokens provided") - # 去重并保持顺序 - unique_tokens = list(dict.fromkeys(tokens)) - - # 最大数量限制 - max_tokens = get_config("performance.usage_max_tokens", DEFAULT_USAGE_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - - truncated = False - original_count = len(unique_tokens) - if len(unique_tokens) > max_tokens: - unique_tokens = unique_tokens[:max_tokens] - truncated = True - logger.warning( - f"Usage refresh: truncated from {original_count} to {max_tokens} tokens" - ) + # 去重并截断 + max_tokens = int(get_config("performance.usage_max_tokens")) + unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "Usage refresh") # 批量执行配置 - max_concurrent = get_config( - "performance.usage_max_concurrent", DEFAULT_USAGE_MAX_CONCURRENT - ) - batch_size = get_config( - "performance.usage_batch_size", DEFAULT_USAGE_BATCH_SIZE - ) + max_concurrent = get_config("performance.usage_max_concurrent") + batch_size = get_config("performance.usage_batch_size") async def _refresh_one(t): return await mgr.sync_usage( @@ -393,40 +384,20 @@ async def refresh_tokens_api(data: dict): ) async def refresh_tokens_api_async(data: dict): """刷新 Token 状态（异步批量 + SSE 进度）""" - from app.services.token.manager import get_token_manager - from app.services.grok.utils.batch import run_in_batches - mgr = await get_token_manager() - tokens: list[str] = [] - if isinstance(data.get("token"), str) and data["token"].strip(): - tokens.append(data["token"].strip()) - if isinstance(data.get("tokens"), list): - tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()]) + tokens = _collect_tokens(data) if not tokens: raise HTTPException(status_code=400, detail="No tokens provided") - unique_tokens = list(dict.fromkeys(tokens)) - - max_tokens = get_config("performance.usage_max_tokens", DEFAULT_USAGE_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - - truncated = False - original_count = len(unique_tokens) - if len(unique_tokens) > max_tokens: - unique_tokens = unique_tokens[:max_tokens] - truncated = True - logger.warning( - f"Usage refresh: truncated from {original_count} to {max_tokens} tokens" - ) + # 去重并截断 + max_tokens = int(get_config("performance.usage_max_tokens")) + unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "Usage refresh") max_concurrent = get_config( - "performance.usage_max_concurrent", DEFAULT_USAGE_MAX_CONCURRENT + "performance.usage_max_concurrent" ) - batch_size = get_config("performance.usage_batch_size", DEFAULT_USAGE_BATCH_SIZE) + batch_size = get_config("performance.usage_batch_size") task = create_task(len(unique_tokens)) @@ -500,19 +471,13 @@ async def refresh_tokens_api_async(data: dict): async def enable_nsfw_api(data: dict): """批量开启 NSFW (Unhinged) 模式""" from app.services.grok.services.nsfw import NSFWService - from app.services.grok.utils.batch import run_in_batches - from app.services.token.manager import get_token_manager try: mgr = await get_token_manager() nsfw_service = NSFWService() # 收集 token 列表 - tokens: list[str] = [] - if isinstance(data.get("token"), str) and data["token"].strip(): - tokens.append(data["token"].strip()) - if isinstance(data.get("tokens"), list): - tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()]) + tokens = _collect_tokens(data) # 若未指定，则使用所有 pool 中的 token if not tokens: @@ -526,30 +491,15 @@ async def enable_nsfw_api(data: dict): if not tokens: raise HTTPException(status_code=400, detail="No tokens available") - # 去重并保持顺序 - unique_tokens = list(dict.fromkeys(tokens)) - - # 限制最大数量（超出时截取前 N 个） - max_tokens = get_config("performance.nsfw_max_tokens", DEFAULT_NSFW_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - - truncated = False - original_count = len(unique_tokens) - if len(unique_tokens) > max_tokens: - unique_tokens = unique_tokens[:max_tokens] - truncated = True - logger.warning( - f"NSFW enable: truncated from {original_count} to {max_tokens} tokens" - ) + # 去重并截断 + max_tokens = int(get_config("performance.nsfw_max_tokens")) + unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "NSFW enable") # 批量执行配置 max_concurrent = get_config( - "performance.nsfw_max_concurrent", DEFAULT_NSFW_MAX_CONCURRENT + "performance.nsfw_max_concurrent" ) - batch_size = get_config("performance.nsfw_batch_size", DEFAULT_NSFW_BATCH_SIZE) + batch_size = get_config("performance.nsfw_batch_size") # 定义 worker async def _enable(token: str): @@ -576,7 +526,7 @@ async def enable_nsfw_api(data: dict): fail_count = 0 for token, res in raw_results.items(): - masked = f"{token[:8]}...{token[-8:]}" if len(token) > 20 else token + masked = _mask_token(token) if res.get("ok") and res.get("data", {}).get("success"): ok_count += 1 results[masked] = res.get("data", {}) @@ -615,17 +565,11 @@ async def enable_nsfw_api(data: dict): async def enable_nsfw_api_async(data: dict): """批量开启 NSFW (Unhinged) 模式（异步批量 + SSE 进度）""" from app.services.grok.services.nsfw import NSFWService - from app.services.grok.utils.batch import run_in_batches - from app.services.token.manager import get_token_manager mgr = await get_token_manager() nsfw_service = NSFWService() - tokens: list[str] = [] - if isinstance(data.get("token"), str) and data["token"].strip(): - tokens.append(data["token"].strip()) - if isinstance(data.get("tokens"), list): - tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()]) + tokens = _collect_tokens(data) if not tokens: for pool_name, pool in mgr.pools.items(): @@ -636,27 +580,12 @@ async def enable_nsfw_api_async(data: dict): if not tokens: raise HTTPException(status_code=400, detail="No tokens available") - unique_tokens = list(dict.fromkeys(tokens)) + # 去重并截断 + max_tokens = int(get_config("performance.nsfw_max_tokens")) + unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "NSFW enable") - max_tokens = get_config("performance.nsfw_max_tokens", DEFAULT_NSFW_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - - truncated = False - original_count = len(unique_tokens) - if len(unique_tokens) > max_tokens: - unique_tokens = unique_tokens[:max_tokens] - truncated = True - logger.warning( - f"NSFW enable: truncated from {original_count} to {max_tokens} tokens" - ) - - max_concurrent = get_config( - "performance.nsfw_max_concurrent", DEFAULT_NSFW_MAX_CONCURRENT - ) - batch_size = get_config("performance.nsfw_batch_size", DEFAULT_NSFW_BATCH_SIZE) + max_concurrent = get_config("performance.nsfw_max_concurrent") + batch_size = get_config("performance.nsfw_batch_size") task = create_task(len(unique_tokens)) @@ -791,28 +720,9 @@ async def get_cache_stats_api(request: Request): } online_details = [] account_map = {a["token"]: a for a in accounts} - max_concurrent = get_config( - "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT - ) - batch_size = get_config( - "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE - ) - try: - max_concurrent = int(max_concurrent) - except Exception: - max_concurrent = 25 - try: - batch_size = int(batch_size) - except Exception: - batch_size = 10 - max_concurrent = max(1, max_concurrent) - batch_size = max(1, batch_size) - - max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 + max_concurrent = max(1, int(get_config("performance.assets_max_concurrent"))) + batch_size = max(1, int(get_config("performance.assets_batch_size"))) + max_tokens = int(get_config("performance.assets_max_tokens")) truncated = False original_count = 0 @@ -855,11 +765,7 @@ async def get_cache_stats_api(request: Request): } if selected_tokens: - selected_tokens = list(dict.fromkeys(selected_tokens)) - original_count = len(selected_tokens) - if len(selected_tokens) > max_tokens: - selected_tokens = selected_tokens[:max_tokens] - truncated = True + selected_tokens, truncated, original_count = _truncate_tokens(selected_tokens, max_tokens, "Assets fetch") total = 0 raw_results = await run_in_batches( selected_tokens, @@ -1027,26 +933,11 @@ async def load_online_cache_api_async(data: dict): else: raise HTTPException(status_code=400, detail="No tokens provided") - selected_tokens = list(dict.fromkeys(selected_tokens)) + max_tokens = int(get_config("performance.assets_max_tokens")) + selected_tokens, truncated, original_count = _truncate_tokens(selected_tokens, max_tokens, "Assets load") - max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - - truncated = False - original_count = len(selected_tokens) - if len(selected_tokens) > max_tokens: - selected_tokens = selected_tokens[:max_tokens] - truncated = True - - max_concurrent = get_config( - "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT - ) - batch_size = get_config( - "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE - ) + max_concurrent = get_config("performance.assets_max_concurrent") + batch_size = get_config("performance.assets_batch_size") task = create_task(len(selected_tokens)) @@ -1220,36 +1111,12 @@ async def clear_online_cache_api(data: dict): token_list = list(dict.fromkeys(token_list)) # 最大数量限制 - max_tokens = get_config( - "performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS - ) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - truncated = False - original_count = len(token_list) - if len(token_list) > max_tokens: - token_list = token_list[:max_tokens] - truncated = True + max_tokens = int(get_config("performance.assets_max_tokens")) + token_list, truncated, original_count = _truncate_tokens(token_list, max_tokens, "Clear online cache") results = {} - max_concurrent = get_config( - "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT - ) - batch_size = get_config( - "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE - ) - try: - max_concurrent = int(max_concurrent) - except Exception: - max_concurrent = 25 - try: - batch_size = int(batch_size) - except Exception: - batch_size = 10 - max_concurrent = max(1, max_concurrent) - batch_size = max(1, batch_size) + max_concurrent = max(1, int(get_config("performance.assets_max_concurrent"))) + batch_size = max(1, int(get_config("performance.assets_batch_size"))) async def _clear_one(t: str): try: @@ -1312,25 +1179,11 @@ async def clear_online_cache_api_async(data: dict): if not token_list: raise HTTPException(status_code=400, detail="No tokens provided") - token_list = list(dict.fromkeys(token_list)) + max_tokens = int(get_config("performance.assets_max_tokens")) + token_list, truncated, original_count = _truncate_tokens(token_list, max_tokens, "Clear online cache async") - max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS) - try: - max_tokens = int(max_tokens) - except Exception: - max_tokens = 1000 - truncated = False - original_count = len(token_list) - if len(token_list) > max_tokens: - token_list = token_list[:max_tokens] - truncated = True - - max_concurrent = get_config( - "performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT - ) - batch_size = get_config( - "performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE - ) + max_concurrent = get_config("performance.assets_max_concurrent") + batch_size = get_config("performance.assets_batch_size") task = create_task(len(token_list)) diff --git a/app/api/v1/image.py b/app/api/v1/image.py index 20c6545..9a382e2 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -7,6 +7,7 @@ import base64 import math import random import re +import time from pathlib import Path from typing import List, Optional, Union @@ -31,12 +32,6 @@ from app.core.config import get_config from app.core.logger import logger -DEFAULT_IMAGE_WS = False -DEFAULT_IMAGE_FORMAT = "url" -DEFAULT_IMAGE_WS_NSFW = True -DEFAULT_GROK_TEMPORARY = True - - router = APIRouter(tags=["Images"]) class ImageGenerationRequest(BaseModel): @@ -97,15 +92,13 @@ def _validate_common_request( if allow_ws_stream: # WS 流式仅支持 b64_json (base64 视为同义) if ( - request.stream - and get_config("grok.image_ws", DEFAULT_IMAGE_WS) + request.stream + and get_config("grok.image_ws") and request.response_format and request.response_format not in {"b64_json", "base64"} ): raise ValidationException( - message=( - "Streaming with image_ws only supports response_format=b64_json/base64" - ), + message="Streaming with image_ws only supports response_format=b64_json/base64", param="response_format", code="invalid_response_format", ) @@ -145,25 +138,22 @@ def validate_generation_request(request: ImageGenerationRequest): def resolve_response_format(response_format: Optional[str]) -> str: - fmt = response_format or get_config("app.image_format", DEFAULT_IMAGE_FORMAT) + """解析响应格式""" + fmt = response_format or get_config("app.image_format") if isinstance(fmt, str): fmt = fmt.lower() if fmt in ("b64_json", "base64", "url"): return fmt - allowed_formats = {"b64_json", "base64", "url"} raise ValidationException( - message=f"response_format must be one of {sorted(allowed_formats)}", + message="response_format must be one of b64_json, base64, url", param="response_format", code="invalid_response_format", ) def response_field_name(response_format: str) -> str: - if response_format == "url": - return "url" - if response_format == "base64": - return "base64" - return "b64_json" + """获取响应字段名""" + return {"url": "url", "base64": "base64"}.get(response_format, "b64_json") def resolve_aspect_ratio(size: str) -> str: @@ -215,6 +205,48 @@ def validate_edit_request(request: ImageEditRequest, images: List[UploadFile]): ) +def _get_effort(model_info) -> EffortType: + """获取模型消耗级别""" + return EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + + +async def _wrap_stream_with_usage(stream, token_mgr, token, model_info): + """包装流式响应，成功完成时记录使用""" + success = False + try: + async for chunk in stream: + yield chunk + success = True + finally: + if success: + try: + await token_mgr.consume(token, _get_effort(model_info)) + except Exception as e: + logger.warning(f"Failed to consume token: {e}") + + +async def _get_token(model: str): + """获取可用 token""" + token_mgr = await get_token_manager() + await token_mgr.reload_if_stale() + + token = None + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name) + if token: + break + + if not token: + raise AppException( + message="No available tokens. Please try again later.", + error_type=ErrorType.RATE_LIMIT.value, + code="rate_limit_exceeded", + status_code=429, + ) + + return token_mgr, token + + async def call_grok( token_mgr, token: str, @@ -237,7 +269,6 @@ async def call_grok( file_attachments=file_attachments, ) - # 收集图片 processor = ImageCollectProcessor( model_info.model_id, token, response_format=response_format ) @@ -249,15 +280,9 @@ async def call_grok( logger.error(f"Grok image call failed: {e}") return [] finally: - # 只在成功时记录使用，失败时不扣费（避免清零 fail_count） if success: try: - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) + await token_mgr.consume(token, _get_effort(model_info)) except Exception as e: logger.warning(f"Failed to consume token: {e}") @@ -291,40 +316,16 @@ async def create_image(request: ImageGenerationRequest): response_format = resolve_response_format(request.response_format) response_field = response_field_name(response_format) - # 获取 token - try: - token_mgr = await get_token_manager() - await token_mgr.reload_if_stale() - token = None - for pool_name in ModelService.pool_candidates_for_model(request.model): - token = token_mgr.get_token(pool_name) - if token: - break - except Exception as e: - logger.error(f"Failed to get token: {e}") - raise AppException( - message="Internal service error obtaining token", - error_type=ErrorType.SERVER.value, - code="internal_error", - ) - - if not token: - raise AppException( - message="No available tokens. Please try again later.", - error_type=ErrorType.RATE_LIMIT.value, - code="rate_limit_exceeded", - status_code=429, - ) - - # 获取模型信息 + # 获取 token 和模型信息 + token_mgr, token = await _get_token(request.model) model_info = ModelService.get(request.model) - use_ws = bool(get_config("grok.image_ws", DEFAULT_IMAGE_WS)) + use_ws = bool(get_config("grok.image_ws")) # 流式模式 if request.stream: if use_ws: aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("grok.image_ws_nsfw", DEFAULT_IMAGE_WS_NSFW)) + enable_nsfw = bool(get_config("grok.image_ws_nsfw")) upstream = image_service.stream( token=token, prompt=request.prompt, @@ -340,26 +341,8 @@ async def create_image(request: ImageGenerationRequest): size=request.size, ) - async def _wrap_stream(stream): - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - if success: - try: - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - except Exception as e: - logger.warning(f"Failed to consume token: {e}") - return StreamingResponse( - _wrap_stream(processor.process(upstream)), + _wrap_stream_with_usage(processor.process(upstream), token_mgr, token, model_info), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) @@ -377,28 +360,8 @@ async def create_image(request: ImageGenerationRequest): model_info.model_id, token, n=request.n, response_format=response_format ) - # 包装流式响应，在成功完成时记录使用 - async def _wrap_stream(stream): - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - # 只在成功完成时扣费 - if success: - try: - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - except Exception as e: - logger.warning(f"Failed to consume token: {e}") - return StreamingResponse( - _wrap_stream(processor.process(response)), + _wrap_stream_with_usage(processor.process(response), token_mgr, token, model_info), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) @@ -409,7 +372,7 @@ async def create_image(request: ImageGenerationRequest): usage_override = None if use_ws: aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("grok.image_ws_nsfw", DEFAULT_IMAGE_WS_NSFW)) + enable_nsfw = bool(get_config("grok.image_ws_nsfw")) all_images = [] seen = set() expected_per_call = 6 @@ -452,12 +415,7 @@ async def create_image(request: ImageGenerationRequest): if len(all_images) >= n: break try: - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) + await token_mgr.consume(token, _get_effort(model_info)) except Exception as e: logger.warning(f"Failed to consume token: {e}") usage_override = { @@ -510,16 +468,11 @@ async def create_image(request: ImageGenerationRequest): selected_images.append("error") # 构建响应 - import time - data = [{response_field: img} for img in selected_images] - usage = usage_override or { - "total_tokens": 0 - * len([img for img in selected_images if img != "error"]), + "total_tokens": 0, "input_tokens": 0, - "output_tokens": 0 - * len([img for img in selected_images if img != "error"]), + "output_tokens": 0, "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, } @@ -627,32 +580,8 @@ async def edit_image( b64 = base64.b64encode(content).decode() images.append(f"data:{mime};base64,{b64}") - # 获取 token - try: - token_mgr = await get_token_manager() - await token_mgr.reload_if_stale() - token = None - for pool_name in ModelService.pool_candidates_for_model(edit_request.model): - token = token_mgr.get_token(pool_name) - if token: - break - except Exception as e: - logger.error(f"Failed to get token: {e}") - raise AppException( - message="Internal service error obtaining token", - error_type=ErrorType.SERVER.value, - code="internal_error", - ) - - if not token: - raise AppException( - message="No available tokens. Please try again later.", - error_type=ErrorType.RATE_LIMIT.value, - code="rate_limit_exceeded", - status_code=429, - ) - - # 获取模型信息 + # 获取 token 和模型信息 + token_mgr, token = await _get_token(edit_request.model) model_info = ModelService.get(edit_request.model) # 上传图片 @@ -714,7 +643,7 @@ async def edit_image( ] = parent_post_id raw_payload = { - "temporary": bool(get_config("grok.temporary", DEFAULT_GROK_TEMPORARY)), + "temporary": bool(get_config("grok.temporary")), "modelName": model_info.grok_model, "message": edit_request.prompt, "enableImageGeneration": True, @@ -749,26 +678,8 @@ async def edit_image( model_info.model_id, token, n=edit_request.n, response_format=response_format ) - async def _wrap_stream(stream): - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - if success: - try: - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - except Exception as e: - logger.warning(f"Failed to consume token: {e}") - return StreamingResponse( - _wrap_stream(processor.process(response)), + _wrap_stream_with_usage(processor.process(response), token_mgr, token, model_info), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) @@ -795,10 +706,7 @@ async def edit_image( if calls_needed == 1: all_images = await _call_edit() else: - tasks = [ - _call_edit() - for _ in range(calls_needed) - ] + tasks = [_call_edit() for _ in range(calls_needed)] results = await asyncio.gather(*tasks, return_exceptions=True) all_images = [] @@ -808,6 +716,7 @@ async def edit_image( elif isinstance(result, list): all_images.extend(result) + # 选择图片 if len(all_images) >= n: selected_images = random.sample(all_images, n) else: @@ -815,8 +724,6 @@ async def edit_image( while len(selected_images) < n: selected_images.append("error") - import time - data = [{response_field: img} for img in selected_images] return JSONResponse( @@ -824,11 +731,9 @@ async def edit_image( "created": int(time.time()), "data": data, "usage": { - "total_tokens": 0 - * len([img for img in selected_images if img != "error"]), + "total_tokens": 0, "input_tokens": 0, - "output_tokens": 0 - * len([img for img in selected_images if img != "error"]), + "output_tokens": 0, "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, }, } diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index 9ac3eb7..91a5270 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -246,6 +246,8 @@ class UploadService(BaseService): else: filename, b64, mime = self.parse_b64(file_input) + logger.debug(f"Upload prepare: filename={filename}, type={mime}, size={len(b64)}") + if not b64: raise ValidationException("Invalid file input: empty content") @@ -353,6 +355,7 @@ class ListService(BaseService): total = 0 async for page_assets in self.iter_assets(token): total += len(page_assets) + logger.debug(f"Asset count: {total}") return total diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index ee7de10..8d2aab7 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -197,7 +197,7 @@ class GrokChatService: proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None timeout = get_config("grok.timeout") - logger.debug(f"Payload: {payload}") + logger.debug(f"Chat request: model={model}, mode={mode}, stream={stream}, attachments={len(file_attachments or [])}") # 建立连接 async def establish_connection(): @@ -220,8 +220,7 @@ class GrokChatService: except Exception: pass - logger.error(f"Chat failed: {response.status_code}", - extra={"status": response.status_code, "token": token[:10] + "..."}) + logger.error(f"Chat failed: status={response.status_code}, token={token[:10]}...") await session.close() raise UpstreamException( @@ -229,6 +228,7 @@ class GrokChatService: details={"status": response.status_code, "body": content}, ) + logger.info(f"Chat connected: model={model}, stream={stream}") return session, response except UpstreamException: @@ -287,6 +287,7 @@ class GrokChatService: # 提取消息和附件 try: message, attachments = MessageExtractor.extract(request.messages, is_video=is_video) + logger.debug(f"Extracted message length={len(message)}, attachments={len(attachments)}") except ValueError as e: raise ValidationException(str(e)) @@ -295,9 +296,10 @@ class GrokChatService: if attachments: upload_service = UploadService() try: - for _, attach_data in attachments: + for attach_type, attach_data in attachments: file_id, _ = await upload_service.upload(attach_data, token) file_ids.append(file_id) + logger.debug(f"Attachment uploaded: type={attach_type}, file_id={file_id}") finally: await upload_service.close() @@ -349,15 +351,18 @@ class ChatService: # 处理响应 if is_stream: + logger.debug(f"Processing stream response: model={model}") processor = StreamProcessor(model_name, token, think) return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) # 非流式 + logger.debug(f"Processing non-stream response: model={model}") result = await CollectProcessor(model_name, token).process(response) try: model_info = ModelService.get(model) effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW await token_mgr.consume(token, effort) + logger.info(f"Chat completed: model={model}, effort={effort.value}") except Exception as e: logger.warning(f"Failed to record usage: {e}") return result diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 06b1dd7..9e51207 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -93,6 +93,7 @@ class ImageService: max_retries: int = None, ) -> AsyncGenerator[Dict[str, object], None]: retries = max(1, max_retries if max_retries is not None else 1) + logger.info(f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}") for attempt in range(retries): try: @@ -199,6 +200,7 @@ class ImageService: if info["is_final"] and not existing.get("is_final"): completed += 1 + logger.debug(f"Final image received: id={image_id}, size={info['blob_size']}") images[image_id] = {"is_final": info["is_final"] or existing.get("is_final")} yield info diff --git a/app/services/grok/services/media.py b/app/services/grok/services/media.py index 8004028..39b0849 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/media.py @@ -183,12 +183,14 @@ class VideoService: ) if response.status_code != 200: - logger.error(f"Video generation failed: {response.status_code}") + logger.error(f"Video generation failed: status={response.status_code}, post_id={post_id}") raise UpstreamException( message=f"Video generation failed: {response.status_code}", details={"status": response.status_code}, ) + logger.info(f"Video generation started: post_id={post_id}") + async def stream_response(): try: async for line in response.aiter_lines(): @@ -217,6 +219,7 @@ class VideoService: preset: str = "normal" ) -> AsyncGenerator[bytes, None]: """生成视频""" + logger.info(f"Video generation: prompt='{prompt[:50]}...', ratio={aspect_ratio}, length={video_length}s, preset={preset}") async with _get_semaphore(): post_id = await self.create_post(token, prompt) return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution_name, preset) @@ -229,6 +232,7 @@ class VideoService: preset: str = "normal" ) -> AsyncGenerator[bytes, None]: """从图片生成视频""" + logger.info(f"Image to video: prompt='{prompt[:50]}...', image={image_url[:80]}") async with _get_semaphore(): post_id = await self.create_image_post(token, image_url) return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution, preset) diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index 8c386f4..bf3d5cb 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -112,10 +112,10 @@ class UsageService: if response.status_code == 200: data = response.json() remaining = data.get("remainingTokens", 0) - logger.info(f"Usage: quota {remaining} remaining") + logger.info(f"Usage sync success: remaining={remaining}, token={token[:10]}...") return data - logger.error(f"Usage failed: {response.status_code}") + logger.error(f"Usage sync failed: status={response.status_code}, token={token[:10]}...") raise UpstreamException( message=f"Failed to get usage stats: {response.status_code}", diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index b887d31..6c231ae 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -35,6 +35,7 @@ class VoiceService: Returns: Dict containing token and livekitUrl """ + logger.debug(f"Voice token request: voice={voice}, personality={personality}, speed={speed}") headers = self._build_headers(token) payload = self._build_payload(voice, personality, speed) @@ -53,16 +54,16 @@ class VoiceService: ) if response.status_code != 200: - logger.error( - f"Voice token failed: {response.status_code}", - extra={"response": response.text[:200]} - ) + body = response.text[:200] + logger.error(f"Voice token failed: status={response.status_code}, body={body}") raise UpstreamException( message=f"Failed to get voice token: {response.status_code}", details={"status": response.status_code, "body": response.text} ) - return response.json() + result = response.json() + logger.info(f"Voice token obtained: voice={voice}") + return result except Exception as e: logger.error(f"Voice service error: {e}") From bcd94069a74b6264b2cb7f4744c611b14beb1c48 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sat, 7 Feb 2026 13:30:50 +0800 Subject: [PATCH 4/6] feat: migrate configuration structure to improve organization and flexibility --- app/api/v1/image.py | 10 +- app/core/config.py | 87 ++++++++++- app/services/grok/defaults.py | 76 ++++++---- .../grok/processors/chat_processors.py | 10 +- .../grok/processors/image_processors.py | 4 +- .../grok/processors/video_processors.py | 6 +- app/services/grok/services/assets.py | 8 +- app/services/grok/services/chat.py | 18 +-- app/services/grok/services/image.py | 12 +- app/services/grok/services/media.py | 12 +- app/services/grok/services/nsfw.py | 14 +- app/services/grok/services/usage.py | 8 +- app/services/grok/services/voice.py | 6 +- app/services/grok/utils/headers.py | 2 +- app/services/grok/utils/retry.py | 12 +- app/services/grok/utils/statsig.py | 2 +- app/static/config/config.html | 2 +- app/static/config/config.js | 106 +++++++------ config.defaults.toml | 139 ++++++++++++++---- data/config.toml | 93 ++++++++---- docs/README.en.md | 54 ++++--- readme.md | 60 ++++---- 22 files changed, 488 insertions(+), 253 deletions(-) diff --git a/app/api/v1/image.py b/app/api/v1/image.py index 9a382e2..fc588ff 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -93,7 +93,7 @@ def _validate_common_request( # WS 流式仅支持 b64_json (base64 视为同义) if ( request.stream - and get_config("grok.image_ws") + and get_config("image.image_ws") and request.response_format and request.response_format not in {"b64_json", "base64"} ): @@ -319,13 +319,13 @@ async def create_image(request: ImageGenerationRequest): # 获取 token 和模型信息 token_mgr, token = await _get_token(request.model) model_info = ModelService.get(request.model) - use_ws = bool(get_config("grok.image_ws")) + use_ws = bool(get_config("image.image_ws")) # 流式模式 if request.stream: if use_ws: aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("grok.image_ws_nsfw")) + enable_nsfw = bool(get_config("image.image_ws_nsfw")) upstream = image_service.stream( token=token, prompt=request.prompt, @@ -372,7 +372,7 @@ async def create_image(request: ImageGenerationRequest): usage_override = None if use_ws: aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("grok.image_ws_nsfw")) + enable_nsfw = bool(get_config("image.image_ws_nsfw")) all_images = [] seen = set() expected_per_call = 6 @@ -643,7 +643,7 @@ async def edit_image( ] = parent_post_id raw_payload = { - "temporary": bool(get_config("grok.temporary")), + "temporary": bool(get_config("chat.temporary")), "modelName": model_info.grok_model, "message": edit_request.prompt, "enableImageGeneration": True, diff --git a/app/core/config.py b/app/core/config.py index d9d3e8c..22c125b 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -32,6 +32,81 @@ def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any return result +def _migrate_deprecated_config(config: Dict[str, Any], valid_sections: set) -> tuple[Dict[str, Any], set]: + """ + 迁移废弃的配置节到新配置结构 + + Returns: + (迁移后的配置, 废弃的配置节集合) + """ + # 配置映射规则：旧配置 -> 新配置 + MIGRATION_MAP = { + # grok.* -> 对应的新配置节 + "grok.temporary": "chat.temporary", + "grok.disable_memory": "chat.disable_memory", + "grok.stream": "chat.stream", + "grok.thinking": "chat.thinking", + "grok.dynamic_statsig": "chat.dynamic_statsig", + "grok.filter_tags": "chat.filter_tags", + + "grok.timeout": "network.timeout", + "grok.base_proxy_url": "network.base_proxy_url", + "grok.asset_proxy_url": "network.asset_proxy_url", + + "grok.cf_clearance": "security.cf_clearance", + "grok.browser": "security.browser", + "grok.user_agent": "security.user_agent", + + "grok.max_retry": "retry.max_retry", + "grok.retry_status_codes": "retry.retry_status_codes", + "grok.retry_backoff_base": "retry.retry_backoff_base", + "grok.retry_backoff_factor": "retry.retry_backoff_factor", + "grok.retry_backoff_max": "retry.retry_backoff_max", + "grok.retry_budget": "retry.retry_budget", + + "grok.stream_idle_timeout": "timeout.stream_idle_timeout", + "grok.video_idle_timeout": "timeout.video_idle_timeout", + + "grok.image_ws": "image.image_ws", + "grok.image_ws_nsfw": "image.image_ws_nsfw", + "grok.image_ws_blocked_seconds": "image.image_ws_blocked_seconds", + "grok.image_ws_final_min_bytes": "image.image_ws_final_min_bytes", + "grok.image_ws_medium_min_bytes": "image.image_ws_medium_min_bytes", + } + + deprecated_sections = set(config.keys()) - valid_sections + if not deprecated_sections: + return config, set() + + result = {k: deepcopy(v) for k, v in config.items() if k in valid_sections} + migrated_count = 0 + + # 处理废弃配置节中的配置项 + for old_section in deprecated_sections: + if old_section not in config or not isinstance(config[old_section], dict): + continue + + for old_key, old_value in config[old_section].items(): + # 查找映射规则 + old_path = f"{old_section}.{old_key}" + new_path = MIGRATION_MAP.get(old_path) + + if new_path: + new_section, new_key = new_path.split(".", 1) + # 确保新配置节存在 + if new_section not in result: + result[new_section] = {} + # 迁移配置项（保留用户的自定义值） + result[new_section][new_key] = old_value + migrated_count += 1 + logger.debug(f"Migrated config: {old_path} -> {new_path} = {old_value}") + + if migrated_count > 0: + logger.info(f"Migrated {migrated_count} config items from deprecated sections") + + return result, deprecated_sections + + def _load_defaults() -> Dict[str, Any]: """加载默认配置文件""" if not DEFAULT_CONFIG_FILE.exists(): @@ -91,10 +166,18 @@ class Config: config_data = {} config_data = config_data or {} + + # 检查是否有废弃的配置节 + valid_sections = set(self._defaults.keys()) + config_data, deprecated_sections = _migrate_deprecated_config(config_data, valid_sections) + if deprecated_sections: + logger.info(f"Cleaned deprecated config sections: {deprecated_sections}") + merged = _deep_merge(self._defaults, config_data) # 自动回填缺失配置到存储 - should_persist = (not from_remote) or (merged != config_data) + # 或迁移了配置后需要更新 + should_persist = (not from_remote) or (merged != config_data) or deprecated_sections if should_persist: async with storage.acquire_lock("config_save", timeout=10): await storage.save_config(merged) @@ -102,6 +185,8 @@ class Config: logger.info( f"Initialized remote storage ({storage.__class__.__name__}) with config baseline." ) + if deprecated_sections: + logger.info("Configuration automatically migrated and cleaned.") self._config = merged except Exception as e: diff --git a/app/services/grok/defaults.py b/app/services/grok/defaults.py index 78ade9a..e23f515 100644 --- a/app/services/grok/defaults.py +++ b/app/services/grok/defaults.py @@ -6,60 +6,84 @@ Grok 服务默认配置 # Grok 服务默认配置 GROK_DEFAULTS = { - "grok": { - # 网络配置 - "browser": "chrome136", + "app": { + "app_url": "", + "app_key": "grok2api", + "api_key": "", + "image_format": "url", + "video_format": "html", + }, + + "network": { "timeout": 120, - "user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36", "base_proxy_url": "", "asset_proxy_url": "", + }, + + "security": { "cf_clearance": "", - - # 对话配置 + "browser": "chrome136", + "user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36", + }, + + "chat": { "temporary": True, "disable_memory": True, "stream": True, "thinking": False, + "dynamic_statsig": True, "filter_tags": ["grok:render", "xaiartifact", "xai:tool_usage_card"], - - # 重试配置 + }, + + "retry": { "max_retry": 3, "retry_status_codes": [401, 429, 403], "retry_backoff_base": 0.5, "retry_backoff_factor": 2.0, "retry_backoff_max": 30.0, "retry_budget": 90.0, - - # 超时配置 + }, + + "timeout": { "stream_idle_timeout": 45.0, "video_idle_timeout": 90.0, - - # 图片配置 + }, + + "image": { + "image_ws": True, + "image_ws_nsfw": True, "image_ws_blocked_seconds": 15, "image_ws_final_min_bytes": 100000, "image_ws_medium_min_bytes": 30000, - - # Statsig - "dynamic_statsig": True, }, - "app": { - "app_url": "", - "image_format": "url", - "video_format": "html", - }, - - "performance": { - "assets_max_concurrent": 25, - "assets_delete_batch_size": 10, - "media_max_concurrent": 50, - "usage_max_concurrent": 25, + "token": { + "auto_refresh": True, + "refresh_interval_hours": 8, + "super_refresh_interval_hours": 2, + "fail_threshold": 5, + "save_delay_ms": 500, + "reload_interval_sec": 30, }, "cache": { "enable_auto_clean": True, "limit_mb": 1024, }, + + "performance": { + "assets_max_concurrent": 25, + "assets_delete_batch_size": 10, + "assets_batch_size": 10, + "assets_max_tokens": 1000, + "media_max_concurrent": 50, + "usage_max_concurrent": 25, + "usage_batch_size": 50, + "usage_max_tokens": 1000, + "nsfw_max_concurrent": 10, + "nsfw_batch_size": 50, + "nsfw_max_tokens": 1000, + }, } diff --git a/app/services/grok/processors/chat_processors.py b/app/services/grok/processors/chat_processors.py index a7bf5d4..d7002ca 100644 --- a/app/services/grok/processors/chat_processors.py +++ b/app/services/grok/processors/chat_processors.py @@ -32,13 +32,13 @@ class StreamProcessor(BaseProcessor): self.fingerprint: str = "" self.think_opened: bool = False self.role_sent: bool = False - self.filter_tags = get_config("grok.filter_tags") + self.filter_tags = get_config("chat.filter_tags") self.image_format = get_config("app.image_format") self._tag_buffer: str = "" self._in_filter_tag: bool = False if think is None: - self.show_think = get_config("grok.thinking") + self.show_think = get_config("chat.thinking") else: self.show_think = think @@ -114,7 +114,7 @@ class StreamProcessor(BaseProcessor): async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: """处理流式响应""" - idle_timeout = get_config("grok.stream_idle_timeout") + idle_timeout = get_config("timeout.stream_idle_timeout") try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -223,7 +223,7 @@ class CollectProcessor(BaseProcessor): def __init__(self, model: str, token: str = ""): super().__init__(model, token) self.image_format = get_config("app.image_format") - self.filter_tags = get_config("grok.filter_tags") + self.filter_tags = get_config("chat.filter_tags") def _filter_content(self, content: str) -> str: """过滤内容中的特殊标签""" @@ -242,7 +242,7 @@ class CollectProcessor(BaseProcessor): response_id = "" fingerprint = "" content = "" - idle_timeout = get_config("grok.stream_idle_timeout") + idle_timeout = get_config("timeout.stream_idle_timeout") try: async for line in _with_idle_timeout(response, idle_timeout, self.model): diff --git a/app/services/grok/processors/image_processors.py b/app/services/grok/processors/image_processors.py index 8550108..42214c7 100644 --- a/app/services/grok/processors/image_processors.py +++ b/app/services/grok/processors/image_processors.py @@ -45,7 +45,7 @@ class ImageStreamProcessor(BaseProcessor): async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: """处理流式响应""" final_images = [] - idle_timeout = get_config("grok.stream_idle_timeout") + idle_timeout = get_config("timeout.stream_idle_timeout") try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -160,7 +160,7 @@ class ImageCollectProcessor(BaseProcessor): async def process(self, response: AsyncIterable[bytes]) -> List[str]: """处理并收集图片""" images = [] - idle_timeout = get_config("grok.stream_idle_timeout") + idle_timeout = get_config("timeout.stream_idle_timeout") try: async for line in _with_idle_timeout(response, idle_timeout, self.model): diff --git a/app/services/grok/processors/video_processors.py b/app/services/grok/processors/video_processors.py index df8750c..ec39aae 100644 --- a/app/services/grok/processors/video_processors.py +++ b/app/services/grok/processors/video_processors.py @@ -32,7 +32,7 @@ class VideoStreamProcessor(BaseProcessor): self.video_format = str(get_config("app.video_format")).lower() if think is None: - self.show_think = get_config("grok.thinking") + self.show_think = get_config("chat.thinking") else: self.show_think = think @@ -68,7 +68,7 @@ class VideoStreamProcessor(BaseProcessor): async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: """处理视频流式响应""" - idle_timeout = get_config("grok.video_idle_timeout") + idle_timeout = get_config("timeout.video_idle_timeout") try: async for line in _with_idle_timeout(response, idle_timeout, self.model): @@ -171,7 +171,7 @@ class VideoCollectProcessor(BaseProcessor): """处理并收集视频响应""" response_id = "" content = "" - idle_timeout = get_config("grok.video_idle_timeout") + idle_timeout = get_config("timeout.video_idle_timeout") try: async for line in _with_idle_timeout(response, idle_timeout, self.model): diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index 91a5270..d576b81 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -111,10 +111,10 @@ class ServiceConfig: @classmethod def from_settings(cls, proxy: Optional[str] = None): return cls( - proxy=proxy or get_config("grok.asset_proxy_url") or get_config("grok.base_proxy_url"), - timeout=get_config("grok.timeout"), - browser=get_config("grok.browser"), - user_agent=get_config("grok.user_agent"), + proxy=proxy or get_config("network.asset_proxy_url") or get_config("network.base_proxy_url"), + timeout=get_config("network.timeout"), + browser=get_config("security.browser"), + user_agent=get_config("security.user_agent"), ) def get_proxies(self) -> Optional[dict]: diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index 8d2aab7..99ea4ce 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -102,7 +102,7 @@ class ChatRequestBuilder: @staticmethod def build_headers(token: str) -> Dict[str, str]: """构造请求头""" - user_agent = get_config("grok.user_agent") + user_agent = get_config("security.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -142,7 +142,7 @@ class ChatRequestBuilder: merged_attachments.extend(image_attachments) payload = { - "temporary": get_config("grok.temporary"), + "temporary": get_config("chat.temporary"), "modelName": model, "message": message, "fileAttachments": merged_attachments, @@ -160,7 +160,7 @@ class ChatRequestBuilder: "modelConfigOverride": {"modelMap": {}}, "requestModelDetails": {"modelId": model}, }, - "disableMemory": get_config("grok.disable_memory"), + "disableMemory": get_config("chat.disable_memory"), "deviceEnvInfo": { "darkModeEnabled": False, "devicePixelRatio": 2, @@ -181,27 +181,27 @@ class GrokChatService: """Grok API 调用服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url") + self.proxy = proxy or get_config("network.base_proxy_url") async def chat(self, token: str, message: str, model: str = "grok-3", mode: str = None, stream: bool = None, file_attachments: List[str] = None, image_attachments: List[str] = None, raw_payload: Dict[str, Any] = None): """发送聊天请求""" if stream is None: - stream = get_config("grok.stream") + stream = get_config("chat.stream") headers = ChatRequestBuilder.build_headers(token) payload = raw_payload if raw_payload is not None else ChatRequestBuilder.build_payload( message, model, mode, file_attachments, image_attachments ) proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None - timeout = get_config("grok.timeout") + timeout = get_config("network.timeout") logger.debug(f"Chat request: model={model}, mode={mode}, stream={stream}, attachments={len(file_attachments or [])}") # 建立连接 async def establish_connection(): - browser = get_config("grok.browser") + browser = get_config("security.browser") session = AsyncSession(impersonate=browser) try: response = await session.post( @@ -303,7 +303,7 @@ class GrokChatService: finally: await upload_service.close() - stream = request.stream if request.stream is not None else get_config("grok.stream") + stream = request.stream if request.stream is not None else get_config("chat.stream") response = await self.chat( token, message, grok_model, mode, stream, @@ -340,7 +340,7 @@ class ChatService: # 解析参数 think = {"enabled": True, "disabled": False}.get(thinking) - is_stream = stream if stream is not None else get_config("grok.stream") + is_stream = stream if stream is not None else get_config("chat.stream") # 构造请求 chat_request = ChatRequest(model=model, messages=messages, stream=is_stream, think=think) diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 9e51207..b94feb6 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -32,7 +32,7 @@ class ImageService: self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)") def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]: - proxy_url = get_config("grok.base_proxy_url") + proxy_url = get_config("network.base_proxy_url") if not proxy_url: return aiohttp.TCPConnector(ssl=self._ssl_context), None @@ -46,7 +46,7 @@ class ImageService: def _get_ws_headers(self, token: str) -> Dict[str, str]: cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent") + user_agent = get_config("security.user_agent") return { "Cookie": cookie, "Origin": "https://grok.com", @@ -61,7 +61,7 @@ class ImageService: return match.group(1) if match else None def _is_final_image(self, url: str, blob_size: int) -> bool: - return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > get_config("grok.image_ws_final_min_bytes") + return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > get_config("image.image_ws_final_min_bytes") def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]: if not url or not blob: @@ -71,7 +71,7 @@ class ImageService: blob_size = len(blob) is_final = self._is_final_image(url, blob_size) - stage = "final" if is_final else ("medium" if blob_size > get_config("grok.image_ws_medium_min_bytes") else "preview") + stage = "final" if is_final else ("medium" if blob_size > get_config("image.image_ws_medium_min_bytes") else "preview") return { "type": "image", @@ -122,8 +122,8 @@ class ImageService: ) -> AsyncGenerator[Dict[str, object], None]: request_id = str(uuid.uuid4()) headers = self._get_ws_headers(token) - timeout = float(get_config("grok.timeout")) - blocked_seconds = float(get_config("grok.image_ws_blocked_seconds")) + timeout = float(get_config("network.timeout")) + blocked_seconds = float(get_config("image.image_ws_blocked_seconds")) try: connector, proxy = self._resolve_proxy() diff --git a/app/services/grok/services/media.py b/app/services/grok/services/media.py index 39b0849..10c4ba4 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/media.py @@ -43,14 +43,14 @@ class VideoService: """视频生成服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url") - self.timeout = get_config("grok.timeout") + self.proxy = proxy or get_config("network.base_proxy_url") + self.timeout = get_config("network.timeout") def _build_headers( self, token: str, referer: str = "https://grok.com/imagine" ) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent") + user_agent = get_config("security.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -103,7 +103,7 @@ class VideoService: CREATE_POST_API, headers=headers, json=payload, - impersonate=get_config("grok.browser"), + impersonate=get_config("security.browser"), timeout=30, proxies=self._build_proxies(), ) @@ -172,7 +172,7 @@ class VideoService: headers = self._build_headers(token) payload = self._build_payload(prompt, post_id, aspect_ratio, video_length, resolution_name, preset) - session = AsyncSession(impersonate=get_config("grok.browser")) + session = AsyncSession(impersonate=get_config("security.browser")) response = await session.post( CHAT_API, headers=headers, @@ -267,7 +267,7 @@ class VideoService: ) think = {"enabled": True, "disabled": False}.get(thinking) - is_stream = stream if stream is not None else get_config("grok.stream") + is_stream = stream if stream is not None else get_config("chat.stream") # 提取内容 from app.services.grok.services.chat import MessageExtractor diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py index 169ed4d..fa02b8f 100644 --- a/app/services/grok/services/nsfw.py +++ b/app/services/grok/services/nsfw.py @@ -37,8 +37,8 @@ class NSFWService: """NSFW 模式服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url") - self.timeout = float(get_config("grok.timeout")) + self.proxy = proxy or get_config("network.base_proxy_url") + self.timeout = float(get_config("network.timeout")) def _build_proxies(self) -> Optional[dict]: """构建代理配置""" @@ -60,7 +60,7 @@ class NSFWService: def _build_headers(self, token: str) -> dict: """构造 gRPC-Web 请求头""" cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent") + user_agent = get_config("security.user_agent") return { "accept": "*/*", "content-type": "application/grpc-web+proto", @@ -75,7 +75,7 @@ class NSFWService: def _build_birth_headers(self, token: str) -> dict: """构造设置出生日期请求头""" cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent") + user_agent = get_config("security.user_agent") return { "accept": "*/*", "content-type": "application/json", @@ -97,9 +97,7 @@ class NSFWService: protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner return encode_grpc_web_payload(protobuf) - async def _set_birth_date( - self, session: AsyncSession, token: str - ) -> tuple[bool, int, Optional[str]]: + async def _set_birth_date(self, session: AsyncSession, token: str) -> tuple[bool, int, Optional[str]]: """设置出生日期""" headers = self._build_birth_headers(token) payload = {"birthDate": self._random_birth_date()} @@ -125,7 +123,7 @@ class NSFWService: logger.debug(f"NSFW payload: len={len(payload)} hex={payload.hex()}") try: - browser = get_config("grok.browser") + browser = get_config("security.browser") async with AsyncSession(impersonate=browser) as session: # 先设置出生日期 ok, birth_status, birth_err = await self._set_birth_date(session, token) diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index bf3d5cb..adb04b3 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -23,12 +23,12 @@ class UsageService: """用量查询服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url") - self.timeout = get_config("grok.timeout") + self.proxy = proxy or get_config("network.base_proxy_url") + self.timeout = get_config("network.timeout") def _build_headers(self, token: str) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent") + user_agent = get_config("security.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -97,7 +97,7 @@ class UsageService: try: headers = self._build_headers(token) payload = {"requestKind": "DEFAULT", "modelName": model_name} - browser = get_config("grok.browser") + browser = get_config("security.browser") async with AsyncSession() as session: response = await session.post( diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 6c231ae..ae003f0 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -18,7 +18,7 @@ class VoiceService: """Voice Mode Service (LiveKit)""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url") + self.proxy = proxy or get_config("network.base_proxy_url") async def get_token( self, @@ -42,8 +42,8 @@ class VoiceService: proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None try: - browser = get_config("grok.browser") - timeout = get_config("grok.timeout") + browser = get_config("security.browser") + timeout = get_config("network.timeout") async with AsyncSession(impersonate=browser) as session: response = await session.post( LIVEKIT_TOKEN_API, diff --git a/app/services/grok/utils/headers.py b/app/services/grok/utils/headers.py index 146e704..7a5e1c2 100644 --- a/app/services/grok/utils/headers.py +++ b/app/services/grok/utils/headers.py @@ -17,7 +17,7 @@ def _normalize_token(token: str) -> str: def build_sso_cookie(token: str, include_rw: bool = False) -> str: token = _normalize_token(token) - cf = get_config("grok.cf_clearance") + cf = get_config("security.cf_clearance") cookie = f"sso={token}" if include_rw: cookie = f"{cookie}; sso-rw={token}" diff --git a/app/services/grok/utils/retry.py b/app/services/grok/utils/retry.py index 2cd52c5..7fa7eba 100644 --- a/app/services/grok/utils/retry.py +++ b/app/services/grok/utils/retry.py @@ -23,17 +23,17 @@ class RetryContext: def __init__(self): self.attempt = 0 - self.max_retry = int(get_config("grok.max_retry")) - self.retry_codes = get_config("grok.retry_status_codes") + self.max_retry = int(get_config("retry.max_retry")) + self.retry_codes = get_config("retry.retry_status_codes") self.last_error = None self.last_status = None self.total_delay = 0.0 - self.retry_budget = float(get_config("grok.retry_budget")) + self.retry_budget = float(get_config("retry.retry_budget")) # 退避参数 - self.backoff_base = float(get_config("grok.retry_backoff_base")) - self.backoff_factor = float(get_config("grok.retry_backoff_factor")) - self.backoff_max = float(get_config("grok.retry_backoff_max")) + self.backoff_base = float(get_config("retry.retry_backoff_base")) + self.backoff_factor = float(get_config("retry.retry_backoff_factor")) + self.backoff_max = float(get_config("retry.retry_backoff_max")) # decorrelated jitter 状态 self._last_delay = self.backoff_base diff --git a/app/services/grok/utils/statsig.py b/app/services/grok/utils/statsig.py index bf78741..c2cd15f 100644 --- a/app/services/grok/utils/statsig.py +++ b/app/services/grok/utils/statsig.py @@ -30,7 +30,7 @@ class StatsigService: Returns: Base64 编码的 ID """ - dynamic = get_config("grok.dynamic_statsig") + dynamic = get_config("chat.dynamic_statsig") if not dynamic: return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" diff --git a/app/static/config/config.html b/app/static/config/config.html index 7fb207f..6fdff9e 100644 --- a/app/static/config/config.html +++ b/app/static/config/config.html @@ -50,7 +50,7 @@ - + diff --git a/app/static/config/config.js b/app/static/config/config.js index bcb8036..f0282a7 100644 --- a/app/static/config/config.js +++ b/app/static/config/config.js @@ -26,6 +26,7 @@ const NUMERIC_FIELDS = new Set([ 'video_idle_timeout', 'image_ws_blocked_seconds', 'image_ws_final_min_bytes', + 'image_ws_medium_min_bytes', 'nsfw_max_concurrent', 'nsfw_batch_size', 'nsfw_max_tokens' @@ -40,44 +41,60 @@ const LOCALE_MAP = { "image_format": { title: "图片格式", desc: "生成的图片格式（url 或 base64）。" }, "video_format": { title: "视频格式", desc: "生成的视频格式（html 或 url，url 为处理后的链接）。" } }, - "grok": { - "label": "Grok 设置", + "network": { + "label": "网络配置", + "timeout": { title: "请求超时", desc: "请求 Grok 服务的超时时间（秒）。" }, + "base_proxy_url": { title: "基础代理 URL", desc: "代理请求到 Grok 官网的基础服务地址。" }, + "asset_proxy_url": { title: "资源代理 URL", desc: "代理请求到 Grok 官网的静态资源（图片/视频）地址。" } + }, + "security": { + "label": "反爬虫验证", + "cf_clearance": { title: "CF Clearance", desc: "Cloudflare Clearance Cookie，用于绕过反爬虫验证。" }, + "browser": { title: "浏览器指纹", desc: "curl_cffi 浏览器指纹标识（如 chrome136）。" }, + "user_agent": { title: "User-Agent", desc: "HTTP 请求的 User-Agent 字符串，需与浏览器指纹匹配。" } + }, + "chat": { + "label": "对话配置", "temporary": { title: "临时对话", desc: "是否启用临时对话模式。" }, - "disable_memory": { title: "禁用记忆", desc: "禁用 Grok 记忆功能，以防止响应中出现不相关上下文" }, + "disable_memory": { title: "禁用记忆", desc: "禁用 Grok 记忆功能，以防止响应中出现不相关上下文。" }, "stream": { title: "流式响应", desc: "是否默认启用流式输出。" }, "thinking": { title: "思维链", desc: "是否启用模型思维链输出。" }, "dynamic_statsig": { title: "动态指纹", desc: "是否启用动态生成 Statsig 值。" }, - "filter_tags": { title: "过滤标签", desc: "自动过滤 Grok 响应中的特殊标签。" }, - "timeout": { title: "超时时间", desc: "请求 Grok 服务的超时时间（秒）。" }, - "base_proxy_url": { title: "基础代理 URL", desc: "代理请求到 Grok 官网的基础服务地址。" }, - "asset_proxy_url": { title: "资源代理 URL", desc: "代理请求到 Grok 官网的静态资源（图片/视频）地址。" }, - "image_ws": { title: "Imagine WebSocket 生成", desc: "启用后 /v1/images/generations 走 WebSocket 直连。" }, - "image_ws_nsfw": { title: "Imagine NSFW", desc: "WS 请求里是否启用 NSFW。" }, - "image_ws_blocked_seconds": { title: "Imagine WebSocket Blocked 阈值", desc: "收到中等图后超过该秒数仍无最终图则判定 blocked。" }, - "image_ws_final_min_bytes": { title: "Imagine WebSocket 最终图最小字节", desc: "判定最终图的最小字节数（通常 JPG > 100KB）。" }, - "max_retry": { title: "最大重试", desc: "请求 Grok 服务失败时的最大重试次数。" }, + "filter_tags": { title: "过滤标签", desc: "自动过滤 Grok 响应中的特殊标签。" } + }, + "retry": { + "label": "重试策略", + "max_retry": { title: "最大重试次数", desc: "请求 Grok 服务失败时的最大重试次数。" }, "retry_status_codes": { title: "重试状态码", desc: "触发重试的 HTTP 状态码列表。" }, "retry_backoff_base": { title: "退避基数", desc: "重试退避的基础延迟（秒）。" }, "retry_backoff_factor": { title: "退避倍率", desc: "重试退避的指数放大系数。" }, "retry_backoff_max": { title: "退避上限", desc: "单次重试等待的最大延迟（秒）。" }, - "retry_budget": { title: "退避预算", desc: "单次请求的最大重试总耗时（秒）。" }, + "retry_budget": { title: "退避预算", desc: "单次请求的最大重试总耗时（秒）。" } + }, + "timeout": { + "label": "超时配置", "stream_idle_timeout": { title: "流空闲超时", desc: "流式响应空闲超时（秒），超过将断开。" }, - "video_idle_timeout": { title: "视频空闲超时", desc: "视频生成空闲超时（秒），超过将断开。" }, - "cf_clearance": { title: "CF Clearance", desc: "" }, - "browser": { title: "curl_cffi", desc: "" }, - "user_agent": { title: "User-Agent", desc: "" } + "video_idle_timeout": { title: "视频空闲超时", desc: "视频生成空闲超时（秒），超过将断开。" } + }, + "image": { + "label": "图片生成", + "image_ws": { title: "WebSocket 生成", desc: "启用后 /v1/images/generations 走 WebSocket 直连。" }, + "image_ws_nsfw": { title: "NSFW 模式", desc: "WebSocket 请求是否启用 NSFW。" }, + "image_ws_blocked_seconds": { title: "Blocked 阈值", desc: "收到中等图后超过该秒数仍无最终图则判定 blocked。" }, + "image_ws_final_min_bytes": { title: "最终图最小字节", desc: "判定最终图的最小字节数（通常 JPG > 100KB）。" }, + "image_ws_medium_min_bytes": { title: "中等图最小字节", desc: "判定中等质量图的最小字节数。" } }, "token": { - "label": "Token 池设置", + "label": "Token 池管理", "auto_refresh": { title: "自动刷新", desc: "是否开启 Token 自动刷新机制。" }, - "refresh_interval_hours": { title: "刷新间隔", desc: "Token 刷新的时间间隔（小时）。" }, + "refresh_interval_hours": { title: "刷新间隔", desc: "普通 Token 刷新的时间间隔（小时）。" }, "super_refresh_interval_hours": { title: "Super 刷新间隔", desc: "Super Token 刷新的时间间隔（小时）。" }, "fail_threshold": { title: "失败阈值", desc: "单个 Token 连续失败多少次后被标记为不可用。" }, "save_delay_ms": { title: "保存延迟", desc: "Token 变更合并写入的延迟（毫秒）。" }, - "reload_interval_sec": { title: "一致性刷新", desc: "多 worker 场景下 Token 状态刷新间隔（秒）。" } + "reload_interval_sec": { title: "同步间隔", desc: "多 worker 场景下 Token 状态刷新间隔（秒）。" } }, "cache": { - "label": "缓存设置", + "label": "缓存管理", "enable_auto_clean": { title: "自动清理", desc: "是否启用缓存自动清理，开启后按上限自动回收。" }, "limit_mb": { title: "清理阈值", desc: "缓存大小阈值（MB），超过阈值会触发清理。" } }, @@ -97,11 +114,9 @@ const LOCALE_MAP = { } }; -// 反爬虫验证字段（从 grok 中单独分组显示） -const ANTIBOT_FIELDS = ['cf_clearance', 'browser', 'user_agent']; -const ANTIBOT_GROUP = { - label: "反爬虫验证", - desc: "以下三个参数用于绕过 Cloudflare 反爬虫验证。配置不正确将导致 403 错误。服务首次请求 Grok 时的 IP 必须与获取 CF Clearance 时的 IP 一致，后续服务器请求 IP 变化不会导致403。" +// 配置部分说明（可选） +const SECTION_DESCRIPTIONS = { + "security": "配置不正确将导致 403 错误。服务首次请求 Grok 时的 IP 必须与获取 CF Clearance 时的 IP 一致，后续服务器请求 IP 变化不会导致 403。" }; const SECTION_ORDER = new Map(Object.keys(LOCALE_MAP).map((key, index) => [key, index])); @@ -250,24 +265,29 @@ function renderConfig(data) { const localeSection = LOCALE_MAP[section]; const keyOrder = localeSection ? new Map(Object.keys(localeSection).map((k, i) => [k, i])) : null; - // 分离反爬虫字段和普通字段 const allKeys = sortByOrder(Object.keys(items), keyOrder); - const normalKeys = section === 'grok' ? allKeys.filter(k => !ANTIBOT_FIELDS.includes(k)) : allKeys; - const antibotKeys = section === 'grok' ? allKeys.filter(k => ANTIBOT_FIELDS.includes(k)) : []; - // 渲染普通字段 - if (normalKeys.length > 0) { + if (allKeys.length > 0) { const card = document.createElement('div'); card.className = 'config-section'; const header = document.createElement('div'); header.innerHTML = `
${getSectionLabel(section)}
`; + + // 添加部分说明（如果有） + if (SECTION_DESCRIPTIONS[section]) { + const descP = document.createElement('p'); + descP.className = 'text-[var(--accents-4)] text-sm mt-1 mb-4'; + descP.textContent = SECTION_DESCRIPTIONS[section]; + header.appendChild(descP); + } + card.appendChild(header); const grid = document.createElement('div'); grid.className = 'config-grid'; - normalKeys.forEach(key => { + allKeys.forEach(key => { const fieldCard = buildFieldCard(section, key, items[key]); grid.appendChild(fieldCard); }); @@ -277,28 +297,6 @@ function renderConfig(data) { fragment.appendChild(card); } } - - // 渲染反爬虫验证分组（仅 grok section） - if (antibotKeys.length > 0) { - const card = document.createElement('div'); - card.className = 'config-section'; - - const header = document.createElement('div'); - header.innerHTML = `
${ANTIBOT_GROUP.label}
-

${ANTIBOT_GROUP.desc}

`; - card.appendChild(header); - - const grid = document.createElement('div'); - grid.className = 'config-grid'; - - antibotKeys.forEach(key => { - const fieldCard = buildFieldCard(section, key, items[key]); - grid.appendChild(fieldCard); - }); - - card.appendChild(grid); - fragment.appendChild(card); - } }); container.appendChild(fragment); diff --git a/config.defaults.toml b/config.defaults.toml index d425cb1..cf95fdf 100644 --- a/config.defaults.toml +++ b/config.defaults.toml @@ -1,58 +1,139 @@ -[grok] -temporary = true -stream = true -thinking = true -dynamic_statsig = true -filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] -timeout = 120 -base_proxy_url = "" -asset_proxy_url = "" -cf_clearance = "" -# 浏览器指纹配置 -browser = "chrome136" -user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" -disable_memory = true -max_retry = 3 -retry_status_codes = [401,429,403] -image_ws = true -image_ws_blocked_seconds = 15 -image_ws_final_min_bytes = 100000 -image_ws_nsfw = true -retry_backoff_base = 0.5 -retry_backoff_factor = 2.0 -retry_backoff_max = 30.0 -retry_budget = 90.0 -stream_idle_timeout = 120.0 -video_idle_timeout = 90.0 - +# ==================== 应用设置 ==================== [app] +# 应用访问地址（用于生成文件链接） app_url = "http://127.0.0.1:8000" +# 后台管理密码 app_key = "grok2api" +# API 调用密钥（可选） api_key = "" +# 生成图片的格式（url 或 base64） image_format = "url" +# 生成视频的格式（html 或 url） video_format = "html" + +# ==================== 网络配置 ==================== +[network] +# 请求超时时间（秒） +timeout = 120 +# 基础代理地址（代理到 Grok 官网） +base_proxy_url = "" +# 资源代理地址（代理静态资源如图片/视频） +asset_proxy_url = "" + + +# ==================== 反爬虫验证 ==================== +[security] +# Cloudflare Clearance Cookie +cf_clearance = "" +# curl_cffi 浏览器指纹 +browser = "chrome136" +# User-Agent 字符串 +user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" + + +# ==================== 对话配置 ==================== +[chat] +# 是否启用临时对话模式 +temporary = true +# 是否禁用 Grok 记忆功能 +disable_memory = true +# 是否默认启用流式响应 +stream = true +# 是否启用思维链输出 +thinking = true +# 是否动态生成 Statsig 指纹 +dynamic_statsig = true +# 过滤的特殊标签列表 +filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] + + +# ==================== 重试策略 ==================== +[retry] +# 最大重试次数 +max_retry = 3 +# 触发重试的 HTTP 状态码 +retry_status_codes = [401,429,403] +# 退避基础延迟（秒） +retry_backoff_base = 0.5 +# 退避倍率 +retry_backoff_factor = 2.0 +# 单次重试最大延迟（秒） +retry_backoff_max = 30.0 +# 总重试预算时间（秒） +retry_budget = 90.0 + + +# ==================== 超时配置 ==================== +[timeout] +# 流式响应空闲超时（秒） +stream_idle_timeout = 120.0 +# 视频生成空闲超时（秒） +video_idle_timeout = 90.0 + + +# ==================== 图片生成 ==================== +[image] +# 是否启用 WebSocket 直连生成图片 +image_ws = true +# WebSocket 生成时是否启用 NSFW +image_ws_nsfw = true +# 收到中等质量图后等待最终图的超时秒数 +image_ws_blocked_seconds = 15 +# 判定为最终图的最小字节数 +image_ws_final_min_bytes = 100000 +# 判定为中等质量图的最小字节数 +image_ws_medium_min_bytes = 30000 + + +# ==================== Token 池管理 ==================== [token] +# 是否启用 Token 自动刷新 auto_refresh = true +# 普通 Token 刷新间隔（小时） refresh_interval_hours = 8 +# Super Token 刷新间隔（小时） super_refresh_interval_hours = 2 +# Token 连续失败阈值 fail_threshold = 5 +# Token 变更保存延迟（毫秒） save_delay_ms = 500 +# 多 worker 状态同步间隔（秒） reload_interval_sec = 30 + +# ==================== 缓存管理 ==================== [cache] +# 是否启用自动清理 enable_auto_clean = true +# 缓存大小上限（MB） limit_mb = 1024 + +# ==================== 并发性能 ==================== [performance] +# Assets 批量处理并发上限 assets_max_concurrent = 25 -media_max_concurrent = 50 -usage_max_concurrent = 25 +# Assets 批量删除大小 assets_delete_batch_size = 10 +# Assets 单批处理数量 assets_batch_size = 10 +# Assets 单次最大处理数量 assets_max_tokens = 1000 + +# Media 生成并发上限 +media_max_concurrent = 50 + +# Token 用量刷新并发上限 +usage_max_concurrent = 25 +# Token 用量刷新批次大小 usage_batch_size = 50 +# Token 用量单次最大数量 usage_max_tokens = 1000 + +# NSFW 批量开启并发上限 nsfw_max_concurrent = 10 +# NSFW 批量开启批次大小 nsfw_batch_size = 50 +# NSFW 单次最大数量 nsfw_max_tokens = 1000 diff --git a/data/config.toml b/data/config.toml index 22296cf..e51f56f 100644 --- a/data/config.toml +++ b/data/config.toml @@ -1,53 +1,46 @@ -[grok] -browser = "chrome136" +[app] +app_url = "http://127.0.0.1:8000" +app_key = "grok2api" +api_key = "" +image_format = "url" +video_format = "html" + +[network] timeout = 120 -user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" base_proxy_url = "" asset_proxy_url = "" + +[security] cf_clearance = "" +browser = "chrome136" +user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" + +[chat] temporary = true disable_memory = true stream = true thinking = true +dynamic_statsig = true filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] + +[retry] max_retry = 3 retry_status_codes = [401,429,403] retry_backoff_base = 0.5 retry_backoff_factor = 2.0 retry_backoff_max = 30.0 retry_budget = 90.0 + +[timeout] stream_idle_timeout = 120.0 video_idle_timeout = 90.0 + +[image] +image_ws = true +image_ws_nsfw = true image_ws_blocked_seconds = 15 image_ws_final_min_bytes = 100000 image_ws_medium_min_bytes = 30000 -dynamic_statsig = true -image_ws = true -image_ws_nsfw = true - -[app] -app_url = "http://127.0.0.1:8000" -image_format = "url" -video_format = "html" -app_key = "grok2api" -api_key = "" - -[performance] -assets_max_concurrent = 25 -media_max_concurrent = 50 -usage_max_concurrent = 25 -assets_delete_batch_size = 10 -assets_batch_size = 10 -assets_max_tokens = 1000 -usage_batch_size = 50 -usage_max_tokens = 1000 -nsfw_max_concurrent = 10 -nsfw_batch_size = 50 -nsfw_max_tokens = 1000 - -[cache] -enable_auto_clean = true -limit_mb = 1024 [token] auto_refresh = true @@ -56,3 +49,43 @@ super_refresh_interval_hours = 2 fail_threshold = 5 save_delay_ms = 200 reload_interval_sec = 30 + +[cache] +enable_auto_clean = true +limit_mb = 1024 + +[performance] +assets_max_concurrent = 25 +assets_delete_batch_size = 10 +assets_batch_size = 10 +assets_max_tokens = 1000 +media_max_concurrent = 50 +usage_max_concurrent = 25 +usage_batch_size = 50 +usage_max_tokens = 1000 +nsfw_max_concurrent = 10 +nsfw_batch_size = 50 +nsfw_max_tokens = 1000 + +[grok] +temporary = true +stream = true +thinking = true +dynamic_statsig = true +filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] +timeout = 120 +base_proxy_url = "" +asset_proxy_url = "" +cf_clearance = "" +max_retry = 3 +retry_status_codes = [401,429,403] +image_ws = true +image_ws_blocked_seconds = 15 +image_ws_final_min_bytes = 100000 +image_ws_nsfw = true +retry_backoff_base = 0.5 +retry_backoff_factor = 2.0 +retry_backoff_max = 30.0 +retry_budget = 90.0 +stream_idle_timeout = 120.0 +video_idle_timeout = 90.0 diff --git a/docs/README.en.md b/docs/README.en.md index 34bc0e5..66d5c0f 100644 --- a/docs/README.en.md +++ b/docs/README.en.md @@ -1,6 +1,6 @@ # Grok2API -[中文](README.md) | **English** +[中文](../README.md) | **English** > [!NOTE] > This project is for learning and research only. You must comply with Grok's Terms of Use and applicable laws. Do not use it for illegal purposes. @@ -194,6 +194,10 @@ Config file: `data/config.toml` > In production or behind a reverse proxy, make sure `app.app_url` is set to the public URL. > Otherwise file links may be incorrect or return 403. +> [!TIP] +> **v2.0 Config Upgrade**: Existing users will have their config **auto-migrated** to the new structure upon update. +> Custom values from the old `[grok]` section will be automatically mapped to the corresponding new sections. + | Module | Field | Key | Description | Default | | :--- | :--- | :--- | :--- | :--- | | **app** | `app_url` | App URL | External access URL for Grok2API (used for file links). | `http://127.0.0.1:8000` | @@ -201,46 +205,50 @@ Config file: `data/config.toml` | | `api_key` | API key | Token for calling Grok2API (optional). | `""` | | | `image_format` | Image format | Output image format (`url` or `base64`). | `url` | | | `video_format` | Video format | Output video format (html tag or processed url). | `html` | -| **grok** | `temporary` | Temporary chat | Enable temporary conversation mode. | `true` | +| **network** | `timeout` | Request timeout | Timeout for Grok requests (seconds). | `120` | +| | `base_proxy_url` | Base proxy URL | Base service address proxying Grok official site. | `""` | +| | `asset_proxy_url` | Asset proxy URL | Proxy URL for Grok static assets (images/videos). | `""` | +| **security** | `cf_clearance` | CF Clearance | Cloudflare clearance cookie for bypassing anti-bot. | `""` | +| | `browser` | Browser fingerprint | curl_cffi browser fingerprint (e.g. chrome136). | `chrome136` | +| | `user_agent` | User-Agent | HTTP User-Agent string. | `Mozilla/5.0 (Macintosh; ...)` | +| **chat** | `temporary` | Temporary chat | Enable temporary conversation mode. | `true` | +| | `disable_memory` | Disable memory | Disable Grok memory to prevent irrelevant context. | `true` | | | `stream` | Streaming | Enable streaming by default. | `true` | | | `thinking` | Thinking chain | Enable model thinking output. | `true` | | | `dynamic_statsig` | Dynamic fingerprint | Enable dynamic Statsig value generation. | `true` | | | `filter_tags` | Filter tags | Auto-filter special tags in Grok responses. | `["xaiartifact", "xai:tool_usage_card", "grok:render"]` | -| | `timeout` | Timeout | Timeout for Grok requests (seconds). | `120` | -| | `base_proxy_url` | Base proxy URL | Base service address proxying Grok official site. | `""` | -| | `asset_proxy_url` | Asset proxy URL | Proxy URL for Grok static assets (images/videos). | `""` | -| | `cf_clearance` | CF Clearance | Cloudflare clearance cookie for verification. | `""` | -| | `max_retry` | Max retries | Max retries on Grok request failure. | `3` | +| **retry** | `max_retry` | Max retries | Max retries on Grok request failure. | `3` | | | `retry_status_codes` | Retry status codes | HTTP status codes that trigger retry. | `[401, 429, 403]` | | | `retry_backoff_base` | Backoff base | Base delay for retry backoff (seconds). | `0.5` | | | `retry_backoff_factor` | Backoff factor | Exponential multiplier for retry backoff. | `2.0` | | | `retry_backoff_max` | Backoff max | Max wait per retry (seconds). | `30.0` | | | `retry_budget` | Backoff budget | Max total retry time per request (seconds). | `90.0` | -| | `stream_idle_timeout` | Stream idle timeout | Idle timeout for streaming responses (seconds). | `45.0` | +| **timeout** | `stream_idle_timeout` | Stream idle timeout | Idle timeout for streaming responses (seconds). | `120.0` | | | `video_idle_timeout` | Video idle timeout | Idle timeout for video generation (seconds). | `90.0` | -| | `image_ws` | Image WS generation | When enabled, `/v1/images/generations` uses WebSocket. | `false` | -| | `image_ws_blocked_seconds` | WS blocked threshold | Mark blocked if no final image after this many seconds post-medium. | `15` | -| | `image_ws_final_min_bytes` | WS final min bytes | Minimum bytes to treat an image as final (JPG usually > 100KB). | `100000` | -| | `image_ws_nsfw` | WS NSFW | Enable NSFW in WS requests. | `true` | +| **image** | `image_ws` | WebSocket generation | When enabled, `/v1/images/generations` uses WebSocket. | `true` | +| | `image_ws_nsfw` | NSFW mode | Enable NSFW in WebSocket requests. | `true` | +| | `image_ws_blocked_seconds` | Blocked threshold | Mark blocked if no final image after this many seconds post-medium. | `15` | +| | `image_ws_final_min_bytes` | Final min bytes | Minimum bytes to treat an image as final (JPG usually > 100KB). | `100000` | +| | `image_ws_medium_min_bytes` | Medium min bytes | Minimum bytes for medium quality image. | `30000` | | **token** | `auto_refresh` | Auto refresh | Enable automatic token refresh. | `true` | -| | `refresh_interval_hours` | Refresh interval | Token refresh interval (hours). | `8` | +| | `refresh_interval_hours` | Refresh interval | Regular token refresh interval (hours). | `8` | | | `super_refresh_interval_hours` | Super refresh interval | Super token refresh interval (hours). | `2` | | | `fail_threshold` | Failure threshold | Consecutive failures before a token is disabled. | `5` | | | `save_delay_ms` | Save delay | Debounced save delay for token changes (ms). | `500` | -| | `reload_interval_sec` | Consistency refresh | Token state refresh interval in multi-worker setups (sec). | `30` | +| | `reload_interval_sec` | Sync interval | Token state refresh interval in multi-worker setups (sec). | `30` | | **cache** | `enable_auto_clean` | Auto clean | Enable cache auto clean; cleanup when exceeding limit. | `true` | | | `limit_mb` | Cleanup threshold | Cache size threshold (MB) that triggers cleanup. | `1024` | -| **performance** | `nsfw_max_concurrent` | NSFW enable concurrency | Concurrency cap for enabling NSFW in batch. Recommended 10. | `10` | +| **performance** | `media_max_concurrent` | Media concurrency | Concurrency cap for video/media generation. Recommended 50. | `50` | +| | `assets_max_concurrent` | Assets concurrency | Concurrency cap for batch asset find/delete. Recommended 25. | `25` | +| | `assets_batch_size` | Assets batch size | Batch size for asset find/delete. Recommended 10. | `10` | +| | `assets_max_tokens` | Assets max tokens | Max tokens per asset find/delete batch. Recommended 1000. | `1000` | +| | `assets_delete_batch_size` | Assets delete batch | Batch concurrency for single-account asset deletion. Recommended 10. | `10` | +| | `usage_max_concurrent` | Token refresh concurrency | Concurrency cap for batch usage refresh. Recommended 25. | `25` | +| | `usage_batch_size` | Token refresh batch size | Batch size for usage refresh. Recommended 50. | `50` | +| | `usage_max_tokens` | Token refresh max tokens | Max tokens per usage refresh batch. Recommended 1000. | `1000` | +| | `nsfw_max_concurrent` | NSFW enable concurrency | Concurrency cap for enabling NSFW in batch. Recommended 10. | `10` | | | `nsfw_batch_size` | NSFW enable batch size | Batch size for enabling NSFW. Recommended 50. | `50` | | | `nsfw_max_tokens` | NSFW enable max tokens | Max tokens per NSFW batch to avoid mistakes. Recommended 1000. | `1000` | -| | `usage_max_concurrent` | Token usage refresh concurrency | Concurrency cap for batch usage refresh. Recommended 25. | `25` | -| | `usage_batch_size` | Token usage refresh batch size | Batch size for usage refresh. Recommended 50. | `50` | -| | `usage_max_tokens` | Token usage refresh max tokens | Max tokens per usage refresh batch. Recommended 1000. | `1000` | -| | `assets_max_concurrent` | Online assets find/delete concurrency | Concurrency cap for online asset find/delete. Recommended 25. | `25` | -| | `assets_batch_size` | Online assets find/delete batch size | Batch size for online asset find/delete. Recommended 10. | `10` | -| | `assets_max_tokens` | Online assets find/delete max tokens | Max tokens per online asset find/delete batch. Recommended 1000. | `1000` | -| | `assets_delete_batch_size` | Online assets delete batch | Batch concurrency for online asset deletion. Recommended 10. | `10` | -| | `media_max_concurrent` | Media concurrency | Concurrency cap for video/media generation. Recommended 50. | `50` |
diff --git a/readme.md b/readme.md index 90ff53a..bf8ea55 100644 --- a/readme.md +++ b/readme.md @@ -1,6 +1,6 @@ # Grok2API -**中文** | [English](README.en.md) +**中文** | [English](docs/README.en.md) > [!NOTE] > 本项目仅供学习与研究，使用者必须在遵循 Grok 的 **使用条款** 以及 **法律法规** 的情况下使用，不得用于非法用途。 @@ -223,6 +223,10 @@ curl http://localhost:8000/v1/images/edits \ > 生产环境或反向代理部署时，请确保 `app.app_url` 配置为对外可访问的完整 URL， > 否则可能出现文件访问链接不正确或 403 等问题。 +> [!TIP] +> **v2.0 配置结构升级**：旧版本用户更新后，配置会**自动迁移**到新结构，无需手动修改。 +> 旧的 `[grok]` 配置节中的自定义值会自动映射到对应的新配置节。 + | 模块 | 字段 | 配置名 | 说明 | 默认值 | | :-------------------- | :--------------------------- | :----------- | :--------------------------------------------------- | :-------------------------------------------------------- | | **app** | `app_url` | 应用地址 | 当前 Grok2API 服务的外部访问 URL，用于文件链接访问。 | `http://127.0.0.1:8000` | @@ -230,46 +234,50 @@ curl http://localhost:8000/v1/images/edits \ | | `api_key` | API 密钥 | 调用 Grok2API 服务的 Token（可选）。 | `""` | | | `image_format` | 图片格式 | 生成的图片格式（url 或 base64）。 | `url` | | | `video_format` | 视频格式 | 生成的视频格式（html 或 url，url 为处理后的链接）。 | `html` | -| **grok** | `temporary` | 临时对话 | 是否启用临时对话模式。 | `true` | +| **network** | `timeout` | 请求超时 | 请求 Grok 服务的超时时间（秒）。 | `120` | +| | `base_proxy_url` | 基础代理 URL | 代理请求到 Grok 官网的基础服务地址。 | `""` | +| | `asset_proxy_url` | 资源代理 URL | 代理请求到 Grok 官网的静态资源（图片/视频）地址。 | `""` | +| **security** | `cf_clearance` | CF Clearance | Cloudflare 验证 Cookie，用于绕过反爬虫验证。 | `""` | +| | `browser` | 浏览器指纹 | curl_cffi 浏览器指纹标识（如 chrome136）。 | `chrome136` | +| | `user_agent` | User-Agent | HTTP 请求的 User-Agent 字符串。 | `Mozilla/5.0 (Macintosh; ...)` | +| **chat** | `temporary` | 临时对话 | 是否启用临时对话模式。 | `true` | +| | `disable_memory` | 禁用记忆 | 禁用 Grok 记忆功能，防止响应中出现不相关上下文。 | `true` | | | `stream` | 流式响应 | 是否默认启用流式输出。 | `true` | | | `thinking` | 思维链 | 是否启用模型思维链输出。 | `true` | | | `dynamic_statsig` | 动态指纹 | 是否启用动态生成 Statsig 值。 | `true` | | | `filter_tags` | 过滤标签 | 自动过滤 Grok 响应中的特殊标签。 | `["xaiartifact", "xai:tool_usage_card", "grok:render"]` | -| | `timeout` | 超时时间 | 请求 Grok 服务的超时时间（秒）。 | `120` | -| | `base_proxy_url` | 基础代理 URL | 代理请求到 Grok 官网的基础服务地址。 | `""` | -| | `asset_proxy_url` | 资源代理 URL | 代理请求到 Grok 官网的静态资源（图片/视频）地址。 | `""` | -| | `cf_clearance` | CF Clearance | Cloudflare 验证 Cookie，用于验证 Cloudflare 的验证。 | `""` | -| | `max_retry` | 最大重试 | 请求 Grok 服务失败时的最大重试次数。 | `3` | +| **retry** | `max_retry` | 最大重试 | 请求 Grok 服务失败时的最大重试次数。 | `3` | | | `retry_status_codes` | 重试状态码 | 触发重试的 HTTP 状态码列表。 | `[401, 429, 403]` | | | `retry_backoff_base` | 退避基数 | 重试退避的基础延迟（秒）。 | `0.5` | | | `retry_backoff_factor` | 退避倍率 | 重试退避的指数放大系数。 | `2.0` | | | `retry_backoff_max` | 退避上限 | 单次重试等待的最大延迟（秒）。 | `30.0` | | | `retry_budget` | 退避预算 | 单次请求的最大重试总耗时（秒）。 | `90.0` | -| | `stream_idle_timeout` | 流空闲超时 | 流式响应空闲超时（秒），超过将断开。 | `45.0` | +| **timeout** | `stream_idle_timeout` | 流空闲超时 | 流式响应空闲超时（秒），超过将断开。 | `120.0` | | | `video_idle_timeout` | 视频空闲超时 | 视频生成空闲超时（秒），超过将断开。 | `90.0` | -| | `image_ws` | 图片 WS 生成 | 启用后 `/v1/images/generations` 走 WebSocket 直连。 | `false` | -| | `image_ws_blocked_seconds` | WS Blocked 阈值 | 收到中等图后超过该秒数仍无最终图则判定 blocked。 | `15` | -| | `image_ws_final_min_bytes` | WS 最终图最小字节 | 判定最终图的最小字节数（通常 JPG > 100KB）。 | `100000` | -| | `image_ws_nsfw` | WS NSFW | WS 请求里是否启用 NSFW。 | `true` | +| **image** | `image_ws` | WebSocket 生成 | 启用后 `/v1/images/generations` 走 WebSocket 直连。| `true` | +| | `image_ws_nsfw` | NSFW 模式 | WebSocket 请求是否启用 NSFW。 | `true` | +| | `image_ws_blocked_seconds` | Blocked 阈值 | 收到中等图后超过该秒数仍无最终图则判定 blocked。 | `15` | +| | `image_ws_final_min_bytes` | 最终图最小字节 | 判定最终图的最小字节数（通常 JPG > 100KB）。 | `100000` | +| | `image_ws_medium_min_bytes`| 中等图最小字节 | 判定中等质量图的最小字节数。 | `30000` | | **token** | `auto_refresh` | 自动刷新 | 是否开启 Token 自动刷新机制。 | `true` | -| | `refresh_interval_hours` | 刷新间隔 | Token 刷新的时间间隔（小时）。 | `8` | -| | `super_refresh_interval_hours` | Super 刷新间隔 | Super Token 刷新的时间间隔（小时）。 | `2` | +| | `refresh_interval_hours` | 刷新间隔 | 普通 Token 刷新的时间间隔（小时）。 | `8` | +| | `super_refresh_interval_hours` | Super 刷新间隔 | Super Token 刷新的时间间隔（小时）。 | `2` | | | `fail_threshold` | 失败阈值 | 单个 Token 连续失败多少次后被标记为不可用。 | `5` | | | `save_delay_ms` | 保存延迟 | Token 变更合并写入的延迟（毫秒）。 | `500` | -| | `reload_interval_sec` | 一致性刷新 | 多 worker 场景下 Token 状态刷新间隔（秒）。 | `30` | +| | `reload_interval_sec` | 同步间隔 | 多 worker 场景下 Token 状态刷新间隔（秒）。 | `30` | | **cache** | `enable_auto_clean` | 自动清理 | 是否启用缓存自动清理，开启后按上限自动回收。 | `true` | | | `limit_mb` | 清理阈值 | 缓存大小阈值（MB），超过阈值会触发清理。 | `1024` | -| **performance** | `nsfw_max_concurrent` | 开启 NSFW 模式并发上限 | 批量开启 NSFW 模式时的并发请求上限。推荐 10。 | `10` | -| | `nsfw_batch_size` | 开启 NSFW 模式批次大小 | 批量开启 NSFW 模式的单批处理数量。推荐 50。 | `50` | -| | `nsfw_max_tokens` | 开启 NSFW 模式最大数量 | 单次批量开启 NSFW 的 Token 数量上限，防止误操作。推荐 1000。 | `1000` | -| | `usage_max_concurrent` | Token 用量刷新并发上限 | 批量刷新 Token 用量时的并发请求上限。推荐 25。 | `25` | -| | `usage_batch_size` | Token 用量刷新批次大小 | 批量刷新 Token 用量的单批处理数量。推荐 50。 | `50` | -| | `usage_max_tokens` | Token 用量刷新最大数量 | 批量刷新 Token 用量的单次处理数量上限。推荐 1000。 | `1000` | -| | `assets_max_concurrent` | 在线资产查找/删除并发上限 | 在线资产查找/删除时的并发请求上限。推荐 25。 | `25` | -| | `assets_batch_size` | 在线资产查找/删除批次大小 | 在线资产查找/删除的单批处理数量。推荐 10。 | `10` | -| | `assets_max_tokens` | 在线资产查找/删除最大数量 | 在线资产查找/删除的单次处理数量上限。推荐 1000。 | `1000` | -| | `assets_delete_batch_size` | 在线资产删除批量 | 在线资产删除单批并发数量。推荐 10。 | `10` | -| | `media_max_concurrent` | 媒体并发上限 | 视频/媒体生成请求的并发上限。推荐 50。 | `50` | +| **performance** | `media_max_concurrent` | Media 并发上限 | 视频/媒体生成请求的并发上限。推荐 50。 | `50` | +| | `assets_max_concurrent` | Assets 并发上限 | 批量查找/删除资产时的并发请求上限。推荐 25。 | `25` | +| | `assets_batch_size` | Assets 批次大小 | 批量查找/删除资产时的单批处理数量。推荐 10。 | `10` | +| | `assets_max_tokens` | Assets 最大数量 | 单次批量查找/删除资产时的处理数量上限。推荐 1000。 | `1000` | +| | `assets_delete_batch_size` | Assets 删除批量 | 单账号批量删除资产时的单批并发数量。推荐 10。 | `10` | +| | `usage_max_concurrent` | Token 刷新并发上限 | 批量刷新 Token 用量时的并发请求上限。推荐 25。 | `25` | +| | `usage_batch_size` | Token 刷新批次大小 | 批量刷新 Token 用量的单批处理数量。推荐 50。 | `50` | +| | `usage_max_tokens` | Token 刷新最大数量 | 单次批量刷新 Token 用量时的处理数量上限。推荐 1000。| `1000` | +| | `nsfw_max_concurrent` | NSFW 开启并发上限 | 批量开启 NSFW 模式时的并发请求上限。推荐 10。 | `10` | +| | `nsfw_batch_size` | NSFW 开启批次大小 | 批量开启 NSFW 模式的单批处理数量。推荐 50。 | `50` | +| | `nsfw_max_tokens` | NSFW 开启最大数量 | 单次批量开启 NSFW 的 Token 数量上限。推荐 1000。 | `1000` |
From 8361e65c31b180daee7b7f3b3339e360de1cdeac Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sat, 7 Feb 2026 13:34:55 +0800 Subject: [PATCH 5/6] refactor: clean up code formatting and improve readability across multiple files --- app/api/v1/admin.py | 70 +++++--- app/api/v1/image.py | 47 ++--- app/core/config.py | 41 +++-- app/services/grok/defaults.py | 9 - app/services/grok/models/model.py | 1 - app/services/grok/processors/base.py | 1 + .../grok/processors/chat_processors.py | 68 +++++-- .../grok/processors/image_processors.py | 39 +++- .../grok/processors/image_ws_processors.py | 55 +++++- .../grok/processors/video_processors.py | 72 ++++++-- app/services/grok/services/assets.py | 169 ++++++++++++------ app/services/grok/services/chat.py | 141 +++++++++++---- app/services/grok/services/image.py | 94 ++++++++-- app/services/grok/services/media.py | 104 ++++++++--- app/services/grok/services/nsfw.py | 14 +- app/services/grok/services/usage.py | 9 +- app/services/grok/services/voice.py | 37 ++-- app/services/grok/utils/retry.py | 1 - app/services/grok/utils/stream.py | 17 +- main.py | 4 +- pyproject.toml | 5 + uv.lock | 33 ++++ 22 files changed, 751 insertions(+), 280 deletions(-) diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 1c8cbcd..ab12b46 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -32,17 +32,21 @@ def _collect_tokens(data: dict) -> list[str]: return tokens -def _truncate_tokens(tokens: list[str], max_tokens: int, operation: str = "operation") -> tuple[list[str], bool, int]: +def _truncate_tokens( + tokens: list[str], max_tokens: int, operation: str = "operation" +) -> tuple[list[str], bool, int]: """去重并截断 token 列表，返回 (unique_tokens, truncated, original_count)""" unique_tokens = list(dict.fromkeys(tokens)) original_count = len(unique_tokens) truncated = False - + if len(unique_tokens) > max_tokens: unique_tokens = unique_tokens[:max_tokens] truncated = True - logger.warning(f"{operation}: truncated from {original_count} to {max_tokens} tokens") - + logger.warning( + f"{operation}: truncated from {original_count} to {max_tokens} tokens" + ) + return unique_tokens, truncated, original_count @@ -50,6 +54,7 @@ def _mask_token(token: str) -> str: """掩码 token 显示""" return f"{token[:8]}...{token[-8:]}" if len(token) > 20 else token + async def render_template(filename: str): """渲染指定模板""" template_path = TEMPLATE_DIR / filename @@ -248,8 +253,11 @@ async def get_storage_info(): storage_type = "redis" elif isinstance(storage, SQLStorage): storage_type = { - "mysql": "mysql", "mariadb": "mysql", - "postgres": "pgsql", "postgresql": "pgsql", "pgsql": "pgsql" + "mysql": "mysql", + "mariadb": "mysql", + "postgres": "pgsql", + "postgresql": "pgsql", + "pgsql": "pgsql", }.get(storage.dialect, storage.dialect) return {"type": storage_type or "local"} @@ -309,7 +317,9 @@ async def update_tokens_api(data: dict): if isinstance(raw_token, str) and raw_token.startswith("sso="): token_data["token"] = raw_token[4:] - base = existing_map.get(pool_name, {}).get(token_data.get("token"), {}) + base = existing_map.get(pool_name, {}).get( + token_data.get("token"), {} + ) merged = dict(base) merged.update(token_data) if merged.get("tags") is None: @@ -338,13 +348,15 @@ async def refresh_tokens_api(data: dict): try: mgr = await get_token_manager() tokens = _collect_tokens(data) - + if not tokens: raise HTTPException(status_code=400, detail="No tokens provided") # 去重并截断 max_tokens = int(get_config("performance.usage_max_tokens")) - unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "Usage refresh") + unique_tokens, truncated, original_count = _truncate_tokens( + tokens, max_tokens, "Usage refresh" + ) # 批量执行配置 max_concurrent = get_config("performance.usage_max_concurrent") @@ -392,11 +404,11 @@ async def refresh_tokens_api_async(data: dict): # 去重并截断 max_tokens = int(get_config("performance.usage_max_tokens")) - unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "Usage refresh") - - max_concurrent = get_config( - "performance.usage_max_concurrent" + unique_tokens, truncated, original_count = _truncate_tokens( + tokens, max_tokens, "Usage refresh" ) + + max_concurrent = get_config("performance.usage_max_concurrent") batch_size = get_config("performance.usage_batch_size") task = create_task(len(unique_tokens)) @@ -493,12 +505,12 @@ async def enable_nsfw_api(data: dict): # 去重并截断 max_tokens = int(get_config("performance.nsfw_max_tokens")) - unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "NSFW enable") + unique_tokens, truncated, original_count = _truncate_tokens( + tokens, max_tokens, "NSFW enable" + ) # 批量执行配置 - max_concurrent = get_config( - "performance.nsfw_max_concurrent" - ) + max_concurrent = get_config("performance.nsfw_max_concurrent") batch_size = get_config("performance.nsfw_batch_size") # 定义 worker @@ -582,7 +594,9 @@ async def enable_nsfw_api_async(data: dict): # 去重并截断 max_tokens = int(get_config("performance.nsfw_max_tokens")) - unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "NSFW enable") + unique_tokens, truncated, original_count = _truncate_tokens( + tokens, max_tokens, "NSFW enable" + ) max_concurrent = get_config("performance.nsfw_max_concurrent") batch_size = get_config("performance.nsfw_batch_size") @@ -765,7 +779,9 @@ async def get_cache_stats_api(request: Request): } if selected_tokens: - selected_tokens, truncated, original_count = _truncate_tokens(selected_tokens, max_tokens, "Assets fetch") + selected_tokens, truncated, original_count = _truncate_tokens( + selected_tokens, max_tokens, "Assets fetch" + ) total = 0 raw_results = await run_in_batches( selected_tokens, @@ -934,7 +950,9 @@ async def load_online_cache_api_async(data: dict): raise HTTPException(status_code=400, detail="No tokens provided") max_tokens = int(get_config("performance.assets_max_tokens")) - selected_tokens, truncated, original_count = _truncate_tokens(selected_tokens, max_tokens, "Assets load") + selected_tokens, truncated, original_count = _truncate_tokens( + selected_tokens, max_tokens, "Assets load" + ) max_concurrent = get_config("performance.assets_max_concurrent") batch_size = get_config("performance.assets_batch_size") @@ -1112,10 +1130,14 @@ async def clear_online_cache_api(data: dict): # 最大数量限制 max_tokens = int(get_config("performance.assets_max_tokens")) - token_list, truncated, original_count = _truncate_tokens(token_list, max_tokens, "Clear online cache") + token_list, truncated, original_count = _truncate_tokens( + token_list, max_tokens, "Clear online cache" + ) results = {} - max_concurrent = max(1, int(get_config("performance.assets_max_concurrent"))) + max_concurrent = max( + 1, int(get_config("performance.assets_max_concurrent")) + ) batch_size = max(1, int(get_config("performance.assets_batch_size"))) async def _clear_one(t: str): @@ -1180,7 +1202,9 @@ async def clear_online_cache_api_async(data: dict): raise HTTPException(status_code=400, detail="No tokens provided") max_tokens = int(get_config("performance.assets_max_tokens")) - token_list, truncated, original_count = _truncate_tokens(token_list, max_tokens, "Clear online cache async") + token_list, truncated, original_count = _truncate_tokens( + token_list, max_tokens, "Clear online cache async" + ) max_concurrent = get_config("performance.assets_max_concurrent") batch_size = get_config("performance.assets_batch_size") diff --git a/app/api/v1/image.py b/app/api/v1/image.py index fc588ff..ec207c1 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -34,6 +34,7 @@ from app.core.logger import logger router = APIRouter(tags=["Images"]) + class ImageGenerationRequest(BaseModel): """图片生成请求 - OpenAI 兼容""" @@ -47,7 +48,6 @@ class ImageGenerationRequest(BaseModel): stream: Optional[bool] = Field(False, description="是否流式输出") - class ImageEditRequest(BaseModel): """图片编辑请求 - OpenAI 兼容""" @@ -62,7 +62,6 @@ class ImageEditRequest(BaseModel): stream: Optional[bool] = Field(False, description="是否流式输出") - def _validate_common_request( request: Union[ImageGenerationRequest, ImageEditRequest], *, @@ -184,9 +183,7 @@ def validate_edit_request(request: ImageEditRequest, images: List[UploadFile]): """验证图片编辑请求参数""" if request.model != "grok-imagine-1.0-edit": raise ValidationException( - message=( - "The model `grok-imagine-1.0-edit` is required for image edits." - ), + message=("The model `grok-imagine-1.0-edit` is required for image edits."), param="model", code="model_not_supported", ) @@ -207,7 +204,11 @@ def validate_edit_request(request: ImageEditRequest, images: List[UploadFile]): def _get_effort(model_info) -> EffortType: """获取模型消耗级别""" - return EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + return ( + EffortType.HIGH + if (model_info and model_info.cost.value == "high") + else EffortType.LOW + ) async def _wrap_stream_with_usage(stream, token_mgr, token, model_info): @@ -229,13 +230,13 @@ async def _get_token(model: str): """获取可用 token""" token_mgr = await get_token_manager() await token_mgr.reload_if_stale() - + token = None for pool_name in ModelService.pool_candidates_for_model(model): token = token_mgr.get_token(pool_name) if token: break - + if not token: raise AppException( message="No available tokens. Please try again later.", @@ -243,7 +244,7 @@ async def _get_token(model: str): code="rate_limit_exceeded", status_code=429, ) - + return token_mgr, token @@ -342,7 +343,9 @@ async def create_image(request: ImageGenerationRequest): ) return StreamingResponse( - _wrap_stream_with_usage(processor.process(upstream), token_mgr, token, model_info), + _wrap_stream_with_usage( + processor.process(upstream), token_mgr, token, model_info + ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) @@ -361,7 +364,9 @@ async def create_image(request: ImageGenerationRequest): ) return StreamingResponse( - _wrap_stream_with_usage(processor.process(response), token_mgr, token, model_info), + _wrap_stream_with_usage( + processor.process(response), token_mgr, token, model_info + ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) @@ -485,7 +490,6 @@ async def create_image(request: ImageGenerationRequest): ) - @router.post("/images/edits") async def edit_image( prompt: str = Form(...), @@ -594,9 +598,7 @@ async def edit_image( if file_uri.startswith("http"): image_urls.append(file_uri) else: - image_urls.append( - f"https://assets.grok.com/{file_uri.lstrip('/')}" - ) + image_urls.append(f"https://assets.grok.com/{file_uri.lstrip('/')}") finally: await upload_service.close() @@ -638,9 +640,9 @@ async def edit_image( } if parent_post_id: - model_config_override["modelMap"]["imageEditModelConfig"][ - "parentPostId" - ] = parent_post_id + model_config_override["modelMap"]["imageEditModelConfig"]["parentPostId"] = ( + parent_post_id + ) raw_payload = { "temporary": bool(get_config("chat.temporary")), @@ -675,11 +677,16 @@ async def edit_image( ) processor = ImageStreamProcessor( - model_info.model_id, token, n=edit_request.n, response_format=response_format + model_info.model_id, + token, + n=edit_request.n, + response_format=response_format, ) return StreamingResponse( - _wrap_stream_with_usage(processor.process(response), token_mgr, token, model_info), + _wrap_stream_with_usage( + processor.process(response), token_mgr, token, model_info + ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) diff --git a/app/core/config.py b/app/core/config.py index 22c125b..bec8714 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -32,10 +32,12 @@ def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any return result -def _migrate_deprecated_config(config: Dict[str, Any], valid_sections: set) -> tuple[Dict[str, Any], set]: +def _migrate_deprecated_config( + config: Dict[str, Any], valid_sections: set +) -> tuple[Dict[str, Any], set]: """ 迁移废弃的配置节到新配置结构 - + Returns: (迁移后的配置, 废弃的配置节集合) """ @@ -48,49 +50,44 @@ def _migrate_deprecated_config(config: Dict[str, Any], valid_sections: set) -> t "grok.thinking": "chat.thinking", "grok.dynamic_statsig": "chat.dynamic_statsig", "grok.filter_tags": "chat.filter_tags", - "grok.timeout": "network.timeout", "grok.base_proxy_url": "network.base_proxy_url", "grok.asset_proxy_url": "network.asset_proxy_url", - "grok.cf_clearance": "security.cf_clearance", "grok.browser": "security.browser", "grok.user_agent": "security.user_agent", - "grok.max_retry": "retry.max_retry", "grok.retry_status_codes": "retry.retry_status_codes", "grok.retry_backoff_base": "retry.retry_backoff_base", "grok.retry_backoff_factor": "retry.retry_backoff_factor", "grok.retry_backoff_max": "retry.retry_backoff_max", "grok.retry_budget": "retry.retry_budget", - "grok.stream_idle_timeout": "timeout.stream_idle_timeout", "grok.video_idle_timeout": "timeout.video_idle_timeout", - "grok.image_ws": "image.image_ws", "grok.image_ws_nsfw": "image.image_ws_nsfw", "grok.image_ws_blocked_seconds": "image.image_ws_blocked_seconds", "grok.image_ws_final_min_bytes": "image.image_ws_final_min_bytes", "grok.image_ws_medium_min_bytes": "image.image_ws_medium_min_bytes", } - + deprecated_sections = set(config.keys()) - valid_sections if not deprecated_sections: return config, set() - + result = {k: deepcopy(v) for k, v in config.items() if k in valid_sections} migrated_count = 0 - + # 处理废弃配置节中的配置项 for old_section in deprecated_sections: if old_section not in config or not isinstance(config[old_section], dict): continue - + for old_key, old_value in config[old_section].items(): # 查找映射规则 old_path = f"{old_section}.{old_key}" new_path = MIGRATION_MAP.get(old_path) - + if new_path: new_section, new_key = new_path.split(".", 1) # 确保新配置节存在 @@ -100,10 +97,10 @@ def _migrate_deprecated_config(config: Dict[str, Any], valid_sections: set) -> t result[new_section][new_key] = old_value migrated_count += 1 logger.debug(f"Migrated config: {old_path} -> {new_path} = {old_value}") - + if migrated_count > 0: logger.info(f"Migrated {migrated_count} config items from deprecated sections") - + return result, deprecated_sections @@ -166,18 +163,24 @@ class Config: config_data = {} config_data = config_data or {} - + # 检查是否有废弃的配置节 valid_sections = set(self._defaults.keys()) - config_data, deprecated_sections = _migrate_deprecated_config(config_data, valid_sections) + config_data, deprecated_sections = _migrate_deprecated_config( + config_data, valid_sections + ) if deprecated_sections: - logger.info(f"Cleaned deprecated config sections: {deprecated_sections}") - + logger.info( + f"Cleaned deprecated config sections: {deprecated_sections}" + ) + merged = _deep_merge(self._defaults, config_data) # 自动回填缺失配置到存储 # 或迁移了配置后需要更新 - should_persist = (not from_remote) or (merged != config_data) or deprecated_sections + should_persist = ( + (not from_remote) or (merged != config_data) or deprecated_sections + ) if should_persist: async with storage.acquire_lock("config_save", timeout=10): await storage.save_config(merged) diff --git a/app/services/grok/defaults.py b/app/services/grok/defaults.py index e23f515..03f1c10 100644 --- a/app/services/grok/defaults.py +++ b/app/services/grok/defaults.py @@ -13,19 +13,16 @@ GROK_DEFAULTS = { "image_format": "url", "video_format": "html", }, - "network": { "timeout": 120, "base_proxy_url": "", "asset_proxy_url": "", }, - "security": { "cf_clearance": "", "browser": "chrome136", "user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36", }, - "chat": { "temporary": True, "disable_memory": True, @@ -34,7 +31,6 @@ GROK_DEFAULTS = { "dynamic_statsig": True, "filter_tags": ["grok:render", "xaiartifact", "xai:tool_usage_card"], }, - "retry": { "max_retry": 3, "retry_status_codes": [401, 429, 403], @@ -43,12 +39,10 @@ GROK_DEFAULTS = { "retry_backoff_max": 30.0, "retry_budget": 90.0, }, - "timeout": { "stream_idle_timeout": 45.0, "video_idle_timeout": 90.0, }, - "image": { "image_ws": True, "image_ws_nsfw": True, @@ -56,7 +50,6 @@ GROK_DEFAULTS = { "image_ws_final_min_bytes": 100000, "image_ws_medium_min_bytes": 30000, }, - "token": { "auto_refresh": True, "refresh_interval_hours": 8, @@ -65,12 +58,10 @@ GROK_DEFAULTS = { "save_delay_ms": 500, "reload_interval_sec": 30, }, - "cache": { "enable_auto_clean": True, "limit_mb": 1024, }, - "performance": { "assets_max_concurrent": 25, "assets_delete_batch_size": 10, diff --git a/app/services/grok/models/model.py b/app/services/grok/models/model.py index 1a8a018..e7cd8d6 100644 --- a/app/services/grok/models/model.py +++ b/app/services/grok/models/model.py @@ -112,7 +112,6 @@ class ModelService: cost=Cost.HIGH, display_name="GROK-4.1-EXPERT", ), - ModelInfo( model_id="grok-4.1-thinking", grok_model="grok-4-1-thinking-1129", diff --git a/app/services/grok/processors/base.py b/app/services/grok/processors/base.py index bfda0d2..76b4838 100644 --- a/app/services/grok/processors/base.py +++ b/app/services/grok/processors/base.py @@ -136,6 +136,7 @@ class BaseProcessor: """处理资产 URL""" if path.startswith("http"): from urllib.parse import urlparse + path = urlparse(path).path if not path.startswith("/"): diff --git a/app/services/grok/processors/chat_processors.py b/app/services/grok/processors/chat_processors.py index d7002ca..8f7c153 100644 --- a/app/services/grok/processors/chat_processors.py +++ b/app/services/grok/processors/chat_processors.py @@ -112,7 +112,9 @@ class StreamProcessor(BaseProcessor): } return f"data: {orjson.dumps(chunk).decode()}\n\n" - async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: """处理流式响应""" idle_timeout = get_config("timeout.stream_idle_timeout") @@ -145,7 +147,9 @@ class StreamProcessor(BaseProcessor): self.think_opened = True idx = img.get("imageIndex", 0) + 1 progress = img.get("progress", 0) - yield self._sse(f"正在生成第{idx}张图片中，当前进度{progress}%\n") + yield self._sse( + f"正在生成第{idx}张图片中，当前进度{progress}%\n" + ) continue # modelResponse @@ -163,7 +167,9 @@ class StreamProcessor(BaseProcessor): if self.image_format == "base64": dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: yield self._sse(f"![{img_id}]({base64_data})\n") else: @@ -173,7 +179,11 @@ class StreamProcessor(BaseProcessor): final_url = await self.process_url(url, "image") yield self._sse(f"![{img_id}]({final_url})\n") - if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + if ( + (meta := mr.get("metadata", {})) + .get("llm_info", {}) + .get("modelHash") + ): self.fingerprint = meta["llm_info"]["modelHash"] continue @@ -194,7 +204,11 @@ class StreamProcessor(BaseProcessor): raise UpstreamException( message=f"Stream idle timeout after {e.idle_seconds}s", status_code=504, - details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}, + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, ) except RequestsError as e: if _is_http2_stream_error(e): @@ -211,7 +225,10 @@ class StreamProcessor(BaseProcessor): details={"error": str(e)}, ) except Exception as e: - logger.error(f"Stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Stream processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) raise finally: await self.close() @@ -271,7 +288,9 @@ class CollectProcessor(BaseProcessor): if self.image_format == "base64": dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: content += f"![{img_id}]({base64_data})\n" else: @@ -281,7 +300,11 @@ class CollectProcessor(BaseProcessor): final_url = await self.process_url(url, "image") content += f"![{img_id}]({final_url})\n" - if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + if ( + (meta := mr.get("metadata", {})) + .get("llm_info", {}) + .get("modelHash") + ): fingerprint = meta["llm_info"]["modelHash"] except asyncio.CancelledError: @@ -290,11 +313,16 @@ class CollectProcessor(BaseProcessor): logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model}) except RequestsError as e: if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in collect: {e}", extra={"model": self.model}) + logger.warning( + f"HTTP/2 stream error in collect: {e}", extra={"model": self.model} + ) else: logger.error(f"Collect request error: {e}", extra={"model": self.model}) except Exception as e: - logger.error(f"Collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Collect processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) finally: await self.close() @@ -309,7 +337,12 @@ class CollectProcessor(BaseProcessor): "choices": [ { "index": 0, - "message": {"role": "assistant", "content": content, "refusal": None, "annotations": []}, + "message": { + "role": "assistant", + "content": content, + "refusal": None, + "annotations": [], + }, "finish_reason": "stop", } ], @@ -317,8 +350,17 @@ class CollectProcessor(BaseProcessor): "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, - "prompt_tokens_details": {"cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, - "completion_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 0, + "audio_tokens": 0, + "image_tokens": 0, + }, + "completion_tokens_details": { + "text_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + }, }, } diff --git a/app/services/grok/processors/image_processors.py b/app/services/grok/processors/image_processors.py index 42214c7..2dac0bd 100644 --- a/app/services/grok/processors/image_processors.py +++ b/app/services/grok/processors/image_processors.py @@ -4,7 +4,7 @@ import asyncio import random -from typing import Any, AsyncGenerator, AsyncIterable, List +from typing import AsyncGenerator, AsyncIterable, List import orjson from curl_cffi.requests.errors import RequestsError @@ -25,7 +25,9 @@ from .base import ( class ImageStreamProcessor(BaseProcessor): """图片生成流式响应处理器""" - def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json"): + def __init__( + self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" + ): super().__init__(model, token) self.partial_index = 0 self.n = n @@ -42,7 +44,9 @@ class ImageStreamProcessor(BaseProcessor): """构建 SSE 响应""" return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" - async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: """处理流式响应""" final_images = [] idle_timeout = get_config("timeout.stream_idle_timeout") @@ -90,7 +94,9 @@ class ImageStreamProcessor(BaseProcessor): final_images.append(processed) continue dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: if "," in base64_data: b64 = base64_data.split(",", 1)[1] @@ -117,7 +123,10 @@ class ImageStreamProcessor(BaseProcessor): "total_tokens": 0, "input_tokens": 0, "output_tokens": 0, - "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + "input_tokens_details": { + "text_tokens": 0, + "image_tokens": 0, + }, }, }, ) @@ -127,7 +136,11 @@ class ImageStreamProcessor(BaseProcessor): raise UpstreamException( message=f"Image stream idle timeout after {e.idle_seconds}s", status_code=504, - details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}, + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, ) except RequestsError as e: if _is_http2_stream_error(e): @@ -144,7 +157,10 @@ class ImageStreamProcessor(BaseProcessor): details={"error": str(e)}, ) except Exception as e: - logger.error(f"Image stream processing error: {e}", extra={"error_type": type(e).__name__}) + logger.error( + f"Image stream processing error: {e}", + extra={"error_type": type(e).__name__}, + ) raise finally: await self.close() @@ -183,7 +199,9 @@ class ImageCollectProcessor(BaseProcessor): images.append(processed) continue dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: if "," in base64_data: b64 = base64_data.split(",", 1)[1] @@ -201,7 +219,10 @@ class ImageCollectProcessor(BaseProcessor): else: logger.error(f"Image collect request error: {e}") except Exception as e: - logger.error(f"Image collect processing error: {e}", extra={"error_type": type(e).__name__}) + logger.error( + f"Image collect processing error: {e}", + extra={"error_type": type(e).__name__}, + ) finally: await self.close() diff --git a/app/services/grok/processors/image_ws_processors.py b/app/services/grok/processors/image_ws_processors.py index 86a6aa9..2074ec2 100644 --- a/app/services/grok/processors/image_ws_processors.py +++ b/app/services/grok/processors/image_ws_processors.py @@ -31,7 +31,12 @@ class ImageWSBaseProcessor(BaseProcessor): def _ensure_image_dir(self) -> Path: if self._image_dir is None: - base_dir = Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" / "image" + base_dir = ( + Path(__file__).parent.parent.parent.parent.parent + / "data" + / "tmp" + / "image" + ) base_dir.mkdir(parents=True, exist_ok=True) self._image_dir = base_dir return self._image_dir @@ -78,7 +83,9 @@ class ImageWSBaseProcessor(BaseProcessor): def _to_output(self, image_id: str, item: Dict) -> str: try: if self.response_format == "url": - return self._save_blob(image_id, item.get("blob", ""), item.get("is_final", False)) + return self._save_blob( + image_id, item.get("blob", ""), item.get("is_final", False) + ) return self._strip_base64(item.get("blob", "")) except Exception as e: logger.warning(f"Image output failed: {e}") @@ -88,7 +95,14 @@ class ImageWSBaseProcessor(BaseProcessor): class ImageWSStreamProcessor(ImageWSBaseProcessor): """WebSocket 图片流式响应处理器""" - def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json", size: str = "1024x1024"): + def __init__( + self, + model: str, + token: str = "", + n: int = 1, + response_format: str = "b64_json", + size: str = "1024x1024", + ): super().__init__(model, token, "b64_json") self.n = n self.size = size @@ -114,7 +128,16 @@ class ImageWSStreamProcessor(ImageWSBaseProcessor): if item.get("type") == "error": message = item.get("error") or "Upstream error" code = item.get("error_code") or "upstream_error" - yield self._sse("error", {"error": {"message": message, "type": "server_error", "code": code}}) + yield self._sse( + "error", + { + "error": { + "message": message, + "type": "server_error", + "code": code, + } + }, + ) return if item.get("type") != "image": continue @@ -159,9 +182,25 @@ class ImageWSStreamProcessor(ImageWSBaseProcessor): if self._target_id and self._target_id in images: selected = [(self._target_id, images[self._target_id])] else: - selected = [max(images.items(), key=lambda x: (x[1].get("is_final", False), x[1].get("blob_size", 0)))] if images else [] + selected = ( + [ + max( + images.items(), + key=lambda x: ( + x[1].get("is_final", False), + x[1].get("blob_size", 0), + ), + ) + ] + if images + else [] + ) else: - selected = [(image_id, images[image_id]) for image_id in self._index_map if image_id in images] + selected = [ + (image_id, images[image_id]) + for image_id in self._index_map + if image_id in images + ] for image_id, item in selected: output = self._strip_base64(item.get("blob", "")) @@ -193,7 +232,9 @@ class ImageWSStreamProcessor(ImageWSBaseProcessor): class ImageWSCollectProcessor(ImageWSBaseProcessor): """WebSocket 图片非流式响应处理器""" - def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json"): + def __init__( + self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" + ): super().__init__(model, token, response_format) self.n = n diff --git a/app/services/grok/processors/video_processors.py b/app/services/grok/processors/video_processors.py index ec39aae..a0ead8c 100644 --- a/app/services/grok/processors/video_processors.py +++ b/app/services/grok/processors/video_processors.py @@ -59,6 +59,7 @@ class VideoStreamProcessor(BaseProcessor): def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: """构建视频 HTML 标签""" import html + safe_video_url = html.escape(video_url) safe_thumbnail_url = html.escape(thumbnail_url) poster_attr = f' poster="{safe_thumbnail_url}"' if safe_thumbnail_url else "" @@ -66,7 +67,9 @@ class VideoStreamProcessor(BaseProcessor): ''' - async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: """处理视频流式响应""" idle_timeout = get_config("timeout.video_idle_timeout") @@ -111,12 +114,16 @@ class VideoStreamProcessor(BaseProcessor): final_video_url = await self.process_url(video_url, "video") final_thumbnail_url = "" if thumbnail_url: - final_thumbnail_url = await self.process_url(thumbnail_url, "image") + final_thumbnail_url = await self.process_url( + thumbnail_url, "image" + ) if self.video_format == "url": yield self._sse(final_video_url) else: - video_html = self._build_video_html(final_video_url, final_thumbnail_url) + video_html = self._build_video_html( + final_video_url, final_thumbnail_url + ) yield self._sse(video_html) logger.info(f"Video generated: {video_url}") @@ -127,29 +134,42 @@ class VideoStreamProcessor(BaseProcessor): yield self._sse(finish="stop") yield "data: [DONE]\n\n" except asyncio.CancelledError: - logger.debug("Video stream cancelled by client", extra={"model": self.model}) + logger.debug( + "Video stream cancelled by client", extra={"model": self.model} + ) except StreamIdleTimeoutError as e: raise UpstreamException( message=f"Video stream idle timeout after {e.idle_seconds}s", status_code=504, - details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds}, + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, ) except RequestsError as e: if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in video: {e}", extra={"model": self.model}) + logger.warning( + f"HTTP/2 stream error in video: {e}", extra={"model": self.model} + ) raise UpstreamException( message="Upstream connection closed unexpectedly", status_code=502, details={"error": str(e), "type": "http2_stream_error"}, ) - logger.error(f"Video stream request error: {e}", extra={"model": self.model}) + logger.error( + f"Video stream request error: {e}", extra={"model": self.model} + ) raise UpstreamException( message=f"Upstream request failed: {e}", status_code=502, details={"error": str(e)}, ) except Exception as e: - logger.error(f"Video stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Video stream processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) finally: await self.close() @@ -195,25 +215,41 @@ class VideoCollectProcessor(BaseProcessor): final_video_url = await self.process_url(video_url, "video") final_thumbnail_url = "" if thumbnail_url: - final_thumbnail_url = await self.process_url(thumbnail_url, "image") + final_thumbnail_url = await self.process_url( + thumbnail_url, "image" + ) if self.video_format == "url": content = final_video_url else: - content = self._build_video_html(final_video_url, final_thumbnail_url) + content = self._build_video_html( + final_video_url, final_thumbnail_url + ) logger.info(f"Video generated: {video_url}") except asyncio.CancelledError: - logger.debug("Video collect cancelled by client", extra={"model": self.model}) + logger.debug( + "Video collect cancelled by client", extra={"model": self.model} + ) except StreamIdleTimeoutError as e: - logger.warning(f"Video collect idle timeout: {e}", extra={"model": self.model}) + logger.warning( + f"Video collect idle timeout: {e}", extra={"model": self.model} + ) except RequestsError as e: if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in video collect: {e}", extra={"model": self.model}) + logger.warning( + f"HTTP/2 stream error in video collect: {e}", + extra={"model": self.model}, + ) else: - logger.error(f"Video collect request error: {e}", extra={"model": self.model}) + logger.error( + f"Video collect request error: {e}", extra={"model": self.model} + ) except Exception as e: - logger.error(f"Video collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Video collect processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) finally: await self.close() @@ -225,7 +261,11 @@ class VideoCollectProcessor(BaseProcessor): "choices": [ { "index": 0, - "message": {"role": "assistant", "content": content, "refusal": None}, + "message": { + "role": "assistant", + "content": content, + "refusal": None, + }, "finish_reason": "stop", } ], diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index d576b81..d19ab4a 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -42,11 +42,24 @@ _ASSETS_SEM_VALUE = None # 常用 MIME 类型（业务数据，非配置） MIME_TYPES = { - ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".gif": "image/gif", - ".webp": "image/webp", ".bmp": "image/bmp", ".pdf": "application/pdf", ".txt": "text/plain", - ".md": "text/markdown", ".csv": "text/csv", ".json": "application/json", ".xml": "application/xml", - ".py": "text/x-python-script", ".js": "application/javascript", ".html": "text/html", - ".css": "text/css", ".mp4": "video/mp4", ".webm": "video/webm", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", + ".pdf": "application/pdf", + ".txt": "text/plain", + ".md": "text/markdown", + ".csv": "text/csv", + ".json": "application/json", + ".xml": "application/xml", + ".py": "text/x-python-script", + ".js": "application/javascript", + ".html": "text/html", + ".css": "text/css", + ".mp4": "video/mp4", + ".webm": "video/webm", } IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} @@ -54,10 +67,11 @@ VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} # ==================== 工具函数 ==================== + def _get_assets_semaphore() -> asyncio.Semaphore: """获取全局并发控制信号量""" value = max(1, int(get_config("performance.assets_max_concurrent"))) - + global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE if _ASSETS_SEMAPHORE is None or value != _ASSETS_SEM_VALUE: _ASSETS_SEM_VALUE = value @@ -71,13 +85,13 @@ async def _file_lock(name: str, timeout: int = 10): if fcntl is None: yield return - + LOCK_DIR.mkdir(parents=True, exist_ok=True) lock_path = LOCK_DIR / f"{name}.lock" fd = None locked = False start = time.monotonic() - + try: fd = open(lock_path, "a+") while True: @@ -103,26 +117,30 @@ async def _file_lock(name: str, timeout: int = 10): @dataclass class ServiceConfig: """服务配置""" + proxy: str timeout: int browser: str user_agent: str - + @classmethod def from_settings(cls, proxy: Optional[str] = None): return cls( - proxy=proxy or get_config("network.asset_proxy_url") or get_config("network.base_proxy_url"), + proxy=proxy + or get_config("network.asset_proxy_url") + or get_config("network.base_proxy_url"), timeout=get_config("network.timeout"), browser=get_config("security.browser"), user_agent=get_config("security.user_agent"), ) - + def get_proxies(self) -> Optional[dict]: return {"http": self.proxy, "https": self.proxy} if self.proxy else None # ==================== 基础服务 ==================== + class BaseService: """基础服务类""" @@ -130,7 +148,9 @@ class BaseService: self.config = ServiceConfig.from_settings(proxy) self._session: Optional[AsyncSession] = None - def _build_headers(self, token: str, referer: str = "https://grok.com/", download: bool = False) -> dict: + def _build_headers( + self, token: str, referer: str = "https://grok.com/", download: bool = False + ) -> dict: """构建请求头""" if download: headers = { @@ -148,7 +168,7 @@ class BaseService: "User-Agent": self.config.user_agent, } apply_statsig(headers) - + headers["Cookie"] = build_sso_cookie(token) return headers @@ -184,11 +204,13 @@ class BaseService: message=f"Failed to fetch: {response.status_code}", details={"url": url, "status": response.status_code}, ) - + filename = url.split("/")[-1].split("?")[0] or "download" - content_type = response.headers.get("content-type", "application/octet-stream").split(";")[0] + content_type = response.headers.get( + "content-type", "application/octet-stream" + ).split(";")[0] b64 = base64.b64encode(response.content).decode() - + logger.debug(f"Fetched: {url}") return filename, b64, content_type except Exception as e: @@ -202,15 +224,15 @@ class BaseService: """解析 Base64 数据""" if not data_uri.startswith("data:"): return "file.bin", data_uri, "application/octet-stream" - + try: header, b64 = data_uri.split(",", 1) except ValueError: return "file.bin", data_uri, "application/octet-stream" - + if ";base64" not in header: return "file.bin", data_uri, "application/octet-stream" - + mime = header[5:].split(";", 1)[0] or "application/octet-stream" b64 = re.sub(r"\s+", "", b64) ext = mime.split("/")[-1] if "/" in mime else "bin" @@ -224,18 +246,21 @@ class BaseService: return f"data:{mime_type};base64,{b64_data}" except Exception as e: logger.error(f"File to base64 failed: {file_path} - {e}") - raise AppException(f"Failed to read file: {file_path}", code="file_read_error") + raise AppException( + f"Failed to read file: {file_path}", code="file_read_error" + ) # ==================== 上传服务 ==================== + class UploadService(BaseService): """文件上传服务""" async def upload(self, file_input: str, token: str) -> Tuple[str, str]: """ 上传文件到 Grok - + Returns: (file_id, file_uri) """ @@ -245,9 +270,11 @@ class UploadService(BaseService): filename, b64, mime = await self.fetch(file_input) else: filename, b64, mime = self.parse_b64(file_input) - - logger.debug(f"Upload prepare: filename={filename}, type={mime}, size={len(b64)}") - + + logger.debug( + f"Upload prepare: filename={filename}, type={mime}, size={len(b64)}" + ) + if not b64: raise ValidationException("Invalid file input: empty content") @@ -274,10 +301,12 @@ class UploadService(BaseService): if response.status_code in (401, 403): logger.warning(f"Upload auth failed: {response.status_code}") try: - await TokenService.record_fail(token, response.status_code, "upload_auth_failed") + await TokenService.record_fail( + token, response.status_code, "upload_auth_failed" + ) except Exception as e: logger.error(f"Failed to record token failure: {e}") - + raise UpstreamException( message=f"Upload authentication failed: {response.status_code}", details={"status": response.status_code, "token_invalidated": True}, @@ -293,6 +322,7 @@ class UploadService(BaseService): # ==================== 列表服务 ==================== + class ListService(BaseService): """文件列表查询服务""" @@ -361,6 +391,7 @@ class ListService(BaseService): # ==================== 删除服务 ==================== + class DeleteService(BaseService): """文件删除服务""" @@ -390,12 +421,12 @@ class DeleteService(BaseService): """删除所有文件""" total = success = failed = 0 list_service = ListService(self.config.proxy) - + try: async for assets in list_service.iter_assets(token): if not assets: continue - + total += len(assets) batch_result = await self._delete_batch(token, assets) success += batch_result["success"] @@ -416,10 +447,13 @@ class DeleteService(BaseService): success = failed = 0 for i in range(0, len(assets), batch_size): - batch = assets[i:i + batch_size] + batch = assets[i : i + batch_size] results = await asyncio.gather( - *[self._delete_one(token, asset, idx) for idx, asset in enumerate(batch)], - return_exceptions=True + *[ + self._delete_one(token, asset, idx) + for idx, asset in enumerate(batch) + ], + return_exceptions=True, ) success += sum(1 for r in results if r is True) failed += sum(1 for r in results if r is not True) @@ -440,12 +474,15 @@ class DeleteService(BaseService): # ==================== 下载服务 ==================== + class DownloadService(BaseService): """文件下载服务""" def __init__(self, proxy: Optional[str] = None): super().__init__(proxy) - self.base_dir = Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" + self.base_dir = ( + Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" + ) self.image_dir = self.base_dir / "image" self.video_dir = self.base_dir / "video" self.image_dir.mkdir(parents=True, exist_ok=True) @@ -461,10 +498,14 @@ class DownloadService(BaseService): def _get_mime(self, cache_path: Path, response=None) -> str: """获取 MIME 类型""" if response: - return response.headers.get("content-type", "application/octet-stream").split(";")[0] + return response.headers.get( + "content-type", "application/octet-stream" + ).split(";")[0] return MIME_TYPES.get(cache_path.suffix.lower(), "application/octet-stream") - async def download(self, file_path: str, token: str, media_type: str = "image") -> Tuple[Optional[Path], str]: + async def download( + self, file_path: str, token: str, media_type: str = "image" + ) -> Tuple[Optional[Path], str]: """下载文件到本地""" async with _get_assets_semaphore(): cache_path = self._cache_path(file_path, media_type) @@ -536,7 +577,9 @@ class DownloadService(BaseService): return self._get_mime(cache_path, response) - async def to_base64(self, file_path: str, token: str, media_type: str = "image") -> str: + async def to_base64( + self, file_path: str, token: str, media_type: str = "image" + ) -> str: """下载并转 base64""" cache_path, mime = await self.download(file_path, token, media_type) if not cache_path or not cache_path.exists(): @@ -558,31 +601,43 @@ class DownloadService(BaseService): cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"count": 0, "size_mb": 0.0} - + allowed = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS - files = [f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed] + files = [ + f + for f in cache_dir.glob("*") + if f.is_file() and f.suffix.lower() in allowed + ] total_size = sum(f.stat().st_size for f in files) return {"count": len(files), "size_mb": round(total_size / 1024 / 1024, 2)} - def list_files(self, media_type: str = "image", page: int = 1, page_size: int = 1000) -> Dict[str, Any]: + def list_files( + self, media_type: str = "image", page: int = 1, page_size: int = 1000 + ) -> Dict[str, Any]: """列出缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"total": 0, "page": page, "page_size": page_size, "items": []} - + allowed = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS - files = [f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed] + files = [ + f + for f in cache_dir.glob("*") + if f.is_file() and f.suffix.lower() in allowed + ] # 构建文件列表 items = [] for f in files: try: stat = f.stat() - items.append({ - "name": f.name, - "size_bytes": stat.st_size, - "mtime_ms": int(stat.st_mtime * 1000), - }) + items.append( + { + "name": f.name, + "size_bytes": stat.st_size, + "mtime_ms": int(stat.st_mtime * 1000), + } + ) except Exception: continue @@ -591,7 +646,7 @@ class DownloadService(BaseService): # 分页 total = len(items) start = max(0, (page - 1) * page_size) - paged = items[start:start + page_size] + paged = items[start : start + page_size] # 添加 URL for item in paged: @@ -603,7 +658,7 @@ class DownloadService(BaseService): """删除缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir file_path = cache_dir / name.replace("/", "-") - + if file_path.exists(): try: file_path.unlink() @@ -636,7 +691,7 @@ class DownloadService(BaseService): """检查并清理缓存""" if self._cleanup_running or not get_config("cache.enable_auto_clean"): return - + self._cleanup_running = True try: async with _file_lock("cache_cleanup", timeout=5): @@ -648,9 +703,11 @@ class DownloadService(BaseService): return # 清理到 80% - logger.info(f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning...") + logger.info( + f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning..." + ) all_files.sort(key=lambda x: x[1]) # 按时间排序 - + deleted_count = 0 deleted_size = 0 target_mb = limit_mb * 0.8 @@ -666,7 +723,9 @@ class DownloadService(BaseService): except Exception: pass - logger.info(f"Cache cleanup: {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)") + logger.info( + f"Cache cleanup: {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)" + ) finally: self._cleanup_running = False @@ -689,4 +748,10 @@ class DownloadService(BaseService): return all_files, total_size -__all__ = ["BaseService", "UploadService", "ListService", "DeleteService", "DownloadService"] +__all__ = [ + "BaseService", + "UploadService", + "ListService", + "DeleteService", + "DownloadService", +] diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index 99ea4ce..dcd36cc 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -3,14 +3,19 @@ Grok Chat 服务 """ import orjson -from typing import Dict, List, Any, AsyncGenerator +from typing import Dict, List, Any from dataclasses import dataclass from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config -from app.core.exceptions import AppException, UpstreamException, ValidationException, ErrorType +from app.core.exceptions import ( + AppException, + UpstreamException, + ValidationException, + ErrorType, +) from app.services.grok.models.model import ModelService from app.services.grok.services.assets import UploadService from app.services.grok.processors import StreamProcessor, CollectProcessor @@ -26,6 +31,7 @@ CHAT_API = "https://grok.com/rest/app-chat/conversations/new" @dataclass class ChatRequest: """聊天请求数据""" + model: str messages: List[Dict[str, Any]] stream: bool = None @@ -36,7 +42,9 @@ class MessageExtractor: """消息内容提取器""" @staticmethod - def extract(messages: List[Dict[str, Any]], is_video: bool = False) -> tuple[str, List[tuple[str, str]]]: + def extract( + messages: List[Dict[str, Any]], is_video: bool = False + ) -> tuple[str, List[tuple[str, str]]]: """从 OpenAI 消息格式提取内容，返回 (text, attachments)""" texts = [] attachments = [] @@ -60,7 +68,11 @@ class MessageExtractor: elif item_type == "image_url": image_data = item.get("image_url", {}) - url = image_data.get("url", "") if isinstance(image_data, dict) else str(image_data) + url = ( + image_data.get("url", "") + if isinstance(image_data, dict) + else str(image_data) + ) if url: attachments.append(("image", url)) @@ -68,7 +80,11 @@ class MessageExtractor: if is_video: raise ValueError("视频模型不支持 input_audio 类型") audio_data = item.get("input_audio", {}) - data = audio_data.get("data", "") if isinstance(audio_data, dict) else str(audio_data) + data = ( + audio_data.get("data", "") + if isinstance(audio_data, dict) + else str(audio_data) + ) if data: attachments.append(("audio", data)) @@ -86,7 +102,14 @@ class MessageExtractor: extracted.append({"role": role, "text": "\n".join(parts)}) # 找到最后一条 user 消息 - last_user_index = next((i for i in range(len(extracted) - 1, -1, -1) if extracted[i]["role"] == "user"), None) + last_user_index = next( + ( + i + for i in range(len(extracted) - 1, -1, -1) + if extracted[i]["role"] == "user" + ), + None, + ) for i, item in enumerate(extracted): role = item["role"] or "user" @@ -132,8 +155,13 @@ class ChatRequestBuilder: return headers @staticmethod - def build_payload(message: str, model: str, mode: str = None, - file_attachments: List[str] = None, image_attachments: List[str] = None) -> Dict[str, Any]: + def build_payload( + message: str, + model: str, + mode: str = None, + file_attachments: List[str] = None, + image_attachments: List[str] = None, + ) -> Dict[str, Any]: """构造请求体""" merged_attachments = [] if file_attachments: @@ -183,22 +211,36 @@ class GrokChatService: def __init__(self, proxy: str = None): self.proxy = proxy or get_config("network.base_proxy_url") - async def chat(self, token: str, message: str, model: str = "grok-3", mode: str = None, - stream: bool = None, file_attachments: List[str] = None, - image_attachments: List[str] = None, raw_payload: Dict[str, Any] = None): + async def chat( + self, + token: str, + message: str, + model: str = "grok-3", + mode: str = None, + stream: bool = None, + file_attachments: List[str] = None, + image_attachments: List[str] = None, + raw_payload: Dict[str, Any] = None, + ): """发送聊天请求""" if stream is None: stream = get_config("chat.stream") headers = ChatRequestBuilder.build_headers(token) - payload = raw_payload if raw_payload is not None else ChatRequestBuilder.build_payload( - message, model, mode, file_attachments, image_attachments + payload = ( + raw_payload + if raw_payload is not None + else ChatRequestBuilder.build_payload( + message, model, mode, file_attachments, image_attachments + ) ) proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None timeout = get_config("network.timeout") - logger.debug(f"Chat request: model={model}, mode={mode}, stream={stream}, attachments={len(file_attachments or [])}") - + logger.debug( + f"Chat request: model={model}, mode={mode}, stream={stream}, attachments={len(file_attachments or [])}" + ) + # 建立连接 async def establish_connection(): browser = get_config("security.browser") @@ -220,8 +262,10 @@ class GrokChatService: except Exception: pass - logger.error(f"Chat failed: status={response.status_code}, token={token[:10]}...") - + logger.error( + f"Chat failed: status={response.status_code}, token={token[:10]}..." + ) + await session.close() raise UpstreamException( message=f"Grok API request failed: {response.status_code}", @@ -250,7 +294,9 @@ class GrokChatService: session = None response = None try: - session, response = await retry_on_status(establish_connection, extract_status=extract_status) + session, response = await retry_on_status( + establish_connection, extract_status=extract_status + ) except Exception as e: status_code = extract_status(e) if status_code: @@ -286,8 +332,12 @@ class GrokChatService: # 提取消息和附件 try: - message, attachments = MessageExtractor.extract(request.messages, is_video=is_video) - logger.debug(f"Extracted message length={len(message)}, attachments={len(attachments)}") + message, attachments = MessageExtractor.extract( + request.messages, is_video=is_video + ) + logger.debug( + f"Extracted message length={len(message)}, attachments={len(attachments)}" + ) except ValueError as e: raise ValidationException(str(e)) @@ -299,15 +349,24 @@ class GrokChatService: for attach_type, attach_data in attachments: file_id, _ = await upload_service.upload(attach_data, token) file_ids.append(file_id) - logger.debug(f"Attachment uploaded: type={attach_type}, file_id={file_id}") + logger.debug( + f"Attachment uploaded: type={attach_type}, file_id={file_id}" + ) finally: await upload_service.close() - stream = request.stream if request.stream is not None else get_config("chat.stream") + stream = ( + request.stream if request.stream is not None else get_config("chat.stream") + ) response = await self.chat( - token, message, grok_model, mode, stream, - file_attachments=file_ids, image_attachments=[] + token, + message, + grok_model, + mode, + stream, + file_attachments=file_ids, + image_attachments=[], ) return response, stream, request.model @@ -317,13 +376,17 @@ class ChatService: """Chat 业务服务""" @staticmethod - async def completions(model: str, messages: List[Dict[str, Any]], - stream: bool = None, thinking: str = None): + async def completions( + model: str, + messages: List[Dict[str, Any]], + stream: bool = None, + thinking: str = None, + ): """Chat Completions 入口""" # 获取 token token_mgr = await get_token_manager() await token_mgr.reload_if_stale() - + token = None for pool_name in ModelService.pool_candidates_for_model(model): token = token_mgr.get_token(pool_name) @@ -343,7 +406,9 @@ class ChatService: is_stream = stream if stream is not None else get_config("chat.stream") # 构造请求 - chat_request = ChatRequest(model=model, messages=messages, stream=is_stream, think=think) + chat_request = ChatRequest( + model=model, messages=messages, stream=is_stream, think=think + ) # 请求 Grok service = GrokChatService() @@ -353,14 +418,20 @@ class ChatService: if is_stream: logger.debug(f"Processing stream response: model={model}") processor = StreamProcessor(model_name, token, think) - return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) - + return wrap_stream_with_usage( + processor.process(response), token_mgr, token, model + ) + # 非流式 logger.debug(f"Processing non-stream response: model={model}") result = await CollectProcessor(model_name, token).process(response) try: model_info = ModelService.get(model) - effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + effort = ( + EffortType.HIGH + if (model_info and model_info.cost.value == "high") + else EffortType.LOW + ) await token_mgr.consume(token, effort) logger.info(f"Chat completed: model={model}, effort={effort.value}") except Exception as e: @@ -368,4 +439,10 @@ class ChatService: return result -__all__ = ["GrokChatService", "ChatRequest", "ChatRequestBuilder", "MessageExtractor", "ChatService"] +__all__ = [ + "GrokChatService", + "ChatRequest", + "ChatRequestBuilder", + "MessageExtractor", + "ChatService", +] diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index b94feb6..089a5f0 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -24,6 +24,7 @@ WS_URL = "wss://grok.com/ws/imagine/listen" class _BlockedError(Exception): pass + class ImageService: """Grok Imagine WebSocket image service.""" @@ -61,7 +62,9 @@ class ImageService: return match.group(1) if match else None def _is_final_image(self, url: str, blob_size: int) -> bool: - return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > get_config("image.image_ws_final_min_bytes") + return (url or "").lower().endswith( + (".jpg", ".jpeg") + ) and blob_size > get_config("image.image_ws_final_min_bytes") def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]: if not url or not blob: @@ -70,8 +73,16 @@ class ImageService: image_id = self._extract_image_id(url) or uuid.uuid4().hex blob_size = len(blob) is_final = self._is_final_image(url, blob_size) - - stage = "final" if is_final else ("medium" if blob_size > get_config("image.image_ws_medium_min_bytes") else "preview") + + stage = ( + "final" + if is_final + else ( + "medium" + if blob_size > get_config("image.image_ws_medium_min_bytes") + else "preview" + ) + ) return { "type": "image", @@ -93,19 +104,27 @@ class ImageService: max_retries: int = None, ) -> AsyncGenerator[Dict[str, object], None]: retries = max(1, max_retries if max_retries is not None else 1) - logger.info(f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}") + logger.info( + f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}" + ) for attempt in range(retries): try: yielded_any = False - async for item in self._stream_once(token, prompt, aspect_ratio, n, enable_nsfw): + async for item in self._stream_once( + token, prompt, aspect_ratio, n, enable_nsfw + ): yielded_any = True yield item return except _BlockedError: if yielded_any or attempt + 1 >= retries: if not yielded_any: - yield {"type": "error", "error_code": "blocked", "error": "blocked_no_final_image"} + yield { + "type": "error", + "error_code": "blocked", + "error": "blocked_no_final_image", + } return logger.warning(f"WebSocket blocked, retry {attempt + 1}/{retries}") except Exception as e: @@ -175,10 +194,17 @@ class ImageService: try: ws_msg = await asyncio.wait_for(ws.receive(), timeout=5.0) except asyncio.TimeoutError: - if medium_received_time and completed == 0 and time.time() - medium_received_time > min(10, blocked_seconds): + if ( + medium_received_time + and completed == 0 + and time.time() - medium_received_time + > min(10, blocked_seconds) + ): raise _BlockedError() if completed > 0 and time.time() - last_activity > 10: - logger.info(f"WebSocket idle timeout, collected {completed} images") + logger.info( + f"WebSocket idle timeout, collected {completed} images" + ) break continue @@ -188,44 +214,74 @@ class ImageService: msg_type = msg.get("type") if msg_type == "image": - info = self._classify_image(msg.get("url", ""), msg.get("blob", "")) + info = self._classify_image( + msg.get("url", ""), msg.get("blob", "") + ) if not info: continue image_id = info["image_id"] existing = images.get(image_id, {}) - - if info["stage"] == "medium" and medium_received_time is None: + + if ( + info["stage"] == "medium" + and medium_received_time is None + ): medium_received_time = time.time() if info["is_final"] and not existing.get("is_final"): completed += 1 - logger.debug(f"Final image received: id={image_id}, size={info['blob_size']}") + logger.debug( + f"Final image received: id={image_id}, size={info['blob_size']}" + ) - images[image_id] = {"is_final": info["is_final"] or existing.get("is_final")} + images[image_id] = { + "is_final": info["is_final"] + or existing.get("is_final") + } yield info elif msg_type == "error": - logger.warning(f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}") - yield {"type": "error", "error_code": msg.get("err_code", ""), "error": msg.get("err_msg", "")} + logger.warning( + f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}" + ) + yield { + "type": "error", + "error_code": msg.get("err_code", ""), + "error": msg.get("err_msg", ""), + } return if completed >= n: - logger.info(f"WebSocket collected {completed} final images") + logger.info( + f"WebSocket collected {completed} final images" + ) break - if medium_received_time and completed == 0 and time.time() - medium_received_time > blocked_seconds: + if ( + medium_received_time + and completed == 0 + and time.time() - medium_received_time > blocked_seconds + ): raise _BlockedError() - elif ws_msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR): + elif ws_msg.type in ( + aiohttp.WSMsgType.CLOSED, + aiohttp.WSMsgType.ERROR, + ): logger.warning(f"WebSocket closed/error: {ws_msg.type}") - yield {"type": "error", "error_code": "ws_closed", "error": f"websocket closed: {ws_msg.type}"} + yield { + "type": "error", + "error_code": "ws_closed", + "error": f"websocket closed: {ws_msg.type}", + } break except aiohttp.ClientError as e: logger.error(f"WebSocket connection error: {e}") yield {"type": "error", "error_code": "connection_failed", "error": str(e)} + image_service = ImageService() __all__ = ["image_service", "ImageService"] diff --git a/app/services/grok/services/media.py b/app/services/grok/services/media.py index 10c4ba4..e3f4917 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/media.py @@ -84,14 +84,16 @@ class VideoService: return {"http": self.proxy, "https": self.proxy} if self.proxy else None async def create_post( - self, token: str, prompt: str, + self, + token: str, + prompt: str, media_type: str = "MEDIA_POST_TYPE_VIDEO", - media_url: str = None + media_url: str = None, ) -> str: """创建媒体帖子，返回 post ID""" try: headers = self._build_headers(token) - + # 根据类型构建不同的载荷 if media_type == "MEDIA_POST_TYPE_IMAGE" and media_url: payload = {"mediaType": media_type, "mediaUrl": media_url} @@ -110,7 +112,9 @@ class VideoService: if response.status_code != 200: logger.error(f"Create post failed: {response.status_code}") - raise UpstreamException(f"Failed to create post: {response.status_code}") + raise UpstreamException( + f"Failed to create post: {response.status_code}" + ) post_id = response.json().get("post", {}).get("id", "") if not post_id: @@ -127,19 +131,27 @@ class VideoService: async def create_image_post(self, token: str, image_url: str) -> str: """创建图片帖子，返回 post ID""" - return await self.create_post(token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url) + return await self.create_post( + token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url + ) def _build_payload( - self, prompt: str, post_id: str, + self, + prompt: str, + post_id: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution_name: str = "480p", - preset: str = "normal" + preset: str = "normal", ) -> dict: """构建视频生成载荷""" - mode_map = {"fun": "--mode=extremely-crazy", "normal": "--mode=normal", "spicy": "--mode=extremely-spicy-or-crazy"} + mode_map = { + "fun": "--mode=extremely-crazy", + "normal": "--mode=normal", + "spicy": "--mode=extremely-spicy-or-crazy", + } mode_flag = mode_map.get(preset, "--mode=custom") - + return { "temporary": True, "modelName": "grok-3", @@ -162,15 +174,22 @@ class VideoService: } async def _generate_internal( - self, token: str, post_id: str, prompt: str, - aspect_ratio: str, video_length: int, - resolution_name: str, preset: str + self, + token: str, + post_id: str, + prompt: str, + aspect_ratio: str, + video_length: int, + resolution_name: str, + preset: str, ) -> AsyncGenerator[bytes, None]: """内部生成逻辑""" session = None try: headers = self._build_headers(token) - payload = self._build_payload(prompt, post_id, aspect_ratio, video_length, resolution_name, preset) + payload = self._build_payload( + prompt, post_id, aspect_ratio, video_length, resolution_name, preset + ) session = AsyncSession(impersonate=get_config("security.browser")) response = await session.post( @@ -183,7 +202,9 @@ class VideoService: ) if response.status_code != 200: - logger.error(f"Video generation failed: status={response.status_code}, post_id={post_id}") + logger.error( + f"Video generation failed: status={response.status_code}, post_id={post_id}" + ) raise UpstreamException( message=f"Video generation failed: {response.status_code}", details={"status": response.status_code}, @@ -212,46 +233,66 @@ class VideoService: raise UpstreamException(f"Video generation error: {str(e)}") async def generate( - self, token: str, prompt: str, + self, + token: str, + prompt: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution_name: str = "480p", - preset: str = "normal" + preset: str = "normal", ) -> AsyncGenerator[bytes, None]: """生成视频""" - logger.info(f"Video generation: prompt='{prompt[:50]}...', ratio={aspect_ratio}, length={video_length}s, preset={preset}") + logger.info( + f"Video generation: prompt='{prompt[:50]}...', ratio={aspect_ratio}, length={video_length}s, preset={preset}" + ) async with _get_semaphore(): post_id = await self.create_post(token, prompt) - return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution_name, preset) + return await self._generate_internal( + token, + post_id, + prompt, + aspect_ratio, + video_length, + resolution_name, + preset, + ) async def generate_from_image( - self, token: str, prompt: str, image_url: str, + self, + token: str, + prompt: str, + image_url: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "480p", - preset: str = "normal" + preset: str = "normal", ) -> AsyncGenerator[bytes, None]: """从图片生成视频""" - logger.info(f"Image to video: prompt='{prompt[:50]}...', image={image_url[:80]}") + logger.info( + f"Image to video: prompt='{prompt[:50]}...', image={image_url[:80]}" + ) async with _get_semaphore(): post_id = await self.create_image_post(token, image_url) - return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution, preset) + return await self._generate_internal( + token, post_id, prompt, aspect_ratio, video_length, resolution, preset + ) @staticmethod async def completions( - model: str, messages: list, + model: str, + messages: list, stream: bool = None, thinking: str = None, aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "480p", - preset: str = "normal" + preset: str = "normal", ): """视频生成入口""" # 获取 token token_mgr = await get_token_manager() await token_mgr.reload_if_stale() - + token = None for pool_name in ModelService.pool_candidates_for_model(model): token = token_mgr.get_token(pool_name) @@ -306,16 +347,23 @@ class VideoService: # 处理响应 if is_stream: processor = VideoStreamProcessor(model, token, think) - return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) - + return wrap_stream_with_usage( + processor.process(response), token_mgr, token, model + ) + result = await VideoCollectProcessor(model, token).process(response) try: model_info = ModelService.get(model) - effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + effort = ( + EffortType.HIGH + if (model_info and model_info.cost.value == "high") + else EffortType.LOW + ) await token_mgr.consume(token, effort) logger.debug(f"Video completed, recorded usage (effort={effort.value})") except Exception as e: logger.warning(f"Failed to record video usage: {e}") return result + __all__ = ["VideoService"] diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py index fa02b8f..26a4f26 100644 --- a/app/services/grok/services/nsfw.py +++ b/app/services/grok/services/nsfw.py @@ -23,6 +23,7 @@ from app.services.grok.utils.headers import build_sso_cookie NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" BIRTH_DATE_API = "https://grok.com/rest/auth/set-birth-date" + @dataclass class NSFWResult: """NSFW 操作结果""" @@ -33,6 +34,7 @@ class NSFWResult: grpc_message: Optional[str] = None error: Optional[str] = None + class NSFWService: """NSFW 模式服务""" @@ -97,11 +99,13 @@ class NSFWService: protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner return encode_grpc_web_payload(protobuf) - async def _set_birth_date(self, session: AsyncSession, token: str) -> tuple[bool, int, Optional[str]]: + async def _set_birth_date( + self, session: AsyncSession, token: str + ) -> tuple[bool, int, Optional[str]]: """设置出生日期""" headers = self._build_birth_headers(token) payload = {"birthDate": self._random_birth_date()} - + try: response = await session.post( BIRTH_DATE_API, @@ -133,7 +137,7 @@ class NSFWService: http_status=birth_status, error=f"Set birth date failed: {birth_err}", ) - + # 开启 NSFW response = await session.post( NSFW_API, @@ -152,8 +156,7 @@ class NSFWService: # 解析 gRPC-Web 响应 _, trailers = parse_grpc_web_response( - response.content, - content_type=response.headers.get("content-type") + response.content, content_type=response.headers.get("content-type") ) grpc_status = get_grpc_status(trailers) @@ -176,4 +179,5 @@ class NSFWService: logger.error(f"NSFW enable failed: {e}") return NSFWResult(success=False, http_status=0, error=str(e)[:100]) + __all__ = ["NSFWService", "NSFWResult"] diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index adb04b3..7550c82 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -112,10 +112,14 @@ class UsageService: if response.status_code == 200: data = response.json() remaining = data.get("remainingTokens", 0) - logger.info(f"Usage sync success: remaining={remaining}, token={token[:10]}...") + logger.info( + f"Usage sync success: remaining={remaining}, token={token[:10]}..." + ) return data - logger.error(f"Usage sync failed: status={response.status_code}, token={token[:10]}...") + logger.error( + f"Usage sync failed: status={response.status_code}, token={token[:10]}..." + ) raise UpstreamException( message=f"Failed to get usage stats: {response.status_code}", @@ -142,4 +146,5 @@ class UsageService: # 最后一次失败已经被记录 raise + __all__ = ["UsageService"] diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index ae003f0..006fc54 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -14,6 +14,7 @@ from app.services.grok.utils.headers import apply_statsig, build_sso_cookie LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens" + class VoiceService: """Voice Mode Service (LiveKit)""" @@ -29,18 +30,20 @@ class VoiceService: ) -> Dict[str, Any]: """ Get LiveKit token - + Args: token: Auth token Returns: Dict containing token and livekitUrl """ - logger.debug(f"Voice token request: voice={voice}, personality={personality}, speed={speed}") + logger.debug( + f"Voice token request: voice={voice}, personality={personality}, speed={speed}" + ) headers = self._build_headers(token) payload = self._build_payload(voice, personality, speed) - + proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None - + try: browser = get_config("security.browser") timeout = get_config("network.timeout") @@ -55,12 +58,14 @@ class VoiceService: if response.status_code != 200: body = response.text[:200] - logger.error(f"Voice token failed: status={response.status_code}, body={body}") + logger.error( + f"Voice token failed: status={response.status_code}, body={body}" + ) raise UpstreamException( message=f"Failed to get voice token: {response.status_code}", - details={"status": response.status_code, "body": response.text} + details={"status": response.status_code, "body": response.text}, ) - + result = response.json() logger.info(f"Voice token obtained: voice={voice}") return result @@ -93,14 +98,16 @@ class VoiceService: ) -> Dict[str, Any]: """Construct payload with voice settings""" return { - "sessionPayload": orjson.dumps({ - "voice": voice, - "personality": personality, - "playback_speed": speed, - "enable_vision": False, - "turn_detection": {"type": "server_vad"} - }).decode(), + "sessionPayload": orjson.dumps( + { + "voice": voice, + "personality": personality, + "playback_speed": speed, + "enable_vision": False, + "turn_detection": {"type": "server_vad"}, + } + ).decode(), "requestAgentDispatch": False, "livekitUrl": "wss://livekit.grok.com", - "params": {"enable_markdown_transcript": "true"} + "params": {"enable_markdown_transcript": "true"}, } diff --git a/app/services/grok/utils/retry.py b/app/services/grok/utils/retry.py index 7fa7eba..162c4f8 100644 --- a/app/services/grok/utils/retry.py +++ b/app/services/grok/utils/retry.py @@ -257,7 +257,6 @@ def with_retry( __all__ = [ - "RetryConfig", "RetryContext", "retry_on_status", "with_retry", diff --git a/app/services/grok/utils/stream.py b/app/services/grok/utils/stream.py index fb107b7..c9e64dd 100644 --- a/app/services/grok/utils/stream.py +++ b/app/services/grok/utils/stream.py @@ -10,14 +10,11 @@ from app.services.token import EffortType async def wrap_stream_with_usage( - stream: AsyncGenerator, - token_mgr, - token: str, - model: str + stream: AsyncGenerator, token_mgr, token: str, model: str ) -> AsyncGenerator: """ 包装流式响应，在完成时记录使用 - + Args: stream: 原始 AsyncGenerator token_mgr: TokenManager 实例 @@ -33,9 +30,15 @@ async def wrap_stream_with_usage( if success: try: model_info = ModelService.get(model) - effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + effort = ( + EffortType.HIGH + if (model_info and model_info.cost.value == "high") + else EffortType.LOW + ) await token_mgr.consume(token, effort) - logger.debug(f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})") + logger.debug( + f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})" + ) except Exception as e: logger.warning(f"Failed to record stream usage: {e}") diff --git a/main.py b/main.py index 97da6ac..8a0e537 100644 --- a/main.py +++ b/main.py @@ -44,9 +44,9 @@ async def lifespan(app: FastAPI): # 1. 注册服务默认配置 from app.core.config import config, register_defaults from app.services.grok.defaults import get_grok_defaults - + register_defaults(get_grok_defaults()) - + # 2. 加载配置 await config.load() diff --git a/pyproject.toml b/pyproject.toml index a49e1eb..64171bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,3 +26,8 @@ dependencies = [ "tomli>=2.4.0", "uvicorn>=0.37.0", ] + +[dependency-groups] +dev = [ + "ruff>=0.15.0", +] diff --git a/uv.lock b/uv.lock index 43bc82d..74d3f7a 100644 --- a/uv.lock +++ b/uv.lock @@ -443,6 +443,11 @@ dependencies = [ { name = "uvicorn" }, ] +[package.dev-dependencies] +dev = [ + { name = "ruff" }, +] + [package.metadata] requires-dist = [ { name = "aiofiles", specifier = ">=25.1.0" }, @@ -467,6 +472,9 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.37.0" }, ] +[package.metadata.requires-dev] +dev = [{ name = "ruff", specifier = ">=0.15.0" }] + [[package]] name = "h11" version = "0.16.0" @@ -970,6 +978,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/89/f0/8956f8a86b20d7bb9d6ac0187cf4cd54d8065bc9a1a09eb8011d4d326596/redis-7.1.0-py3-none-any.whl", hash = "sha256:23c52b208f92b56103e17c5d06bdc1a6c2c0b3106583985a76a18f83b265de2b", size = 354159, upload-time = "2025-11-19T15:54:38.064Z" }, ] +[[package]] +name = "ruff" +version = "0.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c8/39/5cee96809fbca590abea6b46c6d1c586b49663d1d2830a751cc8fc42c666/ruff-0.15.0.tar.gz", hash = "sha256:6bdea47cdbea30d40f8f8d7d69c0854ba7c15420ec75a26f463290949d7f7e9a", size = 4524893, upload-time = "2026-02-03T17:53:35.357Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/88/3fd1b0aa4b6330d6aaa63a285bc96c9f71970351579152d231ed90914586/ruff-0.15.0-py3-none-linux_armv6l.whl", hash = "sha256:aac4ebaa612a82b23d45964586f24ae9bc23ca101919f5590bdb368d74ad5455", size = 10354332, upload-time = "2026-02-03T17:52:54.892Z" }, + { url = "https://files.pythonhosted.org/packages/72/f6/62e173fbb7eb75cc29fe2576a1e20f0a46f671a2587b5f604bfb0eaf5f6f/ruff-0.15.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:dcd4be7cc75cfbbca24a98d04d0b9b36a270d0833241f776b788d59f4142b14d", size = 10767189, upload-time = "2026-02-03T17:53:19.778Z" }, + { url = "https://files.pythonhosted.org/packages/99/e4/968ae17b676d1d2ff101d56dc69cf333e3a4c985e1ec23803df84fc7bf9e/ruff-0.15.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:d747e3319b2bce179c7c1eaad3d884dc0a199b5f4d5187620530adf9105268ce", size = 10075384, upload-time = "2026-02-03T17:53:29.241Z" }, + { url = "https://files.pythonhosted.org/packages/a2/bf/9843c6044ab9e20af879c751487e61333ca79a2c8c3058b15722386b8cae/ruff-0.15.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:650bd9c56ae03102c51a5e4b554d74d825ff3abe4db22b90fd32d816c2e90621", size = 10481363, upload-time = "2026-02-03T17:52:43.332Z" }, + { url = "https://files.pythonhosted.org/packages/55/d9/4ada5ccf4cd1f532db1c8d44b6f664f2208d3d93acbeec18f82315e15193/ruff-0.15.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a6664b7eac559e3048223a2da77769c2f92b43a6dfd4720cef42654299a599c9", size = 10187736, upload-time = "2026-02-03T17:53:00.522Z" }, + { url = "https://files.pythonhosted.org/packages/86/e2/f25eaecd446af7bb132af0a1d5b135a62971a41f5366ff41d06d25e77a91/ruff-0.15.0-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6f811f97b0f092b35320d1556f3353bf238763420ade5d9e62ebd2b73f2ff179", size = 10968415, upload-time = "2026-02-03T17:53:15.705Z" }, + { url = "https://files.pythonhosted.org/packages/e7/dc/f06a8558d06333bf79b497d29a50c3a673d9251214e0d7ec78f90b30aa79/ruff-0.15.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:761ec0a66680fab6454236635a39abaf14198818c8cdf691e036f4bc0f406b2d", size = 11809643, upload-time = "2026-02-03T17:53:23.031Z" }, + { url = "https://files.pythonhosted.org/packages/dd/45/0ece8db2c474ad7df13af3a6d50f76e22a09d078af63078f005057ca59eb/ruff-0.15.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:940f11c2604d317e797b289f4f9f3fa5555ffe4fb574b55ed006c3d9b6f0eb78", size = 11234787, upload-time = "2026-02-03T17:52:46.432Z" }, + { url = "https://files.pythonhosted.org/packages/8a/d9/0e3a81467a120fd265658d127db648e4d3acfe3e4f6f5d4ea79fac47e587/ruff-0.15.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bcbca3d40558789126da91d7ef9a7c87772ee107033db7191edefa34e2c7f1b4", size = 11112797, upload-time = "2026-02-03T17:52:49.274Z" }, + { url = "https://files.pythonhosted.org/packages/b2/cb/8c0b3b0c692683f8ff31351dfb6241047fa873a4481a76df4335a8bff716/ruff-0.15.0-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:9a121a96db1d75fa3eb39c4539e607f628920dd72ff1f7c5ee4f1b768ac62d6e", size = 11033133, upload-time = "2026-02-03T17:53:33.105Z" }, + { url = "https://files.pythonhosted.org/packages/f8/5e/23b87370cf0f9081a8c89a753e69a4e8778805b8802ccfe175cc410e50b9/ruff-0.15.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:5298d518e493061f2eabd4abd067c7e4fb89e2f63291c94332e35631c07c3662", size = 10442646, upload-time = "2026-02-03T17:53:06.278Z" }, + { url = "https://files.pythonhosted.org/packages/e1/9a/3c94de5ce642830167e6d00b5c75aacd73e6347b4c7fc6828699b150a5ee/ruff-0.15.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:afb6e603d6375ff0d6b0cee563fa21ab570fd15e65c852cb24922cef25050cf1", size = 10195750, upload-time = "2026-02-03T17:53:26.084Z" }, + { url = "https://files.pythonhosted.org/packages/30/15/e396325080d600b436acc970848d69df9c13977942fb62bb8722d729bee8/ruff-0.15.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:77e515f6b15f828b94dc17d2b4ace334c9ddb7d9468c54b2f9ed2b9c1593ef16", size = 10676120, upload-time = "2026-02-03T17:53:09.363Z" }, + { url = "https://files.pythonhosted.org/packages/8d/c9/229a23d52a2983de1ad0fb0ee37d36e0257e6f28bfd6b498ee2c76361874/ruff-0.15.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:6f6e80850a01eb13b3e42ee0ebdf6e4497151b48c35051aab51c101266d187a3", size = 11201636, upload-time = "2026-02-03T17:52:57.281Z" }, + { url = "https://files.pythonhosted.org/packages/6f/b0/69adf22f4e24f3677208adb715c578266842e6e6a3cc77483f48dd999ede/ruff-0.15.0-py3-none-win32.whl", hash = "sha256:238a717ef803e501b6d51e0bdd0d2c6e8513fe9eec14002445134d3907cd46c3", size = 10465945, upload-time = "2026-02-03T17:53:12.591Z" }, + { url = "https://files.pythonhosted.org/packages/51/ad/f813b6e2c97e9b4598be25e94a9147b9af7e60523b0cb5d94d307c15229d/ruff-0.15.0-py3-none-win_amd64.whl", hash = "sha256:dd5e4d3301dc01de614da3cdffc33d4b1b96fb89e45721f1598e5532ccf78b18", size = 11564657, upload-time = "2026-02-03T17:52:51.893Z" }, + { url = "https://files.pythonhosted.org/packages/f6/b0/2d823f6e77ebe560f4e397d078487e8d52c1516b331e3521bc75db4272ca/ruff-0.15.0-py3-none-win_arm64.whl", hash = "sha256:c480d632cc0ca3f0727acac8b7d053542d9e114a462a145d0b00e7cd658c515a", size = 10865753, upload-time = "2026-02-03T17:53:03.014Z" }, +] + [[package]] name = "sqlalchemy" version = "2.0.46" From a692725baa95568a90b0f007543544f911aa7ca9 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sat, 7 Feb 2026 13:38:08 +0800 Subject: [PATCH 6/6] chore: remove deprecated grok configuration options to streamline settings --- data/config.toml | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/data/config.toml b/data/config.toml index e51f56f..93016ea 100644 --- a/data/config.toml +++ b/data/config.toml @@ -66,26 +66,3 @@ usage_max_tokens = 1000 nsfw_max_concurrent = 10 nsfw_batch_size = 50 nsfw_max_tokens = 1000 - -[grok] -temporary = true -stream = true -thinking = true -dynamic_statsig = true -filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] -timeout = 120 -base_proxy_url = "" -asset_proxy_url = "" -cf_clearance = "" -max_retry = 3 -retry_status_codes = [401,429,403] -image_ws = true -image_ws_blocked_seconds = 15 -image_ws_final_min_bytes = 100000 -image_ws_nsfw = true -retry_backoff_base = 0.5 -retry_backoff_factor = 2.0 -retry_backoff_max = 30.0 -retry_budget = 90.0 -stream_idle_timeout = 120.0 -video_idle_timeout = 90.0