diff --git a/.dockerignore b/.dockerignore index a3c2840..49a3e57 100644 --- a/.dockerignore +++ b/.dockerignore @@ -80,6 +80,10 @@ logs/ *.log *.log.* +# 数据临时文件 +data/temp/ +data/token.json + # 数据库文件 *.db *.sqlite diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 74602dd..805a9e6 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -21,6 +21,17 @@ jobs: contents: read packages: write id-token: write + + strategy: + fail-fast: false + matrix: + include: + # AMD64 平台 + - platform: linux/amd64 + suffix: amd64 + # ARM64 平台 + - platform: linux/arm64 + suffix: arm64 steps: - name: Checkout repository @@ -42,19 +53,58 @@ jobs: with: images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} tags: | - type=ref,event=branch - type=ref,event=pr - type=semver,pattern={{version}} - type=semver,pattern={{major}}.{{minor}} - type=raw,value=latest,enable={{is_default_branch}} + # 分支推送 + type=ref,event=branch,suffix=-${{ matrix.suffix }} + # 语义化版本:完整版本号 + type=semver,pattern={{version}},suffix=-${{ matrix.suffix }} + # latest标签 + type=raw,value=latest-${{ matrix.suffix }},enable={{is_default_branch}} - name: Build and push Docker image uses: docker/build-push-action@v5 with: context: . - platforms: linux/amd64,linux/arm64 + platforms: ${{ matrix.platform }} push: ${{ github.event_name != 'pull_request' }} tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} - cache-from: type=gha - cache-to: type=gha,mode=max \ No newline at end of file + cache-from: type=gha,scope=${{ matrix.suffix }} + cache-to: type=gha,mode=max,scope=${{ matrix.suffix }} + pull: true + + # 合并多架构镜像为统一标签 + merge-manifests: + runs-on: ubuntu-latest + needs: build-and-push + if: github.event_name != 'pull_request' + permissions: + contents: read + packages: write + + steps: + - name: Log in to Container Registry + uses: docker/login-action@v3 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Extract metadata + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + tags: | + type=ref,event=branch + type=semver,pattern={{version}} + type=raw,value=latest,enable={{is_default_branch}} + + - name: Create and push manifest + run: | + TAGS="${{ steps.meta.outputs.tags }}" + for tag in $TAGS; do + echo "合并标签: $tag" + docker buildx imagetools create -t $tag \ + ${tag}-amd64 \ + ${tag}-arm64 + done \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 1992333..f625296 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,37 +1,55 @@ -FROM python:3.11-slim AS base +# 构建阶段 +FROM python:3.11-slim AS builder -WORKDIR /app +WORKDIR /build -# 安装依赖阶段 -FROM base AS dependencies - -# 安装运行时需要的系统依赖 -RUN apt-get update && \ - apt-get install -y --no-install-recommends \ - gcc \ - && rm -rf /var/lib/apt/lists/* - -# 复制并安装 Python 依赖 +# 安装依赖到独立目录 COPY requirements.txt . -RUN pip install --no-cache-dir --upgrade pip && \ - pip install --no-cache-dir -r requirements.txt +RUN pip install --no-cache-dir --only-binary=:all: --prefix=/install -r requirements.txt && \ + find /install -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true && \ + find /install -type d -name "tests" -exec rm -rf {} + 2>/dev/null || true && \ + find /install -type d -name "test" -exec rm -rf {} + 2>/dev/null || true && \ + find /install -type d -name "*.dist-info" -exec sh -c 'rm -f "$1"/RECORD "$1"/INSTALLER' _ {} \; && \ + find /install -type f -name "*.pyc" -delete && \ + find /install -type f -name "*.pyo" -delete && \ + find /install -name "*.so" -exec strip --strip-unneeded {} \; 2>/dev/null || true -# 最终运行阶段 -FROM python:3.11-slim AS runtime +# 运行阶段 - 使用最小镜像 +FROM python:3.11-slim WORKDIR /app -# 复制必要的 Python 依赖 -COPY --from=dependencies /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages -COPY --from=dependencies /usr/local/bin /usr/local/bin +# 清理基础镜像中的冗余文件 +RUN rm -rf /usr/share/doc/* \ + /usr/share/man/* \ + /usr/share/locale/* \ + /var/cache/apt/* \ + /var/lib/apt/lists/* \ + /tmp/* \ + /var/tmp/* + +# 从构建阶段复制已安装的包 +COPY --from=builder /install /usr/local + +# 创建必要的目录(包括用于挂载的data目录) +RUN mkdir -p /app/logs /app/data/temp/image /app/data/temp/video # 复制应用代码 -COPY . . +COPY app/ ./app/ +COPY main.py . -# 创建必要的目录和文件 -RUN mkdir -p /app/logs /app/data/temp && \ - echo '{"ssoNormal": {}, "ssoSuper": {}}' > /app/data/token.json +# 复制并设置 entrypoint 脚本 +COPY docker-entrypoint.sh /usr/local/bin/ +RUN chmod +x /usr/local/bin/docker-entrypoint.sh + +# 删除 Python 字节码和缓存 +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 EXPOSE 8000 -CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] \ No newline at end of file +# 使用 entrypoint 脚本初始化配置 +ENTRYPOINT ["docker-entrypoint.sh"] + +# 默认启动命令 +CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/app/api/admin/manage.py b/app/api/admin/manage.py index 3c8ca43..cb0fe81 100644 --- a/app/api/admin/manage.py +++ b/app/api/admin/manage.py @@ -1,8 +1,4 @@ -""" -管理接口模块 - -提供Token管理功能,包括登录验证、Token增删查等操作。 -""" +"""管理接口 - Token管理和系统配置""" import secrets from typing import Dict, Any, List, Optional @@ -18,10 +14,9 @@ from app.services.grok.token import token_manager from app.models.grok_models import TokenType -# 创建路由器 router = APIRouter(tags=["管理"]) -# 常量定义 +# 常量 STATIC_DIR = Path(__file__).parents[2] / "template" TEMP_DIR = Path(__file__).parents[3] / "data" / "temp" IMAGE_CACHE_DIR = TEMP_DIR / "image" @@ -30,68 +25,86 @@ SESSION_EXPIRE_HOURS = 24 BYTES_PER_KB = 1024 BYTES_PER_MB = 1024 * 1024 -# 简单的会话存储 +# 会话存储 _sessions: Dict[str, datetime] = {} # === 请求/响应模型 === class LoginRequest(BaseModel): - """登录请求""" username: str password: str class LoginResponse(BaseModel): - """登录响应""" success: bool token: Optional[str] = None message: str class AddTokensRequest(BaseModel): - """批量添加Token请求""" tokens: List[str] - token_type: str # "sso" 或 "ssoSuper" + token_type: str class DeleteTokensRequest(BaseModel): - """批量删除Token请求""" tokens: List[str] - token_type: str # "sso" 或 "ssoSuper" + token_type: str class TokenInfo(BaseModel): - """Token信息""" token: str token_type: str created_time: Optional[int] = None remaining_queries: int heavy_remaining_queries: int - status: str # "未使用"、"限流中"、"失效"、"正常" + status: str + tags: List[str] = [] + note: str = "" class TokenListResponse(BaseModel): - """Token列表响应""" success: bool data: List[TokenInfo] total: int +class UpdateSettingsRequest(BaseModel): + global_config: Optional[Dict[str, Any]] = None + grok_config: Optional[Dict[str, Any]] = None + + +class UpdateTokenTagsRequest(BaseModel): + token: str + token_type: str + tags: List[str] + + +class UpdateTokenNoteRequest(BaseModel): + token: str + token_type: str + note: str + + +class TestTokenRequest(BaseModel): + token: str + token_type: str + + # === 辅助函数 === def validate_token_type(token_type_str: str) -> TokenType: - """验证并转换Token类型字符串为枚举""" + """验证Token类型""" if token_type_str not in ["sso", "ssoSuper"]: raise HTTPException( status_code=400, - detail={"error": "无效的Token类型,必须是 'sso' 或 'ssoSuper'", "code": "INVALID_TYPE"} + detail={"error": "无效的Token类型", "code": "INVALID_TYPE"} ) return TokenType.NORMAL if token_type_str == "sso" else TokenType.SUPER def parse_created_time(created_time) -> Optional[int]: - """解析创建时间,统一处理不同格式""" + """解析创建时间""" if isinstance(created_time, str): return int(created_time) if created_time else None elif isinstance(created_time, int): @@ -100,7 +113,7 @@ def parse_created_time(created_time) -> Optional[int]: def calculate_token_stats(tokens: Dict[str, Any], token_type: str) -> Dict[str, int]: - """计算Token统计信息""" + """计算Token统计""" total = len(tokens) expired = sum(1 for t in tokens.values() if t.get("status") == "expired") @@ -111,7 +124,7 @@ def calculate_token_stats(tokens: Dict[str, Any], token_type: str) -> Dict[str, if t.get("status") != "expired" and t.get("remainingQueries", -1) == 0) active = sum(1 for t in tokens.values() if t.get("status") != "expired" and t.get("remainingQueries", -1) > 0) - else: # super token + else: unused = sum(1 for t in tokens.values() if t.get("status") != "expired" and t.get("remainingQueries", -1) == -1 and t.get("heavyremainingQueries", -1) == -1) @@ -122,69 +135,64 @@ def calculate_token_stats(tokens: Dict[str, Any], token_type: str) -> Dict[str, if t.get("status") != "expired" and (t.get("remainingQueries", -1) > 0 or t.get("heavyremainingQueries", -1) > 0)) - return { - "total": total, - "unused": unused, - "limited": limited, - "expired": expired, - "active": active - } + return {"total": total, "unused": unused, "limited": limited, "expired": expired, "active": active} def verify_admin_session(authorization: Optional[str] = Header(None)) -> bool: """验证管理员会话""" if not authorization or not authorization.startswith("Bearer "): - raise HTTPException( - status_code=401, - detail={"error": "未授权访问", "code": "UNAUTHORIZED"} - ) + raise HTTPException(status_code=401, detail={"error": "未授权访问", "code": "UNAUTHORIZED"}) - token = authorization[7:] # 移除 "Bearer " 前缀 + token = authorization[7:] - # 检查token是否存在且未过期 if token not in _sessions: - raise HTTPException( - status_code=401, - detail={"error": "会话已过期或无效", "code": "SESSION_INVALID"} - ) + raise HTTPException(status_code=401, detail={"error": "会话无效", "code": "SESSION_INVALID"}) - # 检查会话是否过期(24小时) if datetime.now() > _sessions[token]: del _sessions[token] - raise HTTPException( - status_code=401, - detail={"error": "会话已过期", "code": "SESSION_EXPIRED"} - ) + raise HTTPException(status_code=401, detail={"error": "会话已过期", "code": "SESSION_EXPIRED"}) return True def get_token_status(token_data: Dict[str, Any], token_type: str) -> str: """获取Token状态""" - # 首先检查是否失效(来自 token.json 的 status 字段) if token_data.get("status") == "expired": return "失效" - # 获取剩余次数 - remaining_queries = token_data.get("remainingQueries", -1) + remaining = token_data.get("remainingQueries", -1) heavy_remaining = token_data.get("heavyremainingQueries", -1) - # 根据token类型选择正确的字段 - if token_type == "ssoSuper": - # Super token 可能使用 heavy 模型 - relevant_remaining = max(remaining_queries, heavy_remaining) - else: - # 普通token主要看 remaining_queries - relevant_remaining = remaining_queries + relevant = max(remaining, heavy_remaining) if token_type == "ssoSuper" else remaining - if relevant_remaining == -1: + if relevant == -1: return "未使用" - elif relevant_remaining == 0: + elif relevant == 0: return "限流中" else: return "正常" +def _calculate_dir_size(directory: Path) -> int: + """计算目录大小""" + total = 0 + for file_path in directory.iterdir(): + if file_path.is_file(): + try: + total += file_path.stat().st_size + except Exception as e: + logger.warning(f"[Admin] 无法获取文件大小: {file_path.name}, {e}") + return total + + +def _format_size(size_bytes: int) -> str: + """格式化文件大小""" + size_mb = size_bytes / BYTES_PER_MB + if size_mb < 1: + return f"{size_bytes / BYTES_PER_KB:.1f} KB" + return f"{size_mb:.1f} MB" + + # === 页面路由 === @router.get("/login", response_class=HTMLResponse) @@ -209,292 +217,163 @@ async def manage_page(): @router.post("/api/login", response_model=LoginResponse) async def admin_login(request: LoginRequest) -> LoginResponse: - """ - 管理员登录 - - 验证用户名和密码,成功后返回会话token。 - """ + """管理员登录""" try: - logger.debug(f"[Admin] 管理员登录尝试 - 用户名: {request.username}") + logger.debug(f"[Admin] 登录尝试: {request.username}") - # 验证用户名和密码 - expected_username = setting.global_config.get("admin_username", "") - expected_password = setting.global_config.get("admin_password", "") + expected_user = setting.global_config.get("admin_username", "") + expected_pass = setting.global_config.get("admin_password", "") - if request.username != expected_username or request.password != expected_password: - logger.warning(f"[Admin] 登录失败: 用户名或密码错误 - 用户名: {request.username}") - return LoginResponse( - success=False, - message="用户名或密码错误" - ) + if request.username != expected_user or request.password != expected_pass: + logger.warning(f"[Admin] 登录失败: {request.username}") + return LoginResponse(success=False, message="用户名或密码错误") - # 生成会话token session_token = secrets.token_urlsafe(32) + _sessions[session_token] = datetime.now() + timedelta(hours=SESSION_EXPIRE_HOURS) - # 设置会话过期时间 - expire_time = datetime.now() + timedelta(hours=SESSION_EXPIRE_HOURS) - _sessions[session_token] = expire_time - - logger.debug(f"[Admin] 管理员登录成功 - 用户名: {request.username}") - - return LoginResponse( - success=True, - token=session_token, - message="登录成功" - ) + logger.debug(f"[Admin] 登录成功: {request.username}") + return LoginResponse(success=True, token=session_token, message="登录成功") except Exception as e: - logger.error(f"[Admin] 登录处理异常 - 用户名: {request.username}, 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"登录失败: {str(e)}", "code": "LOGIN_ERROR"} - ) + logger.error(f"[Admin] 登录异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"登录失败: {e}", "code": "LOGIN_ERROR"}) @router.post("/api/logout") -async def admin_logout(_: bool = Depends(verify_admin_session), - authorization: Optional[str] = Header(None)) -> Dict[str, Any]: - """ - 管理员登出 - - 清除会话token。 - """ +async def admin_logout(_: bool = Depends(verify_admin_session), authorization: Optional[str] = Header(None)) -> Dict[str, Any]: + """管理员登出""" try: if authorization and authorization.startswith("Bearer "): token = authorization[7:] if token in _sessions: del _sessions[token] - logger.debug("[Admin] 管理员登出成功") + logger.debug("[Admin] 登出成功") return {"success": True, "message": "登出成功"} - logger.warning("[Admin] 登出失败: 无效或缺失的会话Token") + logger.warning("[Admin] 登出失败: 无效会话") return {"success": False, "message": "无效的会话"} except Exception as e: - logger.error(f"[Admin] 登出处理异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"登出失败: {str(e)}", "code": "LOGOUT_ERROR"} - ) + logger.error(f"[Admin] 登出异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"登出失败: {e}", "code": "LOGOUT_ERROR"}) @router.get("/api/tokens", response_model=TokenListResponse) async def list_tokens(_: bool = Depends(verify_admin_session)) -> TokenListResponse: - """ - 获取所有Token列表 - - 返回系统中所有Token及其状态信息。 - """ + """获取Token列表""" try: - logger.debug("[Admin] 开始获取Token列表") + logger.debug("[Admin] 获取Token列表") - all_tokens_data = token_manager.get_tokens() + all_tokens = token_manager.get_tokens() token_list: List[TokenInfo] = [] - # 处理普通Token - normal_tokens = all_tokens_data.get(TokenType.NORMAL.value, {}) - for token, data in normal_tokens.items(): + # 普通Token + for token, data in all_tokens.get(TokenType.NORMAL.value, {}).items(): token_list.append(TokenInfo( token=token, token_type="sso", created_time=parse_created_time(data.get("createdTime")), remaining_queries=data.get("remainingQueries", -1), heavy_remaining_queries=data.get("heavyremainingQueries", -1), - status=get_token_status(data, "sso") + status=get_token_status(data, "sso"), + tags=data.get("tags", []), + note=data.get("note", "") )) - # 处理Super Token - super_tokens = all_tokens_data.get(TokenType.SUPER.value, {}) - for token, data in super_tokens.items(): + # Super Token + for token, data in all_tokens.get(TokenType.SUPER.value, {}).items(): token_list.append(TokenInfo( token=token, token_type="ssoSuper", created_time=parse_created_time(data.get("createdTime")), remaining_queries=data.get("remainingQueries", -1), heavy_remaining_queries=data.get("heavyremainingQueries", -1), - status=get_token_status(data, "ssoSuper") + status=get_token_status(data, "ssoSuper"), + tags=data.get("tags", []), + note=data.get("note", "") )) - normal_count = len(normal_tokens) - super_count = len(super_tokens) - total_count = len(token_list) - - logger.debug(f"[Admin] Token列表获取成功 - 普通Token: {normal_count}, Super Token: {super_count}, 总计: {total_count}") - - return TokenListResponse( - success=True, - data=token_list, - total=total_count - ) + logger.debug(f"[Admin] Token列表获取成功: {len(token_list)}个") + return TokenListResponse(success=True, data=token_list, total=len(token_list)) except Exception as e: - logger.error(f"[Admin] 获取Token列表异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"获取Token列表失败: {str(e)}", "code": "LIST_ERROR"} - ) + logger.error(f"[Admin] 获取Token列表异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "LIST_ERROR"}) @router.post("/api/tokens/add") -async def add_tokens(request: AddTokensRequest, - _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """ - 批量添加Token - - 支持添加普通Token(sso)和Super Token(ssoSuper)。 - """ +async def add_tokens(request: AddTokensRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: + """批量添加Token""" try: - logger.debug(f"[Admin] 批量添加Token - 类型: {request.token_type}, 数量: {len(request.tokens)}") + logger.debug(f"[Admin] 添加Token: {request.token_type}, {len(request.tokens)}个") - # 验证并转换token类型 token_type = validate_token_type(request.token_type) - - # 添加Token await token_manager.add_token(request.tokens, token_type) - logger.debug(f"[Admin] Token添加成功 - 类型: {request.token_type}, 数量: {len(request.tokens)}") - - return { - "success": True, - "message": f"成功添加 {len(request.tokens)} 个Token", - "count": len(request.tokens) - } + logger.debug(f"[Admin] Token添加成功: {len(request.tokens)}个") + return {"success": True, "message": f"成功添加 {len(request.tokens)} 个Token", "count": len(request.tokens)} except HTTPException: raise except Exception as e: - logger.error(f"[Admin] Token添加异常 - 类型: {request.token_type}, 数量: {len(request.tokens)}, 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"添加Token失败: {str(e)}", "code": "ADD_ERROR"} - ) + logger.error(f"[Admin] Token添加异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"添加失败: {e}", "code": "ADD_ERROR"}) @router.post("/api/tokens/delete") -async def delete_tokens(request: DeleteTokensRequest, - _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """ - 批量删除Token - - 支持删除普通Token(sso)和Super Token(ssoSuper)。 - """ +async def delete_tokens(request: DeleteTokensRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: + """批量删除Token""" try: - logger.debug(f"[Admin] 批量删除Token - 类型: {request.token_type}, 数量: {len(request.tokens)}") + logger.debug(f"[Admin] 删除Token: {request.token_type}, {len(request.tokens)}个") - # 验证并转换token类型 token_type = validate_token_type(request.token_type) - - # 删除Token await token_manager.delete_token(request.tokens, token_type) - logger.debug(f"[Admin] Token删除成功 - 类型: {request.token_type}, 数量: {len(request.tokens)}") - - return { - "success": True, - "message": f"成功删除 {len(request.tokens)} 个Token", - "count": len(request.tokens) - } + logger.debug(f"[Admin] Token删除成功: {len(request.tokens)}个") + return {"success": True, "message": f"成功删除 {len(request.tokens)} 个Token", "count": len(request.tokens)} except HTTPException: raise except Exception as e: - logger.error(f"[Admin] Token删除异常 - 类型: {request.token_type}, 数量: {len(request.tokens)}, 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"删除Token失败: {str(e)}", "code": "DELETE_ERROR"} - ) + logger.error(f"[Admin] Token删除异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"删除失败: {e}", "code": "DELETE_ERROR"}) @router.get("/api/settings") async def get_settings(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """获取全局配置""" + """获取配置""" try: - logger.debug("[Admin] 获取全局配置") - return { - "success": True, - "data": { - "global": setting.global_config, - "grok": setting.grok_config - } - } + logger.debug("[Admin] 获取配置") + return {"success": True, "data": {"global": setting.global_config, "grok": setting.grok_config}} except Exception as e: - logger.error(f"[Admin] 获取配置失败: {str(e)}") - raise HTTPException(status_code=500, detail={"error": f"获取配置失败: {str(e)}", "code": "GET_SETTINGS_ERROR"}) - - -class UpdateSettingsRequest(BaseModel): - """更新配置请求""" - global_config: Optional[Dict[str, Any]] = None - grok_config: Optional[Dict[str, Any]] = None - - -class StreamTimeoutSettings(BaseModel): - """流式超时配置""" - stream_chunk_timeout: int = 120 - stream_first_response_timeout: int = 30 - stream_total_timeout: int = 600 + logger.error(f"[Admin] 获取配置失败: {e}") + raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "GET_SETTINGS_ERROR"}) @router.post("/api/settings") async def update_settings(request: UpdateSettingsRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """更新全局配置""" + """更新配置""" try: - logger.debug("[Admin] 更新全局配置") - - # 使用ConfigManager的save方法(支持存储抽象层) - await setting.save( - global_config=request.global_config, - grok_config=request.grok_config - ) - + logger.debug("[Admin] 更新配置") + await setting.save(global_config=request.global_config, grok_config=request.grok_config) logger.debug("[Admin] 配置更新成功") return {"success": True, "message": "配置更新成功"} except Exception as e: - logger.error(f"[Admin] 更新配置失败: {str(e)}") - raise HTTPException(status_code=500, detail={"error": f"更新配置失败: {str(e)}", "code": "UPDATE_SETTINGS_ERROR"}) - - -def _calculate_dir_size(directory: Path) -> int: - """计算目录中所有文件的大小(字节)""" - total_size = 0 - for file_path in directory.iterdir(): - if file_path.is_file(): - try: - total_size += file_path.stat().st_size - except Exception as e: - logger.warning(f"[Admin] 无法获取文件大小: {file_path.name}, 错误: {str(e)}") - return total_size - - -def _format_size(size_bytes: int) -> str: - """格式化字节大小为可读字符串""" - size_mb = size_bytes / BYTES_PER_MB - if size_mb < 1: - size_kb = size_bytes / BYTES_PER_KB - return f"{size_kb:.1f} KB" - return f"{size_mb:.1f} MB" + logger.error(f"[Admin] 更新配置失败: {e}") + raise HTTPException(status_code=500, detail={"error": f"更新失败: {e}", "code": "UPDATE_SETTINGS_ERROR"}) @router.get("/api/cache/size") async def get_cache_size(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: """获取缓存大小""" try: - logger.debug("[Admin] 开始获取缓存大小") + logger.debug("[Admin] 获取缓存大小") - # 计算图片缓存大小 - image_size = 0 - if IMAGE_CACHE_DIR.exists(): - image_size = _calculate_dir_size(IMAGE_CACHE_DIR) - - # 计算视频缓存大小 - video_size = 0 - if VIDEO_CACHE_DIR.exists(): - video_size = _calculate_dir_size(VIDEO_CACHE_DIR) - - # 总大小 + image_size = _calculate_dir_size(IMAGE_CACHE_DIR) if IMAGE_CACHE_DIR.exists() else 0 + video_size = _calculate_dir_size(VIDEO_CACHE_DIR) if VIDEO_CACHE_DIR.exists() else 0 total_size = image_size + video_size - logger.debug(f"[Admin] 缓存大小获取完成 - 图片: {_format_size(image_size)}, 视频: {_format_size(video_size)}, 总计: {_format_size(total_size)}") + logger.debug(f"[Admin] 缓存大小: 图片{_format_size(image_size)}, 视频{_format_size(video_size)}") return { "success": True, @@ -509,213 +388,235 @@ async def get_cache_size(_: bool = Depends(verify_admin_session)) -> Dict[str, A } except Exception as e: - logger.error(f"[Admin] 获取缓存大小异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"获取缓存大小失败: {str(e)}", "code": "CACHE_SIZE_ERROR"} - ) + logger.error(f"[Admin] 获取缓存大小异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "CACHE_SIZE_ERROR"}) @router.post("/api/cache/clear") async def clear_cache(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """清理缓存 - - 删除所有临时文件""" + """清理所有缓存""" try: - logger.debug("[Admin] 开始清理缓存") + logger.debug("[Admin] 清理缓存") - deleted_count = 0 image_count = 0 video_count = 0 - # 清理图片缓存 + # 清理图片 if IMAGE_CACHE_DIR.exists(): for file_path in IMAGE_CACHE_DIR.iterdir(): if file_path.is_file(): try: file_path.unlink() image_count += 1 - logger.debug(f"[Admin] 删除图片缓存: {file_path.name}") except Exception as e: - logger.error(f"[Admin] 删除图片缓存失败: {file_path.name}, 错误: {str(e)}") + logger.error(f"[Admin] 删除失败: {file_path.name}, {e}") - # 清理视频缓存 + # 清理视频 if VIDEO_CACHE_DIR.exists(): for file_path in VIDEO_CACHE_DIR.iterdir(): if file_path.is_file(): try: file_path.unlink() video_count += 1 - logger.debug(f"[Admin] 删除视频缓存: {file_path.name}") except Exception as e: - logger.error(f"[Admin] 删除视频缓存失败: {file_path.name}, 错误: {str(e)}") + logger.error(f"[Admin] 删除失败: {file_path.name}, {e}") - deleted_count = image_count + video_count - logger.debug(f"[Admin] 缓存清理完成 - 图片: {image_count}, 视频: {video_count}, 总计: {deleted_count}") + total = image_count + video_count + logger.debug(f"[Admin] 缓存清理完成: 图片{image_count}, 视频{video_count}") return { "success": True, - "message": f"成功清理缓存,删除图片 {image_count} 个,视频 {video_count} 个,共 {deleted_count} 个文件", - "data": { - "deleted_count": deleted_count, - "image_count": image_count, - "video_count": video_count - } + "message": f"成功清理缓存,删除图片 {image_count} 个,视频 {video_count} 个,共 {total} 个文件", + "data": {"deleted_count": total, "image_count": image_count, "video_count": video_count} } except Exception as e: - logger.error(f"[Admin] 清理缓存异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"清理缓存失败: {str(e)}", "code": "CACHE_CLEAR_ERROR"} - ) + logger.error(f"[Admin] 清理缓存异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"清理失败: {e}", "code": "CACHE_CLEAR_ERROR"}) @router.post("/api/cache/clear/images") async def clear_image_cache(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """清理图片缓存 - - 仅删除图片缓存文件""" + """清理图片缓存""" try: - logger.debug("[Admin] 开始清理图片缓存") + logger.debug("[Admin] 清理图片缓存") - deleted_count = 0 - - # 清理图片缓存 + count = 0 if IMAGE_CACHE_DIR.exists(): for file_path in IMAGE_CACHE_DIR.iterdir(): if file_path.is_file(): try: file_path.unlink() - deleted_count += 1 - logger.debug(f"[Admin] 删除图片缓存: {file_path.name}") + count += 1 except Exception as e: - logger.error(f"[Admin] 删除图片缓存失败: {file_path.name}, 错误: {str(e)}") + logger.error(f"[Admin] 删除失败: {file_path.name}, {e}") - logger.debug(f"[Admin] 图片缓存清理完成 - 删除 {deleted_count} 个文件") - - return { - "success": True, - "message": f"成功清理图片缓存,删除 {deleted_count} 个文件", - "data": { - "deleted_count": deleted_count, - "type": "images" - } - } + logger.debug(f"[Admin] 图片缓存清理完成: {count}个") + return {"success": True, "message": f"成功清理图片缓存,删除 {count} 个文件", "data": {"deleted_count": count, "type": "images"}} except Exception as e: - logger.error(f"[Admin] 清理图片缓存异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"清理图片缓存失败: {str(e)}", "code": "IMAGE_CACHE_CLEAR_ERROR"} - ) + logger.error(f"[Admin] 清理图片缓存异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"清理失败: {e}", "code": "IMAGE_CACHE_CLEAR_ERROR"}) @router.post("/api/cache/clear/videos") async def clear_video_cache(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """清理视频缓存 - - 仅删除视频缓存文件""" + """清理视频缓存""" try: - logger.debug("[Admin] 开始清理视频缓存") + logger.debug("[Admin] 清理视频缓存") - deleted_count = 0 - - # 清理视频缓存 + count = 0 if VIDEO_CACHE_DIR.exists(): for file_path in VIDEO_CACHE_DIR.iterdir(): if file_path.is_file(): try: file_path.unlink() - deleted_count += 1 - logger.debug(f"[Admin] 删除视频缓存: {file_path.name}") + count += 1 except Exception as e: - logger.error(f"[Admin] 删除视频缓存失败: {file_path.name}, 错误: {str(e)}") + logger.error(f"[Admin] 删除失败: {file_path.name}, {e}") - logger.debug(f"[Admin] 视频缓存清理完成 - 删除 {deleted_count} 个文件") - - return { - "success": True, - "message": f"成功清理视频缓存,删除 {deleted_count} 个文件", - "data": { - "deleted_count": deleted_count, - "type": "videos" - } - } + logger.debug(f"[Admin] 视频缓存清理完成: {count}个") + return {"success": True, "message": f"成功清理视频缓存,删除 {count} 个文件", "data": {"deleted_count": count, "type": "videos"}} except Exception as e: - logger.error(f"[Admin] 清理视频缓存异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"清理视频缓存失败: {str(e)}", "code": "VIDEO_CACHE_CLEAR_ERROR"} - ) + logger.error(f"[Admin] 清理视频缓存异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"清理失败: {e}", "code": "VIDEO_CACHE_CLEAR_ERROR"}) @router.get("/api/stats") async def get_stats(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """ - 获取统计信息 - - 返回Token的统计数据。 - """ + """获取统计信息""" try: logger.debug("[Admin] 开始获取统计信息") - all_tokens_data = token_manager.get_tokens() + all_tokens = token_manager.get_tokens() + normal_stats = calculate_token_stats(all_tokens.get(TokenType.NORMAL.value, {}), "normal") + super_stats = calculate_token_stats(all_tokens.get(TokenType.SUPER.value, {}), "super") + total = normal_stats["total"] + super_stats["total"] - # 统计普通Token - normal_tokens = all_tokens_data.get(TokenType.NORMAL.value, {}) - normal_stats = calculate_token_stats(normal_tokens, "normal") - - # 统计Super Token - super_tokens = all_tokens_data.get(TokenType.SUPER.value, {}) - super_stats = calculate_token_stats(super_tokens, "super") - - total_count = normal_stats["total"] + super_stats["total"] - - stats = { - "success": True, - "data": { - "normal": normal_stats, - "super": super_stats, - "total": total_count - } - } - - logger.debug(f"[Admin] 统计信息获取成功 - 普通Token: {normal_stats['total']}, Super Token: {super_stats['total']}, 总计: {total_count}") - return stats + logger.debug(f"[Admin] 统计信息获取成功 - 普通Token: {normal_stats['total']}, Super Token: {super_stats['total']}, 总计: {total}") + return {"success": True, "data": {"normal": normal_stats, "super": super_stats, "total": total}} except Exception as e: - logger.error(f"[Admin] 获取统计信息异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"获取统计信息失败: {str(e)}", "code": "STATS_ERROR"} - ) + logger.error(f"[Admin] 获取统计信息异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "STATS_ERROR"}) @router.get("/api/storage/mode") async def get_storage_mode(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: - """ - 获取当前存储模式 - - 返回当前的存储模式(file/mysql/redis)。 - """ + """获取存储模式""" try: logger.debug("[Admin] 获取存储模式") - import os - storage_mode = os.getenv("STORAGE_MODE", "file").upper() + mode = os.getenv("STORAGE_MODE", "file").upper() + return {"success": True, "data": {"mode": mode}} + except Exception as e: + logger.error(f"[Admin] 获取存储模式异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "STORAGE_MODE_ERROR"}) - return { - "success": True, - "data": { - "mode": storage_mode - } - } + +@router.post("/api/tokens/tags") +async def update_token_tags(request: UpdateTokenTagsRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: + """更新Token标签""" + try: + logger.debug(f"[Admin] 更新Token标签: {request.token[:10]}..., {request.tags}") + + token_type = validate_token_type(request.token_type) + await token_manager.update_token_tags(request.token, token_type, request.tags) + + logger.debug(f"[Admin] Token标签更新成功: {request.token[:10]}...") + return {"success": True, "message": "标签更新成功", "tags": request.tags} + + except HTTPException: + raise + except Exception as e: + logger.error(f"[Admin] Token标签更新异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"更新失败: {e}", "code": "UPDATE_TAGS_ERROR"}) + + +@router.get("/api/tokens/tags/all") +async def get_all_tags(_: bool = Depends(verify_admin_session)) -> Dict[str, Any]: + """获取所有标签""" + try: + logger.debug("[Admin] 获取所有标签") + + all_tokens = token_manager.get_tokens() + tags_set = set() + + for token_type_data in all_tokens.values(): + for token_data in token_type_data.values(): + tags = token_data.get("tags", []) + if isinstance(tags, list): + tags_set.update(tags) + + tags_list = sorted(list(tags_set)) + logger.debug(f"[Admin] 标签获取成功: {len(tags_list)}个") + return {"success": True, "data": tags_list} except Exception as e: - logger.error(f"[Admin] 获取存储模式异常 - 错误: {str(e)}") - raise HTTPException( - status_code=500, - detail={"error": f"获取存储模式失败: {str(e)}", "code": "STORAGE_MODE_ERROR"} - ) + logger.error(f"[Admin] 获取标签异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "GET_TAGS_ERROR"}) + + +@router.post("/api/tokens/note") +async def update_token_note(request: UpdateTokenNoteRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: + """更新Token备注""" + try: + logger.debug(f"[Admin] 更新Token备注: {request.token[:10]}...") + + token_type = validate_token_type(request.token_type) + await token_manager.update_token_note(request.token, token_type, request.note) + + logger.debug(f"[Admin] Token备注更新成功: {request.token[:10]}...") + return {"success": True, "message": "备注更新成功", "note": request.note} + + except HTTPException: + raise + except Exception as e: + logger.error(f"[Admin] Token备注更新异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"更新失败: {e}", "code": "UPDATE_NOTE_ERROR"}) + + +@router.post("/api/tokens/test") +async def test_token(request: TestTokenRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]: + """测试Token可用性""" + try: + logger.debug(f"[Admin] 测试Token: {request.token[:10]}...") + + token_type = validate_token_type(request.token_type) + auth_token = f"sso-rw={request.token};sso={request.token}" + + result = await token_manager.check_limits(auth_token, "grok-4-fast") + + if result: + logger.debug(f"[Admin] Token测试成功: {request.token[:10]}...") + return { + "success": True, + "message": "Token有效", + "data": { + "valid": True, + "remaining_queries": result.get("remainingTokens", -1), + "limit": result.get("limit", -1) + } + } + else: + logger.warning(f"[Admin] Token测试失败: {request.token[:10]}...") + + all_tokens = token_manager.get_tokens() + token_data = all_tokens.get(token_type.value, {}).get(request.token) + + if token_data: + if token_data.get("status") == "expired": + return {"success": False, "message": "Token已失效", "data": {"valid": False, "error_type": "expired", "error_code": 401}} + elif token_data.get("remainingQueries") == 0: + return {"success": False, "message": "Token已被限流", "data": {"valid": False, "error_type": "limited", "error_code": "other"}} + else: + return {"success": False, "message": "服务器被block或网络错误", "data": {"valid": False, "error_type": "blocked", "error_code": 403}} + else: + return {"success": False, "message": "Token数据异常", "data": {"valid": False, "error_type": "unknown", "error_code": "data_error"}} + + except HTTPException: + raise + except Exception as e: + logger.error(f"[Admin] Token测试异常: {e}") + raise HTTPException(status_code=500, detail={"error": f"测试失败: {e}", "code": "TEST_TOKEN_ERROR"}) diff --git a/app/api/v1/chat.py b/app/api/v1/chat.py index 87460e5..d80424d 100644 --- a/app/api/v1/chat.py +++ b/app/api/v1/chat.py @@ -1,9 +1,4 @@ -# -*- coding: utf-8 -*- -""" -聊天API路由模块 - -提供OpenAI兼容的聊天API接口,支持与Grok模型的交互。 -""" +"""聊天API路由 - OpenAI兼容的聊天接口""" from fastapi import APIRouter, Depends, HTTPException from typing import Optional @@ -15,38 +10,20 @@ from app.core.logger import logger from app.services.grok.client import GrokClient from app.models.openai_schema import OpenAIChatRequest -# 聊天路由 + router = APIRouter(prefix="/chat", tags=["聊天"]) @router.post("/completions", response_model=None) -async def chat_completions( - request: OpenAIChatRequest, - _: Optional[str] = Depends(auth_manager.verify) -): - """ - 创建聊天补全 - - 兼容OpenAI聊天API的端点,支持流式和非流式响应。 - - Args: - request: OpenAI格式的聊天请求 - _: 认证依赖(自动验证) - - Returns: - OpenAIChatCompletionResponse: 非流式响应 - StreamingResponse: 流式响应 - - Raises: - HTTPException: 当请求处理失败时 - """ +async def chat_completions(request: OpenAIChatRequest, _: Optional[str] = Depends(auth_manager.verify)): + """创建聊天补全(支持流式和非流式)""" try: - logger.info(f"[Chat] 聊天请求 - 模型: {request.model}") + logger.info("[Chat] 收到聊天请求") - # 调用Grok客户端处理请求 + # 调用Grok客户端 result = await GrokClient.openai_to_grok(request.model_dump()) - # 如果是流式响应,GrokClient已经返回了Iterator,直接包装为StreamingResponse + # 流式响应 if request.stream: return StreamingResponse( content=result, @@ -58,11 +35,11 @@ async def chat_completions( } ) - # 非流式响应直接返回 + # 非流式响应 return result except GrokApiException as e: - logger.error(f"[Chat] Grok API错误: {str(e)} - 详情: {e.details}") + logger.error(f"[Chat] Grok API错误: {e} - 详情: {e.details}") raise HTTPException( status_code=500, detail={ @@ -74,7 +51,7 @@ async def chat_completions( } ) except Exception as e: - logger.error(f"[Chat] 聊天请求处理失败: {str(e)}") + logger.error(f"[Chat] 处理失败: {e}") raise HTTPException( status_code=500, detail={ diff --git a/app/api/v1/images.py b/app/api/v1/images.py index b9c191e..e0679fa 100644 --- a/app/api/v1/images.py +++ b/app/api/v1/images.py @@ -1,4 +1,4 @@ -"""图片服务API路由""" +"""图片服务API - 提供缓存的图片和视频文件""" from fastapi import APIRouter, HTTPException from fastapi.responses import FileResponse @@ -13,31 +13,26 @@ router = APIRouter() @router.get("/images/{img_path:path}") async def get_image(img_path: str): """获取缓存的图片或视频 - + Args: - img_path: 文件路径,格式如 users-xxx-generated-xxx-image.jpg 或 users-xxx-generated-xxx-video.mp4 - - Returns: - 文件响应 + img_path: 文件路径(格式:users-xxx-generated-xxx-image.jpg) """ try: - # 将路径转换回原始格式(短横线转斜杠) + # 转换路径(短横线→斜杠) original_path = "/" + img_path.replace('-', '/') - # 判断是图片还是视频 + # 判断类型 is_video = any(original_path.lower().endswith(ext) for ext in ['.mp4', '.webm', '.mov', '.avi']) if is_video: - # 检查视频缓存 cache_path = video_cache_service.get_cached(original_path) media_type = "video/mp4" else: - # 检查图片缓存 cache_path = image_cache_service.get_cached(original_path) media_type = "image/jpeg" if cache_path and cache_path.exists(): - logger.debug(f"[MediaAPI] 返回缓存文件: {cache_path}") + logger.debug(f"[MediaAPI] 返回缓存: {cache_path}") return FileResponse( path=str(cache_path), media_type=media_type, @@ -48,11 +43,11 @@ async def get_image(img_path: str): ) # 文件不存在 - logger.warning(f"[MediaAPI] 文件未找到: {original_path}") + logger.warning(f"[MediaAPI] 未找到: {original_path}") raise HTTPException(status_code=404, detail="File not found") except HTTPException: raise except Exception as e: - logger.error(f"[MediaAPI] 获取文件失败: {e}") + logger.error(f"[MediaAPI] 获取失败: {e}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/app/api/v1/models.py b/app/api/v1/models.py index 5c18513..5cdcc71 100644 --- a/app/api/v1/models.py +++ b/app/api/v1/models.py @@ -1,8 +1,4 @@ -""" -模型接口模块 - -提供 OpenAI 兼容的 /v1/models 端点,返回系统支持的所有模型列表。 -""" +"""模型接口 - OpenAI兼容的模型列表端点""" import time from typing import Dict, Any, List, Optional @@ -12,43 +8,27 @@ from app.models.grok_models import Models from app.core.auth import auth_manager from app.core.logger import logger -# 配置日志 -# 创建路由器 router = APIRouter(tags=["模型"]) @router.get("/models") async def list_models(_: Optional[str] = Depends(auth_manager.verify)) -> Dict[str, Any]: - """ - 获取可用模型列表 - - 返回 OpenAI 兼容的模型列表格式,包含系统支持的所有 Grok 模型的详细信息。 - - Args: - _: 认证依赖(自动验证) - - Returns: - Dict[str, Any]: 包含模型列表的响应数据 - """ + """获取可用模型列表""" try: - logger.debug("[Models] 请求获取模型列表") + logger.debug("[Models] 请求模型列表") - # 获取当前时间戳 - current_timestamp = int(time.time()) - - # 构建模型数据列表 + timestamp = int(time.time()) model_data: List[Dict[str, Any]] = [] for model in Models: model_id = model.value config = Models.get_model_info(model_id) - # 基础信息 model_info = { "id": model_id, "object": "model", - "created": current_timestamp, + "created": timestamp, "owned_by": "x-ai", "display_name": config.get("display_name", model_id), "description": config.get("description", ""), @@ -61,22 +41,16 @@ async def list_models(_: Optional[str] = Depends(auth_manager.verify)) -> Dict[s model_data.append(model_info) - # 构建响应 - response = { - "object": "list", - "data": model_data - } - - logger.debug(f"[Models] 成功返回 {len(model_data)} 个模型") - return response + logger.debug(f"[Models] 返回 {len(model_data)} 个模型") + return {"object": "list", "data": model_data} except Exception as e: - logger.error(f"[Models] 获取模型列表时发生错误: {str(e)}") + logger.error(f"[Models] 获取列表失败: {e}") raise HTTPException( status_code=500, detail={ "error": { - "message": f"Failed to retrieve models: {str(e)}", + "message": f"Failed to retrieve models: {e}", "type": "internal_error", "code": "model_list_error" } @@ -86,22 +60,13 @@ async def list_models(_: Optional[str] = Depends(auth_manager.verify)) -> Dict[s @router.get("/models/{model_id}") async def get_model(model_id: str, _: Optional[str] = Depends(auth_manager.verify)) -> Dict[str, Any]: - """ - 获取特定模型信息 - - Args: - model_id (str): 模型ID - _: 认证依赖(自动验证) - - Returns: - Dict[str, Any]: 模型详细信息 - """ + """获取特定模型信息""" try: - logger.debug(f"[Models] 请求获取模型信息: {model_id}") + logger.debug(f"[Models] 请求模型: {model_id}") - # 验证模型是否存在 + # 验证模型 if not Models.is_valid_model(model_id): - logger.warning(f"[Models] 请求的模型不存在: {model_id}") + logger.warning(f"[Models] 模型不存在: {model_id}") raise HTTPException( status_code=404, detail={ @@ -113,17 +78,13 @@ async def get_model(model_id: str, _: Optional[str] = Depends(auth_manager.verif } ) - # 获取当前时间戳 - current_timestamp = int(time.time()) - - # 获取模型配置 + timestamp = int(time.time()) config = Models.get_model_info(model_id) - # 构建模型信息 model_info = { "id": model_id, "object": "model", - "created": current_timestamp, + "created": timestamp, "owned_by": "x-ai", "display_name": config.get("display_name", model_id), "description": config.get("description", ""), @@ -134,18 +95,18 @@ async def get_model(model_id: str, _: Optional[str] = Depends(auth_manager.verif "default_top_p": config.get("default_top_p", 0.95) } - logger.debug(f"[Models] 成功返回模型信息: {model_id}") + logger.debug(f"[Models] 返回模型: {model_id}") return model_info except HTTPException: raise except Exception as e: - logger.error(f"[Models] 获取模型信息时发生错误: {str(e)}") + logger.error(f"[Models] 获取模型失败: {e}") raise HTTPException( status_code=500, detail={ "error": { - "message": f"Failed to retrieve model: {str(e)}", + "message": f"Failed to retrieve model: {e}", "type": "internal_error", "code": "model_retrieve_error" } diff --git a/app/core/auth.py b/app/core/auth.py index fffaa9e..9a95f69 100644 --- a/app/core/auth.py +++ b/app/core/auth.py @@ -1,53 +1,58 @@ -"""认证模块""" +"""认证模块 - API令牌验证""" from typing import Optional from fastapi import Depends, HTTPException from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials + from app.core.config import setting from app.core.logger import logger +# Bearer安全方案 security = HTTPBearer(auto_error=False) +def _build_error(message: str, code: str = "invalid_token") -> dict: + """构建认证错误""" + return { + "error": { + "message": message, + "type": "authentication_error", + "code": code + } + } + + class AuthManager: - """认证管理器""" + """认证管理器 - 验证API令牌""" @staticmethod def verify(credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)) -> Optional[str]: - """验证认证令牌""" + """验证令牌""" api_key = setting.grok_config.get("api_key") + # 未设置时跳过 if not api_key: - logger.debug("[Auth] 未设置API_KEY,跳过验证。") + logger.debug("[Auth] 未设置API_KEY,跳过验证") return credentials.credentials if credentials else None + # 检查令牌 if not credentials: raise HTTPException( status_code=401, - detail={ - "error": { - "message": "缺少认证令牌", - "type": "authentication_error", - "code": "missing_token" - } - } + detail=_build_error("缺少认证令牌", "missing_token") ) + # 验证令牌 if credentials.credentials != api_key: raise HTTPException( status_code=401, - detail={ - "error": { - "message": f"令牌无效,长度: {len(credentials.credentials)}", - "type": "authentication_error", - "code": "invalid_token" - } - } + detail=_build_error(f"令牌无效,长度: {len(credentials.credentials)}", "invalid_token") ) - logger.debug("[Auth] 令牌认证成功。") + logger.debug("[Auth] 令牌认证成功") return credentials.credentials +# 全局实例 auth_manager = AuthManager() \ No newline at end of file diff --git a/app/core/config.py b/app/core/config.py index 01226f7..dc68617 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -1,72 +1,193 @@ -"""配置管理器""" +"""配置管理器 - 管理应用配置的读写""" import toml from pathlib import Path -from typing import Dict, Any +from typing import Dict, Any, Optional, Literal + + +# 默认配置 +DEFAULT_GROK = { + "api_key": "", + "proxy_url": "", + "proxy_pool_url": "", + "proxy_pool_interval": 300, + "cache_proxy_url": "", + "cf_clearance": "", + "x_statsig_id": "", + "dynamic_statsig": True, + "filtered_tags": "xaiartifact,xai:tool_usage_card", + "show_thinking": True, + "temporary": False, + "max_upload_concurrency": 20, + "max_request_concurrency": 100, + "stream_first_response_timeout": 30, + "stream_chunk_timeout": 120, + "stream_total_timeout": 600, + "retry_status_codes": [401, 429], # 可重试的HTTP状态码 +} + +DEFAULT_GLOBAL = { + "base_url": "http://localhost:8000", + "log_level": "INFO", + "image_mode": "url", + "admin_password": "admin", + "admin_username": "admin", + "image_cache_max_size_mb": 512, + "video_cache_max_size_mb": 1024, + "max_upload_concurrency": 20, # 最大并发上传数 + "max_request_concurrency": 50, # 最大并发请求数 + "batch_save_interval": 1.0, # 批量保存间隔(秒) + "batch_save_threshold": 10 # 触发批量保存的变更数阈值 +} class ConfigManager: """配置管理器""" def __init__(self) -> None: - """初始化""" - - # 加载环境变量 + """初始化配置""" self.config_path: Path = Path(__file__).parents[2] / "data" / "setting.toml" + self._storage: Optional[Any] = None + self._ensure_exists() self.global_config: Dict[str, Any] = self.load("global") self.grok_config: Dict[str, Any] = self.load("grok") - self._storage = None + + def _ensure_exists(self) -> None: + """确保配置存在""" + if not self.config_path.exists(): + self.config_path.parent.mkdir(parents=True, exist_ok=True) + self._create_default() + + def _create_default(self) -> None: + """创建默认配置""" + default = {"grok": DEFAULT_GROK.copy(), "global": DEFAULT_GLOBAL.copy()} + with open(self.config_path, "w", encoding="utf-8") as f: + toml.dump(default, f) + + def _normalize_proxy(self, proxy: str) -> str: + """标准化代理URL(socks5:// → socks5h://)""" + if proxy and proxy.startswith("socks5://"): + return proxy.replace("socks5://", "socks5h://", 1) + return proxy + + def _normalize_cf(self, cf: str) -> str: + """标准化CF Clearance(自动添加前缀)""" + if cf and not cf.startswith("cf_clearance="): + return f"cf_clearance={cf}" + return cf - def set_storage(self, storage) -> None: + def set_storage(self, storage: Any) -> None: """设置存储实例""" self._storage = storage - def load(self, section: str) -> Dict[str, Any]: - """配置加载器""" + def load(self, section: Literal["global", "grok"]) -> Dict[str, Any]: + """加载配置节""" try: with open(self.config_path, "r", encoding="utf-8") as f: - return toml.load(f)[section] + config = toml.load(f)[section] + + # 标准化Grok配置 + if section == "grok": + if "proxy_url" in config: + config["proxy_url"] = self._normalize_proxy(config["proxy_url"]) + if "cf_clearance" in config: + config["cf_clearance"] = self._normalize_cf(config["cf_clearance"]) + + return config except Exception as e: - raise Exception(f"[Setting] 配置加载失败: {e}") + raise Exception(f"[Setting] 配置加载失败: {e}") from e async def reload(self) -> None: - """重新加载配置(用于从存储同步后)""" + """重新加载配置""" self.global_config = self.load("global") self.grok_config = self.load("grok") - - async def save(self, global_config: Dict[str, Any] = None, grok_config: Dict[str, Any] = None) -> None: - """保存配置到存储""" - if not self._storage: - # 如果没有设置存储,使用传统文件保存方式(向后兼容) - import aiofiles - async with aiofiles.open(self.config_path, "r", encoding="utf-8") as f: - content = await f.read() - config = toml.loads(content) - - if global_config: - config["global"].update(global_config) - if grok_config: - config["grok"].update(grok_config) - - async with aiofiles.open(self.config_path, "w", encoding="utf-8") as f: - await f.write(toml.dumps(config)) - else: - # 使用存储抽象层 - config_data = await self._storage.load_config() - - if global_config: - config_data["global"].update(global_config) - if grok_config: - config_data["grok"].update(grok_config) - - await self._storage.save_config(config_data) + + async def _save_file(self, updates: Dict[str, Dict[str, Any]]) -> None: + """保存到文件""" + import aiofiles + + async with aiofiles.open(self.config_path, "r", encoding="utf-8") as f: + config = toml.loads(await f.read()) + + for section, data in updates.items(): + if section in config: + config[section].update(data) + + async with aiofiles.open(self.config_path, "w", encoding="utf-8") as f: + await f.write(toml.dumps(config)) + + async def _save_storage(self, updates: Dict[str, Dict[str, Any]]) -> None: + """保存到存储""" + config = await self._storage.load_config() + + for section, data in updates.items(): + if section in config: + config[section].update(data) + + await self._storage.save_config(config) + + def _prepare_grok(self, grok: Dict[str, Any]) -> Dict[str, Any]: + """准备Grok配置(移除前缀)""" + processed = grok.copy() + if "cf_clearance" in processed: + cf = processed["cf_clearance"] + if cf and cf.startswith("cf_clearance="): + processed["cf_clearance"] = cf.replace("cf_clearance=", "", 1) + return processed + + async def save(self, global_config: Optional[Dict[str, Any]] = None, grok_config: Optional[Dict[str, Any]] = None) -> None: + """保存配置""" + updates = {} + + if global_config: + updates["global"] = global_config + if grok_config: + updates["grok"] = self._prepare_grok(grok_config) + + # 选择存储方式 + if self._storage: + await self._save_storage(updates) + else: + await self._save_file(updates) - # 重新加载配置 await self.reload() + + async def get_proxy_async(self, proxy_type: Literal["service", "cache"] = "service") -> str: + """异步获取代理URL(支持代理池) + + Args: + proxy_type: 代理类型 + - service: 服务代理(client/upload) + - cache: 缓存代理(cache) + """ + from app.core.proxy_pool import proxy_pool + + if proxy_type == "cache": + cache_proxy = self.grok_config.get("cache_proxy_url", "") + if cache_proxy: + return cache_proxy + + # 从代理池获取 + return await proxy_pool.get_proxy() or "" + + def get_proxy(self, proxy_type: Literal["service", "cache"] = "service") -> str: + """获取代理URL(同步方法,用于向后兼容) + + Args: + proxy_type: 代理类型 + - service: 服务代理(client/upload) + - cache: 缓存代理(cache) + """ + from app.core.proxy_pool import proxy_pool + + if proxy_type == "cache": + cache_proxy = self.grok_config.get("cache_proxy_url", "") + if cache_proxy: + return cache_proxy + + # 返回当前代理(如果是代理池,返回最后一次获取的) + return proxy_pool.get_current_proxy() or self.grok_config.get("proxy_url", "") -# 全局设置 + +# 全局实例 setting = ConfigManager() - -if __name__ == "__main__": - print(setting.global_config) - print(setting.grok_config) \ No newline at end of file diff --git a/app/core/exception.py b/app/core/exception.py index a5c179b..0f13a1d 100644 --- a/app/core/exception.py +++ b/app/core/exception.py @@ -1,4 +1,4 @@ -"""异常处理器""" +"""异常处理器 - OpenAI兼容的错误响应""" from fastapi import Request, status from fastapi.responses import JSONResponse @@ -6,22 +6,60 @@ from fastapi.exceptions import RequestValidationError from starlette.exceptions import HTTPException as StarletteHTTPException -class GrokApiException(Exception): - """Grok API 业务异常""" +# HTTP错误映射 +HTTP_ERROR_MAP = { + 400: ("invalid_request_error", "请求格式错误或缺少必填参数"), + 401: ("invalid_request_error", "令牌认证失败"), + 403: ("permission_error", "没有权限访问此资源"), + 404: ("invalid_request_error", "请求的资源不存在"), + 429: ("rate_limit_error", "请求频率超出限制,请稍后再试"), + 500: ("api_error", "内部服务器错误"), + 503: ("api_error", "服务暂时不可用"), +} - def __init__(self, message: str, error_code: str = None, details: dict = None): +# Grok错误码映射 +GROK_STATUS_MAP = { + "NO_AUTH_TOKEN": status.HTTP_401_UNAUTHORIZED, + "INVALID_TOKEN": status.HTTP_401_UNAUTHORIZED, + "HTTP_ERROR": status.HTTP_502_BAD_GATEWAY, + "NETWORK_ERROR": status.HTTP_503_SERVICE_UNAVAILABLE, + "JSON_ERROR": status.HTTP_502_BAD_GATEWAY, + "API_ERROR": status.HTTP_502_BAD_GATEWAY, + "STREAM_ERROR": status.HTTP_502_BAD_GATEWAY, + "NO_RESPONSE": status.HTTP_502_BAD_GATEWAY, + "TOKEN_SAVE_ERROR": status.HTTP_500_INTERNAL_SERVER_ERROR, + "NO_AVAILABLE_TOKEN": status.HTTP_503_SERVICE_UNAVAILABLE, +} + +GROK_TYPE_MAP = { + "NO_AUTH_TOKEN": "authentication_error", + "INVALID_TOKEN": "authentication_error", + "HTTP_ERROR": "api_error", + "NETWORK_ERROR": "api_error", + "JSON_ERROR": "api_error", + "API_ERROR": "api_error", + "STREAM_ERROR": "api_error", + "NO_RESPONSE": "api_error", + "TOKEN_SAVE_ERROR": "api_error", + "NO_AVAILABLE_TOKEN": "api_error", +} + + +class GrokApiException(Exception): + """Grok API业务异常""" + + def __init__(self, message: str, error_code: str = None, details: dict = None, context: dict = None): self.message = message self.error_code = error_code self.details = details or {} + self.context = context or {} super().__init__(self.message) def build_error_response(message: str, error_type: str, code: str = None, param: str = None) -> dict: """构建OpenAI兼容的错误响应""" - error = { - "message": message, - "type": error_type, - } + error = {"message": message, "type": error_type} + if code: error["code"] = code if param: @@ -32,18 +70,8 @@ def build_error_response(message: str, error_type: str, code: str = None, param: async def http_exception_handler(_: Request, exc: StarletteHTTPException) -> JSONResponse: """处理HTTP异常""" - error_map = { - 400: ("invalid_request_error", "请求格式错误或缺少必填参数。"), - 401: ("invalid_request_error", "令牌认证失败。"), - 403: ("permission_error", "没有权限访问此资源。"), - 404: ("invalid_request_error", "请求的资源不存在。"), - 429: ("rate_limit_error", "请求频率超出限制,请稍后再试。"), - 500: ("api_error", "内部服务器错误。"), - 503: ("api_error", "服务暂时不可用。"), - } - - error_type, default_message = error_map.get(exc.status_code, ("api_error", str(exc.detail))) - message = str(exc.detail) if exc.detail else default_message + error_type, default_msg = HTTP_ERROR_MAP.get(exc.status_code, ("api_error", str(exc.detail))) + message = str(exc.detail) if exc.detail else default_msg return JSONResponse( status_code=exc.status_code, @@ -55,7 +83,7 @@ async def validation_exception_handler(_: Request, exc: RequestValidationError) """处理验证错误""" errors = exc.errors() param = errors[0]["loc"][-1] if errors and errors[0].get("loc") else None - message = errors[0]["msg"] if errors and errors[0].get("msg") else "请求参数错误。" + message = errors[0]["msg"] if errors and errors[0].get("msg") else "请求参数错误" return JSONResponse( status_code=status.HTTP_400_BAD_REQUEST, @@ -64,36 +92,9 @@ async def validation_exception_handler(_: Request, exc: RequestValidationError) async def grok_api_exception_handler(_: Request, exc: GrokApiException) -> JSONResponse: - """处理Grok API业务异常""" - # 根据错误码映射HTTP状态码 - status_code_map = { - "NO_AUTH_TOKEN": status.HTTP_401_UNAUTHORIZED, - "INVALID_TOKEN": status.HTTP_401_UNAUTHORIZED, - "HTTP_ERROR": status.HTTP_502_BAD_GATEWAY, - "NETWORK_ERROR": status.HTTP_503_SERVICE_UNAVAILABLE, - "JSON_ERROR": status.HTTP_502_BAD_GATEWAY, - "API_ERROR": status.HTTP_502_BAD_GATEWAY, - "STREAM_ERROR": status.HTTP_502_BAD_GATEWAY, - "NO_RESPONSE": status.HTTP_502_BAD_GATEWAY, - "TOKEN_SAVE_ERROR": status.HTTP_500_INTERNAL_SERVER_ERROR, - "NO_AVAILABLE_TOKEN": status.HTTP_503_SERVICE_UNAVAILABLE, - } - - http_status = status_code_map.get(exc.error_code, status.HTTP_500_INTERNAL_SERVER_ERROR) - error_type_map = { - "NO_AUTH_TOKEN": "authentication_error", - "INVALID_TOKEN": "authentication_error", - "HTTP_ERROR": "api_error", - "NETWORK_ERROR": "api_error", - "JSON_ERROR": "api_error", - "API_ERROR": "api_error", - "STREAM_ERROR": "api_error", - "NO_RESPONSE": "api_error", - "TOKEN_SAVE_ERROR": "api_error", - "NO_AVAILABLE_TOKEN": "api_error", - } - - error_type = error_type_map.get(exc.error_code, "api_error") + """处理Grok API异常""" + http_status = GROK_STATUS_MAP.get(exc.error_code, status.HTTP_500_INTERNAL_SERVER_ERROR) + error_type = GROK_TYPE_MAP.get(exc.error_code, "api_error") return JSONResponse( status_code=http_status, @@ -105,15 +106,12 @@ async def global_exception_handler(_: Request, exc: Exception) -> JSONResponse: """处理未捕获异常""" return JSONResponse( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - content=build_error_response( - "服务器遇到意外错误,请重试。", - "api_error" - ) + content=build_error_response("服务器遇到意外错误,请重试", "api_error") ) def register_exception_handlers(app) -> None: - """注册OpenAI兼容的异常处理器""" + """注册异常处理器""" app.add_exception_handler(StarletteHTTPException, http_exception_handler) app.add_exception_handler(RequestValidationError, validation_exception_handler) app.add_exception_handler(GrokApiException, grok_api_exception_handler) diff --git a/app/core/logger.py b/app/core/logger.py index 5fded40..8129ff8 100644 --- a/app/core/logger.py +++ b/app/core/logger.py @@ -1,63 +1,109 @@ -"""全局日志模块""" +"""全局日志模块 - 单例模式的日志管理器""" import sys import logging from pathlib import Path from logging.handlers import RotatingFileHandler + from app.core.config import setting -class LoggerManager: - """日志管理器""" +# 过滤模式 +FILTER_PATTERNS = [ + "chunk: b'", # SSE原始字节 + "Got event:", # SSE事件 + "Closing", # SSE关闭 +] + +class MCPLogFilter(logging.Filter): + """MCP日志过滤器 - 过滤大量数据的DEBUG日志""" + + def filter(self, record: logging.LogRecord) -> bool: + """过滤日志""" + # 过滤SSE的DEBUG日志 + if record.name == "sse_starlette.sse" and record.levelno == logging.DEBUG: + msg = record.getMessage() + return not any(p in msg for p in FILTER_PATTERNS) + + # 过滤MCP streamable_http的DEBUG日志 + if "mcp.server.streamable_http" in record.name and record.levelno == logging.DEBUG: + return False + + return True + + +class LoggerManager: + """日志管理器(单例)""" + + _instance = None _initialized = False + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + def __init__(self): - """初始化日志""" + """初始化日志系统""" if LoggerManager._initialized: return - # 日志配置 + # 配置 log_dir = Path(__file__).parents[2] / "logs" log_dir.mkdir(exist_ok=True) log_level = setting.global_config.get("log_level", "INFO").upper() log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" log_file = log_dir / "app.log" - # 配置根日志器 + # 根日志器 self.logger = logging.getLogger() self.logger.setLevel(log_level) - # 避免重复添加处理器 + # 避免重复添加 if self.logger.handlers: return - # 创建格式器 + # 格式器和过滤器 formatter = logging.Formatter(log_format) + mcp_filter = MCPLogFilter() # 控制台处理器 - console_handler = logging.StreamHandler(sys.stdout) - console_handler.setLevel(log_level) - console_handler.setFormatter(formatter) + console = logging.StreamHandler(sys.stdout) + console.setLevel(log_level) + console.setFormatter(formatter) + console.addFilter(mcp_filter) - # 文件处理器 + # 文件处理器(10MB,5个备份) file_handler = RotatingFileHandler( - log_file, maxBytes=10 * 1024 * 1024, backupCount=5, encoding="utf-8" + log_file, maxBytes=10*1024*1024, backupCount=5, encoding="utf-8" ) file_handler.setLevel(log_level) file_handler.setFormatter(formatter) + file_handler.addFilter(mcp_filter) - # 添加处理器到根日志器 - self.logger.addHandler(console_handler) + # 添加处理器 + self.logger.addHandler(console) self.logger.addHandler(file_handler) - # 配置第三方库日志级别,避免过多调试信息 - logging.getLogger("asyncio").setLevel(logging.WARNING) - logging.getLogger("uvicorn").setLevel(logging.INFO) - logging.getLogger("fastapi").setLevel(logging.INFO) - logging.getLogger("aiomysql").setLevel(logging.WARNING) + # 配置第三方库 + self._configure_third_party() LoggerManager._initialized = True + + def _configure_third_party(self): + """配置第三方库日志级别""" + config = { + "asyncio": logging.WARNING, + "uvicorn": logging.INFO, + "fastapi": logging.INFO, + "aiomysql": logging.WARNING, + "mcp": logging.CRITICAL, + "fastmcp": logging.CRITICAL, + } + + for name, level in config.items(): + logging.getLogger(name).setLevel(level) def debug(self, msg: str) -> None: """调试日志""" @@ -80,5 +126,5 @@ class LoggerManager: self.logger.critical(msg) -# 全局日志器实例 +# 全局实例 logger = LoggerManager() diff --git a/app/core/proxy_pool.py b/app/core/proxy_pool.py new file mode 100644 index 0000000..eaea5f5 --- /dev/null +++ b/app/core/proxy_pool.py @@ -0,0 +1,143 @@ +"""代理池管理器 - 从URL动态获取代理IP""" + +import asyncio +import aiohttp +import time +from typing import Optional, List +from app.core.logger import logger + + +class ProxyPool: + """代理池管理器""" + + def __init__(self): + self._pool_url: Optional[str] = None + self._static_proxy: Optional[str] = None + self._current_proxy: Optional[str] = None + self._last_fetch_time: float = 0 + self._fetch_interval: int = 300 # 5分钟刷新一次 + self._enabled: bool = False + self._lock = asyncio.Lock() + + def configure(self, proxy_url: str, proxy_pool_url: str = "", proxy_pool_interval: int = 300): + """配置代理池 + + Args: + proxy_url: 静态代理URL(socks5h://xxx 或 http://xxx) + proxy_pool_url: 代理池API URL,返回单个代理地址 + proxy_pool_interval: 代理池刷新间隔(秒) + """ + self._static_proxy = proxy_url if proxy_url else None + self._pool_url = proxy_pool_url if proxy_pool_url else None + self._fetch_interval = proxy_pool_interval + self._enabled = bool(self._pool_url) + + if self._enabled: + logger.info(f"[ProxyPool] 代理池已启用: {self._pool_url}, 刷新间隔: {self._fetch_interval}s") + elif self._static_proxy: + logger.info(f"[ProxyPool] 使用静态代理: {self._static_proxy}") + self._current_proxy = self._static_proxy + else: + logger.info("[ProxyPool] 未配置代理") + + async def get_proxy(self) -> Optional[str]: + """获取代理地址 + + Returns: + 代理URL或None + """ + # 如果未启用代理池,返回静态代理 + if not self._enabled: + return self._static_proxy + + # 检查是否需要刷新 + now = time.time() + if not self._current_proxy or (now - self._last_fetch_time) >= self._fetch_interval: + async with self._lock: + # 双重检查 + if not self._current_proxy or (now - self._last_fetch_time) >= self._fetch_interval: + await self._fetch_proxy() + + return self._current_proxy + + async def force_refresh(self) -> Optional[str]: + """强制刷新代理(用于403错误重试) + + Returns: + 新的代理URL或None + """ + if not self._enabled: + return self._static_proxy + + async with self._lock: + await self._fetch_proxy() + + return self._current_proxy + + async def _fetch_proxy(self): + """从代理池URL获取新的代理""" + try: + logger.debug(f"[ProxyPool] 正在从代理池获取新代理: {self._pool_url}") + + timeout = aiohttp.ClientTimeout(total=10) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(self._pool_url) as response: + if response.status == 200: + proxy_text = await response.text() + proxy = proxy_text.strip() + + # 验证代理格式 + if self._validate_proxy(proxy): + self._current_proxy = proxy + self._last_fetch_time = time.time() + logger.info(f"[ProxyPool] 成功获取新代理: {proxy}") + else: + logger.error(f"[ProxyPool] 代理格式无效: {proxy}") + # 降级到静态代理 + if not self._current_proxy: + self._current_proxy = self._static_proxy + else: + logger.error(f"[ProxyPool] 获取代理失败: HTTP {response.status}") + # 降级到静态代理 + if not self._current_proxy: + self._current_proxy = self._static_proxy + + except asyncio.TimeoutError: + logger.error("[ProxyPool] 获取代理超时") + if not self._current_proxy: + self._current_proxy = self._static_proxy + + except Exception as e: + logger.error(f"[ProxyPool] 获取代理异常: {e}") + # 降级到静态代理 + if not self._current_proxy: + self._current_proxy = self._static_proxy + + def _validate_proxy(self, proxy: str) -> bool: + """验证代理格式 + + Args: + proxy: 代理URL + + Returns: + 是否有效 + """ + if not proxy: + return False + + # 支持的协议 + valid_protocols = ['http://', 'https://', 'socks5://', 'socks5h://'] + + return any(proxy.startswith(proto) for proto in valid_protocols) + + def get_current_proxy(self) -> Optional[str]: + """获取当前使用的代理(同步方法) + + Returns: + 当前代理URL或None + """ + return self._current_proxy or self._static_proxy + + +# 全局代理池实例 +proxy_pool = ProxyPool() diff --git a/app/core/storage.py b/app/core/storage.py index 942ee3e..6ef93ec 100644 --- a/app/core/storage.py +++ b/app/core/storage.py @@ -1,7 +1,7 @@ """存储抽象层 - 支持文件、MySQL和Redis存储""" import os -import json +import orjson import toml import asyncio import warnings @@ -47,7 +47,7 @@ class BaseStorage(ABC): class FileStorage(BaseStorage): - """文件存储实现""" + """文件存储""" def __init__(self, data_dir: Path): self.data_dir = data_dir @@ -60,89 +60,87 @@ class FileStorage(BaseStorage): """初始化文件存储""" self.data_dir.mkdir(parents=True, exist_ok=True) - # 初始化token文件 if not self.token_file.exists(): - await self._write_file(self.token_file, json.dumps({"sso": {}, "ssoSuper": {}}, indent=2, ensure_ascii=False)) - logger.info("[Storage] 创建新的token文件") + await self._write(self.token_file, orjson.dumps({"sso": {}, "ssoSuper": {}}, option=orjson.OPT_INDENT_2).decode()) + logger.info("[Storage] 创建token文件") - # 初始化配置文件 if not self.config_file.exists(): - default_config = { + default = { "global": {"api_keys": [], "admin_username": "admin", "admin_password": "admin"}, "grok": {"proxy_url": "", "cf_clearance": "", "x_statsig_id": ""} } - await self._write_file(self.config_file, toml.dumps(default_config)) - logger.info("[Storage] 创建新的配置文件") + await self._write(self.config_file, toml.dumps(default)) + logger.info("[Storage] 创建配置文件") - async def _read_file(self, file_path: Path) -> str: - """读取文件内容""" - async with aiofiles.open(file_path, "r", encoding="utf-8") as f: + async def _read(self, path: Path) -> str: + """读取文件""" + async with aiofiles.open(path, "r", encoding="utf-8") as f: return await f.read() - async def _write_file(self, file_path: Path, content: str) -> None: - """写入文件内容""" - async with aiofiles.open(file_path, "w", encoding="utf-8") as f: + async def _write(self, path: Path, content: str) -> None: + """写入文件""" + async with aiofiles.open(path, "w", encoding="utf-8") as f: await f.write(content) - async def _load_json(self, file_path: Path, default: Dict[str, Any], lock: asyncio.Lock) -> Dict[str, Any]: - """加载JSON文件""" + async def _load_json(self, path: Path, default: Dict, lock: asyncio.Lock) -> Dict[str, Any]: + """加载JSON""" try: async with lock: - if not file_path.exists(): + if not path.exists(): return default - return json.loads(await self._read_file(file_path)) + return orjson.loads(await self._read(path)) except Exception as e: - logger.error(f"[Storage] 加载{file_path.name}失败: {e}") + logger.error(f"[Storage] 加载{path.name}失败: {e}") return default - async def _save_json(self, file_path: Path, data: Dict[str, Any], lock: asyncio.Lock) -> None: - """保存JSON文件""" + async def _save_json(self, path: Path, data: Dict, lock: asyncio.Lock) -> None: + """保存JSON""" try: async with lock: - await self._write_file(file_path, json.dumps(data, indent=2, ensure_ascii=False)) + await self._write(path, orjson.dumps(data, option=orjson.OPT_INDENT_2).decode()) except Exception as e: - logger.error(f"[Storage] 保存{file_path.name}失败: {e}") + logger.error(f"[Storage] 保存{path.name}失败: {e}") raise - async def _load_toml(self, file_path: Path, default: Dict[str, Any], lock: asyncio.Lock) -> Dict[str, Any]: - """加载TOML文件""" + async def _load_toml(self, path: Path, default: Dict, lock: asyncio.Lock) -> Dict[str, Any]: + """加载TOML""" try: async with lock: - if not file_path.exists(): + if not path.exists(): return default - return toml.loads(await self._read_file(file_path)) + return toml.loads(await self._read(path)) except Exception as e: - logger.error(f"[Storage] 加载{file_path.name}失败: {e}") + logger.error(f"[Storage] 加载{path.name}失败: {e}") return default - async def _save_toml(self, file_path: Path, data: Dict[str, Any], lock: asyncio.Lock) -> None: - """保存TOML文件""" + async def _save_toml(self, path: Path, data: Dict, lock: asyncio.Lock) -> None: + """保存TOML""" try: async with lock: - await self._write_file(file_path, toml.dumps(data)) + await self._write(path, toml.dumps(data)) except Exception as e: - logger.error(f"[Storage] 保存{file_path.name}失败: {e}") + logger.error(f"[Storage] 保存{path.name}失败: {e}") raise async def load_tokens(self) -> Dict[str, Any]: - """加载token数据""" + """加载token""" return await self._load_json(self.token_file, {"sso": {}, "ssoSuper": {}}, self._token_lock) async def save_tokens(self, data: Dict[str, Any]) -> None: - """保存token数据""" + """保存token""" await self._save_json(self.token_file, data, self._token_lock) async def load_config(self) -> Dict[str, Any]: - """加载配置数据""" + """加载配置""" return await self._load_toml(self.config_file, {"global": {}, "grok": {}}, self._config_lock) async def save_config(self, data: Dict[str, Any]) -> None: - """保存配置数据""" + """保存配置""" await self._save_toml(self.config_file, data, self._config_lock) class MysqlStorage(BaseStorage): - """MySQL存储实现""" + """MySQL存储""" def __init__(self, database_url: str, data_dir: Path): self.database_url = database_url @@ -155,22 +153,15 @@ class MysqlStorage(BaseStorage): try: import aiomysql parsed = self._parse_url(self.database_url) - logger.info(f"[Storage] 解析数据库连接: {parsed['user']}@{parsed['host']}:{parsed['port']}/{parsed['db']}") + logger.info(f"[Storage] MySQL: {parsed['user']}@{parsed['host']}:{parsed['port']}/{parsed['db']}") - # 创建数据库 await self._create_db(parsed) - - # 创建连接池 self._pool = await aiomysql.create_pool( host=parsed['host'], port=parsed['port'], user=parsed['user'], password=parsed['password'], db=parsed['db'], charset="utf8mb4", autocommit=True, maxsize=10 ) - - # 创建表 await self._create_tables() - - # 初始化文件存储和同步数据 await self._file.init_db() await self._sync_data() @@ -181,26 +172,26 @@ class MysqlStorage(BaseStorage): raise def _parse_url(self, url: str) -> Dict[str, Any]: - """解析数据库URL""" - parsed = urlparse(url) + """解析URL""" + p = urlparse(url) return { - 'user': unquote(parsed.username) if parsed.username else "", - 'password': unquote(parsed.password) if parsed.password else "", - 'host': parsed.hostname, - 'port': parsed.port or 3306, - 'db': parsed.path[1:] if parsed.path else "grok2api" + 'user': unquote(p.username) if p.username else "", + 'password': unquote(p.password) if p.password else "", + 'host': p.hostname, + 'port': p.port or 3306, + 'db': p.path[1:] if p.path else "grok2api" } - async def _create_db(self, parsed: Dict[str, Any]) -> None: + async def _create_db(self, parsed: Dict) -> None: """创建数据库""" import aiomysql - temp_pool = await aiomysql.create_pool( + pool = await aiomysql.create_pool( host=parsed['host'], port=parsed['port'], user=parsed['user'], password=parsed['password'], charset="utf8mb4", autocommit=True, maxsize=1 ) try: - async with temp_pool.acquire() as conn: + async with pool.acquire() as conn: async with conn.cursor() as cursor: with warnings.catch_warnings(): warnings.filterwarnings('ignore', message='.*database exists') @@ -208,10 +199,10 @@ class MysqlStorage(BaseStorage): f"CREATE DATABASE IF NOT EXISTS `{parsed['db']}` " f"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci" ) - logger.info(f"[Storage] MySQL数据库 '{parsed['db']}' 已就绪") + logger.info(f"[Storage] 数据库 '{parsed['db']}' 就绪") finally: - temp_pool.close() - await temp_pool.wait_closed() + pool.close() + await pool.wait_closed() async def _create_tables(self) -> None: """创建表""" @@ -240,7 +231,7 @@ class MysqlStorage(BaseStorage): warnings.filterwarnings('ignore', message='.*already exists') for sql in tables.values(): await cursor.execute(sql) - logger.info("[Storage] MySQL表创建/验证成功") + logger.info("[Storage] MySQL表就绪") async def _sync_data(self) -> None: """同步数据""" @@ -252,39 +243,33 @@ class MysqlStorage(BaseStorage): await self._file.save_tokens(data) else: await self._file.save_config(data) - logger.info(f"[Storage] {table.split('_')[1]}数据已从数据库同步到文件") + logger.info(f"[Storage] {table.split('_')[1]}数据已从DB同步") else: - if table == "grok_tokens": - file_data = await self._file.load_tokens() - if file_data.get(key) or file_data.get("ssoSuper"): - await self._save_db(table, file_data) - logger.info("[Storage] Token数据已从文件初始化到数据库") - else: - file_data = await self._file.load_config() - if file_data.get(key) or file_data.get("grok"): - await self._save_db(table, file_data) - logger.info("[Storage] 配置数据已从文件初始化到数据库") + file_data = await (self._file.load_tokens() if table == "grok_tokens" else self._file.load_config()) + if file_data.get(key) or (table == "grok_tokens" and file_data.get("ssoSuper")): + await self._save_db(table, file_data) + logger.info(f"[Storage] {table.split('_')[1]}数据已初始化到DB") except Exception as e: - logger.warning(f"[Storage] 数据同步失败: {e}") + logger.warning(f"[Storage] 同步失败: {e}") - async def _load_db(self, table: str) -> Optional[Dict[str, Any]]: - """从数据库加载数据""" + async def _load_db(self, table: str) -> Optional[Dict]: + """从DB加载""" try: async with self._pool.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute(f"SELECT data FROM {table} ORDER BY id DESC LIMIT 1") result = await cursor.fetchone() - return json.loads(result[0]) if result else None + return orjson.loads(result[0]) if result else None except Exception as e: - logger.error(f"[Storage] 从数据库加载{table}失败: {e}") + logger.error(f"[Storage] 加载{table}失败: {e}") return None - async def _save_db(self, table: str, data: Dict[str, Any]) -> None: - """保存数据到数据库""" + async def _save_db(self, table: str, data: Dict) -> None: + """保存到DB""" try: async with self._pool.acquire() as conn: async with conn.cursor() as cursor: - json_data = json.dumps(data, ensure_ascii=False) + json_data = orjson.dumps(data).decode() await cursor.execute(f"SELECT id FROM {table} ORDER BY id DESC LIMIT 1") result = await cursor.fetchone() @@ -293,24 +278,24 @@ class MysqlStorage(BaseStorage): else: await cursor.execute(f"INSERT INTO {table} (data) VALUES (%s)", (json_data,)) except Exception as e: - logger.error(f"[Storage] 保存数据到{table}失败: {e}") + logger.error(f"[Storage] 保存{table}失败: {e}") raise async def load_tokens(self) -> Dict[str, Any]: - """加载token数据""" + """加载token""" return await self._file.load_tokens() async def save_tokens(self, data: Dict[str, Any]) -> None: - """保存token数据""" + """保存token""" await self._file.save_tokens(data) await self._save_db("grok_tokens", data) async def load_config(self) -> Dict[str, Any]: - """加载配置数据""" + """加载配置""" return await self._file.load_config() async def save_config(self, data: Dict[str, Any]) -> None: - """保存配置数据""" + """保存配置""" await self._file.save_config(data) await self._save_db("grok_settings", data) @@ -319,11 +304,11 @@ class MysqlStorage(BaseStorage): if self._pool: self._pool.close() await self._pool.wait_closed() - logger.info("[Storage] MySQL连接池已关闭") + logger.info("[Storage] MySQL已关闭") class RedisStorage(BaseStorage): - """Redis存储实现""" + """Redis存储""" def __init__(self, redis_url: str, data_dir: Path): self.redis_url = redis_url @@ -334,22 +319,18 @@ class RedisStorage(BaseStorage): async def init_db(self) -> None: """初始化Redis""" try: - import redis.asyncio as redis - parsed = self._parse_url(self.redis_url) - logger.info(f"[Storage] 解析Redis URL: host={parsed['host']}, port={parsed['port']}, db={parsed.get('db', 0)}, username={parsed.get('username')}, password={'***' if parsed.get('password') else None}") + import redis.asyncio as aioredis + parsed = urlparse(self.redis_url) + db = int(parsed.path.lstrip('/')) if parsed.path and parsed.path != '/' else 0 + logger.info(f"[Storage] Redis: {parsed.hostname}:{parsed.port or 6379}/{db}") - # 创建Redis连接 - self._redis = redis.Redis( - host=parsed['host'], port=parsed['port'], password=parsed.get('password'), - username=parsed.get('username'), db=parsed.get('db', 0), - encoding="utf-8", decode_responses=True + self._redis = aioredis.Redis.from_url( + self.redis_url, encoding="utf-8", decode_responses=True ) - # 测试连接 await self._redis.ping() - logger.info(f"[Storage] Redis连接成功: {parsed['host']}:{parsed['port']}/{parsed['db']}") + logger.info(f"[Storage] Redis连接成功") - # 初始化文件存储和同步数据 await self._file.init_db() await self._sync_data() @@ -359,26 +340,6 @@ class RedisStorage(BaseStorage): logger.error(f"[Storage] Redis初始化失败: {e}") raise - def _parse_url(self, url: str) -> Dict[str, Any]: - """解析Redis URL""" - if url.startswith('redis://'): - url = url[8:] - parsed = urlparse(f'//{url}') - - result = { - 'host': parsed.hostname or 'localhost', - 'port': parsed.port or 6379, - 'db': int(parsed.path.lstrip('/')) if parsed.path and parsed.path != '/' else 0, - 'username': unquote(parsed.username) if parsed.username else None, - 'password': unquote(parsed.password) if parsed.password else None - } - - # Redis 6+ 默认用户名 - if result['password'] and not result['username']: - result['username'] = 'default' - - return result - async def _sync_data(self) -> None: """同步数据""" try: @@ -388,44 +349,43 @@ class RedisStorage(BaseStorage): ]: data = await self._redis.get(key) if data: - parsed = json.loads(data) + parsed = orjson.loads(data) if key == "grok:tokens": await self._file.save_tokens(parsed) else: await self._file.save_config(parsed) - logger.info(f"[Storage] {key.split(':')[1]}数据已从Redis同步到文件") + logger.info(f"[Storage] {key.split(':')[1]}数据已从Redis同步") else: file_data = await file_func() if file_data.get(key_name) or (key == "grok:tokens" and file_data.get("ssoSuper")): - json_data = json.dumps(file_data, ensure_ascii=False) - await self._redis.set(key, json_data) - logger.info(f"[Storage] {key.split(':')[1]}数据已从文件初始化到Redis") + await self._redis.set(key, orjson.dumps(file_data).decode()) + logger.info(f"[Storage] {key.split(':')[1]}数据已初始化到Redis") except Exception as e: - logger.warning(f"[Storage] 数据同步失败: {e}") + logger.warning(f"[Storage] 同步失败: {e}") - async def _save_redis(self, key: str, data: Dict[str, Any]) -> None: + async def _save_redis(self, key: str, data: Dict) -> None: """保存到Redis""" try: - await self._redis.set(key, json.dumps(data, ensure_ascii=False)) + await self._redis.set(key, orjson.dumps(data).decode()) except Exception as e: - logger.error(f"[Storage] 保存到Redis失败: {e}") + logger.error(f"[Storage] 保存Redis失败: {e}") raise async def load_tokens(self) -> Dict[str, Any]: - """加载token数据""" + """加载token""" return await self._file.load_tokens() async def save_tokens(self, data: Dict[str, Any]) -> None: - """保存token数据""" + """保存token""" await self._file.save_tokens(data) await self._save_redis("grok:tokens", data) async def load_config(self) -> Dict[str, Any]: - """加载配置数据""" + """加载配置""" return await self._file.load_config() async def save_config(self, data: Dict[str, Any]) -> None: - """保存配置数据""" + """保存配置""" await self._file.save_config(data) await self._save_redis("grok:settings", data) @@ -433,11 +393,11 @@ class RedisStorage(BaseStorage): """关闭连接""" if self._redis: await self._redis.close() - logger.info("[Storage] Redis连接已关闭") + logger.info("[Storage] Redis已关闭") class StorageManager: - """存储管理器""" + """存储管理器(单例)""" _instance: Optional['StorageManager'] = None _storage: Optional[BaseStorage] = None @@ -457,22 +417,17 @@ class StorageManager: url = os.getenv("DATABASE_URL", "") data_dir = Path(__file__).parents[2] / "data" - storage_classes = { - "mysql": MysqlStorage, - "redis": RedisStorage, - "file": FileStorage - } + classes = {"mysql": MysqlStorage, "redis": RedisStorage, "file": FileStorage} if mode in ("mysql", "redis") and not url: - raise ValueError(f"{mode.upper()}模式需要DATABASE_URL环境变量") + raise ValueError(f"{mode.upper()}模式需要DATABASE_URL") - storage_class = storage_classes.get(mode, FileStorage) + storage_class = classes.get(mode, FileStorage) self._storage = storage_class(url, data_dir) if mode != "file" else storage_class(data_dir) await self._storage.init_db() self._initialized = True - logger.info(f"[Storage] 使用{mode}存储模式") - logger.info("[Storage] 存储管理器初始化完成") + logger.info(f"[Storage] 使用{mode}模式") def get_storage(self) -> BaseStorage: """获取存储实例""" @@ -486,5 +441,5 @@ class StorageManager: await self._storage.close() -# 全局存储管理器实例 +# 全局实例 storage_manager = StorageManager() diff --git a/app/models/grok_models.py b/app/models/grok_models.py index 1f99064..b067c54 100644 --- a/app/models/grok_models.py +++ b/app/models/grok_models.py @@ -1,97 +1,130 @@ -from enum import Enum +"""Grok 模型配置和枚举定义""" -# 模型配置字典 -_MODEL_CONFIG = { - "grok-3-fast": { - "grok_model": ("grok-3", "MODEL_MODE_FAST"), - "rate_limit_model": "grok-3", - "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, - "requires_super": False, - "display_name": "Grok 3 Fast", - "description": "Fast and efficient Grok 3 model", - "raw_model_path": "xai/grok-3", - "default_temperature": 1.0, - "default_max_output_tokens": 8192, - "supported_max_output_tokens": 131072, - "default_top_p": 0.95 - }, - "grok-4-fast": { - "grok_model": ("grok-4-mini-thinking-tahoe", "MODEL_MODE_GROK_4_MINI_THINKING"), - "rate_limit_model": "grok-4-mini-thinking-tahoe", - "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, - "requires_super": False, - "display_name": "Grok 4 Fast", - "description": "Fast version of Grok 4 with mini thinking capabilities", - "raw_model_path": "xai/grok-4-mini-thinking-tahoe", - "default_temperature": 1.0, - "default_max_output_tokens": 8192, - "supported_max_output_tokens": 131072, - "default_top_p": 0.95 - }, - "grok-4-fast-expert": { - "grok_model": ("grok-4-mini-thinking-tahoe", "MODEL_MODE_EXPERT"), - "rate_limit_model": "grok-4-mini-thinking-tahoe", - "cost": {"type": "high_cost", "multiplier": 4, "description": "计4次调用"}, - "requires_super": False, - "display_name": "Grok 4 Fast Expert", - "description": "Expert mode of Grok 4 Fast with enhanced reasoning", - "raw_model_path": "xai/grok-4-mini-thinking-tahoe", - "default_temperature": 1.0, - "default_max_output_tokens": 32768, - "supported_max_output_tokens": 131072, - "default_top_p": 0.95 - }, - "grok-4-expert": { - "grok_model": ("grok-4", "MODEL_MODE_EXPERT"), - "rate_limit_model": "grok-4", - "cost": {"type": "high_cost", "multiplier": 4, "description": "计4次调用"}, - "requires_super": False, - "display_name": "Grok 4 Expert", - "description": "Full Grok 4 model with expert mode capabilities", - "raw_model_path": "xai/grok-4", - "default_temperature": 1.0, - "default_max_output_tokens": 32768, - "supported_max_output_tokens": 131072, - "default_top_p": 0.95 - }, - "grok-4-heavy": { - "grok_model": ("grok-4-heavy", "MODEL_MODE_HEAVY"), - "rate_limit_model": "grok-4-heavy", - "cost": {"type": "independent", "multiplier": 1, "description": "独立计费,只有Super用户可用"}, - "requires_super": True, - "display_name": "Grok 4 Heavy", - "description": "Most powerful Grok 4 model with heavy computational capabilities. Requires Super Token for access.", - "raw_model_path": "xai/grok-4-heavy", - "default_temperature": 1.0, - "default_max_output_tokens": 65536, - "supported_max_output_tokens": 131072, - "default_top_p": 0.95 - }, - "grok-imagine-0.9": { - "grok_model": ("grok-3", "MODEL_MODE_FAST"), - "rate_limit_model": "grok-3", - "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, - "requires_super": False, - "display_name": "Grok Imagine 0.9", - "description": "Video generation model powered by Grok", - "raw_model_path": "xai/grok-imagine-0.9", - "default_temperature": 1.0, - "default_max_output_tokens": 8192, - "supported_max_output_tokens": 131072, - "default_top_p": 0.95, - "is_video_model": True - } +from enum import Enum +from typing import Dict, Any, Tuple + + +# 模型配置 +_MODEL_CONFIG: Dict[str, Dict[str, Any]] = { + "grok-3-fast": { + "grok_model": ("grok-3", "MODEL_MODE_FAST"), + "rate_limit_model": "grok-3", + "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, + "requires_super": False, + "display_name": "Grok 3 Fast", + "description": "Fast and efficient Grok 3 model", + "raw_model_path": "xai/grok-3", + "default_temperature": 1.0, + "default_max_output_tokens": 8192, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-4-fast": { + "grok_model": ("grok-4-mini-thinking-tahoe", "MODEL_MODE_GROK_4_MINI_THINKING"), + "rate_limit_model": "grok-4-mini-thinking-tahoe", + "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, + "requires_super": False, + "display_name": "Grok 4 Fast", + "description": "Fast version of Grok 4 with mini thinking capabilities", + "raw_model_path": "xai/grok-4-mini-thinking-tahoe", + "default_temperature": 1.0, + "default_max_output_tokens": 8192, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-4-fast-expert": { + "grok_model": ("grok-4-mini-thinking-tahoe", "MODEL_MODE_EXPERT"), + "rate_limit_model": "grok-4-mini-thinking-tahoe", + "cost": {"type": "high_cost", "multiplier": 4, "description": "计4次调用"}, + "requires_super": False, + "display_name": "Grok 4 Fast Expert", + "description": "Expert mode of Grok 4 Fast with enhanced reasoning", + "raw_model_path": "xai/grok-4-mini-thinking-tahoe", + "default_temperature": 1.0, + "default_max_output_tokens": 32768, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-4-expert": { + "grok_model": ("grok-4", "MODEL_MODE_EXPERT"), + "rate_limit_model": "grok-4", + "cost": {"type": "high_cost", "multiplier": 4, "description": "计4次调用"}, + "requires_super": False, + "display_name": "Grok 4 Expert", + "description": "Full Grok 4 model with expert mode capabilities", + "raw_model_path": "xai/grok-4", + "default_temperature": 1.0, + "default_max_output_tokens": 32768, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-4-heavy": { + "grok_model": ("grok-4-heavy", "MODEL_MODE_HEAVY"), + "rate_limit_model": "grok-4-heavy", + "cost": {"type": "independent", "multiplier": 1, "description": "独立计费,只有Super用户可用"}, + "requires_super": True, + "display_name": "Grok 4 Heavy", + "description": "Most powerful Grok 4 model with heavy computational capabilities. Requires Super Token for access.", + "raw_model_path": "xai/grok-4-heavy", + "default_temperature": 1.0, + "default_max_output_tokens": 65536, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-4.1": { + "grok_model": ("grok-4-1-non-thinking-w-tool", "MODEL_MODE_GROK_4_1"), + "rate_limit_model": "grok-4-1-non-thinking-w-tool", + "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, + "requires_super": False, + "display_name": "Grok 4.1", + "description": "Latest Grok 4.1 model with tool capabilities", + "raw_model_path": "xai/grok-4-1-non-thinking-w-tool", + "default_temperature": 1.0, + "default_max_output_tokens": 8192, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-4.1-thinking": { + "grok_model": ("grok-4-1-thinking-1108b", "MODEL_MODE_AUTO"), + "rate_limit_model": "grok-4-1-thinking-1108b", + "cost": {"type": "high_cost", "multiplier": 1, "description": "计1次调用"}, + "requires_super": False, + "display_name": "Grok 4.1 Thinking", + "description": "Grok 4.1 model with advanced thinking and tool capabilities", + "raw_model_path": "xai/grok-4-1-thinking-1108b", + "default_temperature": 1.0, + "default_max_output_tokens": 32768, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95 + }, + "grok-imagine-0.9": { + "grok_model": ("grok-3", "MODEL_MODE_FAST"), + "rate_limit_model": "grok-3", + "cost": {"type": "low_cost", "multiplier": 1, "description": "计1次调用"}, + "requires_super": False, + "display_name": "Grok Imagine 0.9", + "description": "Image and video generation model. Supports text-to-image and image-to-video generation.", + "raw_model_path": "xai/grok-imagine-0.9", + "default_temperature": 1.0, + "default_max_output_tokens": 8192, + "supported_max_output_tokens": 131072, + "default_top_p": 0.95, + "is_video_model": True } +} + class TokenType(Enum): - """Token类型枚举""" - NORMAL = "ssoNormal" # 普通用户Token - SUPER = "ssoSuper" # 超级用户Token + """Token类型""" + NORMAL = "ssoNormal" + SUPER = "ssoSuper" class Models(Enum): - """支持的模型枚举""" + """支持的模型""" GROK_3_FAST = "grok-3-fast" + GROK_4_1 = "grok-4.1" + GROK_4_1_THINKING = "grok-4.1-thinking" GROK_4_FAST = "grok-4-fast" GROK_4_FAST_EXPERT = "grok-4-fast-expert" GROK_4_EXPERT = "grok-4-expert" @@ -99,8 +132,8 @@ class Models(Enum): GROK_IMAGINE_0_9 = "grok-imagine-0.9" @classmethod - def get_model_info(cls, model: str) -> dict: - """获取模型的完整配置信息""" + def get_model_info(cls, model: str) -> Dict[str, Any]: + """获取模型配置""" return _MODEL_CONFIG.get(model, {}) @classmethod @@ -109,15 +142,22 @@ class Models(Enum): return model in _MODEL_CONFIG @classmethod - def to_grok(cls, model: str) -> tuple[str, str]: - """转换为Grok内部模型名和模式类型""" + def to_grok(cls, model: str) -> Tuple[str, str]: + """转换为Grok内部模型名和模式 + + Returns: + (模型名, 模式类型) 元组 + """ config = _MODEL_CONFIG.get(model) - if config: - return config["grok_model"] - return model, "MODEL_MODE_FAST" + return config["grok_model"] if config else (model, "MODEL_MODE_FAST") @classmethod def to_rate_limit(cls, model: str) -> str: - """转换为速率限制接口模型名""" + """转换为速率限制模型名""" config = _MODEL_CONFIG.get(model) - return config["rate_limit_model"] if config else model \ No newline at end of file + return config["rate_limit_model"] if config else model + + @classmethod + def get_all_model_names(cls) -> list[str]: + """获取所有模型名称""" + return list(_MODEL_CONFIG.keys()) \ No newline at end of file diff --git a/app/models/openai_schema.py b/app/models/openai_schema.py index 37b6eeb..540eed9 100644 --- a/app/models/openai_schema.py +++ b/app/models/openai_schema.py @@ -1,16 +1,18 @@ -"""OpenAI 请求-响应模型""" +"""OpenAI 请求-响应模型定义""" from fastapi import HTTPException from typing import Optional, List, Union, Dict, Any from pydantic import BaseModel, Field, field_validator +from app.models.grok_models import Models + class OpenAIChatRequest(BaseModel): - """OpenAI 聊天请求模型""" + """OpenAI聊天请求""" model: str = Field(..., description="模型名称", min_length=1) messages: List[Dict[str, Any]] = Field(..., description="消息列表", min_length=1) - stream: bool = Field(False, description="启用流式响应") + stream: bool = Field(False, description="流式响应") temperature: Optional[float] = Field(0.7, ge=0, le=2, description="采样温度") max_tokens: Optional[int] = Field(None, ge=1, le=100000, description="最大Token数") top_p: Optional[float] = Field(1.0, ge=0, le=1, description="采样参数") @@ -20,31 +22,19 @@ class OpenAIChatRequest(BaseModel): def validate_messages(cls, v): """验证消息格式""" if not v: - raise HTTPException( - status_code=400, - detail="消息列表不能为空" - ) + raise HTTPException(status_code=400, detail="消息列表不能为空") for msg in v: if not isinstance(msg, dict): - raise HTTPException( - status_code=400, - detail="每个消息必须是一个字典" - ) + raise HTTPException(status_code=400, detail="每个消息必须是字典") if 'role' not in msg: - raise HTTPException( - status_code=400, - detail="消息缺少必填字段 'role'" - ) + raise HTTPException(status_code=400, detail="消息缺少 'role' 字段") if 'content' not in msg: - raise HTTPException( - status_code=400, - detail="消息缺少必填字段 'content'" - ) + raise HTTPException(status_code=400, detail="消息缺少 'content' 字段") if msg['role'] not in ['system', 'user', 'assistant']: raise HTTPException( status_code=400, - detail=f"无效的角色 '{msg['role']}', 必须是 'system', 'user' 或 'assistant'" + detail=f"无效角色 '{msg['role']}', 必须是 system/user/assistant" ) return v @@ -53,29 +43,27 @@ class OpenAIChatRequest(BaseModel): @field_validator('model') def validate_model(cls, v): """验证模型名称""" - allowed_models = [ - 'grok-3-fast', 'grok-4-fast', 'grok-4-fast-expert', - 'grok-4-expert', 'grok-4-heavy', 'grok-imagine-0.9' - ] - if v not in allowed_models: + if not Models.is_valid_model(v): + supported = Models.get_all_model_names() raise HTTPException( status_code=400, - detail=f"不支持的模型 '{v}', 支持的模型: {', '.join(allowed_models)}" + detail=f"不支持的模型 '{v}', 支持: {', '.join(supported)}" ) return v + class OpenAIChatCompletionMessage(BaseModel): """聊天完成消息""" role: str = Field(..., description="角色") - content: str = Field(..., description="消息内容") + content: str = Field(..., description="内容") reference_id: Optional[str] = Field(default=None, description="参考ID") annotations: Optional[List[str]] = Field(default=None, description="注释") class OpenAIChatCompletionChoice(BaseModel): """聊天完成选项""" - index: int = Field(..., description="选项索引") - message: OpenAIChatCompletionMessage = Field(..., description="响应消息") + index: int = Field(..., description="索引") + message: OpenAIChatCompletionMessage = Field(..., description="消息") logprobs: Optional[float] = Field(default=None, description="对数概率") finish_reason: str = Field(default="stop", description="完成原因") @@ -85,20 +73,20 @@ class OpenAIChatCompletionResponse(BaseModel): id: str = Field(..., description="响应ID") object: str = Field("chat.completion", description="对象类型") created: int = Field(..., description="创建时间戳") - model: str = Field(..., description="使用的模型") - choices: List[OpenAIChatCompletionChoice] = Field(..., description="响应选项") + model: str = Field(..., description="模型") + choices: List[OpenAIChatCompletionChoice] = Field(..., description="选项") usage: Optional[Dict[str, Any]] = Field(None, description="令牌使用") class OpenAIChatCompletionChunkMessage(BaseModel): - """流式响应消息片段""" + """流式消息片段""" role: str = Field(..., description="角色") - content: str = Field(..., description="消息内容") + content: str = Field(..., description="内容") class OpenAIChatCompletionChunkChoice(BaseModel): - """流式响应选项""" - index: int = Field(..., description="选项索引") + """流式选项""" + index: int = Field(..., description="索引") delta: Optional[Union[Dict[str, Any], OpenAIChatCompletionChunkMessage]] = Field( None, description="Delta数据" ) @@ -106,10 +94,10 @@ class OpenAIChatCompletionChunkChoice(BaseModel): class OpenAIChatCompletionChunkResponse(BaseModel): - """流式聊天完成响应""" + """流式聊天响应""" id: str = Field(..., description="响应ID") object: str = Field(default="chat.completion.chunk", description="对象类型") created: int = Field(..., description="创建时间戳") - model: str = Field(..., description="使用的模型") + model: str = Field(..., description="模型") system_fingerprint: Optional[str] = Field(default=None, description="系统指纹") - choices: List[OpenAIChatCompletionChunkChoice] = Field(..., description="响应选项") \ No newline at end of file + choices: List[OpenAIChatCompletionChunkChoice] = Field(..., description="选项") \ No newline at end of file diff --git a/app/services/grok/cache.py b/app/services/grok/cache.py index d9d9220..95e86b1 100644 --- a/app/services/grok/cache.py +++ b/app/services/grok/cache.py @@ -1,9 +1,9 @@ -"""缓存服务模块""" +"""缓存服务模块 - 提供图片和视频的下载、缓存和清理功能""" import asyncio import base64 from pathlib import Path -from typing import Optional +from typing import Optional, Tuple from curl_cffi.requests import AsyncSession from app.core.config import setting @@ -11,166 +11,234 @@ from app.core.logger import logger from app.services.grok.statsig import get_dynamic_headers +# 常量定义 +MIME_TYPES = { + '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png', + '.gif': 'image/gif', '.webp': 'image/webp', '.bmp': 'image/bmp', +} +DEFAULT_MIME = 'image/jpeg' +ASSETS_URL = "https://assets.grok.com" + + class CacheService: """缓存服务基类""" - def __init__(self, cache_type: str): - """初始化缓存服务""" + def __init__(self, cache_type: str, timeout: float = 30.0): self.cache_type = cache_type self.cache_dir = Path(f"data/temp/{cache_type}") self.cache_dir.mkdir(parents=True, exist_ok=True) + self.timeout = timeout + self._cleanup_lock = asyncio.Lock() - def _cache_path(self, file_path: str) -> Path: - """获取缓存文件的完整路径""" - filename = file_path.lstrip('/').replace('/', '-') - return self.cache_dir / filename + def _get_path(self, file_path: str) -> Path: + """转换文件路径为缓存路径""" + return self.cache_dir / file_path.lstrip('/').replace('/', '-') - async def download_file(self, file_path: str, auth_token: str, timeout: float = 30.0) -> Optional[Path]: + def _log(self, level: str, msg: str): + """统一日志输出""" + getattr(logger, level)(f"[{self.cache_type.upper()}Cache] {msg}") + + def _build_headers(self, file_path: str, auth_token: str) -> dict: + """构建请求头""" + cf = setting.grok_config.get("cf_clearance", "") + return { + **get_dynamic_headers(pathname=file_path), + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8", + "Sec-Fetch-Dest": "document", + "Sec-Fetch-Mode": "navigate", + "Sec-Fetch-Site": "same-site", + "Sec-Fetch-User": "?1", + "Upgrade-Insecure-Requests": "1", + "Referer": "https://grok.com/", + "Cookie": f"{auth_token};{cf}" if cf else auth_token + } + + async def download(self, file_path: str, auth_token: str, timeout: Optional[float] = None) -> Optional[Path]: """下载并缓存文件""" - cache_path = self._cache_path(file_path) + cache_path = self._get_path(file_path) if cache_path.exists(): - logger.debug(f"[{self.cache_type.upper()}Cache] 文件已缓存: {cache_path}") + self._log("debug", "文件已缓存") return cache_path - try: - cf_clearance = setting.grok_config.get("cf_clearance", "") - headers = { - **get_dynamic_headers(pathname=file_path), - "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", - "Sec-Fetch-Dest": "document", - "Sec-Fetch-Mode": "navigate", - "Sec-Fetch-Site": "same-site", - "Sec-Fetch-User": "?1", - "Upgrade-Insecure-Requests": "1", - "Referer": "https://grok.com/", - "Cookie": f"{auth_token};{cf_clearance}" if cf_clearance else auth_token - } + # 外层重试:可配置状态码(401/429等) + retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) + MAX_OUTER_RETRY = 3 + + for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次 + try: + # 内层重试:403代理池重试(cache使用缓存代理,不支持代理池,所以403只重试一次) + max_403_retries = 5 + retry_403_count = 0 + + while retry_403_count <= max_403_retries: + proxy = setting.get_proxy("cache") + proxies = {"http": proxy, "https": proxy} if proxy else {} + + if proxy and outer_retry == 0 and retry_403_count == 0: + self._log("debug", f"使用代理: {proxy.split('@')[-1] if '@' in proxy else proxy}") - proxy_url = setting.grok_config.get("proxy_url") - proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else {} - - async with AsyncSession() as session: - logger.debug(f"[{self.cache_type.upper()}Cache] 开始下载: https://assets.grok.com{file_path}") - response = await session.get( - f"https://assets.grok.com{file_path}", - headers=headers, - proxies=proxies, - timeout=timeout, - allow_redirects=True, - impersonate="chrome133a" - ) - response.raise_for_status() - cache_path.write_bytes(response.content) - logger.debug(f"[{self.cache_type.upper()}Cache] 文件已缓存: {cache_path} ({len(response.content)} bytes)") - asyncio.create_task(self.cleanup_cache()) - return cache_path - except Exception as e: - logger.error(f"[{self.cache_type.upper()}Cache] 下载文件失败: {e}") - return None + async with AsyncSession() as session: + url = f"{ASSETS_URL}{file_path}" + if outer_retry == 0 and retry_403_count == 0: + self._log("debug", f"下载: {url}") + + response = await session.get( + url, + headers=self._build_headers(file_path, auth_token), + proxies=proxies, + timeout=timeout or self.timeout, + allow_redirects=True, + impersonate="chrome133a" + ) + + # 检查403错误 - 内层重试(cache不使用代理池,所以直接失败) + if response.status_code == 403: + retry_403_count += 1 + + if retry_403_count <= max_403_retries: + self._log("warning", f"遇到403错误,正在重试 ({retry_403_count}/{max_403_retries})...") + await asyncio.sleep(0.5) + continue + + self._log("error", f"403错误,已重试{retry_403_count-1}次,放弃") + return None + + # 检查可配置状态码错误 - 外层重试 + if response.status_code in retry_codes: + if outer_retry < MAX_OUTER_RETRY: + delay = (outer_retry + 1) * 0.1 # 渐进延迟:0.1s, 0.2s, 0.3s + self._log("warning", f"遇到{response.status_code}错误,外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY}),等待{delay}s...") + await asyncio.sleep(delay) + break # 跳出内层循环,进入外层重试 + else: + self._log("error", f"{response.status_code}错误,已重试{outer_retry}次,放弃") + return None + + response.raise_for_status() + + cache_path.write_bytes(response.content) + + if outer_retry > 0 or retry_403_count > 0: + self._log("info", f"重试成功!") + else: + self._log("debug", "缓存成功") + + # 异步清理(带错误处理) + asyncio.create_task(self._safe_cleanup()) + return cache_path + + except Exception as e: + if outer_retry < MAX_OUTER_RETRY - 1: + self._log("warning", f"下载异常: {e},外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})...") + await asyncio.sleep(0.5) + continue + + self._log("error", f"下载失败: {e}(已重试{outer_retry}次)") + return None + + return None def get_cached(self, file_path: str) -> Optional[Path]: - """获取缓存的文件路径""" - cache_path = self._cache_path(file_path) - return cache_path if cache_path.exists() else None + """获取已缓存的文件""" + path = self._get_path(file_path) + return path if path.exists() else None - async def cleanup_cache(self): - """清理缓存目录,确保不超过配置的大小限制""" + async def _safe_cleanup(self): + """安全清理(捕获异常)""" try: - max_size_mb = setting.global_config.get(f"{self.cache_type}_cache_max_size_mb", 500) - max_size_bytes = max_size_mb * 1024 * 1024 - - files = [(fp, (stat := fp.stat()).st_size, stat.st_mtime) - for fp in self.cache_dir.glob("*") if fp.is_file()] - total_size = sum(size for _, size, _ in files) - - if total_size <= max_size_bytes: - logger.debug(f"[{self.cache_type.upper()}Cache] 缓存大小 {total_size / 1024 / 1024:.2f}MB,未超限") - return - - logger.info(f"[{self.cache_type.upper()}Cache] 缓存大小 {total_size / 1024 / 1024:.2f}MB 超过限制 {max_size_mb}MB,开始清理") - files.sort(key=lambda x: x[2]) - - for file_path, size, _ in files: - if total_size <= max_size_bytes: - break - file_path.unlink() - total_size -= size - logger.debug(f"[{self.cache_type.upper()}Cache] 已删除缓存文件: {file_path}") - - logger.info(f"[{self.cache_type.upper()}Cache] 缓存清理完成,当前大小 {total_size / 1024 / 1024:.2f}MB") + await self.cleanup() except Exception as e: - logger.error(f"[{self.cache_type.upper()}Cache] 清理缓存失败: {e}") + self._log("error", f"后台清理失败: {e}") + + async def cleanup(self): + """清理超限缓存""" + if self._cleanup_lock.locked(): + return + + async with self._cleanup_lock: + try: + max_mb = setting.global_config.get(f"{self.cache_type}_cache_max_size_mb", 500) + max_bytes = max_mb * 1024 * 1024 + + # 获取文件信息 (path, size, mtime) + files = [(f, (s := f.stat()).st_size, s.st_mtime) + for f in self.cache_dir.glob("*") if f.is_file()] + total = sum(size for _, size, _ in files) + + if total <= max_bytes: + return + + self._log("info", f"清理缓存 {total/1024/1024:.1f}MB -> {max_mb}MB") + + # 删除最旧的文件 + for path, size, _ in sorted(files, key=lambda x: x[2]): + if total <= max_bytes: + break + path.unlink() + total -= size + + self._log("info", f"清理完成: {total/1024/1024:.1f}MB") + except Exception as e: + self._log("error", f"清理失败: {e}") -class ImageCacheService(CacheService): +class ImageCache(CacheService): """图片缓存服务""" def __init__(self): - super().__init__("image") + super().__init__("image", timeout=30.0) - async def download_image(self, image_path: str, auth_token: str) -> Optional[Path]: - """下载并缓存图片""" - return await self.download_file(image_path, auth_token, timeout=30.0) - - def get_cached(self, image_path: str) -> Optional[Path]: - """获取缓存的图片路径""" - return super().get_cached(image_path) + async def download_image(self, path: str, token: str) -> Optional[Path]: + """下载图片""" + return await self.download(path, token) @staticmethod def to_base64(image_path: Path) -> Optional[str]: - """将图片转换为 base64 编码""" + """图片转base64""" try: if not image_path.exists(): - logger.error(f"[ImageCache] 图片文件不存在: {image_path}") + logger.error(f"[ImageCache] 文件不存在: {image_path}") return None - with open(image_path, "rb") as f: - base64_data = base64.b64encode(f.read()).decode('utf-8') - - mime_type = {'.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png', - '.gif': 'image/gif', '.webp': 'image/webp'}.get(image_path.suffix.lower(), 'image/jpeg') - - return f"data:{mime_type};base64,{base64_data}" + data = base64.b64encode(image_path.read_bytes()).decode() + mime = MIME_TYPES.get(image_path.suffix.lower(), DEFAULT_MIME) + return f"data:{mime};base64,{data}" except Exception as e: - logger.error(f"[ImageCache] 图片转 base64 失败: {e}") + logger.error(f"[ImageCache] 转换失败: {e}") return None - async def download_base64(self, image_path: str, auth_token: str) -> Optional[str]: - """下载图片并转换为 base64 编码(转换后立即删除缓存文件)""" + async def download_base64(self, path: str, token: str) -> Optional[str]: + """下载并转为base64(自动删除临时文件)""" try: - cache_path = await self.download_file(image_path, auth_token, timeout=30.0) + cache_path = await self.download(path, token) if not cache_path: return None - base64_str = self.to_base64(cache_path) - + result = self.to_base64(cache_path) + + # 清理临时文件 try: cache_path.unlink() - logger.debug(f"[ImageCache] 已删除临时文件: {cache_path}") except Exception as e: logger.warning(f"[ImageCache] 删除临时文件失败: {e}") - return base64_str + return result except Exception as e: - logger.error(f"[ImageCache] 下载并转换 base64 失败: {e}") + logger.error(f"[ImageCache] 下载base64失败: {e}") return None -class VideoCacheService(CacheService): +class VideoCache(CacheService): """视频缓存服务""" def __init__(self): - super().__init__("video") + super().__init__("video", timeout=60.0) - async def download_video(self, video_path: str, auth_token: str) -> Optional[Path]: - """下载并缓存视频""" - return await self.download_file(video_path, auth_token, timeout=60.0) - - def get_cached(self, video_path: str) -> Optional[Path]: - """获取缓存的视频路径""" - return super().get_cached(video_path) + async def download_video(self, path: str, token: str) -> Optional[Path]: + """下载视频""" + return await self.download(path, token) # 全局实例 -image_cache_service = ImageCacheService() -video_cache_service = VideoCacheService() +image_cache_service = ImageCache() +video_cache_service = VideoCache() diff --git a/app/services/grok/client.py b/app/services/grok/client.py index b275308..0403de7 100644 --- a/app/services/grok/client.py +++ b/app/services/grok/client.py @@ -1,9 +1,8 @@ -"""Grok API 客户端模块""" +"""Grok API 客户端 - 处理OpenAI到Grok的请求转换和响应处理""" import asyncio -import json -from typing import Dict, List, Tuple, Any - +import orjson +from typing import Dict, List, Tuple, Any, Optional from curl_cffi import requests as curl_requests from app.core.config import setting @@ -13,103 +12,163 @@ from app.services.grok.processer import GrokResponseProcessor from app.services.grok.statsig import get_dynamic_headers from app.services.grok.token import token_manager from app.services.grok.upload import ImageUploadManager +from app.services.grok.create import PostCreateManager from app.core.exception import GrokApiException -# 常量定义 -GROK_API_ENDPOINT = "https://grok.com/rest/app-chat/conversations/new" -REQUEST_TIMEOUT = 120 -IMPERSONATE_BROWSER = "chrome133a" + +# 常量 +API_ENDPOINT = "https://grok.com/rest/app-chat/conversations/new" +TIMEOUT = 120 +BROWSER = "chrome133a" +MAX_RETRY = 3 +MAX_UPLOADS = 20 # 提高并发上传限制以支持更高并发 class GrokClient: """Grok API 客户端""" + + _upload_sem = None # 延迟初始化 @staticmethod - async def openai_to_grok(openai_request: dict): - """转换OpenAI请求为Grok请求并处理响应""" - model = openai_request["model"] - messages = openai_request["messages"] - stream = openai_request.get("stream", False) + def _get_upload_semaphore(): + """获取上传信号量(动态配置)""" + if GrokClient._upload_sem is None: + # 从配置读取,如果不可用则使用默认值 + max_concurrency = setting.global_config.get("max_upload_concurrency", MAX_UPLOADS) + GrokClient._upload_sem = asyncio.Semaphore(max_concurrency) + logger.debug(f"[Client] 初始化上传并发限制: {max_concurrency}") + return GrokClient._upload_sem - logger.debug(f"[Client] 处理请求 - 模型:{model}, 消息数:{len(messages)}, 流式:{stream}") - - # 提取消息内容和图片URL - content, image_urls = GrokClient._extract_content(messages) - - # 获取认证令牌和模型信息 - auth_token = token_manager.get_token(model) - model_name, model_mode = Models.to_grok(model) - - # 检查是否为视频模型 - is_video_model = Models.get_model_info(model).get("is_video_model", False) + @staticmethod + async def openai_to_grok(request: dict): + """转换OpenAI请求为Grok请求""" + model = request["model"] + content, images = GrokClient._extract_content(request["messages"]) + stream = request.get("stream", False) - # 视频模型特殊处理:只允许一张图片 - if is_video_model and len(image_urls) > 1: - logger.warning(f"[Client] 视频模型只允许一张图片,当前有{len(image_urls)}张,只使用第一张") - image_urls = image_urls[:1] + # 获取模型信息 + info = Models.get_model_info(model) + grok_model, mode = Models.to_grok(model) + is_video = info.get("is_video_model", False) - # 上传图片并获取附件ID列表 - image_attachments = await GrokClient._upload_imgs(image_urls, auth_token) + # 视频模型限制 + if is_video and len(images) > 1: + logger.warning(f"[Client] 视频模型仅支持1张图片,已截取前1张") + images = images[:1] + + return await GrokClient._retry(model, content, images, grok_model, mode, is_video, stream) - # 视频模型:文本添加 --mode=custom - if is_video_model: - content = f"{content} --mode=custom" - logger.debug(f"[Client] 视频模型文本处理: {content}") + @staticmethod + async def _retry(model: str, content: str, images: List[str], grok_model: str, mode: str, is_video: bool, stream: bool): + """重试请求""" + last_err = None - # 构建Grok请求载荷 - payload = GrokClient._build_payload(content, model_name, model_mode, image_attachments, is_video_model) + for i in range(MAX_RETRY): + try: + token = token_manager.get_token(model) + img_ids, img_uris = await GrokClient._upload(images, token) - return await GrokClient._send_request(payload, auth_token, model, stream) + # 视频模型创建会话 + post_id = None + if is_video and img_ids and img_uris: + post_id = await GrokClient._create_post(img_ids[0], img_uris[0], token) + + payload = GrokClient._build_payload(content, grok_model, mode, img_ids, img_uris, is_video, post_id) + return await GrokClient._request(payload, token, model, stream, post_id) + + except GrokApiException as e: + last_err = e + # 检查是否可重试 + if e.error_code not in ["HTTP_ERROR", "NO_AVAILABLE_TOKEN"]: + raise + + status = e.context.get("status") if e.context else None + retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) + + if status not in retry_codes: + raise + + if i < MAX_RETRY - 1: + logger.warning(f"[Client] 失败(状态:{status}), 重试 {i+1}/{MAX_RETRY}") + await asyncio.sleep(0.5) + + raise last_err or GrokApiException("请求失败", "REQUEST_ERROR") @staticmethod def _extract_content(messages: List[Dict]) -> Tuple[str, List[str]]: - """提取消息内容和图片URL""" - content_parts = [] - image_urls = [] - + """提取文本和图片""" + texts, images = [], [] + for msg in messages: - msg_content = msg.get("content", "") - - # 处理复杂消息格式(包含文本和图片) - if isinstance(msg_content, list): - for item in msg_content: - item_type = item.get("type") - if item_type == "text": - content_parts.append(item.get("text", "")) - elif item_type == "image_url": - url = item.get("image_url", {}).get("url", "") - if url: - image_urls.append(url) - # 处理纯文本消息 + content = msg.get("content", "") + + if isinstance(content, list): + for item in content: + if item.get("type") == "text": + texts.append(item.get("text", "")) + elif item.get("type") == "image_url": + if url := item.get("image_url", {}).get("url"): + images.append(url) else: - content_parts.append(msg_content) - - return "".join(content_parts), image_urls + texts.append(content) + + return "".join(texts), images @staticmethod - async def _upload_imgs(image_urls: List[str], auth_token: str) -> List[str]: - """上传图片并返回附件ID列表""" - image_attachments = [] - # 并发上传所有图片 - tasks = [ImageUploadManager.upload(url, auth_token) for url in image_urls] - results = await asyncio.gather(*tasks, return_exceptions=True) - - for url, result in zip(image_urls, results): + async def _upload(urls: List[str], token: str) -> Tuple[List[str], List[str]]: + """并发上传图片""" + if not urls: + return [], [] + + async def upload_limited(url): + async with GrokClient._get_upload_semaphore(): + return await ImageUploadManager.upload(url, token) + + results = await asyncio.gather(*[upload_limited(u) for u in urls], return_exceptions=True) + + ids, uris = [], [] + for url, result in zip(urls, results): if isinstance(result, Exception): - logger.warning(f"[Client] 图片上传失败: {url}, 错误: {result}") - elif result: - image_attachments.append(result) - - return image_attachments + logger.warning(f"[Client] 上传失败: {url} - {result}") + elif isinstance(result, tuple) and len(result) == 2: + fid, furi = result + if fid: + ids.append(fid) + uris.append(furi) + + return ids, uris @staticmethod - def _build_payload(content: str, model_name: str, model_mode: str, image_attachments: List[str], is_video_model: bool = False) -> Dict[str, Any]: - """构建Grok API请求载荷""" - payload = { + async def _create_post(file_id: str, file_uri: str, token: str) -> Optional[str]: + """创建视频会话""" + try: + result = await PostCreateManager.create(file_id, file_uri, token) + if result and result.get("success"): + return result.get("post_id") + except Exception as e: + logger.warning(f"[Client] 创建会话失败: {e}") + return None + + @staticmethod + def _build_payload(content: str, model: str, mode: str, img_ids: List[str], img_uris: List[str], is_video: bool = False, post_id: str = None) -> Dict: + """构建请求载荷""" + # 视频模型特殊处理 + if is_video and img_uris: + img_msg = f"https://grok.com/imagine/{post_id}" if post_id else f"https://assets.grok.com/post/{img_uris[0]}" + return { + "temporary": True, + "modelName": "grok-3", + "message": f"{img_msg} {content} --mode=custom", + "fileAttachments": img_ids, + "toolOverrides": {"videoGen": True} + } + + # 标准载荷 + return { "temporary": setting.grok_config.get("temporary", True), - "modelName": model_name, + "modelName": model, "message": content, - "fileAttachments": image_attachments, + "fileAttachments": img_ids, "imageAttachments": [], "disableSearch": False, "enableImageGeneration": True, @@ -124,125 +183,134 @@ class GrokClient: "isReasoning": False, "webpageUrls": [], "disableTextFollowUps": True, - "responseMetadata": {"requestModelDetails": {"modelId": model_name}}, + "responseMetadata": {"requestModelDetails": {"modelId": model}}, "disableMemory": False, "forceSideBySide": False, - "modelMode": model_mode, + "modelMode": mode, "isAsyncChat": False } - - # 视频模型特殊配置 - if is_video_model: - payload["toolOverrides"] = {"videoGen": True} - logger.debug("[Client] 视频模型载荷配置: toolOverrides.videoGen = True") - - return payload @staticmethod - async def _send_request(payload: dict, auth_token: str, model: str, stream: bool): - """发送HTTP请求到Grok API""" - # 验证认证令牌 - if not auth_token: + async def _request(payload: dict, token: str, model: str, stream: bool, post_id: str = None): + """发送请求""" + if not token: raise GrokApiException("认证令牌缺失", "NO_AUTH_TOKEN") - try: - # 构建请求头和代理 - headers = GrokClient._build_headers(auth_token) - proxy_config = GrokClient._get_proxy() - - # 构建请求参数 - request_kwargs = { - "headers": headers, - "data": json.dumps(payload), - "impersonate": IMPERSONATE_BROWSER, - "timeout": REQUEST_TIMEOUT, - "stream": True, - "proxies": proxy_config if proxy_config else None - } - - # 在线程池中执行同步HTTP请求,避免阻塞事件循环 - response = await asyncio.to_thread( - curl_requests.post, - GROK_API_ENDPOINT, - **request_kwargs - ) - - logger.debug(f"[Client] API响应状态码: {response.status_code}") - - # 处理非成功响应 - if response.status_code != 200: - GrokClient._handle_error(response, auth_token) - - # 请求成功,重置失败计数 - asyncio.create_task(token_manager.reset_failure(auth_token)) - - # 处理并返回响应 - return await GrokClient._process_response(response, auth_token, model, stream) - - except curl_requests.RequestsError as e: - logger.error(f"[Client] 网络请求错误: {e}") - raise GrokApiException(f"网络错误: {e}", "NETWORK_ERROR") from e - except json.JSONDecodeError as e: - logger.error(f"[Client] JSON解析错误: {e}") - raise GrokApiException(f"JSON解析错误: {e}", "JSON_ERROR") from e - except Exception as e: - logger.error(f"[Client] 未知请求错误: {type(e).__name__}: {e}") - raise GrokApiException(f"请求处理错误: {e}", "REQUEST_ERROR") from e + # 403重试配置 + max_retries = 5 + retry_count = 0 + + while retry_count <= max_retries: + try: + # 构建请求 + headers = GrokClient._build_headers(token) + if model == "grok-imagine-0.9": + file_attachments = payload.get("fileAttachments", []) + ref_id = post_id or (file_attachments[0] if file_attachments else "") + if ref_id: + headers["Referer"] = f"https://grok.com/imagine/{ref_id}" + + # 异步获取代理 + from app.core.proxy_pool import proxy_pool + + # 如果是重试且使用代理池,强制刷新代理 + if retry_count > 0 and proxy_pool._enabled: + logger.info(f"[Client] 403重试 {retry_count}/{max_retries},刷新代理...") + proxy = await proxy_pool.force_refresh() + else: + proxy = await setting.get_proxy_async("service") + + proxies = {"http": proxy, "https": proxy} if proxy else None + + # 执行请求 + response = await asyncio.to_thread( + curl_requests.post, + API_ENDPOINT, + headers=headers, + data=orjson.dumps(payload), + impersonate=BROWSER, + timeout=TIMEOUT, + stream=True, + proxies=proxies + ) + + # 内层403重试:仅当有代理池时触发 + if response.status_code == 403 and proxy_pool._enabled: + retry_count += 1 + + if retry_count <= max_retries: + logger.warning(f"[Client] 遇到403错误,正在重试 ({retry_count}/{max_retries})...") + await asyncio.sleep(0.5) + continue + + # 内层重试全部失败 + logger.error(f"[Client] 403错误,已重试{retry_count-1}次,放弃") + + # 检查响应状态 + if response.status_code != 200: + GrokClient._handle_error(response, token) + + # 成功 - 重置失败计数 + asyncio.create_task(token_manager.reset_failure(token)) + + # 如果是重试成功,记录日志 + if retry_count > 0: + logger.info(f"[Client] 重试成功!第{retry_count}次重试后请求成功") + + # 处理响应 + result = (GrokResponseProcessor.process_stream(response, token) if stream + else await GrokResponseProcessor.process_normal(response, token, model)) + + asyncio.create_task(GrokClient._update_limits(token, model)) + return result + + except curl_requests.RequestsError as e: + logger.error(f"[Client] 网络错误: {e}") + raise GrokApiException(f"网络错误: {e}", "NETWORK_ERROR") from e + except GrokApiException: + # 重新抛出GrokApiException(包括403错误) + raise + except Exception as e: + logger.error(f"[Client] 请求错误: {e}") + raise GrokApiException(f"请求错误: {e}", "REQUEST_ERROR") from e + + # 理论上不应该到这里,但以防万一 + raise GrokApiException("403错误:已达到最大重试次数", "MAX_RETRIES_EXCEEDED") @staticmethod - def _build_headers(auth_token: str) -> Dict[str, str]: + def _build_headers(token: str) -> Dict[str, str]: """构建请求头""" headers = get_dynamic_headers("/rest/app-chat/conversations/new") - - # 构建Cookie - cf_clearance = setting.grok_config.get("cf_clearance", "") - headers["Cookie"] = f"{auth_token};{cf_clearance}" if cf_clearance else auth_token - + cf = setting.grok_config.get("cf_clearance", "") + headers["Cookie"] = f"{token};{cf}" if cf else token return headers @staticmethod - def _get_proxy() -> Dict[str, str]: - """获取代理配置""" - proxy_url = setting.grok_config.get("proxy_url", "") - if proxy_url: - return {"http": proxy_url, "https": proxy_url} - return {} - - @staticmethod - def _handle_error(response, auth_token: str): - """处理错误响应""" - try: - error_data = response.json() - error_message = str(error_data) - except Exception as e: - error_data = response.text - error_message = error_data[:200] if error_data else e - - # 记录Token失败 - asyncio.create_task(token_manager.record_failure(auth_token, response.status_code, error_message)) - + def _handle_error(response, token: str): + """处理错误""" + if response.status_code == 403: + msg = "您的IP被拦截,请尝试以下方法之一: 1.更换IP 2.使用代理 3.配置CF值" + data = {"cf_blocked": True, "status": 403} + logger.warning(f"[Client] {msg}") + else: + try: + data = response.json() + msg = str(data) + except: + data = response.text + msg = data[:200] if data else "未知错误" + + asyncio.create_task(token_manager.record_failure(token, response.status_code, msg)) raise GrokApiException( - f"请求失败: {response.status_code} - {error_message}", + f"请求失败: {response.status_code} - {msg}", "HTTP_ERROR", - {"status": response.status_code, "data": error_data} + {"status": response.status_code, "data": data} ) @staticmethod - async def _process_response(response, auth_token: str, model: str, stream: bool): - """处理API响应""" - if stream: - result = GrokResponseProcessor.process_stream(response, auth_token) - asyncio.create_task(GrokClient._update_rate_limits(auth_token, model)) - else: - result = await GrokResponseProcessor.process_normal(response, auth_token, model) - asyncio.create_task(GrokClient._update_rate_limits(auth_token, model)) - - return result - - @staticmethod - async def _update_rate_limits(auth_token: str, model: str): - """异步更新速率限制信息""" + async def _update_limits(token: str, model: str): + """更新速率限制""" try: - await token_manager.check_limits(auth_token, model) + await token_manager.check_limits(token, model) except Exception as e: - logger.error(f"[Client] 更新速率限制失败: {e}") \ No newline at end of file + logger.error(f"[Client] 更新限制失败: {e}") \ No newline at end of file diff --git a/app/services/grok/create.py b/app/services/grok/create.py new file mode 100644 index 0000000..3652bc3 --- /dev/null +++ b/app/services/grok/create.py @@ -0,0 +1,140 @@ +"""Post创建管理器 - 用于视频生成前的会话创建""" + +import asyncio +import orjson +from typing import Dict, Any, Optional +from curl_cffi.requests import AsyncSession + +from app.services.grok.statsig import get_dynamic_headers +from app.core.exception import GrokApiException +from app.core.config import setting +from app.core.logger import logger + + +# 常量 +ENDPOINT = "https://grok.com/rest/media/post/create" +TIMEOUT = 30 +BROWSER = "chrome133a" + + +class PostCreateManager: + """会话创建管理器""" + + @staticmethod + async def create(file_id: str, file_uri: str, auth_token: str) -> Optional[Dict[str, Any]]: + """创建会话记录 + + Args: + file_id: 文件ID + file_uri: 文件URI + auth_token: 认证令牌 + + Returns: + 会话信息字典,包含post_id等 + """ + # 参数验证 + if not file_id or not file_uri: + raise GrokApiException("文件ID或URI缺失", "INVALID_PARAMS") + if not auth_token: + raise GrokApiException("认证令牌缺失", "NO_AUTH_TOKEN") + + try: + # 构建请求 + data = { + "media_url": f"https://assets.grok.com/{file_uri}", + "media_type": "MEDIA_POST_TYPE_IMAGE" + } + + cf = setting.grok_config.get("cf_clearance", "") + headers = { + **get_dynamic_headers("/rest/media/post/create"), + "Cookie": f"{auth_token};{cf}" if cf else auth_token + } + + # 外层重试:可配置状态码(401/429等) + retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) + MAX_OUTER_RETRY = 3 + + for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次 + # 内层重试:403代理池重试 + max_403_retries = 5 + retry_403_count = 0 + + while retry_403_count <= max_403_retries: + # 异步获取代理(支持代理池) + from app.core.proxy_pool import proxy_pool + + # 如果是403重试且使用代理池,强制刷新代理 + if retry_403_count > 0 and proxy_pool._enabled: + logger.info(f"[PostCreate] 403重试 {retry_403_count}/{max_403_retries},刷新代理...") + proxy = await proxy_pool.force_refresh() + else: + proxy = await setting.get_proxy_async("service") + + proxies = {"http": proxy, "https": proxy} if proxy else None + + # 发送请求 + async with AsyncSession() as session: + response = await session.post( + ENDPOINT, + headers=headers, + json=data, + impersonate=BROWSER, + timeout=TIMEOUT, + proxies=proxies + ) + + # 内层403重试:仅当有代理池时触发 + if response.status_code == 403 and proxy_pool._enabled: + retry_403_count += 1 + + if retry_403_count <= max_403_retries: + logger.warning(f"[PostCreate] 遇到403错误,正在重试 ({retry_403_count}/{max_403_retries})...") + await asyncio.sleep(0.5) + continue + + # 内层重试全部失败 + logger.error(f"[PostCreate] 403错误,已重试{retry_403_count-1}次,放弃") + + # 检查可配置状态码错误 - 外层重试 + if response.status_code in retry_codes: + if outer_retry < MAX_OUTER_RETRY: + delay = (outer_retry + 1) * 0.1 # 渐进延迟:0.1s, 0.2s, 0.3s + logger.warning(f"[PostCreate] 遇到{response.status_code}错误,外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY}),等待{delay}s...") + await asyncio.sleep(delay) + break # 跳出内层循环,进入外层重试 + else: + logger.error(f"[PostCreate] {response.status_code}错误,已重试{outer_retry}次,放弃") + raise GrokApiException(f"创建失败: {response.status_code}错误", "CREATE_ERROR") + + if response.status_code == 200: + result = response.json() + post_id = result.get("post", {}).get("id", "") + + if outer_retry > 0 or retry_403_count > 0: + logger.info(f"[PostCreate] 重试成功!") + + logger.debug(f"[PostCreate] 成功,会话ID: {post_id}") + return { + "post_id": post_id, + "file_id": file_id, + "file_uri": file_uri, + "success": True, + "data": result + } + + # 其他错误处理 + try: + error = response.json() + msg = f"状态码: {response.status_code}, 详情: {error}" + except: + msg = f"状态码: {response.status_code}, 详情: {response.text[:200]}" + + logger.error(f"[PostCreate] 失败: {msg}") + raise GrokApiException(f"创建失败: {msg}", "CREATE_ERROR") + + except GrokApiException: + raise + except Exception as e: + logger.error(f"[PostCreate] 异常: {e}") + raise GrokApiException(f"创建异常: {e}", "CREATE_ERROR") from e diff --git a/app/services/grok/processer.py b/app/services/grok/processer.py index f671a55..2c51c69 100644 --- a/app/services/grok/processer.py +++ b/app/services/grok/processer.py @@ -1,10 +1,10 @@ -"""Grok API 响应处理器模块""" +"""Grok API 响应处理器 - 处理流式和非流式响应""" -import json +import orjson import uuid import time import asyncio -from typing import AsyncGenerator +from typing import AsyncGenerator, Tuple from app.core.config import setting from app.core.exception import GrokApiException @@ -21,61 +21,43 @@ from app.services.grok.cache import image_cache_service, video_cache_service class StreamTimeoutManager: - """流式响应超时管理器""" + """流式响应超时管理""" - def __init__(self, chunk_timeout: int = 120, first_response_timeout: int = 30, total_timeout: int = 600): - """初始化超时管理器 - - Args: - chunk_timeout: 数据块间隔超时(秒) - first_response_timeout: 首次响应超时(秒) - total_timeout: 总超时限制(秒,0表示不限制) - """ + def __init__(self, chunk_timeout: int = 120, first_timeout: int = 30, total_timeout: int = 600): self.chunk_timeout = chunk_timeout - self.first_response_timeout = first_response_timeout + self.first_timeout = first_timeout self.total_timeout = total_timeout - self.start_time = asyncio.get_event_loop().time() self.last_chunk_time = self.start_time - self.first_chunk_received = False + self.first_received = False - def check_timeout(self) -> tuple[bool, str]: - """检查是否超时 + def check_timeout(self) -> Tuple[bool, str]: + """检查超时""" + now = asyncio.get_event_loop().time() - Returns: - (is_timeout, timeout_message): 是否超时及超时信息 - """ - current_time = asyncio.get_event_loop().time() + if not self.first_received and now - self.start_time > self.first_timeout: + return True, f"首次响应超时({self.first_timeout}秒)" - # 检查首次响应超时 - if not self.first_chunk_received: - if current_time - self.start_time > self.first_response_timeout: - return True, f"首次响应超时 ({self.first_response_timeout}秒未收到首个数据块)" + if self.total_timeout > 0 and now - self.start_time > self.total_timeout: + return True, f"总超时({self.total_timeout}秒)" - # 检查总超时 - if self.total_timeout > 0: - if current_time - self.start_time > self.total_timeout: - return True, f"流式响应总超时 ({self.total_timeout}秒)" - - # 检查数据块间隔超时 - if self.first_chunk_received: - if current_time - self.last_chunk_time > self.chunk_timeout: - return True, f"数据块间隔超时 ({self.chunk_timeout}秒无新数据)" + if self.first_received and now - self.last_chunk_time > self.chunk_timeout: + return True, f"数据块超时({self.chunk_timeout}秒)" return False, "" - def mark_chunk_received(self): - """标记收到数据块""" + def mark_received(self): + """标记收到数据""" self.last_chunk_time = asyncio.get_event_loop().time() - self.first_chunk_received = True + self.first_received = True - def get_total_duration(self) -> float: - """获取总耗时(秒)""" + def duration(self) -> float: + """获取总耗时""" return asyncio.get_event_loop().time() - self.start_time class GrokResponseProcessor: - """Grok API 响应处理器""" + """Grok响应处理器""" @staticmethod async def process_normal(response, auth_token: str, model: str = None) -> OpenAIChatCompletionResponse: @@ -86,7 +68,7 @@ class GrokResponseProcessor: if not chunk: continue - data = json.loads(chunk.decode("utf-8")) + data = orjson.loads(chunk) # 错误检查 if error := data.get("error"): @@ -96,312 +78,190 @@ class GrokResponseProcessor: {"code": error.get("code")} ) - # 提取响应数据 grok_resp = data.get("result", {}).get("response", {}) - # 提取视频数据 + # 视频响应 if video_resp := grok_resp.get("streamingVideoGenerationResponse"): if video_url := video_resp.get("videoUrl"): - logger.debug(f"[Processor] 检测到视频生成: {video_url}") - full_video_url = f"https://assets.grok.com/{video_url}" - - # 下载并缓存视频 - try: - cache_path = await video_cache_service.download_video(f"/{video_url}", auth_token) - if cache_path: - video_path = video_url.replace('/', '-') - base_url = setting.global_config.get("base_url", "") - local_video_url = f"{base_url}/images/{video_path}" if base_url else f"/images/{video_path}" - content = f'\n' - else: - content = f'\n' - except Exception as e: - logger.warning(f"[Processor] 缓存视频失败: {e}") - content = f'\n' - - # 返回视频响应 - result = OpenAIChatCompletionResponse( - id=f"chatcmpl-{uuid.uuid4()}", - object="chat.completion", - created=int(time.time()), - model=model or "grok-imagine-0.9", - choices=[OpenAIChatCompletionChoice( - index=0, - message=OpenAIChatCompletionMessage( - role="assistant", - content=content - ), - finish_reason="stop" - )], - usage=None - ) + content = await GrokResponseProcessor._build_video_content(video_url, auth_token) + result = GrokResponseProcessor._build_response(content, model or "grok-imagine-0.9") response_closed = True response.close() return result - # 提取模型响应 + # 模型响应 model_response = grok_resp.get("modelResponse") if not model_response: continue - # 检查 modelResponse 中的错误 if error_msg := model_response.get("error"): - raise GrokApiException( - f"模型响应错误: {error_msg}", - "MODEL_ERROR" - ) + raise GrokApiException(f"模型错误: {error_msg}", "MODEL_ERROR") - # 构建响应内容 - model_name = model_response.get("model") + # 构建内容 content = model_response.get("message", "") + model_name = model_response.get("model") - # 提取图片数据 + # 处理图片 if images := model_response.get("generatedImageUrls"): - # 获取图片返回模式 - image_mode = setting.global_config.get("image_mode", "url") + content = await GrokResponseProcessor._append_images(content, images, auth_token) - for img in images: - try: - if image_mode == "base64": - # base64 模式:下载并转换为 base64 - base64_str = await image_cache_service.download_base64(f"/{img}", auth_token) - if base64_str: - content += f"\n![Generated Image]({base64_str})" - else: - content += f"\n![Generated Image](https://assets.grok.com/{img})" - else: - # url 模式:缓存并返回链接 - cache_path = await image_cache_service.download_image(f"/{img}", auth_token) - if cache_path: - img_path = img.replace('/', '-') - base_url = setting.global_config.get("base_url", "") - img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" - content += f"\n![Generated Image]({img_url})" - else: - content += f"\n![Generated Image](https://assets.grok.com/{img})" - except Exception as e: - logger.warning(f"[Processor] 处理图片失败: {e}") - content += f"\n![Generated Image](https://assets.grok.com/{img})" - - # 返回 OpenAI 响应格式 - result = OpenAIChatCompletionResponse( - id=f"chatcmpl-{uuid.uuid4()}", - object="chat.completion", - created=int(time.time()), - model=model_name, - choices=[OpenAIChatCompletionChoice( - index=0, - message=OpenAIChatCompletionMessage( - role="assistant", - content=content - ), - finish_reason="stop" - )], - usage=None - ) + result = GrokResponseProcessor._build_response(content, model_name) response_closed = True response.close() return result raise GrokApiException("无响应数据", "NO_RESPONSE") - except json.JSONDecodeError as e: + except orjson.JSONDecodeError as e: logger.error(f"[Processor] JSON解析失败: {e}") raise GrokApiException(f"JSON解析失败: {e}", "JSON_ERROR") from e except Exception as e: - logger.error(f"[Processor] 处理响应时发生未知错误: {type(e).__name__}: {e}") + logger.error(f"[Processor] 处理错误: {type(e).__name__}: {e}") raise GrokApiException(f"响应处理错误: {e}", "PROCESS_ERROR") from e finally: - # 确保响应对象被关闭,避免双重释放 if not response_closed and hasattr(response, 'close'): try: response.close() except Exception as e: - logger.warning(f"[Processor] 关闭响应对象时出错: {e}") + logger.warning(f"[Processor] 关闭响应失败: {e}") @staticmethod async def process_stream(response, auth_token: str) -> AsyncGenerator[str, None]: """处理流式响应""" - # 流式生成状态 + # 状态变量 is_image = False is_thinking = False thinking_finished = False - chunk_index = 0 model = None filtered_tags = setting.grok_config.get("filtered_tags", "").split(",") video_progress_started = False last_video_progress = -1 response_closed = False + show_thinking = setting.grok_config.get("show_thinking", True) - # 初始化超时管理器 - timeout_manager = StreamTimeoutManager( + # 超时管理 + timeout_mgr = StreamTimeoutManager( chunk_timeout=setting.grok_config.get("stream_chunk_timeout", 120), - first_response_timeout=setting.grok_config.get("stream_first_response_timeout", 30), + first_timeout=setting.grok_config.get("stream_first_response_timeout", 30), total_timeout=setting.grok_config.get("stream_total_timeout", 600) ) - def make_chunk(chunk_content: str, finish: str = None): - """生成OpenAI格式的响应块""" + def make_chunk(content: str, finish: str = None): + """生成响应块""" chunk_data = OpenAIChatCompletionChunkResponse( id=f"chatcmpl-{uuid.uuid4()}", created=int(time.time()), model=model or "grok-4-mini-thinking-tahoe", choices=[OpenAIChatCompletionChunkChoice( - index=chunk_index, + index=0, delta=OpenAIChatCompletionChunkMessage( role="assistant", - content=chunk_content - ) if chunk_content else {}, + content=content + ) if content else {}, finish_reason=finish )] - ).model_dump() - # SSE 格式返回 - return f"data: {json.dumps(chunk_data)}\n\n" + ) + return f"data: {chunk_data.model_dump_json()}\n\n" try: for chunk in response.iter_lines(): # 超时检查 - is_timeout, timeout_msg = timeout_manager.check_timeout() + is_timeout, timeout_msg = timeout_mgr.check_timeout() if is_timeout: logger.warning(f"[Processor] {timeout_msg}") yield make_chunk("", "stop") yield "data: [DONE]\n\n" return - logger.debug(f"[Processor] 接收到数据块: {len(chunk)} bytes") + logger.debug(f"[Processor] 收到数据块: {len(chunk)} bytes") if not chunk: continue try: - data = json.loads(chunk.decode("utf-8")) + data = orjson.loads(chunk) # 错误检查 if error := data.get("error"): error_msg = error.get('message', '未知错误') - logger.error(f"[Processor] Grok API返回错误: {error_msg}") + logger.error(f"[Processor] API错误: {error_msg}") yield make_chunk(f"Error: {error_msg}", "stop") yield "data: [DONE]\n\n" return - # 提取响应数据 grok_resp = data.get("result", {}).get("response", {}) - logger.debug(f"[Processor] 解析响应数据: {len(grok_resp)} 字段") + logger.debug(f"[Processor] 解析响应: {len(grok_resp)} bytes") if not grok_resp: continue + + timeout_mgr.mark_received() - # 更新模型名称 + # 更新模型 if user_resp := grok_resp.get("userResponse"): if m := user_resp.get("model"): model = m - # 提取视频数据 + # 视频处理 if video_resp := grok_resp.get("streamingVideoGenerationResponse"): progress = video_resp.get("progress", 0) + v_url = video_resp.get("videoUrl") + # 进度更新 if progress > last_video_progress: last_video_progress = progress - - # 添加 标签 - if not video_progress_started: - content = f"视频已生成{progress}%\n" - video_progress_started = True - elif progress < 100: - content = f"视频已生成{progress}%\n" - else: - # 进度100%时关闭 标签并立即处理视频 - content = f"视频已生成{progress}%\n" - - # 立即下载并缓存视频 - if v_url := video_resp.get("videoUrl"): - logger.debug(f"[Processor] 视频生成完成: {v_url}") - full_video_url = f"https://assets.grok.com/{v_url}" - - try: - cache_path = await video_cache_service.download_video(f"/{v_url}", auth_token) - if cache_path: - video_path = v_url.replace('/', '-') - base_url = setting.global_config.get("base_url", "") - local_video_url = f"{base_url}/images/{video_path}" if base_url else f"/images/{video_path}" - content += f'\n' - else: - content += f'\n' - except Exception as e: - logger.warning(f"[Processor] 缓存视频失败: {e}") - content += f'\n' - - yield make_chunk(content) - timeout_manager.mark_chunk_received() - chunk_index += 1 + if show_thinking: + if not video_progress_started: + content = f"视频已生成{progress}%\n" + video_progress_started = True + elif progress < 100: + content = f"视频已生成{progress}%\n" + else: + content = f"视频已生成{progress}%\n" + yield make_chunk(content) + + # 视频URL + if v_url: + logger.debug("[Processor] 视频生成完成") + video_content = await GrokResponseProcessor._build_video_content(v_url, auth_token) + yield make_chunk(video_content) continue - # 检查生成模式 + # 图片模式 if grok_resp.get("imageAttachmentInfo"): is_image = True - # 获取token token = grok_resp.get("token", "") - # 提取图片数据 + # 图片处理 if is_image: if model_resp := grok_resp.get("modelResponse"): - # 获取图片返回模式 image_mode = setting.global_config.get("image_mode", "url") - - # 初始化内容变量 content = "" - # 生成图片链接并缓存 for img in model_resp.get("generatedImageUrls", []): try: if image_mode == "base64": - # base64 模式:下载并转换为 base64 + # Base64模式 - 分块发送 base64_str = await image_cache_service.download_base64(f"/{img}", auth_token) if base64_str: - # 分块发送 base64 数据,每 8KB 一个 chunk - markdown_prefix = "![Generated Image](data:" - markdown_suffix = ")\n" - - # 提取 data URL 的 mime 和 base64 部分 - if base64_str.startswith("data:"): + # 分块发送大数据 + if not base64_str.startswith("data:"): parts = base64_str.split(",", 1) if len(parts) == 2: - mime_part = parts[0] + "," - b64_data = parts[1] - - # 发送前缀 - yield make_chunk(markdown_prefix + mime_part) - timeout_manager.mark_chunk_received() - chunk_index += 1 - - # 分块发送 base64 数据 - chunk_size = 8192 - for i in range(0, len(b64_data), chunk_size): - chunk_data = b64_data[i:i + chunk_size] - yield make_chunk(chunk_data) - timeout_manager.mark_chunk_received() - chunk_index += 1 - - # 发送后缀 - yield make_chunk(markdown_suffix) - timeout_manager.mark_chunk_received() - chunk_index += 1 + yield make_chunk(f"![Generated Image](data:{parts[0]},") + # 8KB分块 + for i in range(0, len(parts[1]), 8192): + yield make_chunk(parts[1][i:i+8192]) + yield make_chunk(")\n") else: yield make_chunk(f"![Generated Image]({base64_str})\n") - timeout_manager.mark_chunk_received() - chunk_index += 1 else: yield make_chunk(f"![Generated Image]({base64_str})\n") - timeout_manager.mark_chunk_received() - chunk_index += 1 else: yield make_chunk(f"![Generated Image](https://assets.grok.com/{img})\n") - timeout_manager.mark_chunk_received() - chunk_index += 1 else: - # url 模式:缓存并返回链接 + # URL模式 await image_cache_service.download_image(f"/{img}", auth_token) - # 本地图片路径 img_path = img.replace('/', '-') base_url = setting.global_config.get("base_url", "") img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" @@ -410,97 +270,154 @@ class GrokResponseProcessor: logger.warning(f"[Processor] 处理图片失败: {e}") content += f"![Generated Image](https://assets.grok.com/{img})\n" - # 发送内容 yield make_chunk(content.strip(), "stop") - timeout_manager.mark_chunk_received() return elif token: yield make_chunk(token) - timeout_manager.mark_chunk_received() - chunk_index += 1 - # 提取对话数据 + # 对话处理 else: - # 过滤 list 格式的 token if isinstance(token, list): continue - # 过滤特定标签 if any(tag in token for tag in filtered_tags if token): continue - # 获取当前状态 current_is_thinking = grok_resp.get("isThinking", False) message_tag = grok_resp.get("messageTag") - # 跳过后续的 标签 if thinking_finished and current_is_thinking: continue - # 检查 toolUsageCardId + # 搜索结果处理 if grok_resp.get("toolUsageCardId"): if web_search := grok_resp.get("webSearchResults"): if current_is_thinking: - # 封装搜索结果 - for result in web_search.get("results", []): - title = result.get("title", "") - url = result.get("url", "") - preview = result.get("preview", "") - preview_clean = preview.replace("\n", "") if isinstance(preview, str) else "" - token += f'\n- [{title}]({url} "{preview_clean}")' - token += "\n" + if show_thinking: + for result in web_search.get("results", []): + title = result.get("title", "") + url = result.get("url", "") + preview = result.get("preview", "") + preview_clean = preview.replace("\n", "") if isinstance(preview, str) else "" + token += f'\n- [{title}]({url} "{preview_clean}")' + token += "\n" + else: + continue else: - # 有 webSearchResults 但 isThinking 为 false continue else: - # 没有 webSearchResults continue if token: content = token - # header 在 token 后换行 if message_tag == "header": content = f"\n\n{token}\n\n" - # is_thinking 状态切换 + # Thinking状态切换 + should_skip = False if not is_thinking and current_is_thinking: - content = f"\n{content}" + if show_thinking: + content = f"\n{content}" + else: + should_skip = True elif is_thinking and not current_is_thinking: - content = f"\n\n{content}" + if show_thinking: + content = f"\n\n{content}" thinking_finished = True + elif current_is_thinking: + if not show_thinking: + should_skip = True - yield make_chunk(content) - timeout_manager.mark_chunk_received() - chunk_index += 1 + if not should_skip: + yield make_chunk(content) + is_thinking = current_is_thinking - except (json.JSONDecodeError, UnicodeDecodeError) as e: - logger.warning(f"[Processor] 解析chunk失败: {e}") + except (orjson.JSONDecodeError, UnicodeDecodeError) as e: + logger.warning(f"[Processor] 解析失败: {e}") continue except Exception as e: - logger.warning(f"[Processor] 处理chunk出错: {e}") + logger.warning(f"[Processor] 处理出错: {e}") continue - # 发送结束块 yield make_chunk("", "stop") - - # 发送流结束标记 yield "data: [DONE]\n\n" - - # 记录流式响应统计 - logger.info(f"[Processor] 流式响应完成,总耗时: {timeout_manager.get_total_duration():.2f}秒") + logger.info(f"[Processor] 流式完成,耗时: {timeout_mgr.duration():.2f}秒") except Exception as e: - logger.error(f"[Processor] 流式处理严重错误: {e}") + logger.error(f"[Processor] 严重错误: {e}") yield make_chunk(f"处理错误: {e}", "error") - # 发送流结束标记 yield "data: [DONE]\n\n" finally: - # 确保响应对象被关闭 if not response_closed and hasattr(response, 'close'): try: response.close() - logger.debug("[Processor] 流式响应对象已关闭") + logger.debug("[Processor] 响应已关闭") except Exception as e: - logger.warning(f"[Processor] 关闭流式响应对象时出错: {e}") \ No newline at end of file + logger.warning(f"[Processor] 关闭失败: {e}") + + @staticmethod + async def _build_video_content(video_url: str, auth_token: str) -> str: + """构建视频内容""" + logger.debug(f"[Processor] 检测到视频: {video_url}") + full_url = f"https://assets.grok.com/{video_url}" + + try: + cache_path = await video_cache_service.download_video(f"/{video_url}", auth_token) + if cache_path: + video_path = video_url.replace('/', '-') + base_url = setting.global_config.get("base_url", "") + local_url = f"{base_url}/images/{video_path}" if base_url else f"/images/{video_path}" + return f'\n' + except Exception as e: + logger.warning(f"[Processor] 缓存视频失败: {e}") + + return f'\n' + + @staticmethod + async def _append_images(content: str, images: list, auth_token: str) -> str: + """追加图片到内容""" + image_mode = setting.global_config.get("image_mode", "url") + + for img in images: + try: + if image_mode == "base64": + base64_str = await image_cache_service.download_base64(f"/{img}", auth_token) + if base64_str: + content += f"\n![Generated Image]({base64_str})" + else: + content += f"\n![Generated Image](https://assets.grok.com/{img})" + else: + cache_path = await image_cache_service.download_image(f"/{img}", auth_token) + if cache_path: + img_path = img.replace('/', '-') + base_url = setting.global_config.get("base_url", "") + img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" + content += f"\n![Generated Image]({img_url})" + else: + content += f"\n![Generated Image](https://assets.grok.com/{img})" + except Exception as e: + logger.warning(f"[Processor] 处理图片失败: {e}") + content += f"\n![Generated Image](https://assets.grok.com/{img})" + + return content + + @staticmethod + def _build_response(content: str, model: str) -> OpenAIChatCompletionResponse: + """构建响应对象""" + return OpenAIChatCompletionResponse( + id=f"chatcmpl-{uuid.uuid4()}", + object="chat.completion", + created=int(time.time()), + model=model, + choices=[OpenAIChatCompletionChoice( + index=0, + message=OpenAIChatCompletionMessage( + role="assistant", + content=content + ), + finish_reason="stop" + )], + usage=None + ) diff --git a/app/services/grok/statsig.py b/app/services/grok/statsig.py index 2c30601..a3d7017 100644 --- a/app/services/grok/statsig.py +++ b/app/services/grok/statsig.py @@ -1,44 +1,80 @@ -"""Grok 请求头管理模块""" +"""Grok 请求头管理 - 生成动态请求头和Statsig ID""" +import base64 +import random +import string import uuid from typing import Dict +from app.core.logger import logger from app.core.config import setting +# 基础请求头 +BASE_HEADERS = { + "Accept": "*/*", + "Accept-Language": "zh-CN,zh;q=0.9", + "Accept-Encoding": "gzip, deflate, br, zstd", + "Connection": "keep-alive", + "Origin": "https://grok.com", + "Priority": "u=1, i", + "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36", + "Sec-Ch-Ua": '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"', + "Sec-Ch-Ua-Mobile": "?0", + "Sec-Ch-Ua-Platform": '"macOS"', + "Sec-Fetch-Dest": "empty", + "Sec-Fetch-Mode": "cors", + "Sec-Fetch-Site": "same-origin", + "Baggage": "sentry-environment=production,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c", +} + + +def _random_string(length: int, letters_only: bool = True) -> str: + """生成随机字符串""" + chars = string.ascii_lowercase if letters_only else string.ascii_lowercase + string.digits + return ''.join(random.choices(chars, k=length)) + + +def _generate_statsig_id() -> str: + """生成x-statsig-id + + 随机选择两种格式: + 1. e:TypeError: Cannot read properties of null (reading 'children['xxxxx']') + 2. e:TypeError: Cannot read properties of undefined (reading 'xxxxxxxxxx') + """ + if random.choice([True, False]): + rand = _random_string(5, letters_only=False) + msg = f"e:TypeError: Cannot read properties of null (reading 'children['{rand}']')" + else: + rand = _random_string(10) + msg = f"e:TypeError: Cannot read properties of undefined (reading '{rand}')" + + return base64.b64encode(msg.encode()).decode() + + def get_dynamic_headers(pathname: str = "/rest/app-chat/conversations/new") -> Dict[str, str]: """获取请求头 - + Args: pathname: 请求路径 - + Returns: - 请求头字典 + 完整的请求头字典 """ - # 获取配置的 x-statsig-id - statsig_id = setting.grok_config.get("x_statsig_id") - if not statsig_id: - raise ValueError("配置文件中未设置 x_statsig_id") + # 获取或生成statsig-id + if setting.grok_config.get("dynamic_statsig", False): + statsig_id = _generate_statsig_id() + logger.debug(f"[Statsig] 动态生成: {statsig_id}") + else: + statsig_id = setting.grok_config.get("x_statsig_id") + if not statsig_id: + raise ValueError("配置文件中未设置 x_statsig_id") + logger.debug(f"[Statsig] 使用固定值: {statsig_id}") - # 构建基础请求头 - headers = { - "Accept": "*/*", - "Accept-Language": "zh-CN,zh;q=0.9", - "Accept-Encoding": "gzip, deflate, br, zstd", - "Content-Type": "application/json" if "upload-file" not in pathname else "text/plain;charset=UTF-8", - "Connection": "keep-alive", - "Origin": "https://grok.com", - "Priority": "u=1, i", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36", - "Sec-Ch-Ua": '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"', - "Sec-Ch-Ua-Mobile": "?0", - "Sec-Ch-Ua-Platform": '"macOS"', - "Sec-Fetch-Dest": "empty", - "Sec-Fetch-Mode": "cors", - "Sec-Fetch-Site": "same-origin", - "Baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c", - "x-statsig-id": statsig_id, - "x-xai-request-id": str(uuid.uuid4()) - } + # 构建请求头 + headers = BASE_HEADERS.copy() + headers["x-statsig-id"] = statsig_id + headers["x-xai-request-id"] = str(uuid.uuid4()) + headers["Content-Type"] = "text/plain;charset=UTF-8" if "upload-file" in pathname else "application/json" return headers \ No newline at end of file diff --git a/app/services/grok/token.py b/app/services/grok/token.py index a8678b0..80f2e98 100644 --- a/app/services/grok/token.py +++ b/app/services/grok/token.py @@ -1,9 +1,10 @@ -"""Grok Token 管理器模块""" +"""Grok Token 管理器 - 单例模式的Token负载均衡和状态管理""" -import json +import orjson import time import asyncio import aiofiles +import portalocker from pathlib import Path from curl_cffi.requests import AsyncSession from typing import Dict, Any, Optional, Tuple @@ -14,37 +15,28 @@ from app.core.logger import logger from app.core.config import setting from app.services.grok.statsig import get_dynamic_headers -# 常量定义 -RATE_LIMIT_ENDPOINT = "https://grok.com/rest/rate-limits" -REQUEST_TIMEOUT = 30 -IMPERSONATE_BROWSER = "chrome133a" -MAX_FAILURE_COUNT = 3 -TOKEN_INVALID_CODE = 401 # SSO Token失效 -STATSIG_INVALID_CODE = 403 # x-statsig-id失效 + +# 常量 +RATE_LIMIT_API = "https://grok.com/rest/rate-limits" +TIMEOUT = 30 +BROWSER = "chrome133a" +MAX_FAILURES = 3 +TOKEN_INVALID = 401 +STATSIG_INVALID = 403 class GrokTokenManager: - """ - Grok Token管理器 - - 单例模式的Token管理器,负责: - - Token文件的读写操作 - - Token负载均衡 - - Token状态管理 - - 支持普通Token和Super Token - """ + """Token管理器(单例)""" _instance: Optional['GrokTokenManager'] = None _lock = asyncio.Lock() def __new__(cls) -> 'GrokTokenManager': - """单例模式实现""" if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance def __init__(self): - """初始化Token管理器""" if hasattr(self, '_initialized'): return @@ -52,79 +44,129 @@ class GrokTokenManager: self._file_lock = asyncio.Lock() self.token_file.parent.mkdir(parents=True, exist_ok=True) self._storage = None - - # 同步加载初始数据 - self._load_data() + self.token_data = None # 延迟加载 + + # 批量保存队列 + self._save_pending = False # 标记是否有待保存的数据 + self._save_task = None # 后台保存任务 + self._shutdown = False # 关闭标志 + self._initialized = True - - logger.debug(f"[Token] 管理器初始化完成,文件: {self.token_file}") + logger.debug(f"[Token] 初始化完成: {self.token_file}") def set_storage(self, storage) -> None: """设置存储实例""" self._storage = storage - def _load_data(self) -> None: - """同步加载Token数据(仅用于初始化)""" - default_data = { - TokenType.NORMAL.value: {}, - TokenType.SUPER.value: {} - } - + async def _load_data(self) -> None: + """异步加载Token数据(支持多进程)""" + default = {TokenType.NORMAL.value: {}, TokenType.SUPER.value: {}} + try: if self.token_file.exists(): - with open(self.token_file, "r", encoding="utf-8") as f: - self.token_data = json.load(f) + # 使用进程锁读取文件 + async with self._file_lock: + with open(self.token_file, "r", encoding="utf-8") as f: + portalocker.lock(f, portalocker.LOCK_SH) # 共享锁(读) + try: + content = f.read() + self.token_data = orjson.loads(content) + finally: + portalocker.unlock(f) else: - self.token_data = default_data - logger.debug("[Token] 创建新的Token数据文件") - except (json.JSONDecodeError, IOError) as e: - logger.error(f"[Token] 加载Token数据失败: {str(e)}") - self.token_data = default_data + self.token_data = default + logger.debug("[Token] 创建新数据文件") + except Exception as e: + logger.error(f"[Token] 加载失败: {e}") + self.token_data = default async def _save_data(self) -> None: - """异步保存Token数据到存储""" + """保存Token数据(支持多进程)""" try: if not self._storage: - # 如果没有设置存储,使用传统文件保存方式(向后兼容) async with self._file_lock: - async with aiofiles.open(self.token_file, "w", encoding="utf-8") as f: - await f.write(json.dumps(self.token_data, indent=2, ensure_ascii=False)) + # 使用进程锁写入文件 + with open(self.token_file, "w", encoding="utf-8") as f: + portalocker.lock(f, portalocker.LOCK_EX) # 独占锁(写) + try: + content = orjson.dumps(self.token_data, option=orjson.OPT_INDENT_2).decode() + f.write(content) + f.flush() # 确保写入磁盘 + finally: + portalocker.unlock(f) else: - # 使用存储抽象层 await self._storage.save_tokens(self.token_data) except IOError as e: - logger.error(f"[Token] 保存Token数据失败: {str(e)}") - raise GrokApiException( - f"Token数据保存失败: {str(e)}", - "TOKEN_SAVE_ERROR", - {"file_path": str(self.token_file)} - ) + logger.error(f"[Token] 保存失败: {e}") + raise GrokApiException(f"保存失败: {e}", "TOKEN_SAVE_ERROR", {"file": str(self.token_file)}) + + def _mark_dirty(self) -> None: + """标记有待保存的数据""" + self._save_pending = True + + async def _batch_save_worker(self) -> None: + """批量保存后台任务""" + from app.core.config import setting + + interval = setting.global_config.get("batch_save_interval", 1.0) + logger.info(f"[Token] 批量保存任务已启动,间隔: {interval}s") + + while not self._shutdown: + await asyncio.sleep(interval) + + if self._save_pending and not self._shutdown: + try: + await self._save_data() + self._save_pending = False + logger.debug("[Token] 批量保存完成") + except Exception as e: + logger.error(f"[Token] 批量保存失败: {e}") + + async def start_batch_save(self) -> None: + """启动批量保存任务""" + if self._save_task is None: + self._save_task = asyncio.create_task(self._batch_save_worker()) + logger.info("[Token] 批量保存任务已创建") + + async def shutdown(self) -> None: + """关闭并刷新所有待保存数据""" + self._shutdown = True + + if self._save_task: + self._save_task.cancel() + try: + await self._save_task + except asyncio.CancelledError: + pass + + # 最终刷新 + if self._save_pending: + await self._save_data() + logger.info("[Token] 关闭时刷新完成") @staticmethod def _extract_sso(auth_token: str) -> Optional[str]: - """从认证令牌中提取SSO值""" + """提取SSO值""" if "sso=" in auth_token: return auth_token.split("sso=")[1].split(";")[0] - logger.warning("[Token] 无法从认证令牌中提取SSO值") + logger.warning("[Token] 无法提取SSO值") return None - def _find_token(self, sso_value: str) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: - """查找Token数据,返回(token_type, token_data)""" + def _find_token(self, sso: str) -> Tuple[Optional[str], Optional[Dict]]: + """查找Token""" for token_type in [TokenType.NORMAL.value, TokenType.SUPER.value]: - if sso_value in self.token_data[token_type]: - return token_type, self.token_data[token_type][sso_value] + if sso in self.token_data[token_type]: + return token_type, self.token_data[token_type][sso] return None, None async def add_token(self, tokens: list[str], token_type: TokenType) -> None: - """添加Token到管理器""" + """添加Token""" if not tokens: - logger.debug("[Token] 尝试添加空的Token列表") return - added_count = 0 + count = 0 for token in tokens: if not token or not token.strip(): - logger.debug("[Token] 跳过空的Token") continue self.token_data[token_type.value][token] = { @@ -134,265 +176,312 @@ class GrokTokenManager: "status": "active", "failedCount": 0, "lastFailureTime": None, - "lastFailureReason": None + "lastFailureReason": None, + "tags": [], + "note": "" } - added_count += 1 + count += 1 - await self._save_data() - logger.info(f"[Token] 成功添加 {added_count} 个 {token_type.value} Token") + self._mark_dirty() # 批量保存 + logger.info(f"[Token] 添加 {count} 个 {token_type.value} Token") async def delete_token(self, tokens: list[str], token_type: TokenType) -> None: - """删除指定的Token""" + """删除Token""" if not tokens: - logger.debug("[Token] 尝试删除空的Token列表") return - deleted_count = 0 + count = 0 for token in tokens: if token in self.token_data[token_type.value]: del self.token_data[token_type.value][token] - deleted_count += 1 - else: - logger.debug(f"[Token] Token不存在: {token[:10]}...") + count += 1 - await self._save_data() - logger.info(f"[Token] 成功删除 {deleted_count} 个 {token_type.value} Token") + self._mark_dirty() # 批量保存 + logger.info(f"[Token] 删除 {count} 个 {token_type.value} Token") + + async def update_token_tags(self, token: str, token_type: TokenType, tags: list[str]) -> None: + """更新Token标签""" + if token not in self.token_data[token_type.value]: + raise GrokApiException("Token不存在", "TOKEN_NOT_FOUND", {"token": token[:10]}) + + cleaned = [t.strip() for t in tags if t and t.strip()] + self.token_data[token_type.value][token]["tags"] = cleaned + self._mark_dirty() # 批量保存 + logger.info(f"[Token] 更新标签: {token[:10]}... -> {cleaned}") + + async def update_token_note(self, token: str, token_type: TokenType, note: str) -> None: + """更新Token备注""" + if token not in self.token_data[token_type.value]: + raise GrokApiException("Token不存在", "TOKEN_NOT_FOUND", {"token": token[:10]}) + + self.token_data[token_type.value][token]["note"] = note.strip() + self._mark_dirty() # 批量保存 + logger.info(f"[Token] 更新备注: {token[:10]}...") def get_tokens(self) -> Dict[str, Any]: - """获取所有Token数据""" + """获取所有Token""" return self.token_data.copy() + def _reload_if_needed(self) -> None: + """在多进程模式下重新加载数据(同步版本,用于select_token)""" + # 只在文件模式且多进程环境下才重新加载 + if self._storage: + return # 数据库模式不需要 + + try: + if self.token_file.exists(): + with open(self.token_file, "r", encoding="utf-8") as f: + portalocker.lock(f, portalocker.LOCK_SH) + try: + content = f.read() + self.token_data = orjson.loads(content) + finally: + portalocker.unlock(f) + except Exception as e: + logger.warning(f"[Token] 重新加载失败: {e}") + def get_token(self, model: str) -> str: - """获取指定模型的Token""" - jwt_token = self.select_token(model) - return f"sso-rw={jwt_token};sso={jwt_token}" + """获取Token""" + jwt = self.select_token(model) + return f"sso-rw={jwt};sso={jwt}" def select_token(self, model: str) -> str: - """根据模型类型和剩余次数选择最优Token""" - def select_best_token(tokens_dict: Dict[str, Any]) -> Tuple[Optional[str], Optional[int]]: - """从 token 字典中选择最佳 token""" - unused_tokens = [] # remaining = -1 的 token - used_tokens = [] # remaining > 0 的 token + """选择最优Token(多进程安全)""" + # 重新加载最新数据(多进程模式) + self._reload_if_needed() + def select_best(tokens: Dict[str, Any], field: str) -> Tuple[Optional[str], Optional[int]]: + """选择最佳Token""" + unused, used = [], [] - for token_key, token_data in tokens_dict.items(): - # 跳过已失效的Token - if token_data.get("status") == "expired": + for key, data in tokens.items(): + # 跳过已失效的token + if data.get("status") == "expired": + continue + + # 跳过失败次数过多的token(任何错误状态码) + if data.get("failedCount", 0) >= MAX_FAILURES: continue - remaining = int(token_data.get(remaining_field, -1)) - - # 跳过已限流的 token + remaining = int(data.get(field, -1)) if remaining == 0: continue - # 分类存储 if remaining == -1: - unused_tokens.append(token_key) + unused.append(key) elif remaining > 0: - used_tokens.append((token_key, remaining)) - - # 优先返回尚未使用的 token - if unused_tokens: - return unused_tokens[0], -1 - - # 否则返回次数最多的 token - if used_tokens: - used_tokens.sort(key=lambda x: x[1], reverse=True) - return used_tokens[0][0], used_tokens[0][1] + used.append((key, remaining)) + if unused: + return unused[0], -1 + if used: + used.sort(key=lambda x: x[1], reverse=True) + return used[0][0], used[0][1] return None, None - max_token_key = None - max_remaining = None - - # 深拷贝 - token_data_snapshot = { + # 快照 + snapshot = { TokenType.NORMAL.value: self.token_data[TokenType.NORMAL.value].copy(), TokenType.SUPER.value: self.token_data[TokenType.SUPER.value].copy() } + # 选择策略 if model == "grok-4-heavy": - # grok-4-heavy 只能使用Super Token + heavy remaining queries - remaining_field = "heavyremainingQueries" - max_token_key, max_remaining = select_best_token(token_data_snapshot[TokenType.SUPER.value]) + field = "heavyremainingQueries" + token_key, remaining = select_best(snapshot[TokenType.SUPER.value], field) else: - # 其他模型使用 remaining Queries - remaining_field = "remainingQueries" + field = "remainingQueries" + token_key, remaining = select_best(snapshot[TokenType.NORMAL.value], field) + if token_key is None: + token_key, remaining = select_best(snapshot[TokenType.SUPER.value], field) - # 优先使用普通Token - max_token_key, max_remaining = select_best_token(token_data_snapshot[TokenType.NORMAL.value]) - - # 如果普通Token没有可用的,尝试使用Super Token - if max_token_key is None: - max_token_key, max_remaining = select_best_token(token_data_snapshot[TokenType.SUPER.value]) - - if max_token_key is None: + if token_key is None: raise GrokApiException( - f"没有可用Token用于模型 {model}", + f"没有可用Token: {model}", "NO_AVAILABLE_TOKEN", { "model": model, - "normal_count": len(token_data_snapshot[TokenType.NORMAL.value]), - "super_count": len(token_data_snapshot[TokenType.SUPER.value]) + "normal": len(snapshot[TokenType.NORMAL.value]), + "super": len(snapshot[TokenType.SUPER.value]) } ) - status_text = "未使用" if max_remaining == -1 else f"剩余{max_remaining}次" - logger.debug(f"[Token] 为模型 {model} 选择Token ({status_text})") - return max_token_key + status = "未使用" if remaining == -1 else f"剩余{remaining}次" + logger.debug(f"[Token] 分配Token: {model} ({status})") + return token_key async def check_limits(self, auth_token: str, model: str) -> Optional[Dict[str, Any]]: - """检查并更新模型速率限制""" + """检查速率限制""" try: - rate_limit_model_name = Models.to_rate_limit(model) - logger.debug(f"[Token] 检查模型 {model} (接口模型: {rate_limit_model_name}) 的速率限制") - - # 准备请求 - payload = {"requestKind": "DEFAULT", "modelName": rate_limit_model_name} - cf_clearance = setting.grok_config.get("cf_clearance", "") - cookie = f"{auth_token};{cf_clearance}" if cf_clearance else auth_token - + rate_model = Models.to_rate_limit(model) + payload = {"requestKind": "DEFAULT", "modelName": rate_model} + + cf = setting.grok_config.get("cf_clearance", "") headers = get_dynamic_headers("/rest/rate-limits") - headers["Cookie"] = cookie + headers["Cookie"] = f"{auth_token};{cf}" if cf else auth_token - # 获取代理配置 - proxy_url = setting.grok_config.get("proxy_url", "") - proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else None + # 外层重试:可配置状态码(401/429等) + retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) + MAX_OUTER_RETRY = 3 + + for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次 + # 内层重试:403代理池重试 + max_403_retries = 5 + retry_403_count = 0 + + while retry_403_count <= max_403_retries: + # 异步获取代理(支持代理池) + from app.core.proxy_pool import proxy_pool + + # 如果是403重试且使用代理池,强制刷新代理 + if retry_403_count > 0 and proxy_pool._enabled: + logger.info(f"[Token] 403重试 {retry_403_count}/{max_403_retries},刷新代理...") + proxy = await proxy_pool.force_refresh() + else: + proxy = await setting.get_proxy_async("service") + + proxies = {"http": proxy, "https": proxy} if proxy else None + + async with AsyncSession() as session: + response = await session.post( + RATE_LIMIT_API, + headers=headers, + json=payload, + impersonate=BROWSER, + timeout=TIMEOUT, + proxies=proxies + ) - # 发送异步请求 - async with AsyncSession() as session: - response = await session.post( - RATE_LIMIT_ENDPOINT, - headers=headers, - json=payload, - impersonate=IMPERSONATE_BROWSER, - timeout=REQUEST_TIMEOUT, - proxies=proxies - ) + # 内层403重试:仅当有代理池时触发 + if response.status_code == 403 and proxy_pool._enabled: + retry_403_count += 1 + + if retry_403_count <= max_403_retries: + logger.warning(f"[Token] 遇到403错误,正在重试 ({retry_403_count}/{max_403_retries})...") + await asyncio.sleep(0.5) + continue + + # 内层重试全部失败 + logger.error(f"[Token] 403错误,已重试{retry_403_count-1}次,放弃") + sso = self._extract_sso(auth_token) + if sso: + await self.record_failure(auth_token, 403, "服务器被Block") + + # 检查可配置状态码错误 - 外层重试 + if response.status_code in retry_codes: + if outer_retry < MAX_OUTER_RETRY: + delay = (outer_retry + 1) * 0.1 # 渐进延迟:0.1s, 0.2s, 0.3s + logger.warning(f"[Token] 遇到{response.status_code}错误,外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY}),等待{delay}s...") + await asyncio.sleep(delay) + break # 跳出内层循环,进入外层重试 + else: + logger.error(f"[Token] {response.status_code}错误,已重试{outer_retry}次,放弃") + sso = self._extract_sso(auth_token) + if sso: + if response.status_code == 401: + await self.record_failure(auth_token, 401, "Token失效") + else: + await self.record_failure(auth_token, response.status_code, f"错误: {response.status_code}") + return None - if response.status_code == 200: - rate_limit_data = response.json() - logger.debug(f"[Token] 成功获取速率限制信息") - - # 保存速率限制信息 - sso_value = self._extract_sso(auth_token) - if sso_value: - if model == "grok-4-heavy": - await self.update_limits(sso_value, normal=None, heavy=rate_limit_data.get("remainingQueries", -1)) - logger.info(f"[Token] 已更新限制: sso={sso_value[:10]}..., heavy={rate_limit_data.get('remainingQueries', -1)}") + if response.status_code == 200: + data = response.json() + sso = self._extract_sso(auth_token) + + if outer_retry > 0 or retry_403_count > 0: + logger.info(f"[Token] 重试成功!") + + if sso: + if model == "grok-4-heavy": + await self.update_limits(sso, normal=None, heavy=data.get("remainingQueries", -1)) + logger.info(f"[Token] 更新限制: {sso[:10]}..., heavy={data.get('remainingQueries', -1)}") + else: + await self.update_limits(sso, normal=data.get("remainingTokens", -1), heavy=None) + logger.info(f"[Token] 更新限制: {sso[:10]}..., basic={data.get('remainingTokens', -1)}") + + return data else: - await self.update_limits(sso_value, normal=rate_limit_data.get("remainingTokens", -1), heavy=None) - logger.info(f"[Token] 已更新限制: sso={sso_value[:10]}..., 通用={rate_limit_data.get('remainingTokens', -1)}") - - return rate_limit_data - else: - logger.warning(f"[Token] 获取速率限制失败,状态码: {response.status_code}") - return None + # 其他错误 + logger.warning(f"[Token] 获取限制失败: {response.status_code}") + sso = self._extract_sso(auth_token) + if sso: + await self.record_failure(auth_token, response.status_code, f"错误: {response.status_code}") + return None except Exception as e: - logger.error(f"[Token] 检查速率限制时发生错误: {str(e)}") + logger.error(f"[Token] 检查限制错误: {e}") return None - async def update_limits(self, sso_value: str, normal: Optional[int] = None, heavy: Optional[int] = None) -> None: - """更新Token限制信息""" + async def update_limits(self, sso: str, normal: Optional[int] = None, heavy: Optional[int] = None) -> None: + """更新限制""" try: for token_type in [TokenType.NORMAL.value, TokenType.SUPER.value]: - if sso_value in self.token_data[token_type]: + if sso in self.token_data[token_type]: if normal is not None: - self.token_data[token_type][sso_value]["remainingQueries"] = normal + self.token_data[token_type][sso]["remainingQueries"] = normal if heavy is not None: - self.token_data[token_type][sso_value]["heavyremainingQueries"] = heavy - - await self._save_data() - logger.info(f"[Token] 已更新Token {sso_value[:10]}... 的限制信息") + self.token_data[token_type][sso]["heavyremainingQueries"] = heavy + self._mark_dirty() # 批量保存 + logger.info(f"[Token] 更新限制: {sso[:10]}...") return - - logger.warning(f"[Token] 未找到SSO值为 {sso_value[:10]}... 的Token") - + logger.warning(f"[Token] 未找到: {sso[:10]}...") except Exception as e: - logger.error(f"[Token] 更新Token限制时发生错误: {str(e)}") + logger.error(f"[Token] 更新限制错误: {e}") - async def record_failure(self, auth_token: str, status_code: int, error_message: str) -> None: - """记录Token失败信息 - - 错误码说明: - - 401: SSO Token失效,会标记Token为expired - - 403: x-statsig-id失效,不影响Token状态 - - Args: - auth_token: 完整的认证Token (格式: sso-rw=xxx;sso=xxx) - status_code: HTTP状态码 - error_message: 错误信息 - """ + async def record_failure(self, auth_token: str, status: int, msg: str) -> None: + """记录失败""" try: - # 403错误是x-statsig-id失效,不是Token问题 - if status_code == STATSIG_INVALID_CODE: - logger.warning(f"[Token] x-statsig-id失效 (403),需要更新配置文件中的x_statsig_id") + if status == STATSIG_INVALID: + logger.warning("[Token] IP被Block,请: 1.更换IP 2.使用代理 3.配置CF值") return - sso_value = self._extract_sso(auth_token) - if not sso_value: + sso = self._extract_sso(auth_token) + if not sso: return - _, token_data = self._find_token(sso_value) - if not token_data: - logger.warning(f"[Token] 未找到SSO值为 {sso_value[:10]}... 的Token") + _, data = self._find_token(sso) + if not data: + logger.warning(f"[Token] 未找到: {sso[:10]}...") return - # 更新失败计数 - token_data["failedCount"] = token_data.get("failedCount", 0) + 1 - token_data["lastFailureTime"] = int(time.time() * 1000) - token_data["lastFailureReason"] = f"{status_code}: {error_message}" + data["failedCount"] = data.get("failedCount", 0) + 1 + data["lastFailureTime"] = int(time.time() * 1000) + data["lastFailureReason"] = f"{status}: {msg}" logger.warning( - f"[Token] Token {sso_value[:10]}... 失败 (状态码: {status_code}), " - f"失败次数: {token_data['failedCount']}/{MAX_FAILURE_COUNT}, " - f"原因: {error_message}" + f"[Token] 失败: {sso[:10]}... (状态:{status}), " + f"次数: {data['failedCount']}/{MAX_FAILURES}, 原因: {msg}" ) - # 只有401错误(SSO Token失效)且失败次数达到上限时,标记为失效 - if status_code == TOKEN_INVALID_CODE and token_data["failedCount"] >= MAX_FAILURE_COUNT: - token_data["status"] = "expired" - logger.error( - f"[Token] SSO Token {sso_value[:10]}... 已被标记为失效 " - f"(连续401错误{token_data['failedCount']}次)" - ) + if 400 <= status < 500 and data["failedCount"] >= MAX_FAILURES: + data["status"] = "expired" + logger.error(f"[Token] 标记失效: {sso[:10]}... (连续{status}错误{data['failedCount']}次)") - await self._save_data() + self._mark_dirty() # 批量保存 except Exception as e: - logger.error(f"[Token] 记录Token失败信息时发生错误: {str(e)}") + logger.error(f"[Token] 记录失败错误: {e}") async def reset_failure(self, auth_token: str) -> None: - """重置Token失败计数 - - 当Token成功完成请求时调用此方法,用于清除失败记录。 - - Args: - auth_token: 完整的认证Token (格式: sso-rw=xxx;sso=xxx) - """ + """重置失败计数""" try: - sso_value = self._extract_sso(auth_token) - if not sso_value: + sso = self._extract_sso(auth_token) + if not sso: return - _, token_data = self._find_token(sso_value) - if not token_data: - logger.warning(f"[Token] 未找到SSO值为 {sso_value[:10]}... 的Token") + _, data = self._find_token(sso) + if not data: return - # 只有在有失败记录时才重置并保存 - if token_data.get("failedCount", 0) > 0: - token_data["failedCount"] = 0 - token_data["lastFailureTime"] = None - token_data["lastFailureReason"] = None - - await self._save_data() - logger.info(f"[Token] Token {sso_value[:10]}... 失败计数已重置") + if data.get("failedCount", 0) > 0: + data["failedCount"] = 0 + data["lastFailureTime"] = None + data["lastFailureReason"] = None + self._mark_dirty() # 批量保存 + logger.info(f"[Token] 重置失败计数: {sso[:10]}...") except Exception as e: - logger.error(f"[Token] 重置Token失败计数时发生错误: {str(e)}") + logger.error(f"[Token] 重置失败错误: {e}") -# 全局Token管理器实例 +# 全局实例 token_manager = GrokTokenManager() diff --git a/app/services/grok/upload.py b/app/services/grok/upload.py index f5f92b3..4d0c8b4 100644 --- a/app/services/grok/upload.py +++ b/app/services/grok/upload.py @@ -1,10 +1,10 @@ -"""图片上传管理器""" +"""图片上传管理器 - 支持Base64和URL图片上传""" +import asyncio import base64 import re from typing import Tuple, Optional from urllib.parse import urlparse - from curl_cffi.requests import AsyncSession from app.services.grok.statsig import get_dynamic_headers @@ -12,134 +12,198 @@ from app.core.exception import GrokApiException from app.core.config import setting from app.core.logger import logger -# 常量定义 -UPLOAD_ENDPOINT = "https://grok.com/rest/app-chat/upload-file" -REQUEST_TIMEOUT = 30 -IMPERSONATE_BROWSER = "chrome133a" -DEFAULT_MIME_TYPE = "image/jpeg" -DEFAULT_EXTENSION = "jpg" + +# 常量 +UPLOAD_API = "https://grok.com/rest/app-chat/upload-file" +TIMEOUT = 30 +BROWSER = "chrome133a" + +# MIME类型 +MIME_TYPES = { + '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png', + '.gif': 'image/gif', '.webp': 'image/webp', '.bmp': 'image/bmp', +} +DEFAULT_MIME = "image/jpeg" +DEFAULT_EXT = "jpg" class ImageUploadManager: - """ - Grok图片上传管理器 - - 提供图片上传功能,支持: - - Base64格式图片上传 - - URL图片下载并上传 - - 多种图片格式支持 - """ + """图片上传管理器""" @staticmethod - async def upload(image_input: str, auth_token: str) -> str: - """上传图片到Grok,支持Base64或URL""" + async def upload(image_input: str, auth_token: str) -> Tuple[str, str]: + """上传图片(支持Base64或URL) + + Returns: + (file_id, file_uri) 元组 + """ try: + # 判断类型并处理 if ImageUploadManager._is_url(image_input): - # 下载 URL 图片 - image_buffer, mime_type = await ImageUploadManager._download(image_input) - - # 获取图片信息 - file_name, _ = ImageUploadManager._get_info("", mime_type) - + buffer, mime = await ImageUploadManager._download(image_input) + filename, _ = ImageUploadManager._get_info("", mime) else: - # 处理 base64 数据 - image_buffer = image_input.split(",")[1] if "data:image" in image_input else image_input + buffer = image_input.split(",")[1] if "data:image" in image_input else image_input + filename, mime = ImageUploadManager._get_info(image_input) - # 获取图片信息 - file_name, mime_type = ImageUploadManager._get_info(image_input) - - # 构建上传数据 - upload_data = { - "fileName": file_name, - "fileMimeType": mime_type, - "content": image_buffer, + # 构建数据 + data = { + "fileName": filename, + "fileMimeType": mime, + "content": buffer, } - # 获取认证令牌 + if not auth_token: - raise GrokApiException("认证令牌缺失或为空", "NO_AUTH_TOKEN") + raise GrokApiException("认证令牌缺失", "NO_AUTH_TOKEN") - cf_clearance = setting.grok_config.get("cf_clearance", "") - cookie = f"{auth_token};{cf_clearance}" if cf_clearance else auth_token - proxy_url = setting.grok_config.get("proxy_url", "") - proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else None + # 外层重试:可配置状态码(401/429等) + retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) + MAX_OUTER_RETRY = 3 + + for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次 + try: + # 内层重试:403代理池重试 + max_403_retries = 5 + retry_403_count = 0 + + while retry_403_count <= max_403_retries: + # 请求配置 + cf = setting.grok_config.get("cf_clearance", "") + headers = { + **get_dynamic_headers("/rest/app-chat/upload-file"), + "Cookie": f"{auth_token};{cf}" if cf else auth_token, + } + + # 异步获取代理(支持代理池) + from app.core.proxy_pool import proxy_pool + + # 如果是403重试且使用代理池,强制刷新代理 + if retry_403_count > 0 and proxy_pool._enabled: + logger.info(f"[Upload] 403重试 {retry_403_count}/{max_403_retries},刷新代理...") + proxy = await proxy_pool.force_refresh() + else: + proxy = await setting.get_proxy_async("service") + + proxies = {"http": proxy, "https": proxy} if proxy else None - # 发送异步请求 - async with AsyncSession() as session: - response = await session.post( - UPLOAD_ENDPOINT, - headers={ - **get_dynamic_headers("/rest/app-chat/upload-file"), - "Cookie": cookie, - }, - json=upload_data, - impersonate=IMPERSONATE_BROWSER, - timeout=REQUEST_TIMEOUT, - proxies=proxies, - ) + # 上传 + async with AsyncSession() as session: + response = await session.post( + UPLOAD_API, + headers=headers, + json=data, + impersonate=BROWSER, + timeout=TIMEOUT, + proxies=proxies, + ) - # 检查响应 - if response.status_code == 200: - result = response.json() - file_id = result.get("fileMetadataId", "") - logger.debug(f"[Upload] 图片上传成功,文件ID: {file_id}") - return file_id - - return "" + # 内层403重试:仅当有代理池时触发 + if response.status_code == 403 and proxy_pool._enabled: + retry_403_count += 1 + + if retry_403_count <= max_403_retries: + logger.warning(f"[Upload] 遇到403错误,正在重试 ({retry_403_count}/{max_403_retries})...") + await asyncio.sleep(0.5) + continue + + # 内层重试全部失败 + logger.error(f"[Upload] 403错误,已重试{retry_403_count-1}次,放弃") + + # 检查可配置状态码错误 - 外层重试 + if response.status_code in retry_codes: + if outer_retry < MAX_OUTER_RETRY: + delay = (outer_retry + 1) * 0.1 # 渐进延迟:0.1s, 0.2s, 0.3s + logger.warning(f"[Upload] 遇到{response.status_code}错误,外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY}),等待{delay}s...") + await asyncio.sleep(delay) + break # 跳出内层循环,进入外层重试 + else: + logger.error(f"[Upload] {response.status_code}错误,已重试{outer_retry}次,放弃") + return "", "" + + if response.status_code == 200: + result = response.json() + file_id = result.get("fileMetadataId", "") + file_uri = result.get("fileUri", "") + + if outer_retry > 0 or retry_403_count > 0: + logger.info(f"[Upload] 重试成功!") + + logger.debug(f"[Upload] 成功,ID: {file_id}") + return file_id, file_uri + + # 其他错误直接返回 + logger.error(f"[Upload] 失败,状态码: {response._status_code}") + return "", "" + + # 内层循环正常结束(非break),说明403重试全部失败 + return "", "" + + except Exception as e: + if outer_retry < MAX_OUTER_RETRY - 1: + logger.warning(f"[Upload] 异常: {e},外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})...") + await asyncio.sleep(0.5) + continue + + logger.warning(f"[Upload] 失败: {e}") + return "", "" + + return "", "" except Exception as e: - logger.warning(f"[Upload] 上传图片失败: {e}") - return "" + logger.warning(f"[Upload] 失败: {e}") + return "", "" @staticmethod - def _is_url(image_input: str) -> bool: - """检查输入是否为有效的URL""" + def _is_url(input_str: str) -> bool: + """检查是否为URL""" try: - result = urlparse(image_input) + result = urlparse(input_str) return all([result.scheme, result.netloc]) and result.scheme in ['http', 'https'] - except Exception as e: - logger.warning(f"[Upload] URL解析失败: {e}") + except: return False @staticmethod async def _download(url: str) -> Tuple[str, str]: - """下载图片并转换为Base64""" + """下载图片并转Base64 + + Returns: + (base64_string, mime_type) 元组 + """ try: async with AsyncSession() as session: response = await session.get(url, timeout=5) response.raise_for_status() - # 获取内容类型 - content_type = response.headers.get('content-type', DEFAULT_MIME_TYPE) + content_type = response.headers.get('content-type', DEFAULT_MIME) if not content_type.startswith('image/'): - content_type = DEFAULT_MIME_TYPE + content_type = DEFAULT_MIME - # 转换为 Base64 - image_base64 = base64.b64encode(response.content).decode('utf-8') - return image_base64, content_type + b64 = base64.b64encode(response.content).decode() + return b64, content_type except Exception as e: - logger.warning(f"[Upload] 下载图片失败: {e}") + logger.warning(f"[Upload] 下载失败: {e}") return "", "" @staticmethod def _get_info(image_data: str, mime_type: Optional[str] = None) -> Tuple[str, str]: - """获取图片文件名和MIME类型""" - # mime_type 有值,直接使用 + """获取文件名和MIME类型 + + Returns: + (file_name, mime_type) 元组 + """ + # 已提供MIME类型 if mime_type: - extension = mime_type.split("/")[1] if "/" in mime_type else DEFAULT_EXTENSION - file_name = f"image.{extension}" - return file_name, mime_type + ext = mime_type.split("/")[1] if "/" in mime_type else DEFAULT_EXT + return f"image.{ext}", mime_type - # mime_type 没有值,使用默认值 - mime_type = DEFAULT_MIME_TYPE - extension = DEFAULT_EXTENSION + # 从Base64提取 + mime = DEFAULT_MIME + ext = DEFAULT_EXT - # 从 Base64 数据中提取 MIME 类型 if "data:image" in image_data: - match = re.search(r"data:([a-zA-Z0-9]+/[a-zA-Z0-9-.+]+);base64,", image_data) - if match: - mime_type = match.group(1) - extension = mime_type.split("/")[1] + if match := re.search(r"data:([a-zA-Z0-9]+/[a-zA-Z0-9-.+]+);base64,", image_data): + mime = match.group(1) + ext = mime.split("/")[1] - file_name = f"image.{extension}" - return file_name, mime_type \ No newline at end of file + return f"image.{ext}", mime \ No newline at end of file diff --git a/app/services/mcp/__init__.py b/app/services/mcp/__init__.py new file mode 100644 index 0000000..6951deb --- /dev/null +++ b/app/services/mcp/__init__.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 -*- +"""MCP模块初始化""" + +from app.services.mcp.server import mcp + +__all__ = ["mcp"] diff --git a/app/services/mcp/server.py b/app/services/mcp/server.py new file mode 100644 index 0000000..38d7719 --- /dev/null +++ b/app/services/mcp/server.py @@ -0,0 +1,63 @@ +# -*- coding: utf-8 -*- +"""FastMCP服务器实例""" + +from fastmcp import FastMCP +from fastmcp.server.auth.providers.jwt import StaticTokenVerifier +from app.services.mcp.tools import ask_grok_impl +from app.core.config import setting + + +def create_mcp_server() -> FastMCP: + """创建MCP服务器实例,如果配置了API密钥则启用认证""" + # 检查是否配置了API密钥 + api_key = setting.grok_config.get("api_key") + + # 如果配置了API密钥,则启用静态token验证 + auth = None + if api_key: + auth = StaticTokenVerifier( + tokens={ + api_key: { + "client_id": "grok2api-client", + "scopes": ["read", "write", "admin"] + } + }, + required_scopes=["read"] + ) + + # 创建FastMCP实例 + return FastMCP( + name="Grok2API-MCP", + instructions="MCP server providing Grok AI chat capabilities. Use ask_grok tool to interact with Grok AI models.", + auth=auth + ) + + +# 创建全局MCP实例 +mcp = create_mcp_server() + + +# 注册ask_grok工具 +@mcp.tool +async def ask_grok( + query: str, + model: str = "grok-3-fast", + system_prompt: str = None +) -> str: + """ + 调用Grok AI进行对话,尤其适用于当用户询问最新信息,需要调用搜索功能,或是想了解社交平台动态(如Twitter(X)、Reddit等)时。 + + Args: + query: 用户的问题或指令 + model: Grok模型名称,可选值: grok-3-fast(默认), grok-4-fast, grok-4-fast-expert, grok-4-expert, grok-4-heavy + system_prompt: 可选的系统提示词,用于设定AI的角色或行为约束 + + Returns: + Grok AI的完整回复内容,可能包括文本和图片链接(Markdown格式) + + Examples: + - 简单问答: ask_grok("什么是Python?") + - 指定模型: ask_grok("解释量子计算", model="grok-4-fast") + - 带系统提示: ask_grok("写一首诗", system_prompt="你是一位古典诗人") + """ + return await ask_grok_impl(query, model, system_prompt) diff --git a/app/services/mcp/tools.py b/app/services/mcp/tools.py new file mode 100644 index 0000000..abedfae --- /dev/null +++ b/app/services/mcp/tools.py @@ -0,0 +1,77 @@ +# -*- coding: utf-8 -*- +"""MCP Tools - Grok AI 对话工具""" + +import json +from typing import Optional +from app.services.grok.client import GrokClient +from app.core.logger import logger +from app.core.exception import GrokApiException + + +async def ask_grok_impl( + query: str, + model: str = "grok-3-fast", + system_prompt: Optional[str] = None +) -> str: + """ + 内部实现: 调用Grok API并收集完整响应 + + Args: + query: 用户问题 + model: 模型名称 + system_prompt: 系统提示词 + + Returns: + str: 完整的Grok响应内容 + """ + try: + # 构建消息列表 + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": query}) + + # 构建请求 + request_data = { + "model": model, + "messages": messages, + "stream": True + } + + logger.info(f"[MCP] ask_grok 调用, 模型: {model}") + + # 调用Grok客户端(流式) + response_iterator = await GrokClient.openai_to_grok(request_data) + + # 收集所有流式响应块 + content_parts = [] + async for chunk in response_iterator: + if isinstance(chunk, bytes): + chunk = chunk.decode('utf-8') + + # 解析SSE格式 + if chunk.startswith("data: "): + data_str = chunk[6:].strip() + if data_str == "[DONE]": + break + + try: + data = json.loads(data_str) + choices = data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + if content := delta.get("content"): + content_parts.append(content) + except json.JSONDecodeError: + continue + + result = "".join(content_parts) + logger.info(f"[MCP] ask_grok 完成, 响应长度: {len(result)}") + return result + + except GrokApiException as e: + logger.error(f"[MCP] Grok API错误: {str(e)}") + raise Exception(f"Grok API调用失败: {str(e)}") + except Exception as e: + logger.error(f"[MCP] ask_grok异常: {str(e)}", exc_info=True) + raise Exception(f"处理请求时出错: {str(e)}") diff --git a/app/template/admin.html b/app/template/admin.html index 41e850d..0e47cc9 100644 --- a/app/template/admin.html +++ b/app/template/admin.html @@ -1,530 +1,887 @@ + - - - 管理控制台 - Grok2API - - - - + + + 管理控制台 - Grok2API + + + + + - -
-
-
- Grok2API - by @Chenyme -
-
-
-
- - - - - - - FILE - -
- 加载中... -
-
- + +
+
+
+ Grok2API + by @Chenyme +
+
+
+
+ + + + + + + FILE + +
+ 加载中...
+
+
-
+
+
-
- -
- -
- - -
- -
-
-

