Merge branch 'chenyme:main' into main

This commit is contained in:
彼方 authored and GitHub committed 2025-12-19 11:04:41 +08:00
commit 7ea8b7c19a
34 files changed
+4052 -2368

No files matched your search

+4
View File
@@ -80,6 +80,10 @@ logs/
*.log
*.log.*
# 数据临时文件
data/temp/
data/token.json
# 数据库文件
*.db
*.sqlite
+58 -8
View File
@@ -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
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
+42 -24
View File
@@ -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"]
# 使用 entrypoint 脚本初始化配置
ENTRYPOINT ["docker-entrypoint.sh"]
# 默认启动命令
CMD ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
+278 -377
View File
@@ -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"})
+10 -33
View File
@@ -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={
+8 -13
View File
@@ -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))
+18 -57
View File
@@ -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"
}
+24 -19
View File
@@ -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()
+166 -45
View File
@@ -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)
+54 -56
View File
@@ -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)
+67 -21
View File
@@ -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()
+143
View File
@@ -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()
+97 -142
View File
@@ -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()
+135 -95
View File
@@ -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
return config["rate_limit_model"] if config else model
@classmethod
def get_all_model_names(cls) -> list[str]:
"""获取所有模型名称"""
return list(_MODEL_CONFIG.keys())
+26 -38
View File
@@ -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="响应选项")
choices: List[OpenAIChatCompletionChunkChoice] = Field(..., description="选项")
+179 -111
View File
@@ -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()
+243 -175
View File
@@ -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}")
logger.error(f"[Client] 更新限制失败: {e}")
+140
View File
@@ -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
+183 -266
View File
@@ -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'<video src="{local_video_url}" controls="controls" width="500" height="300"></video>\n'
else:
content = f'<video src="{full_video_url}" controls="controls" width="500" height="300"></video>\n'
except Exception as e:
logger.warning(f"[Processor] 缓存视频失败: {e}")
content = f'<video src="{full_video_url}" controls="controls" width="500" height="300"></video>\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
# 添加 <think> 标签
if not video_progress_started:
content = f"<think>视频已生成{progress}%\n"
video_progress_started = True
elif progress < 100:
content = f"视频已生成{progress}%\n"
else:
# 进度100%时关闭 <think> 标签并立即处理视频
content = f"视频已生成{progress}%</think>\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'<video src="{local_video_url}" controls="controls"></video>\n'
else:
content += f'<video src="{full_video_url}" controls="controls"></video>\n'
except Exception as e:
logger.warning(f"[Processor] 缓存视频失败: {e}")
content += f'<video src="{full_video_url}" controls="controls"></video>\n'
yield make_chunk(content)
timeout_manager.mark_chunk_received()
chunk_index += 1
if show_thinking:
if not video_progress_started:
content = f"<think>视频已生成{progress}%\n"
video_progress_started = True
elif progress < 100:
content = f"视频已生成{progress}%\n"
else:
content = f"视频已生成{progress}%</think>\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")
# 跳过后续的 <think> 标签
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"<think>\n{content}"
if show_thinking:
content = f"<think>\n{content}"
else:
should_skip = True
elif is_thinking and not current_is_thinking:
content = f"\n</think>\n{content}"
if show_thinking:
content = f"\n</think>\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}")
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'<video src="{local_url}" controls="controls" width="500" height="300"></video>\n'
except Exception as e:
logger.warning(f"[Processor] 缓存视频失败: {e}")
return f'<video src="{full_url}" controls="controls" width="500" height="300"></video>\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
)
+64 -28
View File
@@ -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
+323 -234
View File
@@ -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()
+155 -91
View File
@@ -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
return f"image.{ext}", mime
+6
View File
@@ -0,0 +1,6 @@
# -*- coding: utf-8 -*-
"""MCP模块初始化"""
from app.services.mcp.server import mcp
__all__ = ["mcp"]
+63
View File
@@ -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)
+77
View File
@@ -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)}")
+873 -516
View File
File diff suppressed because it is too large. Load diff
+8 -2
View File
@@ -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"
+53
View File
@@ -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 "$@"
+83 -13
View File
@@ -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
)
+11 -3
View File
@@ -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 | 获取存储模式信息 | ✅ |
</details>
@@ -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 | ✅ | ❌ | ✅ | ❌ |
<br>
@@ -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 |
<br>
+8 -1
View File
@@ -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
redis==6.4.0
fastmcp==2.12.4
cryptography==46.0.3
orjson==3.11.4
aiohttp==3.13.2
+276
View File
@@ -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}")
+177
View File
@@ -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}"