From 41d5a8d48fc5e6cda35ff04da6090e38ee1f4e69 Mon Sep 17 00:00:00 2001 From: piexian <64474352+piexian@users.noreply.github.com> Date: Tue, 3 Feb 2026 08:00:45 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20NSFW=20=E5=90=AF?= =?UTF-8?q?=E7=94=A8=E5=8A=9F=E8=83=BD=E3=80=81Token=20=E7=AD=9B=E9=80=89?= =?UTF-8?q?=E6=A0=87=E7=AD=BE=E9=A1=B5=E5=8F=8A=E4=BF=AE=E5=A4=8D=20use=5F?= =?UTF-8?q?count=20=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增功能: - 添加 gRPC-Web 协议工具用于 NSFW API 调用 - 添加批量执行工具支持并发控制 - 添加 NSFW 服务及管理 API 端点 POST /api/v1/admin/tokens/nsfw/enable - 添加 Token 状态筛选标签页（全部/正常/限流/失效）支持 ARIA 无障碍 - 添加批量启用 NSFW 按钮到 Token 管理工具栏 缺陷修复: - 修复 use_count 统计：在 chat/video/image 完成后调用 consume() - 修复 use_count 计算：使用 actual_cost 替代 cost 避免过度计数 - 修复 toggleSelectAll 仅选择当前筛选结果 配置变更: - 添加 NSFW 相关配置项到 config.defaults.toml --- app/api/v1/admin.py | 271 +++++++++++++++++++++++++++------- app/api/v1/image.py | 126 ++++++++++------ app/services/grok/batch.py | 68 +++++++++ app/services/grok/chat.py | 242 +++++++++++++++++++----------- app/services/grok/grpc_web.py | 157 ++++++++++++++++++++ app/services/grok/media.py | 249 +++++++++++++++++++------------ app/services/grok/nsfw.py | 112 ++++++++++++++ app/services/token/models.py | 74 +++++----- app/static/token/token.css | 39 +++++ app/static/token/token.html | 25 ++++ app/static/token/token.js | 121 ++++++++++++++- config.defaults.toml | 3 + 12 files changed, 1171 insertions(+), 316 deletions(-) create mode 100644 app/services/grok/batch.py create mode 100644 app/services/grok/grpc_web.py create mode 100644 app/services/grok/nsfw.py diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 8a9b226..7cfdded 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -14,42 +14,49 @@ router = APIRouter() TEMPLATE_DIR = Path(__file__).parent.parent.parent / "static" + async def render_template(filename: str): """渲染指定模板""" template_path = TEMPLATE_DIR / filename if not template_path.exists(): return HTMLResponse(f"Template {filename} not found.", status_code=404) - + async with aiofiles.open(template_path, "r", encoding="utf-8") as f: content = await f.read() return HTMLResponse(content) + @router.get("/admin", response_class=HTMLResponse, include_in_schema=False) async def admin_login_page(): """管理后台登录页""" return await render_template("login/login.html") + @router.get("/admin/config", response_class=HTMLResponse, include_in_schema=False) async def admin_config_page(): """配置管理页""" return await render_template("config/config.html") + @router.get("/admin/token", response_class=HTMLResponse, include_in_schema=False) async def admin_token_page(): """Token 管理页""" return await render_template("token/token.html") + @router.post("/api/v1/admin/login", dependencies=[Depends(verify_app_key)]) async def admin_login_api(): """管理后台登录验证（使用 app_key）""" return {"status": "success", "api_key": get_config("app.api_key", "")} + @router.get("/api/v1/admin/config", dependencies=[Depends(verify_api_key)]) async def get_config_api(): """获取当前配置""" # 暴露原始配置字典 return config._config + @router.post("/api/v1/admin/config", dependencies=[Depends(verify_api_key)]) async def update_config_api(data: dict): """更新配置""" @@ -59,6 +66,7 @@ async def update_config_api(data: dict): except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + @router.get("/api/v1/admin/storage", dependencies=[Depends(verify_api_key)]) async def get_storage_info(): """获取当前存储模式""" @@ -81,6 +89,7 @@ async def get_storage_info(): storage_type = storage.dialect return {"type": storage_type or "local"} + @router.get("/api/v1/admin/tokens", dependencies=[Depends(verify_api_key)]) async def get_tokens_api(): """获取所有 Token""" @@ -88,12 +97,14 @@ async def get_tokens_api(): tokens = await storage.load_tokens() return tokens or {} + @router.post("/api/v1/admin/tokens", dependencies=[Depends(verify_api_key)]) async def update_tokens_api(data: dict): """更新 Token 信息""" storage = get_storage() try: from app.services.token.manager import get_token_manager + async with storage.acquire_lock("tokens_save", timeout=10): await storage.save_tokens(data) mgr = await get_token_manager() @@ -102,11 +113,12 @@ async def update_tokens_api(data: dict): except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + @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 - + try: mgr = await get_token_manager() tokens = [] @@ -114,55 +126,163 @@ async def refresh_tokens_api(data: dict): tokens.append(data["token"]) if "tokens" in data and isinstance(data["tokens"], list): tokens.extend(data["tokens"]) - + if not tokens: - raise HTTPException(status_code=400, detail="No tokens provided") - + raise HTTPException(status_code=400, detail="No tokens provided") + unique_tokens = list(set(tokens)) - + sem = asyncio.Semaphore(10) - + async def _refresh_one(t): async with sem: - return t, await mgr.sync_usage(t, "grok-3", consume_on_fail=False, is_usage=False) - + return t, await mgr.sync_usage( + t, "grok-3", consume_on_fail=False, is_usage=False + ) + results_list = await asyncio.gather(*[_refresh_one(t) for t in unique_tokens]) results = dict(results_list) - + return {"status": "success", "results": results} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + +@router.post("/api/v1/admin/tokens/nsfw/enable", dependencies=[Depends(verify_api_key)]) +async def enable_nsfw_api(data: dict): + """批量开启 NSFW (Unhinged) 模式""" + from app.services.grok.nsfw import NSFWService + from app.services.grok.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()]) + + # 若未指定，则使用所有 pool 中的 token + if not tokens: + for pool_name, pool in mgr.pools.items(): + for info in pool.list(): + raw = ( + info.token[4:] if info.token.startswith("sso=") else info.token + ) + tokens.append(raw) + + if not tokens: + raise HTTPException(status_code=400, detail="No tokens available") + + # 去重并保持顺序 + unique_tokens = list(dict.fromkeys(tokens)) + + # 限制最大数量（防止误操作） + max_tokens = get_config("performance.nsfw_max_tokens", 1000) + try: + max_tokens = int(max_tokens) + except Exception: + max_tokens = 1000 + if len(unique_tokens) > max_tokens: + raise HTTPException( + status_code=400, + detail=f"Too many tokens: {len(unique_tokens)} > {max_tokens}", + ) + + # 批量执行配置 + max_concurrent = get_config("performance.nsfw_max_concurrent", 10) + batch_size = get_config("performance.admin_nsfw_batch_size", 50) + + # 定义 worker + async def _enable(token: str): + result = await nsfw_service.enable(token) + return { + "success": result.success, + "http_status": result.http_status, + "grpc_status": result.grpc_status, + "grpc_message": result.grpc_message, + "error": result.error, + } + + # 执行批量操作 + raw_results = await run_in_batches( + unique_tokens, _enable, max_concurrent=max_concurrent, batch_size=batch_size + ) + + # 构造返回结果（mask token） + results = {} + ok_count = 0 + fail_count = 0 + + for token, res in raw_results.items(): + masked = f"{token[:8]}...{token[-8:]}" if len(token) > 20 else token + if res.get("ok") and res.get("data", {}).get("success"): + ok_count += 1 + results[masked] = res.get("data", {}) + else: + fail_count += 1 + results[masked] = res.get("data") or {"error": res.get("error")} + + return { + "status": "success", + "summary": { + "total": len(unique_tokens), + "ok": ok_count, + "fail": fail_count, + }, + "results": results, + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Enable NSFW failed: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + @router.get("/admin/cache", response_class=HTMLResponse, include_in_schema=False) async def admin_cache_page(): """缓存管理页""" return await render_template("cache/cache.html") + @router.get("/api/v1/admin/cache", dependencies=[Depends(verify_api_key)]) async def get_cache_stats_api(request: Request): """获取缓存统计""" from app.services.grok.assets import DownloadService, ListService from app.services.token.manager import get_token_manager - + try: dl_service = DownloadService() image_stats = dl_service.get_stats("image") video_stats = dl_service.get_stats("video") - + mgr = await get_token_manager() pools = mgr.pools accounts = [] for pool_name, pool in pools.items(): for info in pool.list(): - raw_token = info.token[4:] if info.token.startswith("sso=") else info.token - masked = f"{raw_token[:8]}...{raw_token[-16:]}" if len(raw_token) > 24 else raw_token - accounts.append({ - "token": raw_token, - "token_masked": masked, - "pool": pool_name, - "status": info.status, - "last_asset_clear_at": info.last_asset_clear_at - }) + raw_token = ( + info.token[4:] if info.token.startswith("sso=") else info.token + ) + masked = ( + f"{raw_token[:8]}...{raw_token[-16:]}" + if len(raw_token) > 24 + else raw_token + ) + accounts.append( + { + "token": raw_token, + "token_masked": masked, + "pool": pool_name, + "status": info.status, + "last_asset_clear_at": info.last_asset_clear_at, + } + ) scope = request.query_params.get("scope") selected_token = request.query_params.get("token") @@ -171,7 +291,12 @@ async def get_cache_stats_api(request: Request): if tokens_param: selected_tokens = [t.strip() for t in tokens_param.split(",") if t.strip()] - online_stats = {"count": 0, "status": "unknown", "token": None, "last_asset_clear_at": None} + online_stats = { + "count": 0, + "status": "unknown", + "token": None, + "last_asset_clear_at": None, + } online_details = [] account_map = {a["token"]: a for a in accounts} batch_size = get_config("performance.admin_assets_batch_size", 10) @@ -192,42 +317,66 @@ async def get_cache_stats_api(request: Request): account = account_map.get(token) try: count = await _fetch_assets(token) - return ({ - "token": token, - "token_masked": account["token_masked"] if account else token, - "count": count, - "status": "ok", - "last_asset_clear_at": account["last_asset_clear_at"] if account else None - }, count) + return ( + { + "token": token, + "token_masked": account["token_masked"] if account else token, + "count": count, + "status": "ok", + "last_asset_clear_at": account["last_asset_clear_at"] + if account + else None, + }, + count, + ) except Exception as e: - return ({ - "token": token, - "token_masked": account["token_masked"] if account else token, - "count": 0, - "status": f"error: {str(e)}", - "last_asset_clear_at": account["last_asset_clear_at"] if account else None - }, 0) + return ( + { + "token": token, + "token_masked": account["token_masked"] if account else token, + "count": 0, + "status": f"error: {str(e)}", + "last_asset_clear_at": account["last_asset_clear_at"] + if account + else None, + }, + 0, + ) if selected_tokens: total = 0 for i in range(0, len(selected_tokens), batch_size): - chunk = selected_tokens[i:i + batch_size] - results = await asyncio.gather(*[_fetch_detail(token) for token in chunk]) + chunk = selected_tokens[i : i + batch_size] + results = await asyncio.gather( + *[_fetch_detail(token) for token in chunk] + ) for detail, count in results: online_details.append(detail) total += count - online_stats = {"count": total, "status": "ok" if selected_tokens else "no_token", "token": None, "last_asset_clear_at": None} + online_stats = { + "count": total, + "status": "ok" if selected_tokens else "no_token", + "token": None, + "last_asset_clear_at": None, + } scope = "selected" elif scope == "all": total = 0 tokens = [account["token"] for account in accounts] for i in range(0, len(tokens), batch_size): - chunk = tokens[i:i + batch_size] - results = await asyncio.gather(*[_fetch_detail(token) for token in chunk]) + chunk = tokens[i : i + batch_size] + results = await asyncio.gather( + *[_fetch_detail(token) for token in chunk] + ) for detail, count in results: online_details.append(detail) total += count - online_stats = {"count": total, "status": "ok" if accounts else "no_token", "token": None, "last_asset_clear_at": None} + online_stats = { + "count": total, + "status": "ok" if accounts else "no_token", + "token": None, + "last_asset_clear_at": None, + } else: token = selected_token if token: @@ -239,7 +388,9 @@ async def get_cache_stats_api(request: Request): "status": "ok", "token": token, "token_masked": match["token_masked"] if match else token, - "last_asset_clear_at": match["last_asset_clear_at"] if match else None + "last_asset_clear_at": match["last_asset_clear_at"] + if match + else None, } except Exception as e: match = next((a for a in accounts if a["token"] == token), None) @@ -248,28 +399,37 @@ async def get_cache_stats_api(request: Request): "status": f"error: {str(e)}", "token": token, "token_masked": match["token_masked"] if match else token, - "last_asset_clear_at": match["last_asset_clear_at"] if match else None + "last_asset_clear_at": match["last_asset_clear_at"] + if match + else None, } else: - online_stats = {"count": 0, "status": "not_loaded", "token": None, "last_asset_clear_at": None} - + online_stats = { + "count": 0, + "status": "not_loaded", + "token": None, + "last_asset_clear_at": None, + } + return { "local_image": image_stats, "local_video": video_stats, "online": online_stats, "online_accounts": accounts, "online_scope": scope or "none", - "online_details": online_details + "online_details": online_details, } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + @router.post("/api/v1/admin/cache/clear", dependencies=[Depends(verify_api_key)]) async def clear_local_cache_api(data: dict): """清理本地缓存""" from app.services.grok.assets import DownloadService + cache_type = data.get("type", "image") - + try: dl_service = DownloadService() result = dl_service.clear(cache_type) @@ -277,15 +437,17 @@ async def clear_local_cache_api(data: dict): except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + @router.get("/api/v1/admin/cache/list", dependencies=[Depends(verify_api_key)]) async def list_local_cache_api( cache_type: str = "image", type_: str = Query(default=None, alias="type"), page: int = 1, - page_size: int = 1000 + page_size: int = 1000, ): """列出本地缓存文件""" from app.services.grok.assets import DownloadService + try: if type_: cache_type = type_ @@ -295,10 +457,12 @@ async def list_local_cache_api( except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + @router.post("/api/v1/admin/cache/item/delete", dependencies=[Depends(verify_api_key)]) async def delete_local_cache_item_api(data: dict): """删除单个本地缓存文件""" from app.services.grok.assets import DownloadService + cache_type = data.get("type", "image") name = data.get("name") if not name: @@ -310,12 +474,13 @@ async def delete_local_cache_item_api(data: dict): except Exception as e: raise HTTPException(status_code=500, detail=str(e)) + @router.post("/api/v1/admin/cache/online/clear", dependencies=[Depends(verify_api_key)]) async def clear_online_cache_api(data: dict): """清理在线缓存""" from app.services.grok.assets import DeleteService from app.services.token.manager import get_token_manager - + delete_service = None try: mgr = await get_token_manager() @@ -344,7 +509,7 @@ async def clear_online_cache_api(data: dict): return t, {"status": "error", "error": str(e)} for i in range(0, len(token_list), batch_size): - chunk = token_list[i:i + batch_size] + chunk = token_list[i : i + batch_size] res_list = await asyncio.gather(*[_clear_one(t) for t in chunk]) for t, res in res_list: results[t] = res @@ -353,7 +518,9 @@ async def clear_online_cache_api(data: dict): token = data.get("token") or mgr.get_token() if not token: - raise HTTPException(status_code=400, detail="No available token to perform cleanup") + raise HTTPException( + status_code=400, detail="No available token to perform cleanup" + ) result = await delete_service.delete_all(token) await mgr.mark_asset_clear(token) diff --git a/app/api/v1/image.py b/app/api/v1/image.py index 0aa7f7e..15fb9e1 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field from app.services.grok.chat import GrokChatService from app.services.grok.model import ModelService from app.services.grok.processor import ImageStreamProcessor, ImageCollectProcessor -from app.services.token import get_token_manager +from app.services.token import get_token_manager, EffortType from app.core.exceptions import ValidationException, AppException, ErrorType from app.core.logger import logger @@ -23,6 +23,7 @@ router = APIRouter(tags=["Images"]) class ImageGenerationRequest(BaseModel): """图片生成请求 - OpenAI 兼容""" + prompt: str = Field(..., description="图片描述") model: Optional[str] = Field("grok-imagine-1.0", description="模型名称") n: Optional[int] = Field(1, ge=1, le=10, description="生成数量 (1-10)") @@ -43,40 +44,36 @@ def validate_request(request: ImageGenerationRequest): raise ValidationException( message=f"The model `{request.model}` is not supported for image generation. Supported: {image_models}", param="model", - code="model_not_supported" + code="model_not_supported", ) - + # 验证 prompt if not request.prompt or not request.prompt.strip(): raise ValidationException( - message="Prompt cannot be empty", - param="prompt", - code="empty_prompt" + message="Prompt cannot be empty", param="prompt", code="empty_prompt" ) - + # 验证 n 参数范围 if request.n < 1 or request.n > 10: raise ValidationException( - message="n must be between 1 and 10", - param="n", - code="invalid_n" + message="n must be between 1 and 10", param="n", code="invalid_n" ) - + # 流式只支持 n=1 或 n=2 if request.stream and request.n not in [1, 2]: raise ValidationException( message="Streaming is only supported when n=1 or n=2", param="stream", - code="invalid_stream_n" + code="invalid_stream_n", ) -async def call_grok(token: str, prompt: str, model_info) -> List[str]: +async def call_grok(token_mgr, token: str, prompt: str, model_info) -> List[str]: """ 调用 Grok 获取图片，返回 base64 列表 """ chat_service = GrokChatService() - + try: response = await chat_service.chat( token=token, @@ -84,38 +81,49 @@ async def call_grok(token: str, prompt: str, model_info) -> List[str]: model=model_info.grok_model, mode=model_info.model_mode, think=False, - stream=True + stream=True, ) - + # 收集图片 processor = ImageCollectProcessor(model_info.model_id, token) images = await processor.process(response) return images - + except Exception as e: logger.error(f"Grok image call failed: {e}") return [] + finally: + # 记录使用 + 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}") @router.post("/images/generations") async def create_image(request: ImageGenerationRequest): """ Image Generation API - + 流式响应格式: - event: image_generation.partial_image - event: image_generation.completed - + 非流式响应格式: - {"created": ..., "data": [{"b64_json": "..."}], "usage": {...}} """ # stream 默认为 false if request.stream is None: request.stream = False - + # 参数验证 validate_request(request) - + # 获取 token try: token_mgr = await get_token_manager() @@ -127,7 +135,7 @@ async def create_image(request: ImageGenerationRequest): raise AppException( message="Internal service error obtaining token", error_type=ErrorType.SERVER.value, - code="internal_error" + code="internal_error", ) if not token: @@ -135,12 +143,12 @@ async def create_image(request: ImageGenerationRequest): message="No available tokens. Please try again later.", error_type=ErrorType.RATE_LIMIT.value, code="rate_limit_exceeded", - status_code=429 + status_code=429, ) - + # 获取模型信息 model_info = ModelService.get(request.model) - + # 流式模式 if request.stream: chat_service = GrokChatService() @@ -150,32 +158,49 @@ async def create_image(request: ImageGenerationRequest): model=model_info.grok_model, mode=model_info.model_mode, think=False, - stream=True + stream=True, ) - + processor = ImageStreamProcessor(model_info.model_id, token, n=request.n) + + # 包装流式响应，在完成时记录使用 + async def _wrap_stream(stream): + try: + async for chunk in stream: + yield chunk + finally: + 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( - processor.process(response), + _wrap_stream(processor.process(response)), media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive"} + headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) - + # 非流式模式 n = request.n - - calls_needed = (n + 1) // 2 - + + calls_needed = (n + 1) // 2 + if calls_needed == 1: # 单次调用 - all_images = await call_grok(token, request.prompt, model_info) + all_images = await call_grok(token_mgr, token, request.prompt, model_info) else: # 并发调用 tasks = [ - call_grok(token, request.prompt, model_info) + call_grok(token_mgr, token, request.prompt, model_info) for _ in range(calls_needed) ] results = await asyncio.gather(*tasks, return_exceptions=True) - + # 收集成功的图片 all_images = [] for result in results: @@ -183,7 +208,7 @@ async def create_image(request: ImageGenerationRequest): logger.error(f"Concurrent call failed: {result}") elif isinstance(result, list): all_images.extend(result) - + # 随机选取 n 张图片 if len(all_images) >= n: selected_images = random.sample(all_images, n) @@ -192,21 +217,26 @@ async def create_image(request: ImageGenerationRequest): selected_images = all_images.copy() while len(selected_images) < n: selected_images.append("error") - + # 构建响应 import time + data = [{"b64_json": img} for img in selected_images] - - return JSONResponse(content={ - "created": int(time.time()), - "data": data, - "usage": { - "total_tokens": 0 * len([img for img in selected_images if img != "error"]), - "input_tokens": 0, - "output_tokens": 0 * len([img for img in selected_images if img != "error"]), - "input_tokens_details": {"text_tokens": 0, "image_tokens": 0} + + return JSONResponse( + content={ + "created": int(time.time()), + "data": data, + "usage": { + "total_tokens": 0 + * len([img for img in selected_images if img != "error"]), + "input_tokens": 0, + "output_tokens": 0 + * len([img for img in selected_images if img != "error"]), + "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + }, } - }) + ) __all__ = ["router"] diff --git a/app/services/grok/batch.py b/app/services/grok/batch.py new file mode 100644 index 0000000..899d6f2 --- /dev/null +++ b/app/services/grok/batch.py @@ -0,0 +1,68 @@ +""" +批量执行工具 + +提供分批并发、单项失败隔离的通用批量处理能力。 +""" + +import asyncio +from typing import Any, Awaitable, Callable, Dict, List, TypeVar + +from app.core.logger import logger + +T = TypeVar("T") + + +async def run_in_batches( + items: List[str], + worker: Callable[[str], Awaitable[T]], + *, + max_concurrent: int = 10, + batch_size: int = 50, +) -> Dict[str, Dict[str, Any]]: + """ + 分批并发执行，单项失败不影响整体 + + Args: + items: 待处理项列表 + worker: 异步处理函数 + max_concurrent: 最大并发数 + batch_size: 每批大小 + + Returns: + {item: {"ok": bool, "data": ..., "error": ...}} + """ + try: + max_concurrent = int(max_concurrent) + except Exception: + max_concurrent = 10 + try: + batch_size = int(batch_size) + except Exception: + batch_size = 50 + + max_concurrent = max(1, max_concurrent) + batch_size = max(1, batch_size) + + sem = asyncio.Semaphore(max_concurrent) + + async def _one(item: str) -> tuple[str, dict]: + async with sem: + try: + data = await worker(item) + return item, {"ok": True, "data": data} + except Exception as e: + logger.warning(f"Batch item failed: {item[:16]}... - {e}") + return item, {"ok": False, "error": str(e)} + + results: Dict[str, dict] = {} + + # 分批执行，避免一次性创建所有 task + for i in range(0, len(items), batch_size): + chunk = items[i : i + batch_size] + pairs = await asyncio.gather(*(_one(x) for x in chunk)) + results.update(dict(pairs)) + + return results + + +__all__ = ["run_in_batches"] diff --git a/app/services/grok/chat.py b/app/services/grok/chat.py index 5fcbfea..cca5c8b 100644 --- a/app/services/grok/chat.py +++ b/app/services/grok/chat.py @@ -4,7 +4,7 @@ Grok Chat 服务 import uuid import orjson -from typing import Dict, List, Any +from typing import Dict, List, Any, AsyncGenerator from dataclasses import dataclass from curl_cffi.requests import AsyncSession @@ -12,17 +12,17 @@ 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, + AppException, UpstreamException, ValidationException, - ErrorType + ErrorType, ) from app.services.grok.statsig import StatsigService from app.services.grok.model import ModelService from app.services.grok.assets import UploadService from app.services.grok.processor import StreamProcessor, CollectProcessor from app.services.grok.retry import retry_on_status -from app.services.token import get_token_manager +from app.services.token import get_token_manager, EffortType CHAT_API = "https://grok.com/rest/app-chat/conversations/new" @@ -33,6 +33,7 @@ BROWSER = "chrome136" @dataclass class ChatRequest: """聊天请求数据""" + model: str messages: List[Dict[str, Any]] stream: bool = None @@ -41,24 +42,26 @@ 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[str]]: + def extract( + messages: List[Dict[str, Any]], is_video: bool = False + ) -> tuple[str, List[str]]: """ 从 OpenAI 消息格式提取内容 - + Args: messages: OpenAI 格式消息列表 is_video: 是否为视频模型 - + Returns: (text, attachments): 拼接后的文本和需要上传的附件列表 - + Raises: ValueError: 视频模型遇到不支持的内容类型 """ @@ -92,7 +95,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)) @@ -101,7 +108,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)) @@ -138,7 +149,7 @@ class MessageExtractor: # 换行拼接文本 message = "\n\n".join(texts) return message, attachments - + @staticmethod def extract_text_only(messages: List[Dict[str, Any]]) -> str: """仅提取文本内容""" @@ -148,7 +159,7 @@ class MessageExtractor: class ChatRequestBuilder: """请求构造器""" - + @staticmethod def build_headers(token: str) -> Dict[str, str]: """构造请求头""" @@ -174,30 +185,30 @@ class ChatRequestBuilder: "Sec-Fetch-Site": "same-origin", "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", } - + # Statsig ID headers["x-statsig-id"] = StatsigService.gen_id() headers["x-xai-request-id"] = str(uuid.uuid4()) - + # Cookie token = token[4:] if token.startswith("sso=") else token cf = get_config("grok.cf_clearance", "") headers["Cookie"] = f"sso={token};cf_clearance={cf}" if cf else f"sso={token}" - + return headers - + @staticmethod def build_payload( - message: str, - model: str, - mode: str, + message: str, + model: str, + mode: str, think: bool = None, file_attachments: List[str] = None, - image_attachments: List[str] = None + image_attachments: List[str] = None, ) -> Dict[str, Any]: """ 构造请求体 - + Args: message: 消息文本 model: 模型名称 @@ -209,7 +220,7 @@ class ChatRequestBuilder: temporary = get_config("grok.temporary", True) if think is None: think = get_config("grok.thinking", False) - + return { "temporary": temporary, "modelName": model, @@ -231,7 +242,7 @@ class ChatRequestBuilder: "disableTextFollowUps": False, "responseMetadata": { "modelConfigOverride": {"modelMap": {}}, - "requestModelDetails": {"modelId": model} + "requestModelDetails": {"modelId": model}, }, "disableMemory": False, "forceSideBySide": False, @@ -243,19 +254,20 @@ class ChatRequestBuilder: "screenWidth": 2056, "screenHeight": 1329, "viewportWidth": 2056, - "viewportHeight": 1083 - } + "viewportHeight": 1083, + }, } # ==================== Grok 服务 ==================== + class GrokChatService: """Grok API 调用服务""" - + def __init__(self, proxy: str = None): self.proxy = proxy or get_config("grok.base_proxy_url", "") - + async def chat( self, token: str, @@ -265,11 +277,11 @@ class GrokChatService: think: bool = None, stream: bool = None, file_attachments: List[str] = None, - image_attachments: List[str] = None + image_attachments: List[str] = None, ): """ 发送聊天请求 - + Args: token: 认证 Token message: 消息文本 @@ -279,27 +291,26 @@ class GrokChatService: stream: 是否流式 file_attachments: 文件附件 ID 列表 image_attachments: 图片附件 URL 列表 - + Raises: UpstreamException: 当 Grok API 返回错误且重试耗尽时 """ if stream is None: stream = get_config("grok.stream", True) - + headers = ChatRequestBuilder.build_headers(token) payload = ChatRequestBuilder.build_payload( - message, model, mode, think, - file_attachments, image_attachments + message, model, mode, think, file_attachments, image_attachments ) proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None timeout = get_config("grok.timeout", TIMEOUT) - + # 状态码提取器 def extract_status(e: Exception) -> int | None: if isinstance(e, UpstreamException) and e.details: return e.details.get("status") return None - + # 建立连接函数 async def establish_connection(): """建立连接并返回 response 对象""" @@ -311,19 +322,22 @@ class GrokChatService: data=orjson.dumps(payload), timeout=timeout, stream=True, - proxies=proxies + proxies=proxies, ) - + if response.status_code != 200: try: content = await response.text() - content = content[:1000] # 限制长度避免日志过大 + content = content[:1000] # 限制长度避免日志过大 except: content = "Unable to read response content" logger.error( f"Chat failed: {response.status_code}, {content}", - extra={"status": response.status_code, "token": token[:10] + "..."} + extra={ + "status": response.status_code, + "token": token[:10] + "...", + }, ) # 关闭 session 并抛出异常 try: @@ -332,12 +346,12 @@ class GrokChatService: pass raise UpstreamException( message=f"Grok API request failed: {response.status_code}", - details={"status": response.status_code} + details={"status": response.status_code}, ) - + # 返回 session 和 response return session, response - + except UpstreamException: # 已经处理过的异常，直接抛出 raise @@ -350,16 +364,15 @@ class GrokChatService: pass raise UpstreamException( message=f"Chat connection failed: {str(e)}", - details={"error": str(e)} + details={"error": str(e)}, ) - + # 建立连接 session = None response = None try: session, response = await retry_on_status( - establish_connection, - extract_status=extract_status + establish_connection, extract_status=extract_status ) except Exception as e: # 记录失败 @@ -368,7 +381,7 @@ class GrokChatService: token_mgr = await get_token_manager() await token_mgr.record_fail(token, status_code, str(e)) raise - + # 流式传输 async def stream_response(): try: @@ -377,36 +390,38 @@ class GrokChatService: finally: if session: await session.close() - + return stream_response() - + async def chat_openai(self, token: str, request: ChatRequest): """OpenAI 兼容接口""" model_info = ModelService.get(request.model) if not model_info: raise ValidationException(f"Unknown model: {request.model}") - + grok_model = model_info.grok_model mode = model_info.model_mode is_video = model_info.is_video - + # 提取消息和附件 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 = [] image_ids = [] - + if attachments: upload_service = UploadService() try: for attach_type, attach_data in attachments: # 获取 ID file_id, _ = await upload_service.upload(attach_data, token) - + if attach_type == "image": # 图片 imageAttachments image_ids.append(file_id) @@ -417,40 +432,84 @@ class GrokChatService: logger.debug(f"File uploaded: {file_id}") finally: await upload_service.close() - - stream = request.stream if request.stream is not None else get_config("grok.stream", True) - think = request.think if request.think is not None else get_config("grok.thinking", False) - - response = await self.chat( - token, message, grok_model, mode, think, stream, - file_attachments=file_ids, - image_attachments=image_ids + + stream = ( + request.stream + if request.stream is not None + else get_config("grok.stream", True) ) - + think = ( + request.think + if request.think is not None + else get_config("grok.thinking", False) + ) + + response = await self.chat( + token, + message, + grok_model, + mode, + think, + stream, + file_attachments=file_ids, + image_attachments=image_ids, + ) + 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: 模型名称 + """ + try: + async for chunk in stream: + yield chunk + finally: + # 无论成功完成还是客户端断开，都记录使用 + 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 + thinking: str = None, ): """ Chat Completions 入口 - + Args: model: 模型名称 messages: 消息列表 stream: 是否流式 thinking: 思考模式 - + Returns: AsyncGenerator 或 dict """ @@ -465,7 +524,7 @@ class ChatService: raise AppException( message="Internal service error obtaining token", error_type=ErrorType.SERVER.value, - code="internal_error" + code="internal_error", ) if not token: @@ -473,26 +532,23 @@ class ChatService: message="No available tokens. Please try again later.", error_type=ErrorType.RATE_LIMIT.value, code="rate_limit_exceeded", - status_code=429 + 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", True) - + # 构造请求 chat_request = ChatRequest( - model=model, - messages=messages, - stream=is_stream, - think=think + model=model, messages=messages, stream=is_stream, think=think ) - + # 请求 Grok service = GrokChatService() try: @@ -503,14 +559,32 @@ class ChatService: logger.error(f"Chat service error: {e}") raise UpstreamException( message=f"Service processing failed: {str(e)}", - details={"error": str(e)} + details={"error": str(e)}, ) - + # 处理响应 if is_stream: - return StreamProcessor(model_name, token, think).process(response) + processor = StreamProcessor(model_name, token, think) + return ChatService._wrap_stream( + processor.process(response), token_mgr, token, model + ) else: - return await CollectProcessor(model_name, token).process(response) + 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 __all__ = [ diff --git a/app/services/grok/grpc_web.py b/app/services/grok/grpc_web.py new file mode 100644 index 0000000..0724727 --- /dev/null +++ b/app/services/grok/grpc_web.py @@ -0,0 +1,157 @@ +""" +gRPC-Web 协议工具 + +提供 framing 编码/解码、trailer 解析等通用功能。 +支持 application/grpc-web+proto 和 application/grpc-web-text (base64) 两种格式。 +""" + +from __future__ import annotations + +import base64 +import re +import struct +from dataclasses import dataclass +from typing import Dict, List, Mapping, Tuple +from urllib.parse import unquote + + +_B64_RE = re.compile(rb"^[A-Za-z0-9+/=\r\n]+$") + + +def encode_grpc_web_payload(data: bytes) -> bytes: + """ + 编码 gRPC-Web data frame + + Frame format: + 1-byte flags + 4-byte big-endian length + message bytes + """ + return b"\x00" + struct.pack(">I", len(data)) + data + + +def _maybe_decode_grpc_web_text(body: bytes, content_type: str | None) -> bytes: + """处理 grpc-web-text 模式的 base64 解码""" + ct = (content_type or "").lower() + if "grpc-web-text" in ct: + compact = b"".join(body.split()) + return base64.b64decode(compact, validate=False) + + # 启发式：body 仅包含 base64 字符才尝试解码 + head = body[: min(len(body), 2048)] + if head and _B64_RE.fullmatch(head): + compact = b"".join(body.split()) + try: + return base64.b64decode(compact, validate=True) + except Exception: + return body + return body + + +def _parse_trailer_block(payload: bytes) -> Dict[str, str]: + """解析 trailer frame 内容""" + text = payload.decode("utf-8", errors="replace") + lines = [ln for ln in re.split(r"\r\n|\n", text) if ln] + + trailers: Dict[str, str] = {} + for ln in lines: + if ":" not in ln: + continue + k, v = ln.split(":", 1) + trailers[k.strip().lower()] = v.strip() + + # grpc-message 可能是 percent-encoding + if "grpc-message" in trailers: + trailers["grpc-message"] = unquote(trailers["grpc-message"]) + + return trailers + + +def parse_grpc_web_response( + body: bytes, + content_type: str | None = None, + headers: Mapping[str, str] | None = None, +) -> Tuple[List[bytes], Dict[str, str]]: + """ + 解析 gRPC-Web 响应 + + Returns: + (messages, trailers): data frames 列表和合并后的 trailers + """ + decoded = _maybe_decode_grpc_web_text(body, content_type) + + messages: List[bytes] = [] + trailers: Dict[str, str] = {} + + i = 0 + n = len(decoded) + while i < n: + if n - i < 5: + break + + flag = decoded[i] + length = int.from_bytes(decoded[i + 1 : i + 5], "big") + i += 5 + + if n - i < length: + break + + payload = decoded[i : i + length] + i += length + + if flag & 0x80: # trailer frame + trailers.update(_parse_trailer_block(payload)) + elif flag & 0x01: # compressed (不支持) + raise ValueError("grpc-web compressed flag not supported") + else: + messages.append(payload) + + # 兼容：grpc-status 可能在 response headers 中 + if headers: + lower = {k.lower(): v for k, v in headers.items()} + if "grpc-status" in lower and "grpc-status" not in trailers: + trailers["grpc-status"] = str(lower["grpc-status"]).strip() + if "grpc-message" in lower and "grpc-message" not in trailers: + trailers["grpc-message"] = unquote(str(lower["grpc-message"]).strip()) + + return messages, trailers + + +@dataclass(frozen=True) +class GrpcStatus: + code: int + message: str = "" + + @property + def ok(self) -> bool: + return self.code == 0 + + @property + def http_equiv(self) -> int: + """映射到类 HTTP 状态码""" + mapping = { + 0: 200, # OK + 16: 401, # UNAUTHENTICATED + 7: 403, # PERMISSION_DENIED + 8: 429, # RESOURCE_EXHAUSTED + 4: 504, # DEADLINE_EXCEEDED + 14: 503, # UNAVAILABLE + } + return mapping.get(self.code, 502) + + +def get_grpc_status(trailers: Mapping[str, str]) -> GrpcStatus: + """从 trailers 提取 gRPC 状态""" + raw = str(trailers.get("grpc-status", "")).strip() + msg = str(trailers.get("grpc-message", "")).strip() + try: + code = int(raw) + except Exception: + code = -1 + return GrpcStatus(code=code, message=msg) + + +__all__ = [ + "encode_grpc_web_payload", + "parse_grpc_web_response", + "get_grpc_status", + "GrpcStatus", +] diff --git a/app/services/grok/media.py b/app/services/grok/media.py index 74fa2d2..004e033 100644 --- a/app/services/grok/media.py +++ b/app/services/grok/media.py @@ -11,10 +11,15 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config -from app.core.exceptions import UpstreamException, AppException, ValidationException, ErrorType +from app.core.exceptions import ( + UpstreamException, + AppException, + ValidationException, + ErrorType, +) from app.services.grok.statsig import StatsigService from app.services.grok.model import ModelService -from app.services.token import get_token_manager +from app.services.token import get_token_manager, EffortType from app.services.grok.processor import VideoStreamProcessor, VideoCollectProcessor # API 端点 @@ -28,6 +33,7 @@ DEFAULT_MAX_CONCURRENT = 50 _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) @@ -44,12 +50,14 @@ def _get_media_semaphore() -> asyncio.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) - - def _build_headers(self, token: str, referer: str = "https://grok.com/imagine") -> dict: + + def _build_headers( + self, token: str, referer: str = "https://grok.com/imagine" + ) -> dict: """构建请求头""" headers = { "Accept": "*/*", @@ -73,50 +81,50 @@ class VideoService: "Sec-Fetch-Site": "same-origin", "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", } - + # Statsig ID headers["x-statsig-id"] = StatsigService.gen_id() headers["x-xai-request-id"] = str(uuid.uuid4()) - + # Cookie token = token[4:] if token.startswith("sso=") else token cf = get_config("grok.cf_clearance", "") headers["Cookie"] = f"sso={token};cf_clearance={cf}" if cf else f"sso={token}" - + return headers - + def _build_proxies(self) -> Optional[dict]: """构建代理""" return {"http": self.proxy, "https": self.proxy} if self.proxy else None - - async def create_post(self, token: str, prompt: str, media_type: str = "MEDIA_POST_TYPE_VIDEO", media_url: str = None) -> str: + + async def create_post( + self, + token: str, + prompt: str, + media_type: str = "MEDIA_POST_TYPE_VIDEO", + 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 """ try: headers = self._build_headers(token) - + # 根据类型构建不同的载荷 if media_type == "MEDIA_POST_TYPE_IMAGE" and media_url: - payload = { - "mediaType": media_type, - "mediaUrl": media_url - } + payload = {"mediaType": media_type, "mediaUrl": media_url} else: - payload = { - "mediaType": media_type, - "prompt": prompt - } - + payload = {"mediaType": media_type, "prompt": prompt} + async with AsyncSession() as session: response = await session.post( CREATE_POST_API, @@ -124,46 +132,45 @@ class VideoService: json=payload, impersonate=BROWSER, timeout=30, - proxies=self._build_proxies() + 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}") - + raise UpstreamException( + f"Failed to create post: {response.status_code}" + ) + data = response.json() post_id = data.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 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 + token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url ) - + def _build_payload( self, prompt: str, @@ -171,7 +178,7 @@ class VideoService: aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "SD", - preset: str = "normal" + preset: str = "normal", ) -> dict: """构建视频生成载荷""" mode_flag = "--mode=custom" @@ -181,9 +188,9 @@ class VideoService: mode_flag = "--mode=normal" elif preset == "spicy": mode_flag = "--mode=extremely-spicy-or-crazy" - + full_prompt = f"{prompt} {mode_flag}" - + return { "temporary": True, "modelName": "grok-3", @@ -198,13 +205,13 @@ class VideoService: "parentPostId": post_id, "aspectRatio": aspect_ratio, "videoLength": video_length, - "videoResolution": resolution + "videoResolution": resolution, } } - } - } + }, + }, } - + async def generate( self, token: str, @@ -213,11 +220,11 @@ class VideoService: video_length: int = 6, resolution: str = "SD", stream: bool = True, - preset: str = "normal" + preset: str = "normal", ) -> AsyncGenerator[bytes, None]: """ 生成视频 - + Args: token: 认证 Token prompt: 视频描述 @@ -226,10 +233,10 @@ class VideoService: resolution: 分辨率 stream: 是否流式 preset: 预设 - + Returns: AsyncGenerator，流式传输 - + Raises: UpstreamException: 连接失败时 """ @@ -238,11 +245,13 @@ class VideoService: 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, preset) - + payload = self._build_payload( + prompt, post_id, aspect_ratio, video_length, resolution, preset + ) + session = AsyncSession(impersonate=BROWSER) response = await session.post( CHAT_API, @@ -250,9 +259,9 @@ class VideoService: data=orjson.dumps(payload), timeout=self.timeout, stream=True, - proxies=self._build_proxies() + proxies=self._build_proxies(), ) - + if response.status_code != 200: logger.error(f"Video generation failed: {response.status_code}") try: @@ -261,9 +270,9 @@ class VideoService: pass raise UpstreamException( message=f"Video generation failed: {response.status_code}", - details={"status": response.status_code} + details={"status": response.status_code}, ) - + # Step 3: 流式传输 async def stream_response(): try: @@ -272,9 +281,9 @@ class VideoService: finally: if session: await session.close() - + return stream_response() - + except Exception as e: if session: try: @@ -285,7 +294,7 @@ class VideoService: if isinstance(e, AppException): raise e raise UpstreamException(f"Video generation error: {str(e)}") - + async def generate_from_image( self, token: str, @@ -295,11 +304,11 @@ class VideoService: video_length: int = 6, resolution: str = "SD", stream: bool = True, - preset: str = "normal" + preset: str = "normal", ) -> AsyncGenerator[bytes, None]: """ 从图片生成视频 - + Args: token: 认证 Token prompt: 视频描述 @@ -309,7 +318,7 @@ class VideoService: resolution: 分辨率 stream: 是否流式 preset: 预设 - + Returns: AsyncGenerator，流式传输 """ @@ -318,11 +327,13 @@ class VideoService: 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) - + payload = self._build_payload( + prompt, post_id, aspect_ratio, video_length, resolution, preset + ) + session = AsyncSession(impersonate=BROWSER) response = await session.post( CHAT_API, @@ -330,9 +341,9 @@ class VideoService: data=orjson.dumps(payload), timeout=self.timeout, stream=True, - proxies=self._build_proxies() + proxies=self._build_proxies(), ) - + if response.status_code != 200: logger.error(f"Video from image failed: {response.status_code}") try: @@ -341,9 +352,9 @@ class VideoService: pass raise UpstreamException( message=f"Video from image failed: {response.status_code}", - details={"status": response.status_code} + details={"status": response.status_code}, ) - + # Step 3: 流式传输 async def stream_response(): try: @@ -352,9 +363,9 @@ class VideoService: finally: if session: await session.close() - + return stream_response() - + except Exception as e: if session: try: @@ -365,7 +376,36 @@ class VideoService: 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: 模型名称 + """ + try: + async for chunk in stream: + yield chunk + finally: + 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}") + @staticmethod async def completions( model: str, @@ -375,11 +415,11 @@ class VideoService: aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "SD", - preset: str = "normal" + preset: str = "normal", ): """ 视频生成入口 - + Args: model: 模型名称 messages: 消息列表 @@ -389,7 +429,7 @@ class VideoService: video_length: 视频时长 resolution: 分辨率 preset: 预设模式 - + Returns: AsyncGenerator (流式) 或 dict (非流式) """ @@ -404,35 +444,35 @@ class VideoService: raise AppException( message="Internal service error obtaining token", error_type=ErrorType.SERVER.value, - code="internal_error" + 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 - ) - + 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", True) - + # 提取内容 from app.services.grok.chat import MessageExtractor from app.services.grok.assets import UploadService - + try: prompt, attachments = MessageExtractor.extract(messages, is_video=True) except ValueError as e: raise ValidationException(str(e)) - + # 处理图片附件 image_url = None if attachments: @@ -447,27 +487,50 @@ class VideoService: 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, stream, preset + token, + prompt, + image_url, + aspect_ratio, + video_length, + resolution, + stream, + preset, ) else: response = await service.generate( - token, prompt, - aspect_ratio, video_length, resolution, stream, preset + token, prompt, aspect_ratio, video_length, resolution, stream, preset ) - + # 处理响应 if is_stream: - return VideoStreamProcessor(model, token, think).process(response) + processor = VideoStreamProcessor(model, token, think) + return VideoService._wrap_stream( + processor.process(response), token_mgr, token, model + ) else: - return await VideoCollectProcessor(model, token).process(response) + 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 __all__ = ["VideoService"] diff --git a/app/services/grok/nsfw.py b/app/services/grok/nsfw.py new file mode 100644 index 0000000..507b2af --- /dev/null +++ b/app/services/grok/nsfw.py @@ -0,0 +1,112 @@ +""" +NSFW (Unhinged) 模式服务 + +使用 gRPC-Web 协议开启账号的 NSFW 功能。 +""" + +from dataclasses import dataclass +from typing import Optional + +from curl_cffi.requests import AsyncSession + +from app.core.config import get_config +from app.core.logger import logger +from app.services.grok.grpc_web import ( + encode_grpc_web_payload, + parse_grpc_web_response, + get_grpc_status, +) + + +NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" +BROWSER = "chrome136" +TIMEOUT = 30 + + +@dataclass +class NSFWResult: + """NSFW 操作结果""" + + success: bool + http_status: int + grpc_status: Optional[int] = None + 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", "") + + def _build_headers(self, token: str) -> dict: + """构造 gRPC-Web 请求头""" + token = token[4:] if token.startswith("sso=") else token + cf = get_config("grok.cf_clearance", "") + cookie = f"sso={token}; sso-rw={token}" + if cf: + cookie += f"; cf_clearance={cf}" + + return { + "accept": "*/*", + "content-type": "application/grpc-web+proto", + "origin": "https://grok.com", + "referer": "https://grok.com/", + "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36", + "x-grpc-web": "1", + "x-user-agent": "connect-es/2.1.1", + "cookie": cookie, + } + + @staticmethod + def _build_payload() -> bytes: + """构造请求 payload""" + # protobuf: enable_unhinged=true (field1=1, field2=1) + protobuf = bytes([0x08, 0x01, 0x10, 0x01]) + return encode_grpc_web_payload(protobuf) + + async def enable(self, token: str) -> NSFWResult: + """为单个 token 开启 NSFW 模式""" + headers = self._build_headers(token) + payload = self._build_payload() + proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None + + try: + async with AsyncSession(impersonate=BROWSER) as session: + response = await session.post( + NSFW_API, + data=payload, + headers=headers, + timeout=TIMEOUT, + proxies=proxies, + ) + + if response.status_code != 200: + return NSFWResult( + success=False, + http_status=response.status_code, + error=f"HTTP {response.status_code}", + ) + + # 解析 gRPC-Web 响应 + content_type = response.headers.get("content-type") + _, trailers = parse_grpc_web_response( + response.content, content_type=content_type + ) + + grpc_status = get_grpc_status(trailers) + + return NSFWResult( + success=grpc_status.ok, + http_status=response.status_code, + grpc_status=grpc_status.code, + grpc_message=grpc_status.message or None, + ) + + except Exception as e: + 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/token/models.py b/app/services/token/models.py index 5ed9751..6c06299 100644 --- a/app/services/token/models.py +++ b/app/services/token/models.py @@ -22,6 +22,7 @@ FAIL_THRESHOLD = 5 class TokenStatus(str, Enum): """Token 状态""" + ACTIVE = "active" DISABLED = "disabled" EXPIRED = "expired" @@ -30,7 +31,8 @@ class TokenStatus(str, Enum): class EffortType(str, Enum): """请求消耗类型""" - LOW = "low" # 扣 1 + + LOW = "low" # 扣 1 HIGH = "high" # 扣 4 @@ -42,122 +44,127 @@ EFFORT_COST = { class TokenInfo(BaseModel): """Token 信息""" - + token: str status: TokenStatus = TokenStatus.ACTIVE quota: int = DEFAULT_QUOTA - + # 统计 - created_at: int = Field(default_factory=lambda: int(datetime.now().timestamp() * 1000)) + created_at: int = Field( + default_factory=lambda: int(datetime.now().timestamp() * 1000) + ) last_used_at: Optional[int] = None use_count: int = 0 - + # 失败追踪 fail_count: int = 0 last_fail_at: Optional[int] = None last_fail_reason: Optional[str] = None - + # 冷却管理 last_sync_at: Optional[int] = None # 上次同步时间 - + # 扩展 tags: List[str] = Field(default_factory=list) note: str = "" last_asset_clear_at: Optional[int] = None - + def is_available(self) -> bool: """检查是否可用（状态正常且配额 > 0）""" return self.status == TokenStatus.ACTIVE and self.quota > 0 - + def consume(self, effort: EffortType = EffortType.LOW) -> int: """ 消耗配额 - + Args: - effort: LOW 扣 1，HIGH 扣 4 - + effort: LOW 扣 1 配额并计 1 次，HIGH 扣 4 配额并计 4 次 + Returns: 实际扣除的配额 """ cost = EFFORT_COST[effort] actual_cost = min(cost, self.quota) - + self.last_used_at = int(datetime.now().timestamp() * 1000) - self.use_count += 1 - self.quota = max(0, self.quota - cost) - + self.use_count += actual_cost # 使用 actual_cost 避免配额不足时过度计数 + self.quota = max(0, self.quota - actual_cost) + # 成功消耗后清空失败计数 self.fail_count = 0 self.last_fail_reason = None - + if self.quota == 0: self.status = TokenStatus.COOLING elif self.status in [TokenStatus.COOLING, TokenStatus.EXPIRED]: self.status = TokenStatus.ACTIVE - + return actual_cost - + def update_quota(self, new_quota: int): """ 更新配额（用于 API 同步） - + Args: new_quota: 新的配额值 """ self.quota = max(0, new_quota) - + if self.quota == 0: self.status = TokenStatus.COOLING - elif self.quota > 0 and self.status in [TokenStatus.COOLING, TokenStatus.EXPIRED]: + elif self.quota > 0 and self.status in [ + TokenStatus.COOLING, + TokenStatus.EXPIRED, + ]: self.status = TokenStatus.ACTIVE - + def reset(self): """重置配额到默认值""" self.quota = DEFAULT_QUOTA self.status = TokenStatus.ACTIVE self.fail_count = 0 self.last_fail_reason = None - + def record_fail(self, status_code: int = 401, reason: str = ""): """记录失败，达到阈值后自动标记为 expired""" # 仅 401 错误才计入失败 if status_code != 401: return - + self.fail_count += 1 self.last_fail_at = int(datetime.now().timestamp() * 1000) self.last_fail_reason = reason - + if self.fail_count >= FAIL_THRESHOLD: self.status = TokenStatus.EXPIRED - + def record_success(self, is_usage: bool = True): """记录成功，清空失败计数并根据配额更新状态""" self.fail_count = 0 self.last_fail_at = None self.last_fail_reason = None - + if is_usage: self.use_count += 1 self.last_used_at = int(datetime.now().timestamp() * 1000) - + if self.quota == 0: self.status = TokenStatus.COOLING else: self.status = TokenStatus.ACTIVE - + def need_refresh(self, interval_hours: int = 8) -> bool: """检查是否需要刷新配额""" if self.status != TokenStatus.COOLING: return False - + if self.last_sync_at is None: return True - + now = int(datetime.now().timestamp() * 1000) interval_ms = interval_hours * 3600 * 1000 return (now - self.last_sync_at) >= interval_ms - + def mark_synced(self): """标记已同步""" self.last_sync_at = int(datetime.now().timestamp() * 1000) @@ -165,6 +172,7 @@ class TokenInfo(BaseModel): class TokenPoolStats(BaseModel): """Token 池统计""" + total: int = 0 active: int = 0 disabled: int = 0 diff --git a/app/static/token/token.css b/app/static/token/token.css index a36f1f6..e657b26 100644 --- a/app/static/token/token.css +++ b/app/static/token/token.css @@ -278,3 +278,42 @@ .confirm-dialog .modal-content { padding: 0; } + + /* Tab Styles (Geist Design System) */ + .tab-item { + position: relative; + padding: 12px 0; + margin-right: 20px; + font-size: 14px; + font-weight: 500; + color: var(--accents-5); + background: transparent; + border: none; + box-shadow: inset 0 -2px 0 0 transparent; + cursor: pointer; + transition: color 0.15s ease, box-shadow 0.15s ease; + white-space: nowrap; + } + + .tab-item:hover { + color: var(--geist-foreground, #000); + } + + .tab-item.active { + color: var(--geist-foreground, #000); + box-shadow: inset 0 -2px 0 0 var(--geist-foreground, #000); + } + + .tab-item:disabled { + cursor: not-allowed; + opacity: 0.5; + } + + /* Hide scrollbar for tabs on mobile */ + .hide-scrollbar { + -ms-overflow-style: none; + scrollbar-width: none; + } + .hide-scrollbar::-webkit-scrollbar { + display: none; + } diff --git a/app/static/token/token.html b/app/static/token/token.html index 342fff6..7aad8f2 100644 --- a/app/static/token/token.html +++ b/app/static/token/token.html @@ -89,6 +89,27 @@ + +
+ + + + +
+
@@ -137,6 +158,10 @@ class="geist-button-outline text-xs px-3 gap-1 border-0 hover:bg-gray-100 justify-center"> 刷新 + diff --git a/app/static/token/token.js b/app/static/token/token.js index f4a7bc9..5de918d 100644 --- a/app/static/token/token.js +++ b/app/static/token/token.js @@ -7,6 +7,7 @@ let batchQueue = []; let batchTotal = 0; let batchProcessed = 0; let currentBatchAction = null; +let currentFilter = 'all'; const BATCH_SIZE = 50; async function init() { @@ -107,19 +108,27 @@ function renderTable() { tbody.innerHTML = ''; loading.classList.add('hidden'); - if (flatTokens.length === 0) { + // 更新 Tab 计数 + updateTabCounts(); + + // 获取筛选后的列表 + const filteredTokens = getFilteredTokens(); + + if (filteredTokens.length === 0) { emptyState.classList.remove('hidden'); return; } emptyState.classList.add('hidden'); - flatTokens.forEach((item, index) => { + filteredTokens.forEach((item) => { + // 获取原始索引用于操作 + const originalIndex = flatTokens.indexOf(item); const tr = document.createElement('tr'); // Checkbox (Center) const tdCheck = document.createElement('td'); tdCheck.className = 'text-center'; - tdCheck.innerHTML = ``; + tdCheck.innerHTML = ``; // Token (Left) const tdToken = document.createElement('td'); @@ -168,10 +177,10 @@ function renderTable() { - - @@ -195,7 +204,8 @@ function renderTable() { function toggleSelectAll() { const checkbox = document.getElementById('select-all'); const checked = checkbox.checked; - flatTokens.forEach(t => t._selected = checked); + // 只选择当前筛选后可见的 Token + getFilteredTokens().forEach(t => t._selected = checked); renderTable(); } @@ -734,6 +744,105 @@ function escapeHtml(text) { .replace(/'/g, "'"); } +// ========== Tab 筛选功能 ========== + +function filterByStatus(status) { + currentFilter = status; + + // 更新 Tab 样式和 ARIA + document.querySelectorAll('.tab-item').forEach(tab => { + const isActive = tab.dataset.filter === status; + tab.classList.toggle('active', isActive); + tab.setAttribute('aria-selected', isActive ? 'true' : 'false'); + }); + + renderTable(); +} + +function getFilteredTokens() { + if (currentFilter === 'all') return flatTokens; + + return flatTokens.filter(t => { + if (currentFilter === 'active') return t.status === 'active'; + if (currentFilter === 'cooling') return t.status === 'cooling'; + if (currentFilter === 'expired') return t.status !== 'active' && t.status !== 'cooling'; + return true; + }); +} + +function updateTabCounts() { + const counts = { + all: flatTokens.length, + active: flatTokens.filter(t => t.status === 'active').length, + cooling: flatTokens.filter(t => t.status === 'cooling').length, + expired: flatTokens.filter(t => t.status !== 'active' && t.status !== 'cooling').length + }; + + Object.entries(counts).forEach(([key, count]) => { + const el = document.getElementById(`tab-count-${key}`); + if (el) el.textContent = count; + }); +} + +// ========== NSFW 批量开启 ========== + +async function batchEnableNSFW() { + if (isBatchProcessing) { + showToast('当前有任务进行中', 'info'); + return; + } + + const selected = flatTokens.filter(t => t._selected); + const targetCount = selected.length || flatTokens.length; + const msg = selected.length === 0 + ? `是否对全部 ${targetCount} 个 Token 开启 NSFW 模式？` + : `是否为选中的 ${selected.length} 个 Token 开启 NSFW 模式？`; + + const ok = await confirmAction(msg, { okText: '开启 NSFW' }); + if (!ok) return; + + // 禁用按钮 + const btn = document.getElementById('btn-batch-nsfw'); + if (btn) btn.disabled = true; + + isBatchProcessing = true; + currentBatchAction = 'nsfw'; + updateBatchProgress(); + + try { + const tokens = selected.length > 0 ? selected.map(t => t.token) : null; + + const res = await fetch('/api/v1/admin/tokens/nsfw/enable', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...buildAuthHeaders(apiKey) + }, + body: JSON.stringify({ tokens }) + }); + + const data = await res.json(); + + if (res.ok && data.status === 'success') { + const { ok: okCount, fail } = data.summary; + showToast( + `NSFW 开启完成：成功 ${okCount}，失败 ${fail}`, + fail > 0 ? 'warning' : 'success' + ); + } else { + showToast('开启失败: ' + (data.detail || '未知错误'), 'error'); + } + } catch (e) { + showToast('请求错误: ' + e.message, 'error'); + } finally { + isBatchProcessing = false; + currentBatchAction = null; + if (btn) btn.disabled = false; + updateBatchProgress(); + setActionButtonsState(); + } +} + window.onload = init; diff --git a/config.defaults.toml b/config.defaults.toml index 58cb2cf..8f299fc 100644 --- a/config.defaults.toml +++ b/config.defaults.toml @@ -35,3 +35,6 @@ media_max_concurrent = 50 usage_max_concurrent = 25 assets_delete_batch_size = 10 admin_assets_batch_size = 10 +nsfw_max_concurrent = 10 +admin_nsfw_batch_size = 50 +nsfw_max_tokens = 1000