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}")