diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 2e8651d..5163485 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -78,7 +78,9 @@ async def stream_batch(task_id: str, request: Request): return StreamingResponse(event_stream(), media_type="text/event-stream") -@router.post("/api/v1/admin/batch/{task_id}/cancel", dependencies=[Depends(verify_api_key)]) +@router.post( + "/api/v1/admin/batch/{task_id}/cancel", dependencies=[Depends(verify_api_key)] +) async def cancel_batch(task_id: str): task = get_task(task_id) if not task: @@ -221,7 +223,10 @@ async def refresh_tokens_api(data: dict): ) raw_results = await run_in_batches( - unique_tokens, _refresh_one, max_concurrent=max_concurrent, batch_size=batch_size + unique_tokens, + _refresh_one, + max_concurrent=max_concurrent, + batch_size=batch_size, ) results = {} @@ -283,6 +288,7 @@ async def refresh_tokens_api_async(data: dict): async def _run(): try: + async def _refresh_one(t: str): return await mgr.sync_usage( t, "grok-3", consume_on_fail=False, is_usage=False @@ -444,7 +450,9 @@ async def enable_nsfw_api(data: dict): # 添加截断提示 if truncated: - response["warning"] = f"数量超出限制，仅处理前 {max_tokens} 个（共 {original_count} 个）" + response["warning"] = ( + f"数量超出限制，仅处理前 {max_tokens} 个（共 {original_count} 个）" + ) return response @@ -506,6 +514,7 @@ async def enable_nsfw_api_async(data: dict): async def _run(): try: + async def _enable(token: str): result = await nsfw_service.enable(token) if result.success: @@ -1162,6 +1171,7 @@ async def clear_online_cache_api_async(data: dict): async def _run(): delete_service = DeleteService() try: + async def _clear_one(t: str): try: result = await delete_service.delete_all(t) diff --git a/app/api/v1/chat.py b/app/api/v1/chat.py index 6aa5fef..4866663 100644 --- a/app/api/v1/chat.py +++ b/app/api/v1/chat.py @@ -24,6 +24,7 @@ USER_CONTENT_TYPES = ["text", "image_url", "input_audio", "file"] class MessageItem(BaseModel): """消息项""" + role: str content: Union[str, List[Dict[str, Any]]] tool_call_id: Optional[str] = None # tool 角色需要的字段 @@ -43,11 +44,14 @@ class MessageItem(BaseModel): class VideoConfig(BaseModel): """视频生成配置""" - aspect_ratio: Optional[str] = Field("3:2", description="视频比例: 3:2, 16:9, 1:1 等") + + aspect_ratio: Optional[str] = Field( + "3:2", description="视频比例: 3:2, 16:9, 1:1 等" + ) video_length: Optional[int] = Field(6, description="视频时长(秒): 5-15") resolution: Optional[str] = Field("SD", description="视频分辨率: SD, HD") preset: Optional[str] = Field("custom", description="风格预设: fun, normal, spicy") - + @field_validator("aspect_ratio") @classmethod def validate_aspect_ratio(cls, v): @@ -56,10 +60,10 @@ class VideoConfig(BaseModel): raise ValidationException( message=f"aspect_ratio must be one of {allowed}", param="video_config.aspect_ratio", - code="invalid_aspect_ratio" + code="invalid_aspect_ratio", ) return v - + @field_validator("video_length") @classmethod def validate_video_length(cls, v): @@ -68,7 +72,7 @@ class VideoConfig(BaseModel): raise ValidationException( message="video_length must be between 5 and 15 seconds", param="video_config.video_length", - code="invalid_video_length" + code="invalid_video_length", ) return v @@ -80,10 +84,10 @@ class VideoConfig(BaseModel): raise ValidationException( message=f"resolution must be one of {allowed}", param="video_config.resolution", - code="invalid_resolution" + code="invalid_resolution", ) return v - + @field_validator("preset") @classmethod def validate_preset(cls, v): @@ -92,16 +96,17 @@ class VideoConfig(BaseModel): return "custom" allowed = ["fun", "normal", "spicy", "custom"] if v not in allowed: - raise ValidationException( + raise ValidationException( message=f"preset must be one of {allowed}", param="video_config.preset", - code="invalid_preset" - ) + code="invalid_preset", + ) return v class ChatCompletionRequest(BaseModel): """Chat Completions 请求""" + model: str = Field(..., description="模型名称") messages: List[MessageItem] = Field(..., description="消息数组") stream: Optional[bool] = Field(None, description="是否流式输出") @@ -132,9 +137,7 @@ class ChatCompletionRequest(BaseModel): f"Invalid stream value type '{type(v).__name__}'. Must be a boolean or string." ) - model_config = { - "extra": "ignore" - } + model_config = {"extra": "ignore"} def validate_request(request: ChatCompletionRequest): @@ -144,65 +147,69 @@ def validate_request(request: ChatCompletionRequest): raise ValidationException( message=f"The model `{request.model}` does not exist or you do not have access to it.", param="model", - code="model_not_found" + code="model_not_found", ) - + # 验证消息 for idx, msg in enumerate(request.messages): content = msg.content - + # 字符串内容 if isinstance(content, str): if not content.strip(): raise ValidationException( message="Message content cannot be empty", param=f"messages.{idx}.content", - code="empty_content" + code="empty_content", ) - + # 列表内容 elif isinstance(content, list): if not content: raise ValidationException( message="Message content cannot be an empty array", param=f"messages.{idx}.content", - code="empty_content" + code="empty_content", ) - + for block_idx, block in enumerate(content): # 检查空对象 if not block: raise ValidationException( message="Content block cannot be empty", param=f"messages.{idx}.content.{block_idx}", - code="empty_block" + code="empty_block", ) - + # 检查 type 字段 if "type" not in block: raise ValidationException( message="Content block must have a 'type' field", param=f"messages.{idx}.content.{block_idx}", - code="missing_type" + code="missing_type", ) - + block_type = block.get("type") - + # 检查 type 空值 - if not block_type or not isinstance(block_type, str) or not block_type.strip(): + if ( + not block_type + or not isinstance(block_type, str) + or not block_type.strip() + ): raise ValidationException( message="Content block 'type' cannot be empty", param=f"messages.{idx}.content.{block_idx}.type", - code="empty_type" + code="empty_type", ) - + # 验证 type 有效性 if msg.role == "user": if block_type not in USER_CONTENT_TYPES: raise ValidationException( message=f"Invalid content block type: '{block_type}'", param=f"messages.{idx}.content.{block_idx}.type", - code="invalid_type" + code="invalid_type", ) elif msg.role in ("tool", "function"): # tool/function 角色只支持 text 类型，但内容可以是 JSON 字符串 @@ -210,15 +217,15 @@ def validate_request(request: ChatCompletionRequest): raise ValidationException( message=f"The `{msg.role}` role only supports 'text' type, got '{block_type}'", param=f"messages.{idx}.content.{block_idx}.type", - code="invalid_type" + code="invalid_type", ) elif block_type != "text": raise ValidationException( message=f"The `{msg.role}` role only supports 'text' type, got '{block_type}'", param=f"messages.{idx}.content.{block_idx}.type", - code="invalid_type" + code="invalid_type", ) - + # 验证字段是否存在 & 非空 if block_type == "text": text = block.get("text", "") @@ -226,15 +233,17 @@ def validate_request(request: ChatCompletionRequest): raise ValidationException( message="Text content cannot be empty", param=f"messages.{idx}.content.{block_idx}.text", - code="empty_text" + code="empty_text", ) elif block_type == "image_url": image_url = block.get("image_url") - if not image_url or not (isinstance(image_url, dict) and image_url.get("url")): + if not image_url or not ( + isinstance(image_url, dict) and image_url.get("url") + ): raise ValidationException( message="image_url must have a 'url' field", param=f"messages.{idx}.content.{block_idx}.image_url", - code="missing_url" + code="missing_url", ) @@ -252,10 +261,10 @@ async def chat_completions(request: ChatCompletionRequest): model_info = ModelService.get(request.model) if model_info and model_info.is_video: from app.services.grok.media import VideoService - + # 提取视频配置 (默认值在 Pydantic 模型中处理) v_conf = request.video_config or VideoConfig() - + result = await VideoService.completions( model=request.model, messages=[msg.model_dump() for msg in request.messages], @@ -264,23 +273,23 @@ async def chat_completions(request: ChatCompletionRequest): aspect_ratio=v_conf.aspect_ratio, video_length=v_conf.video_length, resolution=v_conf.resolution, - preset=v_conf.preset + preset=v_conf.preset, ) else: result = await ChatService.completions( model=request.model, messages=[msg.model_dump() for msg in request.messages], stream=request.stream, - thinking=request.thinking + thinking=request.thinking, ) - + if isinstance(result, dict): return JSONResponse(content=result) else: return StreamingResponse( result, media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive"} + headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) diff --git a/app/api/v1/files.py b/app/api/v1/files.py index 9a984bb..46e7213 100644 --- a/app/api/v1/files.py +++ b/app/api/v1/files.py @@ -24,9 +24,9 @@ async def get_image(filename: str): """ if "/" in filename: filename = filename.replace("/", "-") - + file_path = IMAGE_DIR / filename - + if await aiofiles.os.path.exists(file_path): if await aiofiles.os.path.isfile(file_path): content_type = "image/jpeg" @@ -34,14 +34,12 @@ async def get_image(filename: str): content_type = "image/png" elif file_path.suffix.lower() == ".webp": content_type = "image/webp" - + # 增加缓存头，支持高并发场景下的浏览器/CDN缓存 return FileResponse( - file_path, + file_path, media_type=content_type, - headers={ - "Cache-Control": "public, max-age=31536000, immutable" - } + headers={"Cache-Control": "public, max-age=31536000, immutable"}, ) logger.warning(f"Image not found: {filename}") @@ -55,17 +53,15 @@ async def get_video(filename: str): """ if "/" in filename: filename = filename.replace("/", "-") - + file_path = VIDEO_DIR / filename - + if await aiofiles.os.path.exists(file_path): if await aiofiles.os.path.isfile(file_path): return FileResponse( - file_path, + file_path, media_type="video/mp4", - headers={ - "Cache-Control": "public, max-age=31536000, immutable" - } + headers={"Cache-Control": "public, max-age=31536000, immutable"}, ) logger.warning(f"Video not found: {filename}") diff --git a/app/api/v1/video.py b/app/api/v1/video.py index 95699e1..c16e7a0 100644 --- a/app/api/v1/video.py +++ b/app/api/v1/video.py @@ -1,3 +1,3 @@ """ TODO：Video Generation API 路由 -""" \ No newline at end of file +""" diff --git a/app/core/auth.py b/app/core/auth.py index 8dd5b4e..9a07f5a 100644 --- a/app/core/auth.py +++ b/app/core/auth.py @@ -1,8 +1,9 @@ """ API 认证模块 """ + from typing import Optional -from fastapi import Depends, HTTPException, status, Security +from fastapi import HTTPException, status, Security from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from app.core.config import get_config @@ -11,48 +12,48 @@ from app.core.config import get_config security = HTTPBearer( auto_error=False, scheme_name="API Key", - description="Enter your API Key in the format: Bearer " + description="Enter your API Key in the format: Bearer ", ) async def verify_api_key( - auth: Optional[HTTPAuthorizationCredentials] = Security(security) + auth: Optional[HTTPAuthorizationCredentials] = Security(security), ) -> Optional[str]: """ 验证 Bearer Token - + 如果 config.toml 中未配置 api_key，则跳过验证。 如果配置了 api_key，则必须提供正确的 Authorization: Bearer 。 """ api_key = get_config("app.api_key", "") - + # 如果未配置 API Key，直接放行 if not api_key: return None - + if not auth: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing authentication token", headers={"WWW-Authenticate": "Bearer"}, ) - + if auth.credentials != api_key: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authentication token", headers={"WWW-Authenticate": "Bearer"}, ) - + return auth.credentials async def verify_app_key( - auth: Optional[HTTPAuthorizationCredentials] = Security(security) + auth: Optional[HTTPAuthorizationCredentials] = Security(security), ) -> Optional[str]: """ 验证后台登录密钥（app_key）。 - + 如果未配置 app_key，则跳过验证。 """ app_key = get_config("app.app_key", "") diff --git a/app/core/batch_tasks.py b/app/core/batch_tasks.py index 07038ff..ff564ff 100644 --- a/app/core/batch_tasks.py +++ b/app/core/batch_tasks.py @@ -54,7 +54,9 @@ class BatchTask: # Drop if queue is full or closed pass - def record(self, ok: bool, *, item: Any = None, detail: Any = None, error: str = "") -> None: + def record( + self, ok: bool, *, item: Any = None, detail: Any = None, error: str = "" + ) -> None: self.processed += 1 if ok: self.ok += 1 @@ -148,4 +150,3 @@ def delete_task(task_id: str) -> None: async def expire_task(task_id: str, delay: int = 300) -> None: await asyncio.sleep(delay) delete_task(task_id) - diff --git a/app/core/config.py b/app/core/config.py index 401d512..641b1cb 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -46,10 +46,10 @@ def _load_defaults() -> Dict[str, Any]: class Config: """配置管理器""" - + _instance = None _config = {} - + def __init__(self): self._config = {} self._defaults = {} @@ -65,13 +65,13 @@ class Config: """显式加载配置""" try: from app.core.storage import get_storage, LocalStorage - + self._ensure_defaults() storage = get_storage() config_data = await storage.load_config() from_remote = True - + # 从本地 data/config.toml 初始化后端 if config_data is None: local_storage = LocalStorage() @@ -104,7 +104,7 @@ class Config: def get(self, key: str, default: Any = None) -> Any: """ 获取配置值 - + Args: key: 配置键，格式 "section.key" default: 默认值 @@ -115,12 +115,13 @@ class Config: return self._config.get(section, {}).get(attr, default) except (ValueError, AttributeError): return default - + return self._config.get(key, default) async def update(self, new_config: dict): """更新配置""" from app.core.storage import get_storage + storage = get_storage() async with storage.acquire_lock("config_save", timeout=10): self._ensure_defaults() diff --git a/app/core/exceptions.py b/app/core/exceptions.py index a9ac7c8..9ceff00 100644 --- a/app/core/exceptions.py +++ b/app/core/exceptions.py @@ -2,7 +2,7 @@ 全局异常处理 - OpenAI 兼容错误格式 """ -from typing import Any, Optional +from typing import Any from enum import Enum from fastapi import Request, HTTPException from fastapi.responses import JSONResponse @@ -13,8 +13,10 @@ from app.core.logger import logger # ============= 错误类型 ============= + class ErrorType(str, Enum): """OpenAI 错误类型""" + INVALID_REQUEST = "invalid_request_error" AUTHENTICATION = "authentication_error" PERMISSION = "permission_error" @@ -26,35 +28,32 @@ class ErrorType(str, Enum): # ============= 辅助函数 ============= + def error_response( message: str, error_type: str = ErrorType.INVALID_REQUEST.value, param: str = None, - code: str = None + code: str = None, ) -> dict: """构建 OpenAI 错误响应""" return { - "error": { - "message": message, - "type": error_type, - "param": param, - "code": code - } + "error": {"message": message, "type": error_type, "param": param, "code": code} } # ============= 异常类 ============= + class AppException(Exception): """应用基础异常""" - + def __init__( self, message: str, error_type: str = ErrorType.SERVER.value, code: str = None, param: str = None, - status_code: int = 500 + status_code: int = 500, ): self.message = message self.error_type = error_type @@ -66,56 +65,57 @@ class AppException(Exception): class ValidationException(AppException): """验证错误""" - + def __init__(self, message: str, param: str = None, code: str = None): super().__init__( message=message, error_type=ErrorType.INVALID_REQUEST.value, code=code or "invalid_value", param=param, - status_code=400 + status_code=400, ) class AuthenticationException(AppException): """认证错误""" - + def __init__(self, message: str = "Invalid API key"): super().__init__( message=message, error_type=ErrorType.AUTHENTICATION.value, code="invalid_api_key", - status_code=401 + status_code=401, ) class UpstreamException(AppException): """上游服务错误""" - + def __init__(self, message: str, details: Any = None): super().__init__( message=message, error_type=ErrorType.SERVER.value, code="upstream_error", - status_code=502 + status_code=502, ) self.details = details # ============= 异常处理器 ============= + async def app_exception_handler(request: Request, exc: AppException) -> JSONResponse: """处理应用异常""" logger.warning(f"AppException: {exc.error_type} - {exc.message}") - + return JSONResponse( status_code=exc.status_code, content=error_response( message=exc.message, error_type=exc.error_type, param=exc.param, - code=exc.code - ) + code=exc.code, + ), ) @@ -129,7 +129,7 @@ async def http_exception_handler(request: Request, exc: HTTPException) -> JSONRe 429: ErrorType.RATE_LIMIT.value, } error_type = type_map.get(exc.status_code, ErrorType.SERVER.value) - + # 默认 code 映射 code_map = { 401: "invalid_api_key", @@ -138,69 +138,72 @@ async def http_exception_handler(request: Request, exc: HTTPException) -> JSONRe 429: "rate_limit_exceeded", } code = code_map.get(exc.status_code, None) - + logger.warning(f"HTTPException: {exc.status_code} - {exc.detail}") - + return JSONResponse( status_code=exc.status_code, content=error_response( - message=str(exc.detail), - error_type=error_type, - code=code - ) + message=str(exc.detail), error_type=error_type, code=code + ), ) -async def validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse: +async def validation_exception_handler( + request: Request, exc: RequestValidationError +) -> JSONResponse: """处理验证错误""" errors = exc.errors() - + if errors: first = errors[0] loc = first.get("loc", []) msg = first.get("msg", "Invalid request") code = first.get("type", "invalid_value") - + # JSON 解析错误 if code == "json_invalid" or "JSON" in msg: message = "Invalid JSON in request body. Please check for trailing commas or syntax errors." param = "body" else: - param_parts = [str(x) for x in loc if not (isinstance(x, int) or str(x).isdigit())] + param_parts = [ + str(x) for x in loc if not (isinstance(x, int) or str(x).isdigit()) + ] param = ".".join(param_parts) if param_parts else None message = msg else: param, message, code = None, "Invalid request", "invalid_value" - + logger.warning(f"ValidationError: {param} - {message}") - + return JSONResponse( status_code=400, content=error_response( message=message, error_type=ErrorType.INVALID_REQUEST.value, param=param, - code=code - ) + code=code, + ), ) async def generic_exception_handler(request: Request, exc: Exception) -> JSONResponse: """处理未捕获异常""" logger.exception(f"Unhandled: {type(exc).__name__}: {str(exc)}") - + return JSONResponse( status_code=500, content=error_response( message="Internal server error", error_type=ErrorType.SERVER.value, - code="internal_error" - ) + code="internal_error", + ), ) # ============= 注册 ============= + def register_exception_handlers(app): """注册异常处理器""" app.add_exception_handler(AppException, app_exception_handler) diff --git a/app/core/logger.py b/app/core/logger.py index bb705d6..41828f7 100644 --- a/app/core/logger.py +++ b/app/core/logger.py @@ -20,42 +20,46 @@ def _format_json(record) -> str: tz = record["time"].strftime("%z") if tz: time_str += tz[:3] + ":" + tz[3:] - + log_entry = { "time": time_str, "level": record["level"].name.lower(), "msg": record["message"], "caller": f"{record['file'].name}:{record['line']}", } - + # trace 上下文 extra = record["extra"] if extra.get("traceID"): log_entry["traceID"] = extra["traceID"] if extra.get("spanID"): log_entry["spanID"] = extra["spanID"] - + # 其他 extra 字段 for key, value in extra.items(): if key not in ("traceID", "spanID") and not key.startswith("_"): log_entry[key] = value - + # 错误及以上级别添加堆栈跟踪 if record["level"].no >= 40 and record["exception"]: - log_entry["stacktrace"] = "".join(traceback.format_exception( - record["exception"].type, - record["exception"].value, - record["exception"].traceback - )) - + log_entry["stacktrace"] = "".join( + traceback.format_exception( + record["exception"].type, + record["exception"].value, + record["exception"].traceback, + ) + ) + return json.dumps(log_entry, ensure_ascii=False) def _make_json_sink(output): """创建 JSON sink""" + def sink(message): json_str = _format_json(message.record) print(json_str, file=output, flush=True) + return sink @@ -75,7 +79,7 @@ def setup_logging( ): """设置日志配置""" logger.remove() - + # 控制台输出 if json_console: logger.add( @@ -91,7 +95,7 @@ def setup_logging( format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {file.name}:{line} - {message}", colorize=True, ) - + # 文件输出 if file_logging: logger.add( @@ -100,7 +104,7 @@ def setup_logging( format="{message}", enqueue=True, ) - + return logger diff --git a/app/core/response_middleware.py b/app/core/response_middleware.py index 64a4711..2cfa8b3 100644 --- a/app/core/response_middleware.py +++ b/app/core/response_middleware.py @@ -9,39 +9,39 @@ import time import uuid from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request -from starlette.types import ASGIApp from app.core.logger import logger + class ResponseLoggerMiddleware(BaseHTTPMiddleware): """ 请求日志/响应追踪中间件 Request Logging and Response Tracking Middleware """ - + async def dispatch(self, request: Request, call_next): # 生成请求 ID trace_id = str(uuid.uuid4()) request.state.trace_id = trace_id - + start_time = time.time() - + # 记录请求信息 logger.info( f"Request: {request.method} {request.url.path}", extra={ "traceID": trace_id, "method": request.method, - "path": request.url.path - } + "path": request.url.path, + }, ) - + try: response = await call_next(request) - + # 计算耗时 duration = (time.time() - start_time) * 1000 - + # 记录响应信息 logger.info( f"Response: {request.method} {request.url.path} - {response.status_code} ({duration:.2f}ms)", @@ -50,12 +50,12 @@ class ResponseLoggerMiddleware(BaseHTTPMiddleware): "method": request.method, "path": request.url.path, "status": response.status_code, - "duration_ms": round(duration, 2) - } + "duration_ms": round(duration, 2), + }, ) - + return response - + except Exception as e: duration = (time.time() - start_time) * 1000 logger.error( @@ -65,7 +65,7 @@ class ResponseLoggerMiddleware(BaseHTTPMiddleware): "method": request.method, "path": request.url.path, "duration_ms": round(duration, 2), - "error": str(e) - } + "error": str(e), + }, ) raise e diff --git a/app/core/storage.py b/app/core/storage.py index 0cd3590..cb23aff 100644 --- a/app/core/storage.py +++ b/app/core/storage.py @@ -12,13 +12,13 @@ import abc import os import asyncio -import os import hashlib import time import tomllib from typing import Any, Dict, Optional from pathlib import Path from enum import Enum + try: import fcntl except ImportError: # pragma: no cover - non-posix platforms @@ -34,17 +34,22 @@ CONFIG_FILE = Path(__file__).parent.parent.parent / "data" / "config.toml" TOKEN_FILE = Path(__file__).parent.parent.parent / "data" / "token.json" LOCK_DIR = Path(__file__).parent.parent.parent / "data" / ".locks" + # JSON 序列化优化助手函数 def json_dumps(obj: Any) -> str: return orjson.dumps(obj).decode("utf-8") + def json_loads(obj: str | bytes) -> Any: return orjson.loads(obj) + class StorageError(Exception): """存储服务基础异常""" + pass + class BaseStorage(abc.ABC): """存储基类""" @@ -72,13 +77,13 @@ class BaseStorage(abc.ABC): async def close(self): """关闭资源""" pass - + @asynccontextmanager async def acquire_lock(self, name: str, timeout: int = 10): """ 获取锁 (互斥访问) 用于读写操作的临界区保护 - + Args: name: 锁名称 timeout: 超时时间 (秒) @@ -98,10 +103,10 @@ class LocalStorage(BaseStorage): - 使用 asyncio.Lock 进行进程内并发控制 - 如果需要多进程安全，需要系统级文件锁 (fcntl) """ - + def __init__(self): self._lock = asyncio.Lock() - + @asynccontextmanager async def acquire_lock(self, name: str, timeout: int = 10): if fcntl is None: @@ -163,7 +168,8 @@ class LocalStorage(BaseStorage): try: lines = [] for section, items in data.items(): - if not isinstance(items, dict): continue + if not isinstance(items, dict): + continue lines.append(f"[{section}]") for key, val in items.items(): if isinstance(val, bool): @@ -179,9 +185,9 @@ class LocalStorage(BaseStorage): val_str = f'"{str(val)}"' lines.append(f"{key} = {val_str}") lines.append("") - + content = "\n".join(lines) - + CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True) async with aiofiles.open(CONFIG_FILE, "w", encoding="utf-8") as f: await f.write(content) @@ -203,15 +209,15 @@ class LocalStorage(BaseStorage): async def save_tokens(self, data: Dict[str, Any]): try: TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True) - temp_path = TOKEN_FILE.with_suffix('.tmp') - + temp_path = TOKEN_FILE.with_suffix(".tmp") + # 原子写操作: 写入临时文件 -> 重命名 async with aiofiles.open(temp_path, "wb") as f: await f.write(orjson.dumps(data, option=orjson.OPT_INDENT_2)) - + # 使用 os.replace 保证原子性 os.replace(temp_path, TOKEN_FILE) - + except Exception as e: logger.error(f"LocalStorage: 保存 Token 失败: {e}") raise StorageError(f"保存 Token 失败: {e}") @@ -227,25 +233,22 @@ class RedisStorage(BaseStorage): - 支持分布式锁 (redis.lock) - 扁平化数据结构优化性能 """ - + def __init__(self, url: str): try: from redis import asyncio as aioredis - from redis.asyncio.lock import Lock except ImportError: raise ImportError("需要安装 redis 包: pip install redis") - + # 显式配置连接池 # 使用 decode_responses=True 简化字符串处理，但在处理复杂对象时使用 orjson self.redis = aioredis.from_url( - url, - decode_responses=True, - health_check_interval=30 + url, decode_responses=True, health_check_interval=30 ) - self.config_key = "grok2api:config" # Hash: section.key -> value_json - self.key_pools = "grok2api:pools" # Set: pool_names - self.prefix_pool_set = "grok2api:pool:" # Set: pool -> token_ids - self.prefix_token_hash = "grok2api:token:"# Hash: token_id -> token_data + self.config_key = "grok2api:config" # Hash: section.key -> value_json + self.key_pools = "grok2api:pools" # Set: pool_names + self.prefix_pool_set = "grok2api:pool:" # Set: pool -> token_ids + self.prefix_token_hash = "grok2api:token:" # Hash: token_id -> token_data self.lock_prefix = "grok2api:lock:" @asynccontextmanager @@ -279,17 +282,19 @@ class RedisStorage(BaseStorage): raw_data = await self.redis.hgetall(self.config_key) if not raw_data: return None - + config = {} for composite_key, val_str in raw_data.items(): - if "." not in composite_key: continue + if "." not in composite_key: + continue section, key = composite_key.split(".", 1) - - if section not in config: config[section] = {} - + + if section not in config: + config[section] = {} + try: val = json_loads(val_str) - except: + except Exception: val = val_str config[section][key] = val return config @@ -299,15 +304,17 @@ class RedisStorage(BaseStorage): async def save_config(self, data: Dict[str, Any]): """保存配置到 Redis Hash""" - if not data: return + if not data: + return try: mapping = {} for section, items in data.items(): - if not isinstance(items, dict): continue + if not isinstance(items, dict): + continue for key, val in items.items(): composite_key = f"{section}.{key}" mapping[composite_key] = json_dumps(val) - + if mapping: await self.redis.hset(self.config_key, mapping=mapping) except Exception as e: @@ -318,24 +325,25 @@ class RedisStorage(BaseStorage): """加载所有 Token""" try: pool_names = await self.redis.smembers(self.key_pools) - if not pool_names: return None - + if not pool_names: + return None + pools = {} async with self.redis.pipeline() as pipe: for pool_name in pool_names: # 获取该池下所有 Token ID pipe.smembers(f"{self.prefix_pool_set}{pool_name}") pool_tokens_res = await pipe.execute() - + # 收集所有 Token ID 以便批量查询 all_token_ids = [] - pool_map = {} # pool_name -> list[token_id] - + pool_map = {} # pool_name -> list[token_id] + for i, pool_name in enumerate(pool_names): tids = list(pool_tokens_res[i]) pool_map[pool_name] = tids all_token_ids.extend(tids) - + if not all_token_ids: return {name: [] for name in pool_names} @@ -344,24 +352,37 @@ class RedisStorage(BaseStorage): for tid in all_token_ids: pipe.hgetall(f"{self.prefix_token_hash}{tid}") token_data_list = await pipe.execute() - + # 重组数据结构 token_lookup = {} for i, tid in enumerate(all_token_ids): t_data = token_data_list[i] - if not t_data: continue - + if not t_data: + continue + # 恢复 tags (JSON -> List) if "tags" in t_data: - try: t_data["tags"] = json_loads(t_data["tags"]) - except: t_data["tags"] = [] - + try: + t_data["tags"] = json_loads(t_data["tags"]) + except Exception: + t_data["tags"] = [] + # 类型转换 (Redis 返回全 string) - for int_field in ["quota", "created_at", "use_count", "fail_count", "last_used_at", "last_fail_at", "last_sync_at"]: + for int_field in [ + "quota", + "created_at", + "use_count", + "fail_count", + "last_used_at", + "last_fail_at", + "last_sync_at", + ]: if t_data.get(int_field) and t_data[int_field] != "None": - try: t_data[int_field] = int(t_data[int_field]) - except: pass - + try: + t_data[int_field] = int(t_data[int_field]) + except Exception: + pass + token_lookup[tid] = t_data # 按 Pool 分组返回 @@ -370,9 +391,9 @@ class RedisStorage(BaseStorage): for tid in pool_map[pool_name]: if tid in token_lookup: pools[pool_name].append(token_lookup[tid]) - + return pools - + except Exception as e: logger.error(f"RedisStorage: 加载 Token 失败: {e}") return None @@ -438,15 +459,19 @@ class RedisStorage(BaseStorage): if "tags" in t_flat: t_flat["tags"] = json_dumps(t_flat["tags"]) status = t_flat.get("status") - if isinstance(status, str) and status.startswith("TokenStatus."): + if isinstance(status, str) and status.startswith( + "TokenStatus." + ): t_flat["status"] = status.split(".", 1)[1].lower() elif isinstance(status, Enum): t_flat["status"] = status.value t_flat = {k: str(v) for k, v in t_flat.items() if v is not None} - pipe.hset(f"{self.prefix_token_hash}{token_str}", mapping=t_flat) + pipe.hset( + f"{self.prefix_token_hash}{token_str}", mapping=t_flat + ) await pipe.execute() - + except Exception as e: logger.error(f"RedisStorage: 保存 Token 失败: {e}") raise @@ -455,7 +480,7 @@ class RedisStorage(BaseStorage): try: await self.redis.close() except (RuntimeError, asyncio.CancelledError, Exception): - # 忽略关闭时的 Event loop is closed 错误 + # 忽略关闭时的 Event loop is closed 错误 pass @@ -469,13 +494,14 @@ class SQLStorage(BaseStorage): def __init__(self, url: str): try: - from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker - from sqlalchemy import text, MetaData + from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker except ImportError: - raise ImportError("需要安装 sqlalchemy 和 async 驱动: pip install sqlalchemy[asyncio]") + raise ImportError( + "需要安装 sqlalchemy 和 async 驱动: pip install sqlalchemy[asyncio]" + ) self.dialect = url.split(":", 1)[0].split("+", 1)[0].lower() - + # 配置 robust 的连接池 self.engine = create_async_engine( url, @@ -483,55 +509,70 @@ class SQLStorage(BaseStorage): pool_size=20, max_overflow=10, pool_recycle=3600, - pool_pre_ping=True + pool_pre_ping=True, ) self.async_session = async_sessionmaker(self.engine, expire_on_commit=False) - self._initialized = False - + self._initialized = False + async def _ensure_schema(self): """确保数据库表存在""" - if self._initialized: return + if self._initialized: + return try: async with self.engine.begin() as conn: from sqlalchemy import text - + # Tokens 表 (通用 SQL) - await conn.execute(text(""" + await conn.execute( + text(""" CREATE TABLE IF NOT EXISTS tokens ( token VARCHAR(512) PRIMARY KEY, pool_name VARCHAR(64) NOT NULL, data TEXT, updated_at BIGINT ) - """)) - + """) + ) + # 配置表 - await conn.execute(text(""" + await conn.execute( + text(""" CREATE TABLE IF NOT EXISTS app_config ( section VARCHAR(64) NOT NULL, key_name VARCHAR(64) NOT NULL, value TEXT, PRIMARY KEY (section, key_name) ) - """)) - + """) + ) + # 索引 try: - await conn.execute(text("CREATE INDEX idx_tokens_pool ON tokens (pool_name)")) + await conn.execute( + text("CREATE INDEX idx_tokens_pool ON tokens (pool_name)") + ) except Exception: pass # 尝试兼容旧表结构 try: if self.dialect in ("mysql", "mariadb"): - await conn.execute(text("ALTER TABLE tokens MODIFY token VARCHAR(512)")) + await conn.execute( + text("ALTER TABLE tokens MODIFY token VARCHAR(512)") + ) await conn.execute(text("ALTER TABLE tokens MODIFY data TEXT")) elif self.dialect in ("postgres", "postgresql", "pgsql"): - await conn.execute(text("ALTER TABLE tokens ALTER COLUMN token TYPE VARCHAR(512)")) - await conn.execute(text("ALTER TABLE tokens ALTER COLUMN data TYPE TEXT")) + await conn.execute( + text( + "ALTER TABLE tokens ALTER COLUMN token TYPE VARCHAR(512)" + ) + ) + await conn.execute( + text("ALTER TABLE tokens ALTER COLUMN data TYPE TEXT") + ) except Exception: pass - + self._initialized = True except Exception as e: logger.error(f"SQLStorage: Schema 初始化失败: {e}") @@ -541,12 +582,13 @@ class SQLStorage(BaseStorage): async def acquire_lock(self, name: str, timeout: int = 10): # SQL 分布式锁: MySQL GET_LOCK / PG advisory_lock from sqlalchemy import text + lock_name = f"g2a:{hashlib.sha1(name.encode('utf-8')).hexdigest()[:24]}" if self.dialect in ("mysql", "mariadb"): async with self.async_session() as session: res = await session.execute( text("SELECT GET_LOCK(:name, :timeout)"), - {"name": lock_name, "timeout": timeout} + {"name": lock_name, "timeout": timeout}, ) got = res.scalar() if got != 1: @@ -555,18 +597,21 @@ class SQLStorage(BaseStorage): yield finally: try: - await session.execute(text("SELECT RELEASE_LOCK(:name)"), {"name": lock_name}) + await session.execute( + text("SELECT RELEASE_LOCK(:name)"), {"name": lock_name} + ) await session.commit() except Exception: pass elif self.dialect in ("postgres", "postgresql", "pgsql"): - lock_key = int.from_bytes(hashlib.sha256(name.encode("utf-8")).digest()[:8], "big", signed=False) + lock_key = int.from_bytes( + hashlib.sha256(name.encode("utf-8")).digest()[:8], "big", signed=False + ) async with self.async_session() as session: start = time.monotonic() while True: res = await session.execute( - text("SELECT pg_try_advisory_lock(:key)"), - {"key": lock_key} + text("SELECT pg_try_advisory_lock(:key)"), {"key": lock_key} ) if res.scalar(): break @@ -577,7 +622,9 @@ class SQLStorage(BaseStorage): yield finally: try: - await session.execute(text("SELECT pg_advisory_unlock(:key)"), {"key": lock_key}) + await session.execute( + text("SELECT pg_advisory_unlock(:key)"), {"key": lock_key} + ) await session.commit() except Exception: pass @@ -587,18 +634,23 @@ class SQLStorage(BaseStorage): async def load_config(self) -> Dict[str, Any]: await self._ensure_schema() from sqlalchemy import text + try: async with self.async_session() as session: - res = await session.execute(text("SELECT section, key_name, value FROM app_config")) + res = await session.execute( + text("SELECT section, key_name, value FROM app_config") + ) rows = res.fetchall() - if not rows: return None - + if not rows: + return None + config = {} for section, key, val_str in rows: - if section not in config: config[section] = {} + if section not in config: + config[section] = {} try: val = json_loads(val_str) - except: + except Exception: val = val_str config[section][key] = val return config @@ -609,21 +661,27 @@ class SQLStorage(BaseStorage): async def save_config(self, data: Dict[str, Any]): await self._ensure_schema() from sqlalchemy import text + try: async with self.async_session() as session: for section, items in data.items(): - if not isinstance(items, dict): continue + if not isinstance(items, dict): + continue for key, val in items.items(): val_str = json_dumps(val) - + # Upsert 逻辑 (简单实现: Delete + Insert) await session.execute( - text("DELETE FROM app_config WHERE section=:s AND key_name=:k"), - {"s": section, "k": key} + text( + "DELETE FROM app_config WHERE section=:s AND key_name=:k" + ), + {"s": section, "k": key}, ) await session.execute( - text("INSERT INTO app_config (section, key_name, value) VALUES (:s, :k, :v)"), - {"s": section, "k": key, "v": val_str} + text( + "INSERT INTO app_config (section, key_name, value) VALUES (:s, :k, :v)" + ), + {"s": section, "k": key, "v": val_str}, ) await session.commit() except Exception as e: @@ -633,23 +691,26 @@ class SQLStorage(BaseStorage): async def load_tokens(self) -> Dict[str, Any]: await self._ensure_schema() from sqlalchemy import text + try: async with self.async_session() as session: res = await session.execute(text("SELECT pool_name, data FROM tokens")) rows = res.fetchall() - if not rows: return None - + if not rows: + return None + pools = {} for pool_name, data_json in rows: - if pool_name not in pools: pools[pool_name] = [] - + if pool_name not in pools: + pools[pool_name] = [] + try: if isinstance(data_json, str): t_data = json_loads(data_json) else: t_data = data_json pools[pool_name].append(t_data) - except: + except Exception: pass return pools except Exception as e: @@ -659,25 +720,30 @@ class SQLStorage(BaseStorage): async def save_tokens(self, data: Dict[str, Any]): await self._ensure_schema() from sqlalchemy import text + try: async with self.async_session() as session: - await session.execute(text("DELETE FROM tokens")) - + await session.execute(text("DELETE FROM tokens")) + params = [] for pool_name, tokens in data.items(): for t in tokens: - params.append({ - "token": t.get("token"), - "pool_name": pool_name, - "data": json_dumps(t), - "updated_at": 0 - }) - + params.append( + { + "token": t.get("token"), + "pool_name": pool_name, + "data": json_dumps(t), + "updated_at": 0, + } + ) + if params: # 批量插入 await session.execute( - text("INSERT INTO tokens (token, pool_name, data, updated_at) VALUES (:token, :pool_name, :data, :updated_at)"), - params + text( + "INSERT INTO tokens (token, pool_name, data, updated_at) VALUES (:token, :pool_name, :data, :updated_at)" + ), + params, ) await session.commit() except Exception as e: @@ -690,6 +756,7 @@ class SQLStorage(BaseStorage): class StorageFactory: """存储后端工厂""" + _instance: Optional[BaseStorage] = None @staticmethod @@ -698,42 +765,45 @@ class StorageFactory: return url if storage_type == "mysql": if url.startswith("mysql://"): - return f"mysql+aiomysql://{url[len('mysql://'):]}" + return f"mysql+aiomysql://{url[len('mysql://') :]}" if url.startswith("mariadb://"): - return f"mariadb+aiomysql://{url[len('mariadb://'):]}" + return f"mariadb+aiomysql://{url[len('mariadb://') :]}" if storage_type == "pgsql": if url.startswith("postgres://"): - return f"postgresql+asyncpg://{url[len('postgres://'):]}" + return f"postgresql+asyncpg://{url[len('postgres://') :]}" if url.startswith("postgresql://"): - return f"postgresql+asyncpg://{url[len('postgresql://'):]}" + return f"postgresql+asyncpg://{url[len('postgresql://') :]}" if url.startswith("pgsql://"): - return f"postgresql+asyncpg://{url[len('pgsql://'):]}" + return f"postgresql+asyncpg://{url[len('pgsql://') :]}" return url - + @classmethod def get_storage(cls) -> BaseStorage: """获取全局存储实例 (单例)""" if cls._instance: return cls._instance - + storage_type = os.getenv("SERVER_STORAGE_TYPE", "local").lower() storage_url = os.getenv("SERVER_STORAGE_URL", "") - + logger.info(f"StorageFactory: 初始化存储后端: {storage_type}") - + if storage_type == "redis": - if not storage_url: raise ValueError("Redis 存储需要设置 SERVER_STORAGE_URL") + if not storage_url: + raise ValueError("Redis 存储需要设置 SERVER_STORAGE_URL") cls._instance = RedisStorage(storage_url) - + elif storage_type in ("mysql", "pgsql"): - if not storage_url: raise ValueError("SQL 存储需要设置 SERVER_STORAGE_URL") + if not storage_url: + raise ValueError("SQL 存储需要设置 SERVER_STORAGE_URL") storage_url = cls._normalize_sql_url(storage_type, storage_url) cls._instance = SQLStorage(storage_url) - + else: cls._instance = LocalStorage() - + return cls._instance + def get_storage() -> BaseStorage: return StorageFactory.get_storage() diff --git a/app/services/grok/assets.py b/app/services/grok/assets.py index 718e9a5..8775f9f 100644 --- a/app/services/grok/assets.py +++ b/app/services/grok/assets.py @@ -11,6 +11,7 @@ import re import uuid from pathlib import Path from contextlib import asynccontextmanager + try: import fcntl except ImportError: # pragma: no cover - non-posix platforms @@ -23,11 +24,7 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config -from app.core.exceptions import ( - AppException, - UpstreamException, - ValidationException -) +from app.core.exceptions import AppException, UpstreamException, ValidationException from app.services.grok.statsig import StatsigService from app.services.token.service import TokenService @@ -50,6 +47,7 @@ DEFAULT_DELETE_BATCH_SIZE = 10 _ASSETS_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) _ASSETS_SEM_VALUE = DEFAULT_MAX_CONCURRENT + def _get_assets_semaphore() -> asyncio.Semaphore: global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) @@ -63,14 +61,18 @@ def _get_assets_semaphore() -> asyncio.Semaphore: _ASSETS_SEMAPHORE = asyncio.Semaphore(value) return _ASSETS_SEMAPHORE + def _get_delete_batch_size() -> int: - value = get_config("performance.assets_delete_batch_size", DEFAULT_DELETE_BATCH_SIZE) + value = get_config( + "performance.assets_delete_batch_size", DEFAULT_DELETE_BATCH_SIZE + ) try: value = int(value) except Exception: value = DEFAULT_DELETE_BATCH_SIZE return max(1, value) + @asynccontextmanager async def _file_lock(name: str, timeout: int = 10): if fcntl is None: @@ -105,53 +107,77 @@ async def _file_lock(name: str, timeout: int = 10): except Exception: pass + MIME_TYPES = { # 图片 - '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png', - '.gif': 'image/gif', '.webp': 'image/webp', '.bmp': 'image/bmp', - + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", # 文档 - '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown', - '.doc': 'application/msword', - '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', - '.rtf': 'application/rtf', - + ".pdf": "application/pdf", + ".txt": "text/plain", + ".md": "text/markdown", + ".doc": "application/msword", + ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ".rtf": "application/rtf", # 表格 - '.csv': 'text/csv', - '.xls': 'application/vnd.ms-excel', - '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', - + ".csv": "text/csv", + ".xls": "application/vnd.ms-excel", + ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", # 代码 - '.py': 'text/x-python-script', '.js': 'application/javascript', '.ts': 'application/typescript', - '.java': 'text/x-java', '.cpp': 'text/x-c++', '.c': 'text/x-c', - '.go': 'text/x-go', '.rs': 'text/x-rust', '.rb': 'text/x-ruby', - '.php': 'text/x-php', '.sh': 'application/x-sh', '.html': 'text/html', - '.css': 'text/css', '.sql': 'application/sql', - + ".py": "text/x-python-script", + ".js": "application/javascript", + ".ts": "application/typescript", + ".java": "text/x-java", + ".cpp": "text/x-c++", + ".c": "text/x-c", + ".go": "text/x-go", + ".rs": "text/x-rust", + ".rb": "text/x-ruby", + ".php": "text/x-php", + ".sh": "application/x-sh", + ".html": "text/html", + ".css": "text/css", + ".sql": "application/sql", # 数据 - '.json': 'application/json', '.xml': 'application/xml', '.yaml': 'application/x-yaml', - '.yml': 'application/x-yaml', '.toml': 'application/toml', '.ini': 'text/plain', - '.log': 'text/plain', '.tmp': 'application/octet-stream', - + ".json": "application/json", + ".xml": "application/xml", + ".yaml": "application/x-yaml", + ".yml": "application/x-yaml", + ".toml": "application/toml", + ".ini": "text/plain", + ".log": "text/plain", + ".tmp": "application/octet-stream", # 其他 - '.graphql': 'application/graphql', '.proto': 'application/x-protobuf', - '.latex': 'application/x-latex', '.wiki': 'text/plain', '.rst': 'text/x-rst', + ".graphql": "application/graphql", + ".proto": "application/x-protobuf", + ".latex": "application/x-latex", + ".wiki": "text/plain", + ".rst": "text/x-rst", } -IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp'} -VIDEO_EXTS = {'.mp4', '.mov', '.m4v', '.webm', '.avi', '.mkv'} +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} +VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} # ==================== 基础服务 ==================== + class BaseService: """基础服务类""" - + def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.asset_proxy_url") or get_config("grok.base_proxy_url", "") + self.proxy = ( + proxy + or get_config("grok.asset_proxy_url") + or get_config("grok.base_proxy_url", "") + ) self.timeout = get_config("grok.timeout", TIMEOUT) self._session: Optional[AsyncSession] = None - + def _headers(self, token: str, referer: str = "https://grok.com/") -> dict: """构建请求头""" headers = { @@ -176,22 +202,22 @@ class BaseService: "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 _proxies(self) -> Optional[dict]: """构建代理配置""" return {"http": self.proxy, "https": self.proxy} if self.proxy else None - + def _dl_headers(self, token: str, file_path: str) -> dict: """构建下载请求头""" headers = { @@ -203,40 +229,43 @@ class BaseService: "Upgrade-Insecure-Requests": "1", "Referer": "https://grok.com/", } - + # 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 - + async def _get_session(self) -> AsyncSession: """获取复用 Session""" if self._session is None: self._session = AsyncSession() return self._session - + async def close(self): """关闭 Session""" if self._session: await self._session.close() self._session = None - + @staticmethod def is_url(input_str: str) -> bool: """检查是否为 URL""" try: result = urlparse(input_str) - return all([result.scheme, result.netloc]) and result.scheme in ['http', 'https'] - except: + return all([result.scheme, result.netloc]) and result.scheme in [ + "http", + "https", + ] + except Exception: return False - + @staticmethod async def fetch(url: str) -> Tuple[str, str, str]: """ 获取远程资源并转 Base64 - + Raises: UpstreamException: 当获取失败时 """ @@ -246,21 +275,25 @@ class BaseService: if response.status_code >= 400: raise UpstreamException( message=f"Failed to fetch resource: {response.status_code}", - details={"url": url, "status": response.status_code} + details={"url": url, "status": response.status_code}, ) - - filename = url.split('/')[-1].split('?')[0] or 'download' - content_type = response.headers.get('content-type', DEFAULT_MIME).split(';')[0] + + filename = url.split("/")[-1].split("?")[0] or "download" + content_type = response.headers.get("content-type", DEFAULT_MIME).split( + ";" + )[0] b64 = base64.b64encode(response.content).decode() - + logger.debug(f"Fetched: {url} -> {filename}") return filename, b64, content_type except Exception as e: logger.error(f"Fetch failed: {url} - {e}") if isinstance(e, AppException): raise e - raise UpstreamException(f"Resource fetch failed: {str(e)}", details={"url": url}) - + raise UpstreamException( + f"Resource fetch failed: {str(e)}", details={"url": url} + ) + @staticmethod def parse_b64(data_uri: str) -> Tuple[str, str, str]: """解析 Base64 数据""" @@ -269,10 +302,10 @@ class BaseService: if match: mime = match.group(1) b64 = match.group(2) - ext = mime.split('/')[-1] if '/' in mime else 'bin' + ext = mime.split("/")[-1] if "/" in mime else "bin" return f"file.{ext}", b64, mime return "file.bin", data_uri, DEFAULT_MIME - + @staticmethod def to_b64(file_path: Path, mime_type: str) -> str: """将本地文件转为 base64 data URI""" @@ -281,21 +314,24 @@ class BaseService: return f"data:{mime_type};base64,{b64_data}" except Exception as e: logger.error(f"File to base64 failed: {file_path} - {e}") - raise AppException(f"Failed to read file: {file_path}", code="file_read_error") + raise AppException( + f"Failed to read file: {file_path}", code="file_read_error" + ) # ==================== 上传服务 ==================== + class UploadService(BaseService): """文件上传服务""" - + async def upload(self, file_input: str, token: str) -> Tuple[str, str]: """ 上传文件到 Grok - + Returns: (file_id, file_uri) - + Raises: ValidationException: 输入无效 UpstreamException: 上传失败 @@ -307,10 +343,10 @@ class UploadService(BaseService): filename, b64, mime = await self.fetch(file_input) else: filename, b64, mime = self.parse_b64(file_input) - + if not b64: raise ValidationException("Invalid file input: empty content") - + # 构建请求 headers = self._headers(token) payload = { @@ -318,7 +354,7 @@ class UploadService(BaseService): "fileMimeType": mime, "content": b64, } - + # 执行上传 session = await self._get_session() response = await session.post( @@ -329,25 +365,28 @@ class UploadService(BaseService): timeout=self.timeout, proxies=self._proxies(), ) - + if response.status_code == 200: result = response.json() file_id = result.get("fileMetadataId", "") file_uri = result.get("fileUri", "") - logger.info(f"Upload success: {filename} -> {file_id}", extra={"file_id": file_id}) + logger.info( + f"Upload success: {filename} -> {file_id}", + extra={"file_id": file_id}, + ) return file_id, file_uri # 401/403 表示 token 失效，记录失败并标记冷却 if response.status_code in (401, 403): logger.warning( f"Upload token invalid: {response.status_code}, recording failure", - extra={"response": response.text[:200]} + extra={"response": response.text[:200]}, ) try: await TokenService.record_fail( token, status_code=response.status_code, - reason=f"upload_auth_failed" + reason="upload_auth_failed", ) except Exception as mark_err: logger.error(f"Failed to record token failure: {mark_err}") @@ -357,19 +396,22 @@ class UploadService(BaseService): details={ "status": response.status_code, "response": response.text[:200], - "token_invalidated": True - } + "token_invalidated": True, + }, ) logger.error( f"Upload failed: {filename} - {response.status_code}", - extra={"response": response.text[:200]} + extra={"response": response.text[:200]}, ) raise UpstreamException( message=f"Upload failed with status {response.status_code}", - details={"status": response.status_code, "response": response.text[:200]} + details={ + "status": response.status_code, + "response": response.text[:200], + }, ) - + except Exception as e: logger.error(f"Upload error: {e}") if isinstance(e, AppException): @@ -379,9 +421,10 @@ class UploadService(BaseService): # ==================== 列表服务 ==================== + class ListService(BaseService): """文件列表查询服务""" - + async def iter_assets(self, token: str): """ 分页迭代资产列表 @@ -401,7 +444,9 @@ class ListService(BaseService): params = dict(base_params) if page_token: if page_token in seen_tokens: - logger.warning("List pagination stopped due to repeated page token") + logger.warning( + "List pagination stopped due to repeated page token" + ) break seen_tokens.add(page_token) params["pageToken"] = page_token @@ -419,7 +464,7 @@ class ListService(BaseService): logger.error(f"List failed: {response.status_code}") raise UpstreamException( message=f"List assets failed: {response.status_code}", - details={"status": response.status_code} + details={"status": response.status_code}, ) result = response.json() @@ -433,7 +478,7 @@ class ListService(BaseService): async def list(self, token: str) -> List[Dict]: """ 查询文件列表 - + Raises: UpstreamException: 查询失败 """ @@ -444,7 +489,7 @@ class ListService(BaseService): logger.info(f"List success: {len(assets)} files") return assets - + except Exception as e: logger.error(f"List error: {e}") if isinstance(e, AppException): @@ -469,13 +514,14 @@ class ListService(BaseService): # ==================== 删除服务 ==================== + class DeleteService(BaseService): """文件删除服务""" - + async def delete(self, token: str, asset_id: str) -> bool: """ 删除单个文件 - + Raises: UpstreamException: 删除失败 """ @@ -483,7 +529,7 @@ class DeleteService(BaseService): try: headers = self._headers(token, referer="https://grok.com/files") url = f"{DELETE_API}/{asset_id}" - + session = await self._get_session() response = await session.delete( url, @@ -492,26 +538,26 @@ class DeleteService(BaseService): timeout=self.timeout, proxies=self._proxies(), ) - + if response.status_code == 200: logger.debug(f"Delete success: {asset_id}") return True - + logger.error(f"Delete failed: {asset_id} - {response.status_code}") - #: Note: Returning False or raising Exception? + #: Note: Returning False or raising Exception? #: Assuming caller handles Exception for stricter control, or False for loose control. #: Given "optimization" and "standardization", raising exceptions is better for API feedback. raise UpstreamException( message=f"Delete failed: {asset_id}", - details={"status": response.status_code} + details={"status": response.status_code}, ) - + except Exception as e: logger.error(f"Delete error: {asset_id} - {e}") if isinstance(e, AppException): raise e raise UpstreamException(f"Delete error: {str(e)}") - + async def delete_all(self, token: str) -> Dict[str, int]: """ 删除所有文件 @@ -533,16 +579,16 @@ class DeleteService(BaseService): if asset_id: try: return await self.delete(token, asset_id) - except: + except Exception: return False return False batch_size = _get_delete_batch_size() for i in range(0, len(assets), batch_size): - batch = assets[i:i + batch_size] - results = await asyncio.gather(*[ - _delete_one(asset, idx) for idx, asset in enumerate(batch) - ]) + batch = assets[i : i + batch_size] + results = await asyncio.gather( + *[_delete_one(asset, idx) for idx, asset in enumerate(batch)] + ) success += sum(results) failed += len(batch) - sum(results) @@ -561,9 +607,10 @@ class DeleteService(BaseService): # ==================== 下载服务 ==================== + class DownloadService(BaseService): """文件下载服务""" - + def __init__(self, proxy: str = None): super().__init__(proxy) # 创建缓存目录 @@ -573,24 +620,26 @@ class DownloadService(BaseService): self.image_dir.mkdir(parents=True, exist_ok=True) self.video_dir.mkdir(parents=True, exist_ok=True) self._cleanup_running = False - + def _cache_path(self, file_path: str, media_type: str) -> Path: """获取缓存路径""" cache_dir = self.image_dir if media_type == "image" else self.video_dir - filename = file_path.lstrip('/').replace('/', '-') + filename = file_path.lstrip("/").replace("/", "-") return cache_dir / filename - - async def download(self, file_path: str, token: str, media_type: str = "image") -> Tuple[Optional[Path], str]: + + async def download( + self, file_path: str, token: str, media_type: str = "image" + ) -> Tuple[Optional[Path], str]: """ 下载文件到本地 - + Raises: UpstreamException: 下载失败 """ async with _get_assets_semaphore(): try: cache_path = self._cache_path(file_path, media_type) - + # 如果已缓存 if cache_path.exists(): logger.debug(f"Cache hit: {cache_path}") @@ -602,16 +651,18 @@ class DownloadService(BaseService): # Double-check after lock if cache_path.exists(): logger.debug(f"Cache hit after lock: {cache_path}") - mime_type = MIME_TYPES.get(cache_path.suffix.lower(), DEFAULT_MIME) + mime_type = MIME_TYPES.get( + cache_path.suffix.lower(), DEFAULT_MIME + ) return cache_path, mime_type # 下载文件 if not file_path.startswith("/"): file_path = f"/{file_path}" - + url = f"{DOWNLOAD_API}{file_path}" headers = self._dl_headers(token, file_path) - + session = await self._get_session() response = await session.get( url, @@ -622,13 +673,13 @@ class DownloadService(BaseService): impersonate=BROWSER, stream=True, ) - + if response.status_code != 200: raise UpstreamException( message=f"Download failed: {response.status_code}", - details={"path": file_path, "status": response.status_code} + details={"path": file_path, "status": response.status_code}, ) - + # 保存文件（分块写入，避免大文件占用内存） tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp") try: @@ -654,48 +705,47 @@ class DownloadService(BaseService): tmp_path.unlink() except Exception: pass - mime_type = response.headers.get('content-type', DEFAULT_MIME).split(';')[0] - + mime_type = response.headers.get( + "content-type", DEFAULT_MIME + ).split(";")[0] + logger.info(f"Download success: {file_path}") - + # 检查缓存限制 asyncio.create_task(self.check_limit()) - + return cache_path, mime_type - + except Exception as e: logger.error(f"Download failed: {file_path} - {e}") if isinstance(e, AppException): raise e raise UpstreamException(f"Download error: {str(e)}") - + async def to_base64( - self, - file_path: str, - token: str, - media_type: str = "image" + self, file_path: str, token: str, media_type: str = "image" ) -> str: """ 下载文件并转为 base64 """ try: cache_path, mime_type = await self.download(file_path, token, media_type) - + if not cache_path or not cache_path.exists(): raise AppException("File download returned invalid path") - + # 使用基础服务的工具方法转换 data_uri = self.to_b64(cache_path, mime_type) - + # 删除临时文件 if data_uri: try: cache_path.unlink() except Exception as e: logger.warning(f"Delete temp file failed: {e}") - + return data_uri - + except Exception as e: logger.error(f"To base64 failed: {file_path} - {e}") if isinstance(e, AppException): @@ -707,18 +757,17 @@ class DownloadService(BaseService): cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"count": 0, "size_mb": 0.0} - + files = [f for f in cache_dir.glob("*") if f.is_file()] allowed_exts = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS files = [f for f in files if f.suffix.lower() in allowed_exts] total_size = sum(f.stat().st_size for f in files) - - return { - "count": len(files), - "size_mb": round(total_size / 1024 / 1024, 2) - } - def list_files(self, media_type: str = "image", page: int = 1, page_size: int = 1000) -> Dict[str, Any]: + return {"count": len(files), "size_mb": round(total_size / 1024 / 1024, 2)} + + def list_files( + self, media_type: str = "image", page: int = 1, page_size: int = 1000 + ) -> Dict[str, Any]: """列出本地缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): @@ -731,11 +780,13 @@ class DownloadService(BaseService): for f in files: try: stat = f.stat() - items.append({ - "name": f.name, - "size_bytes": stat.st_size, - "mtime_ms": int(stat.st_mtime * 1000), - }) + items.append( + { + "name": f.name, + "size_bytes": stat.st_size, + "mtime_ms": int(stat.st_mtime * 1000), + } + ) except Exception: continue @@ -774,29 +825,26 @@ class DownloadService(BaseService): return {"deleted": True} except Exception: return {"deleted": False} - + def clear(self, media_type: str = "image") -> Dict[str, Any]: """清空缓存""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"count": 0, "size_mb": 0.0} - + files = list(cache_dir.glob("*")) total_size = sum(f.stat().st_size for f in files) count = 0 - + for f in files: try: f.unlink() count += 1 except Exception as e: logger.error(f"Failed to delete {f}: {e}") - - return { - "count": count, - "size_mb": round(total_size / 1024 / 1024, 2) - } - + + return {"count": count, "size_mb": round(total_size / 1024 / 1024, 2)} + async def check_limit(self): """检查并清理缓存限制""" if self._cleanup_running: @@ -812,7 +860,7 @@ class DownloadService(BaseService): # 统计总大小 total_size = 0 all_files = [] - + for d in [self.image_dir, self.video_dir]: if d.exists(): for f in d.glob("*"): @@ -820,52 +868,56 @@ class DownloadService(BaseService): stat = f.stat() total_size += stat.st_size all_files.append((f, stat.st_mtime, stat.st_size)) - except: + except Exception: pass - + current_mb = total_size / 1024 / 1024 if current_mb <= limit_mb: return - + # 需要清理 - logger.info(f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning up...") - + logger.info( + f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning up..." + ) + # 按时间排序 all_files.sort(key=lambda x: x[1]) - + deleted_count = 0 deleted_size = 0 target_mb = limit_mb * 0.8 # 清理到 80% - + for f, _, size in all_files: try: f.unlink() deleted_count += 1 deleted_size += size total_size -= size - + if (total_size / 1024 / 1024) <= target_mb: break except Exception as e: logger.error(f"Cleanup failed for {f}: {e}") - - logger.info(f"Cache cleanup: deleted {deleted_count} files ({deleted_size/1024/1024:.2f}MB)") + + logger.info( + f"Cache cleanup: deleted {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)" + ) finally: self._cleanup_running = False def get_public_url(self, file_path: str) -> str: """ 获取文件的公共访问 URL - + 如果配置了 app_url，则返回自托管 URL，否则返回 Grok 原始 URL """ app_url = get_config("app.app_url", "") if not app_url: return f"{DOWNLOAD_API}{file_path if file_path.startswith('/') else '/' + file_path}" - + if not file_path.startswith("/"): file_path = f"/{file_path}" - + # 自动添加 /v1/files 前缀 return f"{app_url.rstrip('/')}/v1/files{file_path}" diff --git a/app/services/grok/chat.py b/app/services/grok/chat.py index add405a..d1141af 100644 --- a/app/services/grok/chat.py +++ b/app/services/grok/chat.py @@ -329,7 +329,7 @@ class GrokChatService: try: content = await response.text() content = content[:1000] # 限制长度避免日志过大 - except: + except Exception: content = "Unable to read response content" logger.error( @@ -342,7 +342,7 @@ class GrokChatService: # 关闭 session 并抛出异常 try: await session.close() - except: + except Exception: pass raise UpstreamException( message=f"Grok API request failed: {response.status_code}", @@ -360,7 +360,7 @@ class GrokChatService: logger.error(f"Chat request error: {e}") try: await session.close() - except: + except Exception: pass raise UpstreamException( message=f"Chat connection failed: {str(e)}", diff --git a/app/services/grok/media.py b/app/services/grok/media.py index 079ca4f..8c5159e 100644 --- a/app/services/grok/media.py +++ b/app/services/grok/media.py @@ -266,7 +266,7 @@ class VideoService: logger.error(f"Video generation failed: {response.status_code}") try: await session.close() - except: + except Exception: pass raise UpstreamException( message=f"Video generation failed: {response.status_code}", @@ -288,7 +288,7 @@ class VideoService: if session: try: await session.close() - except: + except Exception: pass logger.error(f"Video generation error: {e}") if isinstance(e, AppException): @@ -348,7 +348,7 @@ class VideoService: logger.error(f"Video from image failed: {response.status_code}") try: await session.close() - except: + except Exception: pass raise UpstreamException( message=f"Video from image failed: {response.status_code}", @@ -370,7 +370,7 @@ class VideoService: if session: try: await session.close() - except: + except Exception: pass logger.error(f"Video from image error: {e}") if isinstance(e, AppException): diff --git a/app/services/grok/model.py b/app/services/grok/model.py index a2d4498..9f5475c 100644 --- a/app/services/grok/model.py +++ b/app/services/grok/model.py @@ -11,18 +11,21 @@ from app.core.exceptions import ValidationException class Tier(str, Enum): """模型档位""" + BASIC = "basic" SUPER = "super" class Cost(str, Enum): """计费类型""" + LOW = "low" HIGH = "high" class ModelInfo(BaseModel): """模型信息""" + model_id: str grok_model: str model_mode: str @@ -36,42 +39,42 @@ class ModelInfo(BaseModel): class ModelService: """模型管理服务""" - + MODELS = [ ModelInfo( model_id="grok-3", grok_model="grok-3", model_mode="MODEL_MODE_AUTO", cost=Cost.LOW, - display_name="Grok 3" + display_name="Grok 3", ), ModelInfo( model_id="grok-3-fast", grok_model="grok-3", cost=Cost.LOW, model_mode="MODEL_MODE_FAST", - display_name="Grok 3 Fast" + display_name="Grok 3 Fast", ), ModelInfo( model_id="grok-4", grok_model="grok-4", model_mode="MODEL_MODE_AUTO", cost=Cost.LOW, - display_name="Grok 4" + display_name="Grok 4", ), ModelInfo( model_id="grok-4-mini", grok_model="grok-4-mini-thinking-tahoe", model_mode="MODEL_MODE_GROK_4_MINI_THINKING", cost=Cost.LOW, - display_name="Grok 4 Mini" + display_name="Grok 4 Mini", ), ModelInfo( model_id="grok-4-fast", grok_model="grok-4", model_mode="MODEL_MODE_FAST", cost=Cost.LOW, - display_name="Grok 4 Fast" + display_name="Grok 4 Fast", ), ModelInfo( model_id="grok-4-heavy", @@ -79,21 +82,21 @@ class ModelService: model_mode="MODEL_MODE_HEAVY", cost=Cost.HIGH, tier=Tier.SUPER, - display_name="Grok 4 Heavy" + display_name="Grok 4 Heavy", ), ModelInfo( model_id="grok-4.1", grok_model="grok-4-1-thinking-1129", model_mode="MODEL_MODE_AUTO", cost=Cost.LOW, - display_name="Grok 4.1" + display_name="Grok 4.1", ), ModelInfo( model_id="grok-4.1-thinking", grok_model="grok-4-1-thinking-1129", model_mode="MODEL_MODE_GROK_4_1_THINKING", - cost=Cost.HIGH, - display_name="Grok 4.1 Thinking" + cost=Cost.HIGH, + display_name="Grok 4.1 Thinking", ), ModelInfo( model_id="grok-imagine-1.0", @@ -102,7 +105,7 @@ class ModelService: cost=Cost.HIGH, display_name="Grok Image", description="Image generation model", - is_image=True + is_image=True, ), ModelInfo( model_id="grok-imagine-1.0-video", @@ -111,22 +114,22 @@ class ModelService: cost=Cost.HIGH, display_name="Grok Video", description="Video generation model", - is_video=True + is_video=True, ), ] - + _map = {m.model_id: m for m in MODELS} - + @classmethod def get(cls, model_id: str) -> Optional[ModelInfo]: """获取模型信息""" return cls._map.get(model_id) - + @classmethod def list(cls) -> list[ModelInfo]: """获取所有模型""" return list(cls._map.values()) - + @classmethod def valid(cls, model_id: str) -> bool: """模型是否有效""" diff --git a/app/services/grok/processor.py b/app/services/grok/processor.py index 5493ea1..6650549 100644 --- a/app/services/grok/processor.py +++ b/app/services/grok/processor.py @@ -1,6 +1,7 @@ """ OpenAI 响应格式处理器 """ + import asyncio import time import uuid @@ -27,15 +28,14 @@ T = TypeVar("T") class StreamIdleTimeoutError(Exception): """流空闲超时错误""" + def __init__(self, idle_seconds: float): self.idle_seconds = idle_seconds super().__init__(f"Stream idle timeout after {idle_seconds}s") async def _with_idle_timeout( - iterable: AsyncIterable[T], - idle_timeout: float, - model: str = "" + iterable: AsyncIterable[T], idle_timeout: float, model: str = "" ) -> AsyncGenerator[T, None]: """ 包装异步迭代器，添加空闲超时检测 @@ -59,15 +59,12 @@ async def _with_idle_timeout( iterator = iterable.__aiter__() while True: try: - item = await asyncio.wait_for( - iterator.__anext__(), - timeout=idle_timeout - ) + item = await asyncio.wait_for(iterator.__anext__(), timeout=idle_timeout) yield item except asyncio.TimeoutError: logger.warning( f"Stream idle timeout after {idle_timeout}s", - extra={"model": model, "idle_timeout": idle_timeout} + extra={"model": model, "idle_timeout": idle_timeout}, ) raise StreamIdleTimeoutError(idle_timeout) except StopAsyncIteration: @@ -79,7 +76,7 @@ ASSET_URL = "https://assets.grok.com/" class BaseProcessor: """基础处理器""" - + def __init__(self, model: str, token: str = ""): self.model = model self.token = token @@ -104,39 +101,44 @@ class BaseProcessor: # 处理可能的绝对路径 if path.startswith("http"): from urllib.parse import urlparse + path = urlparse(path).path - + if not path.startswith("/"): path = f"/{path}" - + if self.app_url: dl_service = self._get_dl() await dl_service.download(path, self.token, media_type) return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}" else: return f"{ASSET_URL.rstrip('/')}{path}" - + def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: """构建 SSE 响应 (StreamProcessor 通用)""" - if not hasattr(self, 'response_id'): + if not hasattr(self, "response_id"): self.response_id = None - if not hasattr(self, 'fingerprint'): + if not hasattr(self, "fingerprint"): self.fingerprint = "" - + delta = {} if role: delta["role"] = role delta["content"] = "" elif content: delta["content"] = content - + chunk = { "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", "object": "chat.completion.chunk", "created": self.created, "model": self.model, - "system_fingerprint": self.fingerprint if hasattr(self, 'fingerprint') else "", - "choices": [{"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}] + "system_fingerprint": self.fingerprint + if hasattr(self, "fingerprint") + else "", + "choices": [ + {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} + ], } return f"data: {orjson.dumps(chunk).decode()}\n\n" @@ -178,15 +180,15 @@ class StreamProcessor(BaseProcessor): if self._in_filter_tag: self._tag_buffer += char # 检查是否到达结束标签 - if char == '>': + if char == ">": # 检查是否是自闭合标签 - if '/>' in self._tag_buffer: + if "/>" in self._tag_buffer: self._in_filter_tag = False self._tag_buffer = "" else: # 检查是否是结束标签 for tag in self.filter_tags: - if f'' in self._tag_buffer: + if f"" in self._tag_buffer: self._in_filter_tag = False self._tag_buffer = "" break @@ -196,18 +198,18 @@ class StreamProcessor(BaseProcessor): continue # 检查是否开始一个过滤标签 - if char == '<': + if char == "<": # 查看后续字符 remaining = token[i:] tag_started = False for tag in self.filter_tags: - if remaining.startswith(f'<{tag}'): + if remaining.startswith(f"<{tag}"): tag_started = True break # 部分匹配（可能跨 token） if len(remaining) < len(tag) + 1: for j in range(1, len(remaining) + 1): - if f'<{tag}'.startswith(remaining[:j]): + if f"<{tag}".startswith(remaining[:j]): tag_started = True break @@ -220,9 +222,11 @@ class StreamProcessor(BaseProcessor): result.append(char) i += 1 - return ''.join(result) - - async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + return "".join(result) + + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: """处理流式响应""" # 获取空闲超时配置 idle_timeout = get_config("grok.stream_idle_timeout", 45.0) @@ -235,31 +239,33 @@ class StreamProcessor(BaseProcessor): data = orjson.loads(line) except orjson.JSONDecodeError: continue - + resp = data.get("result", {}).get("response", {}) - + # 元数据 if (llm := resp.get("llmInfo")) and not self.fingerprint: self.fingerprint = llm.get("modelHash", "") if rid := resp.get("responseId"): self.response_id = rid - + # 首次发送 role if not self.role_sent: yield self._sse(role="assistant") self.role_sent = True - + # 图像生成进度 if img := resp.get("streamingImageGenerationResponse"): if self.show_think: if not self.think_opened: yield self._sse("\n") self.think_opened = True - idx = img.get('imageIndex', 0) + 1 - progress = img.get('progress', 0) - yield self._sse(f"正在生成第{idx}张图片中，当前进度{progress}%\n") + idx = img.get("imageIndex", 0) + 1 + progress = img.get("progress", 0) + yield self._sse( + f"正在生成第{idx}张图片中，当前进度{progress}%\n" + ) continue - + # modelResponse if mr := resp.get("modelResponse"): if self.think_opened and self.show_think: @@ -267,15 +273,17 @@ class StreamProcessor(BaseProcessor): yield self._sse(msg + "\n") yield self._sse("\n") self.think_opened = False - + # 处理生成的图片 for url in mr.get("generatedImageUrls", []): parts = url.split("/") img_id = parts[-2] if len(parts) >= 2 else "image" - + if self.image_format == "base64": dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: yield self._sse(f"![{img_id}]({base64_data})\n") else: @@ -284,18 +292,22 @@ class StreamProcessor(BaseProcessor): else: final_url = await self.process_url(url, "image") yield self._sse(f"![{img_id}]({final_url})\n") - - if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + + if ( + (meta := mr.get("metadata", {})) + .get("llm_info", {}) + .get("modelHash") + ): self.fingerprint = meta["llm_info"]["modelHash"] continue - + # 普通 token if (token := resp.get("token")) is not None: if token: filtered = self._filter_token(token) if filtered: yield self._sse(filtered) - + if self.think_opened: yield self._sse("\n") yield self._sse(finish="stop") @@ -308,7 +320,11 @@ class StreamProcessor(BaseProcessor): raise UpstreamException( message=f"Stream idle timeout after {e.idle_seconds}s", status_code=504, - details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds} + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, ) except RequestsError as e: # HTTP/2 流错误转换为 UpstreamException @@ -317,16 +333,19 @@ class StreamProcessor(BaseProcessor): raise UpstreamException( message="Upstream connection closed unexpectedly", status_code=502, - details={"error": str(e), "type": "http2_stream_error"} + details={"error": str(e), "type": "http2_stream_error"}, ) logger.error(f"Stream request error: {e}", extra={"model": self.model}) raise UpstreamException( message=f"Upstream request failed: {e}", status_code=502, - details={"error": str(e)} + details={"error": str(e)}, ) except Exception as e: - logger.error(f"Stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Stream processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) raise finally: await self.close() @@ -346,20 +365,21 @@ class CollectProcessor(BaseProcessor): def _filter_content(self, content: str) -> str: """过滤内容中的特殊标签""" import re + if not content or not self.filter_tags: return content result = content for tag in self.filter_tags: # 匹配 ... 或 - pattern = rf'<{re.escape(tag)}[^>]*>.*?|<{re.escape(tag)}[^>]*/>' - result = re.sub(pattern, '', result, flags=re.DOTALL) + pattern = rf"<{re.escape(tag)}[^>]*>.*?|<{re.escape(tag)}[^>]*/>" + result = re.sub(pattern, "", result, flags=re.DOTALL) # 匹配 ... 跨行 - pattern2 = rf'<{re.escape(tag)}[^>]*>[\s\S]*?' - result = re.sub(pattern2, '', result, flags=re.DOTALL) + pattern2 = rf"<{re.escape(tag)}[^>]*>[\s\S]*?" + result = re.sub(pattern2, "", result, flags=re.DOTALL) return result - + async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: """处理并收集完整响应""" response_id = "" @@ -376,25 +396,27 @@ class CollectProcessor(BaseProcessor): data = orjson.loads(line) except orjson.JSONDecodeError: continue - + resp = data.get("result", {}).get("response", {}) - + if (llm := resp.get("llmInfo")) and not fingerprint: fingerprint = llm.get("modelHash", "") - + if mr := resp.get("modelResponse"): response_id = mr.get("responseId", "") content = mr.get("message", "") - + if urls := mr.get("generatedImageUrls"): content += "\n" for url in urls: parts = url.split("/") img_id = parts[-2] if len(parts) >= 2 else "image" - + if self.image_format == "base64": dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: content += f"![{img_id}]({base64_data})\n" else: @@ -403,8 +425,12 @@ class CollectProcessor(BaseProcessor): else: final_url = await self.process_url(url, "image") content += f"![{img_id}]({final_url})\n" - - if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + + if ( + (meta := mr.get("metadata", {})) + .get("llm_info", {}) + .get("modelHash") + ): fingerprint = meta["llm_info"]["modelHash"] except asyncio.CancelledError: @@ -414,11 +440,16 @@ class CollectProcessor(BaseProcessor): # 非流式模式下，超时后返回已收集的内容 except RequestsError as e: if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in collect: {e}", extra={"model": self.model}) + logger.warning( + f"HTTP/2 stream error in collect: {e}", extra={"model": self.model} + ) else: logger.error(f"Collect request error: {e}", extra={"model": self.model}) except Exception as e: - logger.error(f"Collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Collect processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) finally: await self.close() @@ -431,42 +462,62 @@ class CollectProcessor(BaseProcessor): "created": self.created, "model": self.model, "system_fingerprint": fingerprint, - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": content, "refusal": None, "annotations": []}, - "finish_reason": "stop" - }], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": content, + "refusal": None, + "annotations": [], + }, + "finish_reason": "stop", + } + ], "usage": { - "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, - "prompt_tokens_details": {"cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, - "completion_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0} - } + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 0, + "audio_tokens": 0, + "image_tokens": 0, + }, + "completion_tokens_details": { + "text_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + }, + }, } class VideoStreamProcessor(BaseProcessor): """视频流式响应处理器""" - + def __init__(self, model: str, token: str = "", think: bool = None): super().__init__(model, token) self.response_id: Optional[str] = None self.think_opened: bool = False self.role_sent: bool = False self.video_format = get_config("app.video_format", "url") - + if think is None: self.show_think = get_config("grok.thinking", False) else: self.show_think = think - + def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: """构建视频 HTML 标签""" poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" return f'''''' - - async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: """处理视频流式响应""" # 视频生成使用更长的空闲超时 idle_timeout = get_config("grok.video_idle_timeout", 90.0) @@ -479,92 +530,109 @@ class VideoStreamProcessor(BaseProcessor): data = orjson.loads(line) except orjson.JSONDecodeError: continue - + resp = data.get("result", {}).get("response", {}) - + if rid := resp.get("responseId"): self.response_id = rid - + # 首次发送 role if not self.role_sent: yield self._sse(role="assistant") self.role_sent = True - + # 视频生成进度 if video_resp := resp.get("streamingVideoGenerationResponse"): progress = video_resp.get("progress", 0) - + if self.show_think: if not self.think_opened: yield self._sse("\n") self.think_opened = True yield self._sse(f"正在生成视频中，当前进度{progress}%\n") - + if progress == 100: video_url = video_resp.get("videoUrl", "") thumbnail_url = video_resp.get("thumbnailImageUrl", "") - + if self.think_opened and self.show_think: yield self._sse("\n") self.think_opened = False - + if video_url: final_video_url = await self.process_url(video_url, "video") final_thumbnail_url = "" if thumbnail_url: - final_thumbnail_url = await self.process_url(thumbnail_url, "image") - - video_html = self._build_video_html(final_video_url, final_thumbnail_url) + final_thumbnail_url = await self.process_url( + thumbnail_url, "image" + ) + + video_html = self._build_video_html( + final_video_url, final_thumbnail_url + ) yield self._sse(video_html) - + logger.info(f"Video generated: {video_url}") continue - + if self.think_opened: yield self._sse("\n") yield self._sse(finish="stop") yield "data: [DONE]\n\n" except asyncio.CancelledError: - logger.debug("Video stream cancelled by client", extra={"model": self.model}) + logger.debug( + "Video stream cancelled by client", extra={"model": self.model} + ) except StreamIdleTimeoutError as e: raise UpstreamException( message=f"Video stream idle timeout after {e.idle_seconds}s", status_code=504, - details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds} + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, ) except RequestsError as e: if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in video: {e}", extra={"model": self.model}) + logger.warning( + f"HTTP/2 stream error in video: {e}", extra={"model": self.model} + ) raise UpstreamException( message="Upstream connection closed unexpectedly", status_code=502, - details={"error": str(e), "type": "http2_stream_error"} + details={"error": str(e), "type": "http2_stream_error"}, ) - logger.error(f"Video stream request error: {e}", extra={"model": self.model}) + logger.error( + f"Video stream request error: {e}", extra={"model": self.model} + ) raise UpstreamException( message=f"Upstream request failed: {e}", status_code=502, - details={"error": str(e)} + details={"error": str(e)}, ) except Exception as e: - logger.error(f"Video stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Video stream processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) finally: await self.close() class VideoCollectProcessor(BaseProcessor): """视频非流式响应处理器""" - + def __init__(self, model: str, token: str = ""): super().__init__(model, token) self.video_format = get_config("app.video_format", "url") - + def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" return f'''''' - + async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: """处理并收集视频响应""" response_id = "" @@ -580,66 +648,90 @@ class VideoCollectProcessor(BaseProcessor): data = orjson.loads(line) except orjson.JSONDecodeError: continue - + resp = data.get("result", {}).get("response", {}) - + if video_resp := resp.get("streamingVideoGenerationResponse"): if video_resp.get("progress") == 100: response_id = resp.get("responseId", "") video_url = video_resp.get("videoUrl", "") thumbnail_url = video_resp.get("thumbnailImageUrl", "") - + if video_url: final_video_url = await self.process_url(video_url, "video") final_thumbnail_url = "" if thumbnail_url: - final_thumbnail_url = await self.process_url(thumbnail_url, "image") - - content = self._build_video_html(final_video_url, final_thumbnail_url) + final_thumbnail_url = await self.process_url( + thumbnail_url, "image" + ) + + content = self._build_video_html( + final_video_url, final_thumbnail_url + ) logger.info(f"Video generated: {video_url}") - + except asyncio.CancelledError: - logger.debug("Video collect cancelled by client", extra={"model": self.model}) + logger.debug( + "Video collect cancelled by client", extra={"model": self.model} + ) except StreamIdleTimeoutError as e: - logger.warning(f"Video collect idle timeout: {e}", extra={"model": self.model}) + logger.warning( + f"Video collect idle timeout: {e}", extra={"model": self.model} + ) except RequestsError as e: if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in video collect: {e}", extra={"model": self.model}) + logger.warning( + f"HTTP/2 stream error in video collect: {e}", + extra={"model": self.model}, + ) else: - logger.error(f"Video collect request error: {e}", extra={"model": self.model}) + logger.error( + f"Video collect request error: {e}", extra={"model": self.model} + ) except Exception as e: - logger.error(f"Video collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + logger.error( + f"Video collect processing error: {e}", + extra={"model": self.model, "error_type": type(e).__name__}, + ) finally: await self.close() - + return { "id": response_id, "object": "chat.completion", "created": self.created, "model": self.model, - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": content, "refusal": None}, - "finish_reason": "stop" - }], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": content, + "refusal": None, + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, } class ImageStreamProcessor(BaseProcessor): """图片生成流式响应处理器""" - + def __init__(self, model: str, token: str = "", n: int = 1): super().__init__(model, token) self.partial_index = 0 self.n = n self.target_index = random.randint(0, 1) if n == 1 else None - + def _sse(self, event: str, data: dict) -> str: """构建 SSE 响应 (覆盖基类)""" return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" - - async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: """处理流式响应""" final_images = [] # 图片生成使用标准空闲超时 @@ -653,33 +745,38 @@ class ImageStreamProcessor(BaseProcessor): data = orjson.loads(line) except orjson.JSONDecodeError: continue - + resp = data.get("result", {}).get("response", {}) - + # 图片生成进度 if img := resp.get("streamingImageGenerationResponse"): image_index = img.get("imageIndex", 0) progress = img.get("progress", 0) - + if self.n == 1 and image_index != self.target_index: continue - + out_index = 0 if self.n == 1 else image_index - - yield self._sse("image_generation.partial_image", { - "type": "image_generation.partial_image", - "b64_json": "", - "index": out_index, - "progress": progress - }) + + yield self._sse( + "image_generation.partial_image", + { + "type": "image_generation.partial_image", + "b64_json": "", + "index": out_index, + "progress": progress, + }, + ) continue - + # modelResponse if mr := resp.get("modelResponse"): if urls := mr.get("generatedImageUrls"): for url in urls: dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: if "," in base64_data: b64 = base64_data.split(",", 1)[1] @@ -687,7 +784,7 @@ class ImageStreamProcessor(BaseProcessor): b64 = base64_data final_images.append(b64) continue - + for index, b64 in enumerate(final_images): if self.n == 1: if index != self.target_index: @@ -695,25 +792,35 @@ class ImageStreamProcessor(BaseProcessor): out_index = 0 else: out_index = index - - yield self._sse("image_generation.completed", { - "type": "image_generation.completed", - "b64_json": b64, - "index": out_index, - "usage": { - "total_tokens": 50, - "input_tokens": 25, - "output_tokens": 25, - "input_tokens_details": {"text_tokens": 5, "image_tokens": 20} - } - }) + + yield self._sse( + "image_generation.completed", + { + "type": "image_generation.completed", + "b64_json": b64, + "index": out_index, + "usage": { + "total_tokens": 50, + "input_tokens": 25, + "output_tokens": 25, + "input_tokens_details": { + "text_tokens": 5, + "image_tokens": 20, + }, + }, + }, + ) except asyncio.CancelledError: logger.debug("Image stream cancelled by client") except StreamIdleTimeoutError as e: raise UpstreamException( message=f"Image stream idle timeout after {e.idle_seconds}s", status_code=504, - details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds} + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, ) except RequestsError as e: if _is_http2_stream_error(e): @@ -721,16 +828,19 @@ class ImageStreamProcessor(BaseProcessor): raise UpstreamException( message="Upstream connection closed unexpectedly", status_code=502, - details={"error": str(e), "type": "http2_stream_error"} + details={"error": str(e), "type": "http2_stream_error"}, ) logger.error(f"Image stream request error: {e}") raise UpstreamException( message=f"Upstream request failed: {e}", status_code=502, - details={"error": str(e)} + details={"error": str(e)}, ) except Exception as e: - logger.error(f"Image stream processing error: {e}", extra={"error_type": type(e).__name__}) + logger.error( + f"Image stream processing error: {e}", + extra={"error_type": type(e).__name__}, + ) raise finally: await self.close() @@ -738,10 +848,10 @@ class ImageStreamProcessor(BaseProcessor): class ImageCollectProcessor(BaseProcessor): """图片生成非流式响应处理器""" - + def __init__(self, model: str, token: str = ""): super().__init__(model, token) - + async def process(self, response: AsyncIterable[bytes]) -> List[str]: """处理并收集图片""" images = [] @@ -756,21 +866,23 @@ class ImageCollectProcessor(BaseProcessor): data = orjson.loads(line) except orjson.JSONDecodeError: continue - + resp = data.get("result", {}).get("response", {}) - + if mr := resp.get("modelResponse"): if urls := mr.get("generatedImageUrls"): for url in urls: dl_service = self._get_dl() - base64_data = await dl_service.to_base64(url, self.token, "image") + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) if base64_data: if "," in base64_data: b64 = base64_data.split(",", 1)[1] else: b64 = base64_data images.append(b64) - + except asyncio.CancelledError: logger.debug("Image collect cancelled by client") except StreamIdleTimeoutError as e: @@ -781,10 +893,13 @@ class ImageCollectProcessor(BaseProcessor): else: logger.error(f"Image collect request error: {e}") except Exception as e: - logger.error(f"Image collect processing error: {e}", extra={"error_type": type(e).__name__}) + logger.error( + f"Image collect processing error: {e}", + extra={"error_type": type(e).__name__}, + ) finally: await self.close() - + return images diff --git a/app/services/grok/retry.py b/app/services/grok/retry.py index c9d6862..69299da 100644 --- a/app/services/grok/retry.py +++ b/app/services/grok/retry.py @@ -88,7 +88,9 @@ class RetryContext: self.last_error = error self.attempt += 1 - def calculate_delay(self, status_code: int, retry_after: Optional[float] = None) -> float: + def calculate_delay( + self, status_code: int, retry_after: Optional[float] = None + ) -> float: """ 计算退避延迟时间 @@ -114,7 +116,7 @@ class RetryContext: return delay # 其他状态码使用指数退避 + full jitter - exp_delay = self.backoff_base * (self.backoff_factor ** self.attempt) + exp_delay = self.backoff_base * (self.backoff_factor**self.attempt) delay = random.uniform(0, min(exp_delay, self.backoff_max)) return delay @@ -164,7 +166,7 @@ async def retry_on_status( *args, extract_status: Callable[[Exception], Optional[int]] = None, on_retry: Callable[[int, int, Exception, float], None] = None, - **kwargs + **kwargs, ) -> Any: """ 通用重试函数 @@ -186,6 +188,7 @@ async def retry_on_status( # 状态码提取器 if extract_status is None: + def extract_status(e: Exception) -> Optional[int]: if isinstance(e, UpstreamException): # 优先从 details 获取，回退到 status_code 属性 @@ -256,9 +259,7 @@ async def retry_on_status( f"last status: {status_code}, total delay: {ctx.total_delay:.2f}s" ) else: - logger.error( - f"Non-retryable status code: {status_code}" - ) + logger.error(f"Non-retryable status code: {status_code}") # 抛出最后一次的错误 raise @@ -266,7 +267,7 @@ async def retry_on_status( def with_retry( extract_status: Callable[[Exception], Optional[int]] = None, - on_retry: Callable[[int, int, Exception, float], None] = None + on_retry: Callable[[int, int, Exception, float], None] = None, ): """ 重试装饰器 @@ -276,17 +277,16 @@ def with_retry( async def my_api_call(): ... """ + def decorator(func: Callable): @wraps(func) async def wrapper(*args, **kwargs): return await retry_on_status( - func, - *args, - extract_status=extract_status, - on_retry=on_retry, - **kwargs + func, *args, extract_status=extract_status, on_retry=on_retry, **kwargs ) + return wrapper + return decorator diff --git a/app/services/grok/statsig.py b/app/services/grok/statsig.py index fcfdeeb..7dd1b80 100644 --- a/app/services/grok/statsig.py +++ b/app/services/grok/statsig.py @@ -11,35 +11,41 @@ from app.core.config import get_config class StatsigService: """Statsig ID 生成服务""" - + @staticmethod def _rand(length: int, alphanumeric: bool = False) -> str: """生成随机字符串""" - chars = string.ascii_lowercase + string.digits if alphanumeric else string.ascii_lowercase + chars = ( + string.ascii_lowercase + string.digits + if alphanumeric + else string.ascii_lowercase + ) return "".join(random.choices(chars, k=length)) - + @staticmethod def gen_id() -> str: """ 生成 Statsig ID - + Returns: Base64 编码的 ID """ # 读取配置 dynamic = get_config("grok.dynamic_statsig", True) - + if not dynamic: return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" - + # 随机格式 if random.choice([True, False]): rand = StatsigService._rand(5, alphanumeric=True) message = f"e:TypeError: Cannot read properties of null (reading 'children['{rand}']')" else: rand = StatsigService._rand(10) - message = f"e:TypeError: Cannot read properties of undefined (reading '{rand}')" - + message = ( + f"e:TypeError: Cannot read properties of undefined (reading '{rand}')" + ) + return base64.b64encode(message.encode()).decode() diff --git a/app/services/grok/usage.py b/app/services/grok/usage.py index 17388fa..5618587 100644 --- a/app/services/grok/usage.py +++ b/app/services/grok/usage.py @@ -6,12 +6,11 @@ import asyncio import uuid from typing import Dict -import orjson 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 +from app.core.exceptions import UpstreamException from app.services.grok.statsig import StatsigService from app.services.grok.retry import retry_on_status @@ -23,6 +22,7 @@ DEFAULT_MAX_CONCURRENT = 25 _USAGE_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) _USAGE_SEM_VALUE = DEFAULT_MAX_CONCURRENT + def _get_usage_semaphore() -> asyncio.Semaphore: global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE value = get_config("performance.usage_max_concurrent", DEFAULT_MAX_CONCURRENT) @@ -39,11 +39,11 @@ def _get_usage_semaphore() -> asyncio.Semaphore: class UsageService: """用量查询服务""" - + def __init__(self, proxy: str = None): self.proxy = proxy or get_config("grok.base_proxy_url", "") self.timeout = get_config("grok.timeout", TIMEOUT) - + def _build_headers(self, token: str) -> dict: """构建请求头""" headers = { @@ -68,33 +68,33 @@ class UsageService: "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) -> dict: """构建代理配置""" return {"http": self.proxy, "https": self.proxy} if self.proxy else None - + async def get(self, token: str, model_name: str = "grok-4-1-thinking-1129") -> Dict: """ 获取速率限制信息 - + Args: token: 认证 Token model_name: 模型名称 - + Returns: 响应数据 - + Raises: UpstreamException: 当获取失败且重试耗尽时 """ @@ -104,16 +104,13 @@ class UsageService: if isinstance(e, UpstreamException) and e.details: return e.details.get("status") return None - + # 定义实际的请求函数 async def do_request(): try: headers = self._build_headers(token) - payload = { - "requestKind": "DEFAULT", - "modelName": model_name - } - + payload = {"requestKind": "DEFAULT", "modelName": model_name} + async with AsyncSession() as session: response = await session.post( LIMITS_API, @@ -121,40 +118,39 @@ class UsageService: json=payload, impersonate=BROWSER, timeout=self.timeout, - proxies=self._build_proxies() + proxies=self._build_proxies(), ) - + if response.status_code == 200: data = response.json() - remaining = data.get('remainingTokens', 0) + remaining = data.get("remainingTokens", 0) logger.info(f"Usage: quota {remaining} remaining") return data - + logger.error(f"Usage failed: {response.status_code}") raise UpstreamException( message=f"Failed to get usage stats: {response.status_code}", - details={"status": response.status_code} + details={"status": response.status_code}, ) - + except Exception as e: if isinstance(e, UpstreamException): raise logger.error(f"Usage error: {e}") raise UpstreamException( message=f"Usage service error: {str(e)}", - details={"error": str(e)} + details={"error": str(e)}, ) - + # 带重试的执行 try: result = await retry_on_status( - do_request, - extract_status=extract_status + do_request, extract_status=extract_status ) return result - - except Exception as e: + + except Exception: # 最后一次失败已经被记录 raise diff --git a/app/services/token/__init__.py b/app/services/token/__init__.py index 1bc1339..945e8f1 100644 --- a/app/services/token/__init__.py +++ b/app/services/token/__init__.py @@ -6,7 +6,7 @@ from app.services.token.models import ( TokenPoolStats, EffortType, DEFAULT_QUOTA, - EFFORT_COST + EFFORT_COST, ) from app.services.token.pool import TokenPool from app.services.token.manager import TokenManager, get_token_manager @@ -16,20 +16,17 @@ from app.services.token.scheduler import TokenRefreshScheduler, get_scheduler __all__ = [ # Models "TokenInfo", - "TokenStatus", - "TokenPoolStats", - "EffortType", + "TokenStatus", + "TokenPoolStats", + "EffortType", "DEFAULT_QUOTA", "EFFORT_COST", - # Core "TokenPool", "TokenManager", - # API "TokenService", "get_token_manager", - # Scheduler "TokenRefreshScheduler", "get_scheduler", diff --git a/app/services/token/pool.py b/app/services/token/pool.py index 5f38748..405a84b 100644 --- a/app/services/token/pool.py +++ b/app/services/token/pool.py @@ -8,67 +8,68 @@ from app.services.token.models import TokenInfo, TokenStatus, TokenPoolStats class TokenPool: """Token 池（管理一组 Token）""" - + def __init__(self, name: str): self.name = name self._tokens: Dict[str, TokenInfo] = {} - + def add(self, token: TokenInfo): """添加 Token""" self._tokens[token.token] = token - + def remove(self, token_str: str) -> bool: """删除 Token""" if token_str in self._tokens: del self._tokens[token_str] return True return False - + def get(self, token_str: str) -> Optional[TokenInfo]: """获取 Token""" return self._tokens.get(token_str) - + def select(self) -> Optional[TokenInfo]: """ 选择一个可用 Token - 策略: + 策略: 1. 选择 active 状态且有配额的 token 2. 优先选择剩余额度最多的 3. 如果额度相同，随机选择（避免并发冲突） """ # 选择 token available = [ - t for t in self._tokens.values() + t + for t in self._tokens.values() if t.status == TokenStatus.ACTIVE and t.quota > 0 ] - + if not available: return None - + # 找到最大额度 max_quota = max(t.quota for t in available) - + # 筛选最大额度 candidates = [t for t in available if t.quota == max_quota] - + # 随机选择 return random.choice(candidates) - + def count(self) -> int: """Token 数量""" return len(self._tokens) - + def list(self) -> List[TokenInfo]: """获取所有 Token""" return list(self._tokens.values()) - + def get_stats(self) -> TokenPoolStats: """获取池统计信息""" stats = TokenPoolStats(total=len(self._tokens)) - + for token in self._tokens.values(): stats.total_quota += token.quota - + if token.status == TokenStatus.ACTIVE: stats.active += 1 elif token.status == TokenStatus.DISABLED: @@ -77,16 +78,16 @@ class TokenPool: stats.expired += 1 elif token.status == TokenStatus.COOLING: stats.cooling += 1 - + if stats.total > 0: stats.avg_quota = stats.total_quota / stats.total - + return stats - + def _rebuild_index(self): """重建索引（预留接口，用于加载时调用）""" pass - + def __iter__(self) -> Iterator[TokenInfo]: return iter(self._tokens.values()) diff --git a/app/services/token/scheduler.py b/app/services/token/scheduler.py index 7f5e46e..ad4da7a 100644 --- a/app/services/token/scheduler.py +++ b/app/services/token/scheduler.py @@ -10,17 +10,17 @@ from app.services.token.manager import get_token_manager class TokenRefreshScheduler: """Token 自动刷新调度器""" - + def __init__(self, interval_hours: int = 8): self.interval_hours = interval_hours self.interval_seconds = interval_hours * 3600 self._task: Optional[asyncio.Task] = None self._running = False - + async def _refresh_loop(self): """刷新循环""" logger.info(f"Scheduler: started (interval: {self.interval_hours}h)") - + while self._running: try: await asyncio.sleep(self.interval_seconds) @@ -31,7 +31,9 @@ class TokenRefreshScheduler: if isinstance(storage, RedisStorage): # Redis: non-blocking lock to avoid multi-worker duplication lock_key = "grok2api:lock:token_refresh" - lock = storage.redis.lock(lock_key, timeout=self.interval_seconds + 60, blocking_timeout=0) + lock = storage.redis.lock( + lock_key, timeout=self.interval_seconds + 60, blocking_timeout=0 + ) lock_acquired = await lock.acquire(blocking=False) else: try: @@ -48,7 +50,7 @@ class TokenRefreshScheduler: logger.info("Scheduler: starting token refresh...") manager = await get_token_manager() result = await manager.refresh_cooling_tokens() - + logger.info( f"Scheduler: refresh completed - " f"checked={result['checked']}, " @@ -62,27 +64,27 @@ class TokenRefreshScheduler: await lock.release() except Exception: pass - + except asyncio.CancelledError: break except Exception as e: logger.error(f"Scheduler: refresh error - {e}") - + def start(self): """启动调度器""" if self._running: logger.warning("Scheduler: already running") return - + self._running = True self._task = asyncio.create_task(self._refresh_loop()) logger.info("Scheduler: enabled") - + def stop(self): """停止调度器""" if not self._running: return - + self._running = False if self._task: self._task.cancel() diff --git a/app/services/token/service.py b/app/services/token/service.py index 51ace82..75b23a6 100644 --- a/app/services/token/service.py +++ b/app/services/token/service.py @@ -9,143 +9,141 @@ from app.services.token.models import TokenInfo, EffortType class TokenService: """ Token 服务外观 - + 提供简化的 API，隐藏内部实现细节 """ - + @staticmethod async def get_token(pool_name: str = "ssoBasic") -> Optional[str]: """ 获取可用 Token - + Args: pool_name: Token 池名称 - + Returns: Token 字符串（不含 sso= 前缀）或 None """ manager = await get_token_manager() return manager.get_token(pool_name) - + @staticmethod async def consume(token: str, effort: EffortType = EffortType.LOW) -> bool: """ 消耗 Token 配额（本地预估） - + Args: token: Token 字符串 effort: 消耗力度 - + Returns: 是否成功 """ manager = await get_token_manager() return await manager.consume(token, effort) - + @staticmethod async def sync_usage( - token: str, - model: str, - effort: EffortType = EffortType.LOW + token: str, model: str, effort: EffortType = EffortType.LOW ) -> bool: """ 同步 Token 使用量（优先 API，降级本地） - + Args: token: Token 字符串 model: 模型名称 effort: 降级时的消耗力度 - + Returns: 是否成功 """ manager = await get_token_manager() return await manager.sync_usage(token, model, effort) - + @staticmethod async def record_fail(token: str, status_code: int = 401, reason: str = "") -> bool: """ 记录 Token 失败 - + Args: token: Token 字符串 status_code: HTTP 状态码 reason: 失败原因 - + Returns: 是否成功 """ manager = await get_token_manager() return await manager.record_fail(token, status_code, reason) - + @staticmethod async def add_token(token: str, pool_name: str = "ssoBasic") -> bool: """ 添加 Token - + Args: token: Token 字符串 pool: Token 池名称 - + Returns: 是否成功 """ manager = await get_token_manager() return await manager.add(token, pool_name) - + @staticmethod async def remove_token(token: str) -> bool: """ 删除 Token - + Args: token: Token 字符串 - + Returns: 是否成功 """ manager = await get_token_manager() return await manager.remove(token) - + @staticmethod async def reset_token(token: str) -> bool: """ 重置单个 Token - + Args: token: Token 字符串 - + Returns: 是否成功 """ manager = await get_token_manager() return await manager.reset_token(token) - + @staticmethod async def reset_all(): """重置所有 Token""" manager = await get_token_manager() await manager.reset_all() - + @staticmethod async def get_stats() -> Dict[str, dict]: """ 获取统计信息 - + Returns: 各池的统计信息 """ manager = await get_token_manager() return manager.get_stats() - + @staticmethod async def list_tokens(pool_name: str = "ssoBasic") -> List[TokenInfo]: """ 获取指定池的所有 Token - + Args: pool_name: Token 池名称 - + Returns: Token 列表 """ diff --git a/main.py b/main.py index c20fd28..7d8d19d 100644 --- a/main.py +++ b/main.py @@ -8,8 +8,6 @@ from contextlib import asynccontextmanager import os import platform import sys -import time -import uuid from pathlib import Path from dotenv import load_dotenv @@ -18,58 +16,59 @@ env_file = Path(__file__).parent / ".env" if env_file.exists(): load_dotenv(env_file) -from fastapi import FastAPI, Request -from fastapi.middleware.cors import CORSMiddleware -from fastapi import Depends +from fastapi import FastAPI # noqa: E402 +from fastapi.middleware.cors import CORSMiddleware # noqa: E402 +from fastapi import Depends # noqa: E402 -from app.core.auth import verify_api_key -from app.core.config import get_config -from app.core.logger import logger, setup_logging -from app.core.exceptions import register_exception_handlers -from app.core.response_middleware import ResponseLoggerMiddleware -from app.api.v1.chat import router as chat_router -from app.api.v1.image import router as image_router -from app.api.v1.files import router as files_router -from app.api.v1.models import router as models_router -from app.services.token import get_scheduler +from app.core.auth import verify_api_key # noqa: E402 +from app.core.config import get_config # noqa: E402 +from app.core.logger import logger, setup_logging # noqa: E402 +from app.core.exceptions import register_exception_handlers # noqa: E402 +from app.core.response_middleware import ResponseLoggerMiddleware # noqa: E402 +from app.api.v1.chat import router as chat_router # noqa: E402 +from app.api.v1.image import router as image_router # noqa: E402 +from app.api.v1.files import router as files_router # noqa: E402 +from app.api.v1.models import router as models_router # noqa: E402 +from app.services.token import get_scheduler # noqa: E402 # 初始化日志 setup_logging( - level=os.getenv("LOG_LEVEL", "INFO"), - json_console=False, - file_logging=True + level=os.getenv("LOG_LEVEL", "INFO"), json_console=False, file_logging=True ) + @asynccontextmanager async def lifespan(app: FastAPI): """应用生命周期管理""" # 1. 加载配置 from app.core.config import config + await config.load() - + # 2. 启动服务显示 logger.info("Starting Grok2API...") logger.info(f"Platform: {platform.system()} {platform.release()}") logger.info(f"Python: {sys.version.split()[0]}") - + # 3. 启动 Token 刷新调度器 refresh_enabled = get_config("token.auto_refresh", True) if refresh_enabled: interval = get_config("token.refresh_interval_hours", 8) scheduler = get_scheduler(interval) scheduler.start() - + logger.info("Application startup complete.") yield - + # 关闭 logger.info("Shutting down Grok2API...") - + from app.core.storage import StorageFactory + if StorageFactory._instance: await StorageFactory._instance.close() - + if refresh_enabled: scheduler = get_scheduler() scheduler.stop() @@ -81,7 +80,7 @@ def create_app() -> FastAPI: title="Grok2API", lifespan=lifespan, ) - + # CORS 配置 app.add_middleware( CORSMiddleware, @@ -90,30 +89,37 @@ def create_app() -> FastAPI: allow_methods=["*"], allow_headers=["*"], ) - + # 请求日志和 ID 中间件 - from app.core.response_middleware import ResponseLoggerMiddleware app.add_middleware(ResponseLoggerMiddleware) - + # 注册异常处理器 register_exception_handlers(app) - + # 注册路由 - app.include_router(chat_router, prefix="/v1", dependencies=[Depends(verify_api_key)]) - app.include_router(image_router, prefix="/v1", dependencies=[Depends(verify_api_key)]) - app.include_router(models_router, prefix="/v1", dependencies=[Depends(verify_api_key)]) + app.include_router( + chat_router, prefix="/v1", dependencies=[Depends(verify_api_key)] + ) + app.include_router( + image_router, prefix="/v1", dependencies=[Depends(verify_api_key)] + ) + app.include_router( + models_router, prefix="/v1", dependencies=[Depends(verify_api_key)] + ) app.include_router(files_router, prefix="/v1/files") - + # 静态文件服务 from fastapi.staticfiles import StaticFiles + static_dir = Path(__file__).parent / "app" / "static" if static_dir.exists(): app.mount("/static", StaticFiles(directory=static_dir), name="static") - + # 注册管理路由 from app.api.v1.admin import router as admin_router + app.include_router(admin_router) - + return app @@ -122,15 +128,14 @@ app = create_app() if __name__ == "__main__": import uvicorn - import warnings - + host = os.getenv("SERVER_HOST", "0.0.0.0") port = int(os.getenv("SERVER_PORT", "8000")) workers = int(os.getenv("SERVER_WORKERS", "1")) - + # 平台检查 is_windows = platform.system() == "Windows" - + # 自动降级 if is_windows and workers > 1: logger.warning( @@ -138,11 +143,11 @@ if __name__ == "__main__": "Using single worker instead." ) workers = 1 - + uvicorn.run( "main:app", host=host, port=port, workers=workers, - log_level=os.getenv("LOG_LEVEL", "INFO").lower() + log_level=os.getenv("LOG_LEVEL", "INFO").lower(), )