Merge branch 'chenyme:main' into main
This commit is contained in:
34 files changed
+4052
-2368
No files matched your search
@@ -80,6 +80,10 @@ logs/
|
||||
*.log
|
||||
*.log.*
|
||||
|
||||
# 数据临时文件
|
||||
data/temp/
|
||||
data/token.json
|
||||
|
||||
# 数据库文件
|
||||
*.db
|
||||
*.sqlite
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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()
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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}")
|
||||
@@ -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
@@ -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"
|
||||
else:
|
||||
content += f"\n"
|
||||
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"
|
||||
else:
|
||||
content += f"\n"
|
||||
except Exception as e:
|
||||
logger.warning(f"[Processor] 处理图片失败: {e}")
|
||||
content += f"\n"
|
||||
|
||||
# 返回 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 = "\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"
|
||||
# 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"\n")
|
||||
timeout_manager.mark_chunk_received()
|
||||
chunk_index += 1
|
||||
else:
|
||||
yield make_chunk(f"\n")
|
||||
timeout_manager.mark_chunk_received()
|
||||
chunk_index += 1
|
||||
else:
|
||||
yield make_chunk(f"\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"\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"
|
||||
else:
|
||||
content += f"\n"
|
||||
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"
|
||||
else:
|
||||
content += f"\n"
|
||||
except Exception as e:
|
||||
logger.warning(f"[Processor] 处理图片失败: {e}")
|
||||
content += f"\n"
|
||||
|
||||
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
|
||||
)
|
||||
@@ -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
@@ -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
@@ -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
|
||||
@@ -0,0 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""MCP模块初始化"""
|
||||
|
||||
from app.services.mcp.server import mcp
|
||||
|
||||
__all__ = ["mcp"]
|
||||
@@ -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)
|
||||
@@ -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
File diff suppressed because it is too large.
Load diff
+8
-2
@@ -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"
|
||||
|
||||
BIN
Binary file not shown.
@@ -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 "$@"
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
@@ -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
|
||||
Executable
+276
@@ -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}")
|
||||
Executable
+177
@@ -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}"
|
||||
Reference in new issue
Block a user