Token 总数

-

-

-
- -
-

未使用

-

-

-
- -
-

限流中

-

-

-
- -
-

失效

-

-

-
- -
-

正常

-

-

-
- -
-

总剩余

-

-

-
- -
-

普通剩余

-

-

-
- -
-

高级剩余

-

-

-
-
- - -
- -
-
-
- - -
-
- -
- - - -
-
- - -
- - - - - - - - - - - - - - - - -
- - Token类型状态普通剩余高级剩余创建时间操作
-
- - -
-
- - - -
- - - + + +
+ +
+
+
+ + + +
+
+ +
+ + + +
+
+ + +
+ + + + + + + + + + + + + + + + + + +
+ + Token类型状态普通高级标签备注创建时间 + 操作
+
+ + +
+ + + + + + + + + + + + + + \ No newline at end of file diff --git a/data/setting.toml b/data/setting.toml index 5ec8757..6c9fc17 100644 --- a/data/setting.toml +++ b/data/setting.toml @@ -1,17 +1,23 @@ [grok] api_key = "" proxy_url = "" -temporary = true +cache_proxy_url = "" cf_clearance = "" x_statsig_id = "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" filtered_tags = "xaiartifact,xai:tool_usage_card,grok:render" stream_chunk_timeout = 120 stream_total_timeout = 600 stream_first_response_timeout = 30 +temporary = true +show_thinking = true +dynamic_statsig = true +proxy_pool_url = "" +proxy_pool_interval = 300 +retry_status_codes = [ 401, 429,] [global] base_url = "" -log_level = "DEBUG" +log_level = "INFO" image_mode = "url" admin_password = "admin" admin_username = "admin" diff --git a/data/temp/video/users-8522ce45-679b-4e0e-a0f7-bb18f434eb6b-generated-15f7113f-5d16-4ff1-bdaa-a2eabd66671c-generated_video.mp4 b/data/temp/video/users-8522ce45-679b-4e0e-a0f7-bb18f434eb6b-generated-15f7113f-5d16-4ff1-bdaa-a2eabd66671c-generated_video.mp4 new file mode 100644 index 0000000..195951e Binary files /dev/null and b/data/temp/video/users-8522ce45-679b-4e0e-a0f7-bb18f434eb6b-generated-15f7113f-5d16-4ff1-bdaa-a2eabd66671c-generated_video.mp4 differ diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh new file mode 100644 index 0000000..b98e9f5 --- /dev/null +++ b/docker-entrypoint.sh @@ -0,0 +1,53 @@ +#!/bin/sh +set -e + +# 初始化配置文件(如果不存在) +echo "[Grok2API] 检查配置文件..." + +# 确保数据目录存在 +mkdir -p /app/data/temp/image /app/data/temp/video /app/logs + +# 如果 setting.toml 不存在,创建默认配置 +if [ ! -f /app/data/setting.toml ]; then + echo "[Grok2API] 初始化 setting.toml..." + cat > /app/data/setting.toml << 'EOF' +[global] +base_url = "http://localhost:8000" +log_level = "INFO" +image_mode = "url" +admin_password = "admin" +admin_username = "admin" +image_cache_max_size_mb = 512 +video_cache_max_size_mb = 1024 +max_upload_concurrency = 20 +max_request_concurrency = 50 +batch_save_interval = 1.0 +batch_save_threshold = 10 + +[grok] +api_key = "" +proxy_url = "" +cache_proxy_url = "" +cf_clearance = "" +x_statsig_id = "" +dynamic_statsig = true +filtered_tags = "xaiartifact,xai:tool_usage_card,grok:render" +stream_chunk_timeout = 120 +stream_total_timeout = 600 +stream_first_response_timeout = 30 +temporary = true +show_thinking = true +EOF +fi + +# 如果 token.json 不存在,创建空token文件 +if [ ! -f /app/data/token.json ]; then + echo "[Grok2API] 初始化 token.json..." + echo '{"ssoNormal": {}, "ssoSuper": {}}' > /app/data/token.json +fi + +echo "[Grok2API] 配置文件检查完成" +echo "[Grok2API] 启动应用..." + +# 执行传入的命令 +exec "$@" diff --git a/main.py b/main.py index 3f091dc..cd76815 100644 --- a/main.py +++ b/main.py @@ -13,28 +13,77 @@ from app.api.v1.models import router as models_router from app.api.v1.images import router as images_router from app.api.admin.manage import router as admin_router +# 导入MCP服务器(认证配置在server.py中完成) +from app.services.mcp import mcp +# 创建MCP的FastAPI应用实例 +# 使用流式HTTP传输,支持高效的双向流式通信 +mcp_app = mcp.http_app(stateless_http=True, transport="streamable-http") + +# 2. 定义应用生命周期 @asynccontextmanager async def lifespan(app: FastAPI): - """应用生命周期管理""" - # 初始化存储管理器 - await storage_manager.init() + """ + 启动顺序: + 1. 初始化核心服务 (storage, settings, token_manager) + 2. 异步加载 token 数据 + 3. 启动批量保存任务 + 4. 启动MCP服务生命周期 + 关闭顺序 (LIFO): + 1. 关闭MCP服务生命周期 + 2. 关闭批量保存任务并刷新数据 + 3. 关闭核心服务 + """ + # --- 启动过程 --- + # 1. 初始化核心服务 + await storage_manager.init() + # 设置存储到配置和token管理器 storage = storage_manager.get_storage() setting.set_storage(storage) token_manager.set_storage(storage) - # 重新加载配置和token数据(从存储同步) + # 2. 重新加载配置 await setting.reload() - token_manager._load_data() + logger.info("[Grok2API] 核心服务初始化完成") + # 2.5. 初始化代理池 + from app.core.proxy_pool import proxy_pool + proxy_url = setting.grok_config.get("proxy_url", "") + proxy_pool_url = setting.grok_config.get("proxy_pool_url", "") + proxy_pool_interval = setting.grok_config.get("proxy_pool_interval", 300) + proxy_pool.configure(proxy_url, proxy_pool_url, proxy_pool_interval) + + # 3. 异步加载 token 数据 + await token_manager._load_data() + logger.info("[Grok2API] Token数据加载完成") + + # 4. 启动批量保存任务 + await token_manager.start_batch_save() + + # 5. 管理MCP服务的生命周期 + mcp_lifespan_context = mcp_app.lifespan(app) + await mcp_lifespan_context.__aenter__() + logger.info("[MCP] MCP服务初始化完成") + logger.info("[Grok2API] 应用启动成功") - yield - # 关闭存储连接 - await storage_manager.close() - logger.info("[Grok2API] 应用关闭成功") + try: + yield + finally: + # --- 关闭过程 --- + # 1. 退出MCP服务的生命周期 + await mcp_lifespan_context.__aexit__(None, None, None) + logger.info("[MCP] MCP服务已关闭") + + # 2. 关闭批量保存任务并刷新数据 + await token_manager.shutdown() + logger.info("[Token] Token管理器已关闭") + + # 3. 关闭核心服务 + await storage_manager.close() + logger.info("[Grok2API] 应用关闭成功") # 初始化日志 @@ -44,7 +93,7 @@ logger.info("[Grok2API] 应用正在启动...") app = FastAPI( title="Grok2API", description="Grok API 转换服务", - version="1.0.3", + version="1.3.1", lifespan=lifespan ) @@ -57,12 +106,12 @@ app.include_router(models_router, prefix="/v1") app.include_router(images_router) app.include_router(admin_router) -# 挂载静态文件(注意:这个应该在API路由之后,避免拦截API请求) +# 挂载静态文件 app.mount("/static", StaticFiles(directory="app/template"), name="template") @app.get("/") async def root(): - """根路径 - 重定向到登录页""" + """根路径""" from fastapi.responses import RedirectResponse return RedirectResponse(url="/login") @@ -76,7 +125,28 @@ async def health_check(): "version": "1.0.3" } +# 挂载MCP服务器 +app.mount("", mcp_app) + if __name__ == "__main__": import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8001) + import os + + # 从环境变量读取 worker 数量,默认为 1 + workers = int(os.getenv("WORKERS", "1")) + + # 提示多进程模式 + if workers > 1: + logger.info( + f"[Grok2API] 多进程模式已启用 (workers={workers})。" + f"建议使用 Redis/MySQL 存储以获得最佳性能。" + ) + + uvicorn.run( + "main:app", # 使用字符串以支持多worker + host="0.0.0.0", + port=8001, + workers=workers, + loop="uvloop" if workers == 1 else "auto" # 单进程用uvloop,多进程用auto + ) \ No newline at end of file diff --git a/readme.md b/readme.md index a9d7b16..2e1e76e 100644 --- a/readme.md +++ b/readme.md @@ -78,7 +78,7 @@ services: # - DATABASE_URL=mysql://user:password@host:3306/grok2api ## MySQL格式: mysql://user:password@host:port/database - ## Redis格式: redis://host:port/db 或 redis://user:password@host:port/db + ## Redis格式: redis://host:port/db 或 redis://user:password@host:port/db (SSL: rediss://) volumes: grok_data: @@ -129,6 +129,11 @@ volumes: | POST | /api/cache/clear/images | 清理图片缓存 | ✅ | | POST | /api/cache/clear/videos | 清理视频缓存 | ✅ | | GET | /api/stats | 获取统计信息 | ✅ | +| POST | /api/tokens/tags | 更新 Token 标签 | ✅ | +| POST | /api/tokens/note | 更新 Token 备注 | ✅ | +| POST | /api/tokens/test | 测试 Token 可用性 | ✅ | +| GET | /api/tokens/tags/all | 获取所有标签列表 | ✅ | +| GET | /api/storage/mode | 获取存储模式信息 | ✅ | @@ -138,12 +143,14 @@ volumes: | 模型名称 | 计次 | 账户类型 | 图像生成/编辑 | 深度思考 | 联网搜索 | 视频生成 | |------------------------|--------|--------------|--------------|----------|----------|----------| -| `grok-3-fast` | 1 | Basic/Super | ✅ | ❌ | ✅ | ❌ | +| `grok-4.1` | 1 | Basic/Super | ✅ | ✅ | ✅ | ❌ | +| `grok-4.1-thinking` | 1 | Basic/Super | ✅ | ✅ | ✅ | ❌ | +| `grok-imagine-0.9` | - | Basic/Super | ✅ | ❌ | ❌ | ✅ | | `grok-4-fast` | 1 | Basic/Super | ✅ | ✅ | ✅ | ❌ | | `grok-4-fast-expert` | 4 | Basic/Super | ✅ | ✅ | ✅ | ❌ | | `grok-4-expert` | 4 | Basic/Super | ✅ | ✅ | ✅ | ❌ | | `grok-4-heavy` | 1 | Super | ✅ | ✅ | ✅ | ❌ | -| `grok-imagine-0.9` | - | Basic/Super | ✅ | ❌ | ❌ | ✅ | +| `grok-3-fast` | 1 | Basic/Super | ✅ | ❌ | ✅ | ❌ |
@@ -168,6 +175,7 @@ volumes: | cf_clearance | grok | 否 | Cloudflare安全令牌 | "" | | x_statsig_id | grok | 是 | 反机器人唯一标识符 | "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" | | filtered_tags | grok | 否 | 过滤响应标签(逗号分隔) | "xaiartifact,xai:tool_usage_card,grok:render" | +| show_thinking | grok | 否 | 显示思考过程 true(显示)/false(隐藏) | true | | temporary | grok | 否 | 会话模式 true(临时)/false | true |
diff --git a/requirements.txt b/requirements.txt index 7b97964..9e51d00 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,18 @@ + toml==0.10.2 fastapi==0.119.0 uvicorn==0.37.0 +uvloop==0.21.0 python-dotenv==1.1.1 curl_cffi==0.13.0 requests==2.32.5 starlette==0.48.0 pydantic==2.12.2 aiofiles==25.1.0 +portalocker==3.0.0 aiomysql==0.2.0 -redis==6.4.0 \ No newline at end of file +redis==6.4.0 +fastmcp==2.12.4 +cryptography==46.0.3 +orjson==3.11.4 +aiohttp==3.13.2 \ No newline at end of file diff --git a/test_concurrency.py b/test_concurrency.py new file mode 100755 index 0000000..c0089b1 --- /dev/null +++ b/test_concurrency.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python3 +""" +Grok2API 并发性能测试脚本 + +测试不同并发级别下的API性能表现 +""" + +import asyncio +import aiohttp +import time +import statistics +import argparse +from datetime import datetime +from typing import List, Dict, Any +import json + + +class ConcurrencyTester: + """并发测试器""" + + def __init__(self, base_url: str, api_key: str = None): + self.base_url = base_url.rstrip('/') + self.api_key = api_key + self.results: List[Dict[str, Any]] = [] + + async def test_request(self, session: aiohttp.ClientSession, request_id: int) -> Dict[str, Any]: + """发送单个测试请求""" + url = f"{self.base_url}/v1/chat/completions" + + headers = { + "Content-Type": "application/json" + } + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + payload = { + "model": "grok-3-fast", + "messages": [ + {"role": "user", "content": f"测试请求 #{request_id},请简短回复OK"} + ], + "stream": False, + "max_tokens": 10 + } + + start_time = time.time() + + try: + async with session.post(url, json=payload, headers=headers, timeout=30) as response: + status = response.status + + if status == 200: + data = await response.json() + elapsed = time.time() - start_time + + return { + "id": request_id, + "status": "success", + "http_status": status, + "elapsed": elapsed, + "response_length": len(json.dumps(data)) + } + else: + elapsed = time.time() - start_time + error_text = await response.text() + + return { + "id": request_id, + "status": "error", + "http_status": status, + "elapsed": elapsed, + "error": error_text[:200] + } + + except asyncio.TimeoutError: + elapsed = time.time() - start_time + return { + "id": request_id, + "status": "timeout", + "elapsed": elapsed, + "error": "Request timeout" + } + + except Exception as e: + elapsed = time.time() - start_time + return { + "id": request_id, + "status": "exception", + "elapsed": elapsed, + "error": str(e) + } + + async def run_concurrent_test(self, concurrency: int, total_requests: int): + """运行并发测试""" + print(f"\n{'='*60}") + print(f"📊 测试配置:并发数 {concurrency}, 总请求数 {total_requests}") + print(f"{'='*60}") + + connector = aiohttp.TCPConnector(limit=concurrency, limit_per_host=concurrency) + timeout = aiohttp.ClientTimeout(total=60) + + async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: + # 预热 + print("🔥 预热中...") + await self.test_request(session, 0) + + # 开始测试 + print(f"🚀 开始并发测试...") + start_time = time.time() + + # 创建任务 + tasks = [] + for i in range(1, total_requests + 1): + task = asyncio.create_task(self.test_request(session, i)) + tasks.append(task) + + # 控制并发数 + if len(tasks) >= concurrency: + results = await asyncio.gather(*tasks) + self.results.extend(results) + tasks = [] + + # 显示进度 + print(f" 进度: {i}/{total_requests} ({i/total_requests*100:.1f}%)", end='\r') + + # 处理剩余任务 + if tasks: + results = await asyncio.gather(*tasks) + self.results.extend(results) + + total_time = time.time() - start_time + + # 统计和输出 + self.print_statistics(concurrency, total_requests, total_time) + + def print_statistics(self, concurrency: int, total_requests: int, total_time: float): + """打印统计信息""" + success_results = [r for r in self.results if r["status"] == "success"] + error_results = [r for r in self.results if r["status"] != "success"] + + success_count = len(success_results) + error_count = len(error_results) + + if success_results: + latencies = [r["elapsed"] for r in success_results] + avg_latency = statistics.mean(latencies) + min_latency = min(latencies) + max_latency = max(latencies) + p50_latency = statistics.median(latencies) + p95_latency = sorted(latencies)[int(len(latencies) * 0.95)] if len(latencies) > 1 else latencies[0] + p99_latency = sorted(latencies)[int(len(latencies) * 0.99)] if len(latencies) > 1 else latencies[0] + else: + avg_latency = min_latency = max_latency = p50_latency = p95_latency = p99_latency = 0 + + throughput = total_requests / total_time if total_time > 0 else 0 + + print(f"\n\n{'='*60}") + print(f"📈 测试结果统计") + print(f"{'='*60}") + print(f" 测试时间: {total_time:.2f}s") + print(f" 总请求数: {total_requests}") + print(f" 并发数: {concurrency}") + print(f"") + print(f" 成功请求: {success_count} ({success_count/total_requests*100:.1f}%)") + print(f" 失败请求: {error_count} ({error_count/total_requests*100:.1f}%)") + print(f"") + print(f" 吞吐量: {throughput:.2f} req/s") + print(f"") + print(f" 延迟统计:") + print(f" 最小: {min_latency*1000:.0f}ms") + print(f" 平均: {avg_latency*1000:.0f}ms") + print(f" 最大: {max_latency*1000:.0f}ms") + print(f" P50: {p50_latency*1000:.0f}ms") + print(f" P95: {p95_latency*1000:.0f}ms") + print(f" P99: {p99_latency*1000:.0f}ms") + + # 错误详情 + if error_results: + print(f"\n ⚠️ 错误详情:") + error_types = {} + for r in error_results: + error_type = r.get("status", "unknown") + error_types[error_type] = error_types.get(error_type, 0) + 1 + + for error_type, count in error_types.items(): + print(f" {error_type}: {count}") + + print(f"{'='*60}\n") + + # 性能评级 + self.print_performance_rating(throughput, avg_latency) + + def print_performance_rating(self, throughput: float, avg_latency: float): + """打印性能评级""" + print(f"🎯 性能评级:") + + # 吞吐量评级 + if throughput >= 100: + rating = "⭐⭐⭐⭐⭐ 优秀" + elif throughput >= 60: + rating = "⭐⭐⭐⭐ 良好" + elif throughput >= 30: + rating = "⭐⭐⭐ 中等" + elif throughput >= 10: + rating = "⭐⭐ 较低" + else: + rating = "⭐ 需优化" + + print(f" 吞吐量 ({throughput:.1f} req/s): {rating}") + + # 延迟评级 + if avg_latency < 0.5: + rating = "⭐⭐⭐⭐⭐ 优秀" + elif avg_latency < 1.0: + rating = "⭐⭐⭐⭐ 良好" + elif avg_latency < 2.0: + rating = "⭐⭐⭐ 中等" + elif avg_latency < 5.0: + rating = "⭐⭐ 较高" + else: + rating = "⭐ 需优化" + + print(f" 平均延迟 ({avg_latency*1000:.0f}ms): {rating}") + print() + + +async def main(): + """主函数""" + parser = argparse.ArgumentParser(description='Grok2API 并发性能测试') + parser.add_argument('--url', default='http://localhost:8001', help='API 基础URL') + parser.add_argument('--key', default='', help='API Key(可选)') + parser.add_argument('-c', '--concurrency', type=int, default=10, help='并发数') + parser.add_argument('-n', '--requests', type=int, default=50, help='总请求数') + parser.add_argument('--multi-test', action='store_true', help='运行多级并发测试') + + args = parser.parse_args() + + print(f""" +╔══════════════════════════════════════════════════════════╗ +║ Grok2API 并发性能测试工具 ║ +╚══════════════════════════════════════════════════════════╝ + +🔗 测试目标: {args.url} +🔑 API Key: {'已设置' if args.key else '未设置'} +⏰ 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} +""") + + tester = ConcurrencyTester(args.url, args.key) + + if args.multi_test: + # 多级并发测试 + test_configs = [ + (5, 20), # 5并发,20请求 + (10, 50), # 10并发,50请求 + (20, 100), # 20并发,100请求 + (50, 200), # 50并发,200请求 + ] + + for concurrency, requests in test_configs: + tester.results = [] # 清空结果 + await tester.run_concurrent_test(concurrency, requests) + await asyncio.sleep(2) # 间隔2秒 + else: + # 单次测试 + await tester.run_concurrent_test(args.concurrency, args.requests) + + print(f"\n✅ 测试完成!") + print(f"⏰ 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + print("\n\n⚠️ 测试被用户中断") + except Exception as e: + print(f"\n\n❌ 测试失败: {e}") diff --git a/test_concurrency.sh b/test_concurrency.sh new file mode 100755 index 0000000..b100e2c --- /dev/null +++ b/test_concurrency.sh @@ -0,0 +1,177 @@ +#!/bin/bash + +# Grok2API 并发测试脚本(Shell版本) +# 使用 curl 和 GNU parallel 进行并发测试 + +set -e + +# 配置 +BASE_URL="${BASE_URL:-http://localhost:8001}" +API_KEY="${API_KEY:-}" +CONCURRENCY="${CONCURRENCY:-10}" +TOTAL_REQUESTS="${TOTAL_REQUESTS:-50}" + +# 颜色 +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' # No Color + +echo -e "${BLUE}╔══════════════════════════════════════════════════════════╗${NC}" +echo -e "${BLUE}║ Grok2API 并发性能测试工具 (Shell版) ║${NC}" +echo -e "${BLUE}╚══════════════════════════════════════════════════════════╝${NC}" +echo "" +echo -e "${GREEN}🔗 测试目标:${NC} $BASE_URL" +echo -e "${GREEN}🔑 API Key:${NC} ${API_KEY:-(未设置)}" +echo -e "${GREEN}📊 并发数:${NC} $CONCURRENCY" +echo -e "${GREEN}📈 总请求数:${NC} $TOTAL_REQUESTS" +echo "" + +# 检查依赖 +if ! command -v curl &> /dev/null; then + echo -e "${RED}❌ 错误: 需要安装 curl${NC}" + exit 1 +fi + +# 创建临时目录 +TMP_DIR=$(mktemp -d) +trap "rm -rf $TMP_DIR" EXIT + +# 单个请求函数 +test_request() { + local request_id=$1 + local start_time=$(date +%s.%N) + + # 构建请求 + local headers="Content-Type: application/json" + if [ -n "$API_KEY" ]; then + headers="${headers}\nAuthorization: Bearer ${API_KEY}" + fi + + local response=$(curl -s -w "\n%{http_code}\n%{time_total}" \ + -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + ${API_KEY:+-H "Authorization: Bearer $API_KEY"} \ + -d "{ + \"model\": \"grok-3-fast\", + \"messages\": [{\"role\": \"user\", \"content\": \"测试请求 #${request_id},请简短回复OK\"}], + \"stream\": false, + \"max_tokens\": 10 + }" 2>&1) + + local http_code=$(echo "$response" | tail -n 2 | head -n 1) + local time_total=$(echo "$response" | tail -n 1) + + # 记录结果 + echo "${request_id},${http_code},${time_total}" >> "$TMP_DIR/results.csv" + + # 显示进度 + echo -ne "\r 进度: ${request_id}/${TOTAL_REQUESTS}" +} + +# 导出函数供 parallel 使用 +export -f test_request +export BASE_URL API_KEY TMP_DIR + +# 清空结果文件 +echo "id,status,time" > "$TMP_DIR/results.csv" + +echo -e "${YELLOW}🚀 开始并发测试...${NC}" +START_TIME=$(date +%s.%N) + +# 使用 GNU parallel(如果可用),否则使用简单循环 +if command -v parallel &> /dev/null; then + seq 1 $TOTAL_REQUESTS | parallel -j $CONCURRENCY test_request {} +else + # 简单的后台任务并发 + for i in $(seq 1 $TOTAL_REQUESTS); do + test_request $i & + + # 控制并发数 + if (( i % CONCURRENCY == 0 )); then + wait + fi + done + wait +fi + +END_TIME=$(date +%s.%N) +TOTAL_TIME=$(echo "$END_TIME - $START_TIME" | bc) + +echo -e "\n" + +# 统计结果 +echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}" +echo -e "${BLUE}📈 测试结果统计${NC}" +echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}" + +# 统计成功/失败 +SUCCESS_COUNT=$(awk -F',' '$2 == 200 {count++} END {print count+0}' "$TMP_DIR/results.csv") +ERROR_COUNT=$((TOTAL_REQUESTS - SUCCESS_COUNT)) + +echo -e " 测试时间: ${TOTAL_TIME}s" +echo -e " 总请求数: ${TOTAL_REQUESTS}" +echo -e " 并发数: ${CONCURRENCY}" +echo "" +echo -e " 成功请求: ${GREEN}${SUCCESS_COUNT}${NC} ($(echo "scale=1; $SUCCESS_COUNT * 100 / $TOTAL_REQUESTS" | bc)%)" +echo -e " 失败请求: ${RED}${ERROR_COUNT}${NC} ($(echo "scale=1; $ERROR_COUNT * 100 / $TOTAL_REQUESTS" | bc)%)" +echo "" + +# 计算吞吐量 +THROUGHPUT=$(echo "scale=2; $TOTAL_REQUESTS / $TOTAL_TIME" | bc) +echo -e " 吞吐量: ${GREEN}${THROUGHPUT}${NC} req/s" +echo "" + +# 延迟统计(只统计成功的请求) +if [ $SUCCESS_COUNT -gt 0 ]; then + echo -e " 延迟统计:" + + # 提取成功请求的延迟时间 + awk -F',' '$2 == 200 {print $3}' "$TMP_DIR/results.csv" | sort -n > "$TMP_DIR/latencies.txt" + + MIN=$(head -n 1 "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + MAX=$(tail -n 1 "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + AVG=$(awk '{sum+=$1; count++} END {printf "%.0f", sum/count*1000}' "$TMP_DIR/latencies.txt") + + # P50 + P50_LINE=$((SUCCESS_COUNT / 2)) + P50=$(sed -n "${P50_LINE}p" "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + + # P95 + P95_LINE=$(echo "scale=0; $SUCCESS_COUNT * 0.95 / 1" | bc) + P95=$(sed -n "${P95_LINE}p" "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + + # P99 + P99_LINE=$(echo "scale=0; $SUCCESS_COUNT * 0.99 / 1" | bc) + P99=$(sed -n "${P99_LINE}p" "$TMP_DIR/latencies.txt" | awk '{printf "%.0f", $1*1000}') + + echo -e " 最小: ${MIN}ms" + echo -e " 平均: ${AVG}ms" + echo -e " 最大: ${MAX}ms" + echo -e " P50: ${P50}ms" + echo -e " P95: ${P95}ms" + echo -e " P99: ${P99}ms" +fi + +echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}" + +# 性能评级 +echo -e "${YELLOW}🎯 性能评级:${NC}" + +if (( $(echo "$THROUGHPUT >= 100" | bc -l) )); then + RATING="⭐⭐⭐⭐⭐ 优秀" +elif (( $(echo "$THROUGHPUT >= 60" | bc -l) )); then + RATING="⭐⭐⭐⭐ 良好" +elif (( $(echo "$THROUGHPUT >= 30" | bc -l) )); then + RATING="⭐⭐⭐ 中等" +elif (( $(echo "$THROUGHPUT >= 10" | bc -l) )); then + RATING="⭐⭐ 较低" +else + RATING="⭐ 需优化" +fi + +echo -e " 吞吐量 (${THROUGHPUT} req/s): ${RATING}" + +echo "" +echo -e "${GREEN}✅ 测试完成!${NC}"