diff --git a/app/api/admin/manage.py b/app/api/admin/manage.py
index 25445ec..cb0fe81 100644
--- a/app/api/admin/manage.py
+++ b/app/api/admin/manage.py
@@ -1,8 +1,4 @@
-"""
-管理接口模块
-
-提供Token管理功能，包括登录验证、Token增删查等操作。
-"""
+"""管理接口 - Token管理和系统配置"""
import secrets
from typing import Dict, Any, List, Optional
@@ -18,10 +14,9 @@ from app.services.grok.token import token_manager
from app.models.grok_models import TokenType
-# 创建路由器
router = APIRouter(tags=["管理"])
-# 常量定义
+# 常量
STATIC_DIR = Path(__file__).parents[2] / "template"
TEMP_DIR = Path(__file__).parents[3] / "data" / "temp"
IMAGE_CACHE_DIR = TEMP_DIR / "image"
@@ -30,70 +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 # "未使用"、"限流中"、"失效"、"正常"
- tags: List[str] = [] # 标签列表
- note: 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):
@@ -102,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")
@@ -113,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)
@@ -124,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)
@@ -211,91 +217,58 @@ 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",
@@ -303,13 +276,12 @@ async def list_tokens(_: bool = Depends(verify_admin_session)) -> TokenListRespo
remaining_queries=data.get("remainingQueries", -1),
heavy_remaining_queries=data.get("heavyremainingQueries", -1),
status=get_token_status(data, "sso"),
- tags=data.get("tags", []), # 向后兼容，如果没有tags字段则返回空列表
- note=data.get("note", "") # 向后兼容，如果没有note字段则返回空字符串
+ 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",
@@ -317,190 +289,91 @@ async def list_tokens(_: bool = Depends(verify_admin_session)) -> TokenListRespo
remaining_queries=data.get("remainingQueries", -1),
heavy_remaining_queries=data.get("heavyremainingQueries", -1),
status=get_token_status(data, "ssoSuper"),
- tags=data.get("tags", []), # 向后兼容，如果没有tags字段则返回空列表
- note=data.get("note", "") # 向后兼容，如果没有note字段则返回空字符串
+ 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,
@@ -515,376 +388,208 @@ 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()
-
- return {
- "success": True,
- "data": {
- "mode": storage_mode
- }
- }
-
+ mode = os.getenv("STORAGE_MODE", "file").upper()
+ return {"success": True, "data": {"mode": mode}}
except Exception as e:
- logger.error(f"[Admin] 获取存储模式异常 - 错误: {str(e)}")
- raise HTTPException(
- status_code=500,
- detail={"error": f"获取存储模式失败: {str(e)}", "code": "STORAGE_MODE_ERROR"}
- )
-
-
-class UpdateTokenTagsRequest(BaseModel):
- """更新Token标签请求"""
- token: str
- token_type: str
- tags: List[str]
+ logger.error(f"[Admin] 获取存储模式异常: {e}")
+ raise HTTPException(status_code=500, detail={"error": f"获取失败: {e}", "code": "STORAGE_MODE_ERROR"})
@router.post("/api/tokens/tags")
-async def update_token_tags(
- request: UpdateTokenTagsRequest,
- _: bool = Depends(verify_admin_session)
-) -> Dict[str, Any]:
- """
- 更新Token标签
-
- 为指定Token添加或修改标签。
- """
+async def update_token_tags(request: UpdateTokenTagsRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]:
+ """更新Token标签"""
try:
- logger.debug(f"[Admin] 更新Token标签 - Token: {request.token[:10]}..., Tags: {request.tags}")
+ logger.debug(f"[Admin] 更新Token标签: {request.token[:10]}..., {request.tags}")
- # 验证并转换token类型
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标签更新成功 - Token: {request.token[:10]}..., Tags: {request.tags}")
-
- return {
- "success": True,
- "message": "标签更新成功",
- "tags": 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标签更新异常 - Token: {request.token[:10]}..., 错误: {str(e)}")
- raise HTTPException(
- status_code=500,
- detail={"error": f"更新标签失败: {str(e)}", "code": "UPDATE_TAGS_ERROR"}
- )
+ 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]:
- """
- 获取所有标签
-
- 返回系统中所有Token使用的标签列表（去重）。
- """
+ """获取所有标签"""
try:
logger.debug("[Admin] 获取所有标签")
- all_tokens_data = token_manager.get_tokens()
+ all_tokens = token_manager.get_tokens()
tags_set = set()
- # 收集所有标签
- for token_type_data in all_tokens_data.values():
+ 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
- }
+ 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": "GET_TAGS_ERROR"}
- )
-
-
-class UpdateTokenNoteRequest(BaseModel):
- """更新Token备注请求"""
- token: str
- token_type: str
- note: str
+ 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备注
-
- 为指定Token添加或修改备注信息。
- """
+async def update_token_note(request: UpdateTokenNoteRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]:
+ """更新Token备注"""
try:
- logger.debug(f"[Admin] 更新Token备注 - Token: {request.token[:10]}...")
+ logger.debug(f"[Admin] 更新Token备注: {request.token[:10]}...")
- # 验证并转换token类型
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备注更新成功 - Token: {request.token[:10]}...")
-
- return {
- "success": True,
- "message": "备注更新成功",
- "note": 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备注更新异常 - Token: {request.token[:10]}..., 错误: {str(e)}")
- raise HTTPException(
- status_code=500,
- detail={"error": f"更新备注失败: {str(e)}", "code": "UPDATE_NOTE_ERROR"}
- )
-
-
-class TestTokenRequest(BaseModel):
- """测试Token请求"""
- token: str
- token_type: str
+ 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可用性
-
- 通过发送速率限制检查请求来验证Token是否有效。
- 根据不同的HTTP状态码进行相应的处理：
- - 401: Token失效
- - 403: 服务器被block，不改变Token状态
- - 其他错误: 设置为限流状态
- """
+async def test_token(request: TestTokenRequest, _: bool = Depends(verify_admin_session)) -> Dict[str, Any]:
+ """测试Token可用性"""
try:
- logger.debug(f"[Admin] 测试Token - Token: {request.token[:10]}...")
+ logger.debug(f"[Admin] 测试Token: {request.token[:10]}...")
- # 验证并转换token类型
token_type = validate_token_type(request.token_type)
-
- # 构造完整的auth token
auth_token = f"sso-rw={request.token};sso={request.token}"
- # 使用check_limits方法测试token
result = await token_manager.check_limits(auth_token, "grok-4-fast")
if result:
- logger.debug(f"[Admin] Token测试成功 - Token: {request.token[:10]}...")
+ logger.debug(f"[Admin] Token测试成功: {request.token[:10]}...")
return {
"success": True,
"message": "Token有效",
@@ -895,65 +600,23 @@ async def test_token(
}
}
else:
- # 测试失败，check_limits方法已经调用了record_failure处理错误
- # 现在检查token的状态来判断错误类型
- logger.warning(f"[Admin] Token测试失败 - Token: {request.token[:10]}...")
+ logger.warning(f"[Admin] Token测试失败: {request.token[:10]}...")
- # 检查token当前状态
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":
- # Token被标记为失效（401错误）
- return {
- "success": False,
- "message": "Token已失效",
- "data": {
- "valid": False,
- "error_type": "expired",
- "error_code": 401
- }
- }
+ return {"success": False, "message": "Token已失效", "data": {"valid": False, "error_type": "expired", "error_code": 401}}
elif token_data.get("remainingQueries") == 0:
- # Token被设置为限流状态（其他错误）
- return {
- "success": False,
- "message": "Token已被限流",
- "data": {
- "valid": False,
- "error_type": "limited",
- "error_code": "other"
- }
- }
+ return {"success": False, "message": "Token已被限流", "data": {"valid": False, "error_type": "limited", "error_code": "other"}}
else:
- # 可能是403错误或其他网络问题，token状态未变
- return {
- "success": False,
- "message": "服务器被block或网络错误",
- "data": {
- "valid": False,
- "error_type": "blocked",
- "error_code": 403
- }
- }
+ return {"success": False, "message": "服务器被block或网络错误", "data": {"valid": False, "error_type": "blocked", "error_code": 403}}
else:
- # 找不到token数据
- return {
- "success": False,
- "message": "Token数据异常",
- "data": {
- "valid": False,
- "error_type": "unknown",
- "error_code": "data_error"
- }
- }
+ 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测试异常 - Token: {request.token[:10]}..., 错误: {str(e)}")
- raise HTTPException(
- status_code=500,
- detail={"error": f"测试Token失败: {str(e)}", "code": "TEST_TOKEN_ERROR"}
- )
+ logger.error(f"[Admin] Token测试异常: {e}")
+ raise HTTPException(status_code=500, detail={"error": f"测试失败: {e}", "code": "TEST_TOKEN_ERROR"})
diff --git a/app/api/v1/chat.py b/app/api/v1/chat.py
index 72dddeb..d80424d 100644
--- a/app/api/v1/chat.py
+++ b/app/api/v1/chat.py
@@ -1,9 +1,4 @@
-# -*- coding: utf-8 -*-
-"""
-聊天API路由模块
-
-提供OpenAI兼容的聊天API接口，支持与Grok模型的交互。
-"""
+"""聊天API路由 - OpenAI兼容的聊天接口"""
from fastapi import APIRouter, Depends, HTTPException
from typing import Optional
@@ -15,38 +10,20 @@ from app.core.logger import logger
from app.services.grok.client import GrokClient
from app.models.openai_schema import OpenAIChatRequest
-# 聊天路由
+
router = APIRouter(prefix="/chat", tags=["聊天"])
@router.post("/completions", response_model=None)
-async def chat_completions(
- request: OpenAIChatRequest,
- _: Optional[str] = Depends(auth_manager.verify)
-):
- """
- 创建聊天补全
-
- 兼容OpenAI聊天API的端点，支持流式和非流式响应。
-
- Args:
- request: OpenAI格式的聊天请求
- _: 认证依赖（自动验证）
-
- Returns:
- OpenAIChatCompletionResponse: 非流式响应
- StreamingResponse: 流式响应
-
- Raises:
- HTTPException: 当请求处理失败时
- """
+async def chat_completions(request: OpenAIChatRequest, _: Optional[str] = Depends(auth_manager.verify)):
+ """创建聊天补全（支持流式和非流式）"""
try:
- logger.info(f"[Chat] 收到聊天请求")
+ logger.info("[Chat] 收到聊天请求")
- # 调用Grok客户端处理请求
+ # 调用Grok客户端
result = await GrokClient.openai_to_grok(request.model_dump())
- # 如果是流式响应，GrokClient已经返回了Iterator，直接包装为StreamingResponse
+ # 流式响应
if request.stream:
return StreamingResponse(
content=result,
@@ -58,11 +35,11 @@ async def chat_completions(
}
)
- # 非流式响应直接返回
+ # 非流式响应
return result
except GrokApiException as e:
- logger.error(f"[Chat] Grok API错误: {str(e)} - 详情: {e.details}")
+ logger.error(f"[Chat] Grok API错误: {e} - 详情: {e.details}")
raise HTTPException(
status_code=500,
detail={
@@ -74,7 +51,7 @@ async def chat_completions(
}
)
except Exception as e:
- logger.error(f"[Chat] 聊天请求处理失败: {str(e)}")
+ logger.error(f"[Chat] 处理失败: {e}")
raise HTTPException(
status_code=500,
detail={
diff --git a/app/api/v1/images.py b/app/api/v1/images.py
index b9c191e..e0679fa 100644
--- a/app/api/v1/images.py
+++ b/app/api/v1/images.py
@@ -1,4 +1,4 @@
-"""图片服务API路由"""
+"""图片服务API - 提供缓存的图片和视频文件"""
from fastapi import APIRouter, HTTPException
from fastapi.responses import FileResponse
@@ -13,31 +13,26 @@ router = APIRouter()
@router.get("/images/{img_path:path}")
async def get_image(img_path: str):
"""获取缓存的图片或视频
-
+
Args:
- img_path: 文件路径，格式如 users-xxx-generated-xxx-image.jpg 或 users-xxx-generated-xxx-video.mp4
-
- Returns:
- 文件响应
+ img_path: 文件路径（格式：users-xxx-generated-xxx-image.jpg）
"""
try:
- # 将路径转换回原始格式（短横线转斜杠）
+ # 转换路径（短横线→斜杠）
original_path = "/" + img_path.replace('-', '/')
- # 判断是图片还是视频
+ # 判断类型
is_video = any(original_path.lower().endswith(ext) for ext in ['.mp4', '.webm', '.mov', '.avi'])
if is_video:
- # 检查视频缓存
cache_path = video_cache_service.get_cached(original_path)
media_type = "video/mp4"
else:
- # 检查图片缓存
cache_path = image_cache_service.get_cached(original_path)
media_type = "image/jpeg"
if cache_path and cache_path.exists():
- logger.debug(f"[MediaAPI] 返回缓存文件: {cache_path}")
+ logger.debug(f"[MediaAPI] 返回缓存: {cache_path}")
return FileResponse(
path=str(cache_path),
media_type=media_type,
@@ -48,11 +43,11 @@ async def get_image(img_path: str):
)
# 文件不存在
- logger.warning(f"[MediaAPI] 文件未找到: {original_path}")
+ logger.warning(f"[MediaAPI] 未找到: {original_path}")
raise HTTPException(status_code=404, detail="File not found")
except HTTPException:
raise
except Exception as e:
- logger.error(f"[MediaAPI] 获取文件失败: {e}")
+ logger.error(f"[MediaAPI] 获取失败: {e}")
raise HTTPException(status_code=500, detail=str(e))
diff --git a/app/api/v1/models.py b/app/api/v1/models.py
index 5c18513..5cdcc71 100644
--- a/app/api/v1/models.py
+++ b/app/api/v1/models.py
@@ -1,8 +1,4 @@
-"""
-模型接口模块
-
-提供 OpenAI 兼容的 /v1/models 端点，返回系统支持的所有模型列表。
-"""
+"""模型接口 - OpenAI兼容的模型列表端点"""
import time
from typing import Dict, Any, List, Optional
@@ -12,43 +8,27 @@ from app.models.grok_models import Models
from app.core.auth import auth_manager
from app.core.logger import logger
-# 配置日志
-# 创建路由器
router = APIRouter(tags=["模型"])
@router.get("/models")
async def list_models(_: Optional[str] = Depends(auth_manager.verify)) -> Dict[str, Any]:
- """
- 获取可用模型列表
-
- 返回 OpenAI 兼容的模型列表格式，包含系统支持的所有 Grok 模型的详细信息。
-
- Args:
- _: 认证依赖（自动验证）
-
- Returns:
- Dict[str, Any]: 包含模型列表的响应数据
- """
+ """获取可用模型列表"""
try:
- logger.debug("[Models] 请求获取模型列表")
+ logger.debug("[Models] 请求模型列表")
- # 获取当前时间戳
- current_timestamp = int(time.time())
-
- # 构建模型数据列表
+ timestamp = int(time.time())
model_data: List[Dict[str, Any]] = []
for model in Models:
model_id = model.value
config = Models.get_model_info(model_id)
- # 基础信息
model_info = {
"id": model_id,
"object": "model",
- "created": current_timestamp,
+ "created": timestamp,
"owned_by": "x-ai",
"display_name": config.get("display_name", model_id),
"description": config.get("description", ""),
@@ -61,22 +41,16 @@ async def list_models(_: Optional[str] = Depends(auth_manager.verify)) -> Dict[s
model_data.append(model_info)
- # 构建响应
- response = {
- "object": "list",
- "data": model_data
- }
-
- logger.debug(f"[Models] 成功返回 {len(model_data)} 个模型")
- return response
+ logger.debug(f"[Models] 返回 {len(model_data)} 个模型")
+ return {"object": "list", "data": model_data}
except Exception as e:
- logger.error(f"[Models] 获取模型列表时发生错误: {str(e)}")
+ logger.error(f"[Models] 获取列表失败: {e}")
raise HTTPException(
status_code=500,
detail={
"error": {
- "message": f"Failed to retrieve models: {str(e)}",
+ "message": f"Failed to retrieve models: {e}",
"type": "internal_error",
"code": "model_list_error"
}
@@ -86,22 +60,13 @@ async def list_models(_: Optional[str] = Depends(auth_manager.verify)) -> Dict[s
@router.get("/models/{model_id}")
async def get_model(model_id: str, _: Optional[str] = Depends(auth_manager.verify)) -> Dict[str, Any]:
- """
- 获取特定模型信息
-
- Args:
- model_id (str): 模型ID
- _: 认证依赖（自动验证）
-
- Returns:
- Dict[str, Any]: 模型详细信息
- """
+ """获取特定模型信息"""
try:
- logger.debug(f"[Models] 请求获取模型信息: {model_id}")
+ logger.debug(f"[Models] 请求模型: {model_id}")
- # 验证模型是否存在
+ # 验证模型
if not Models.is_valid_model(model_id):
- logger.warning(f"[Models] 请求的模型不存在: {model_id}")
+ logger.warning(f"[Models] 模型不存在: {model_id}")
raise HTTPException(
status_code=404,
detail={
@@ -113,17 +78,13 @@ async def get_model(model_id: str, _: Optional[str] = Depends(auth_manager.verif
}
)
- # 获取当前时间戳
- current_timestamp = int(time.time())
-
- # 获取模型配置
+ timestamp = int(time.time())
config = Models.get_model_info(model_id)
- # 构建模型信息
model_info = {
"id": model_id,
"object": "model",
- "created": current_timestamp,
+ "created": timestamp,
"owned_by": "x-ai",
"display_name": config.get("display_name", model_id),
"description": config.get("description", ""),
@@ -134,18 +95,18 @@ async def get_model(model_id: str, _: Optional[str] = Depends(auth_manager.verif
"default_top_p": config.get("default_top_p", 0.95)
}
- logger.debug(f"[Models] 成功返回模型信息: {model_id}")
+ logger.debug(f"[Models] 返回模型: {model_id}")
return model_info
except HTTPException:
raise
except Exception as e:
- logger.error(f"[Models] 获取模型信息时发生错误: {str(e)}")
+ logger.error(f"[Models] 获取模型失败: {e}")
raise HTTPException(
status_code=500,
detail={
"error": {
- "message": f"Failed to retrieve model: {str(e)}",
+ "message": f"Failed to retrieve model: {e}",
"type": "internal_error",
"code": "model_retrieve_error"
}
diff --git a/app/core/auth.py b/app/core/auth.py
index d497e4d..9a95f69 100644
--- a/app/core/auth.py
+++ b/app/core/auth.py
@@ -1,53 +1,58 @@
-"""认证模块"""
+"""认证模块 - API令牌验证"""
from typing import Optional
from fastapi import Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
from app.core.config import setting
from app.core.logger import logger
+# Bearer安全方案
security = HTTPBearer(auto_error=False)
+def _build_error(message: str, code: str = "invalid_token") -> dict:
+ """构建认证错误"""
+ return {
+ "error": {
+ "message": message,
+ "type": "authentication_error",
+ "code": code
+ }
+ }
+
+
class AuthManager:
- """认证管理器"""
+ """认证管理器 - 验证API令牌"""
@staticmethod
def verify(credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)) -> Optional[str]:
- """验证认证令牌"""
+ """验证令牌"""
api_key = setting.grok_config.get("api_key")
+ # 未设置时跳过
if not api_key:
- logger.debug("[Auth] 未设置API_KEY，跳过验证。")
+ logger.debug("[Auth] 未设置API_KEY，跳过验证")
return credentials.credentials if credentials else None
+ # 检查令牌
if not credentials:
raise HTTPException(
status_code=401,
- detail={
- "error": {
- "message": "缺少认证令牌",
- "type": "authentication_error",
- "code": "missing_token"
- }
- }
+ detail=_build_error("缺少认证令牌", "missing_token")
)
+ # 验证令牌
if credentials.credentials != api_key:
raise HTTPException(
status_code=401,
- detail={
- "error": {
- "message": f"令牌无效，长度: {len(credentials.credentials)}",
- "type": "authentication_error",
- "code": "invalid_token"
- }
- }
+ detail=_build_error(f"令牌无效，长度: {len(credentials.credentials)}", "invalid_token")
)
logger.debug("[Auth] 令牌认证成功")
return credentials.credentials
+# 全局实例
auth_manager = AuthManager()
\ No newline at end of file
diff --git a/app/core/config.py b/app/core/config.py
index 9ebbbbd..8036266 100644
--- a/app/core/config.py
+++ b/app/core/config.py
@@ -1,164 +1,163 @@
-"""配置管理器"""
+"""配置管理器 - 管理应用配置的读写"""
import toml
from pathlib import Path
-from typing import Dict, Any
+from typing import Dict, Any, Optional, Literal
+
+
+# 默认配置
+DEFAULT_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
+}
+
+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
+}
class ConfigManager:
"""配置管理器"""
def __init__(self) -> None:
- """初始化"""
-
- # 配置文件路径
+ """初始化配置"""
self.config_path: Path = Path(__file__).parents[2] / "data" / "setting.toml"
-
- # 确保配置文件存在(首次启动时创建默认配置)
- self._ensure_config_exists()
-
- # 加载配置
+ 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_config_exists(self) -> None:
- """确保配置文件存在,如果不存在则创建默认配置"""
+ def _ensure_exists(self) -> None:
+ """确保配置存在"""
if not self.config_path.exists():
- # 创建 data 目录
self.config_path.parent.mkdir(parents=True, exist_ok=True)
-
- # 创建默认配置文件
- self._create_default_config()
+ self._create_default()
- def _create_default_config(self) -> None:
- """创建默认配置文件"""
- default_config = {
- "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": False
- },
- "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
- }
- }
-
+ 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_config, 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:
config = toml.load(f)[section]
- # 自动将 SOCKS5 转换为 SOCKS5H
- if section == "grok" and "proxy_url" in config:
- proxy_url = config["proxy_url"]
- if proxy_url and proxy_url.startswith("socks5://"):
- config["proxy_url"] = proxy_url.replace("socks5://", "socks5h://", 1)
+ # 标准化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"])
- # 自动为 CF Clearance 添加前缀
- if section == "grok" and "cf_clearance" in config:
- cf_clearance = config["cf_clearance"]
- if cf_clearance and not cf_clearance.startswith("cf_clearance="):
- config["cf_clearance"] = f"cf_clearance={cf_clearance}"
-
- return config
+ 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:
- # 处理 cf_clearance，移除前缀后保存
- processed_grok_config = grok_config.copy()
- if "cf_clearance" in processed_grok_config:
- cf_clearance = processed_grok_config["cf_clearance"]
- if cf_clearance and cf_clearance.startswith("cf_clearance="):
- processed_grok_config["cf_clearance"] = cf_clearance.replace("cf_clearance=", "", 1)
- config["grok"].update(processed_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:
- # 处理 cf_clearance，移除前缀后保存
- processed_grok_config = grok_config.copy()
- if "cf_clearance" in processed_grok_config:
- cf_clearance = processed_grok_config["cf_clearance"]
- if cf_clearance and cf_clearance.startswith("cf_clearance="):
- processed_grok_config["cf_clearance"] = cf_clearance.replace("cf_clearance=", "", 1)
- config_data["grok"].update(processed_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()
- def get_service_proxy(self) -> str:
- """获取服务代理URL（用于 client 和 upload）"""
- return self.grok_config.get("proxy_url", "")
-
- def get_cache_proxy(self) -> str:
- """获取缓存代理URL（用于 cache）
+ def get_proxy(self, proxy_type: Literal["service", "cache"] = "service") -> str:
+ """获取代理URL
- 逻辑：
- - 如果只设置了 proxy_url，缓存和服务都使用 proxy_url
- - 如果同时设置了 proxy_url 和 cache_proxy_url，缓存使用 cache_proxy_url
+ Args:
+ proxy_type: 代理类型
+ - service: 服务代理（client/upload）
+ - cache: 缓存代理（cache）
"""
- cache_proxy = self.grok_config.get("cache_proxy_url", "")
- service_proxy = self.grok_config.get("proxy_url", "")
+ if proxy_type == "cache":
+ cache_proxy = self.grok_config.get("cache_proxy_url", "")
+ if cache_proxy:
+ return cache_proxy
- # 如果设置了 cache_proxy_url，优先使用
- if cache_proxy:
- return cache_proxy
-
- # 否则使用 proxy_url（服务代理）
- return service_proxy
+ return self.grok_config.get("proxy_url", "")
-# 全局设置
+
+# 全局实例
setting = ConfigManager()
-
-if __name__ == "__main__":
- print(setting.global_config)
- print(setting.grok_config)
\ No newline at end of file
diff --git a/app/core/exception.py b/app/core/exception.py
index a5c179b..0f13a1d 100644
--- a/app/core/exception.py
+++ b/app/core/exception.py
@@ -1,4 +1,4 @@
-"""异常处理器"""
+"""异常处理器 - OpenAI兼容的错误响应"""
from fastapi import Request, status
from fastapi.responses import JSONResponse
@@ -6,22 +6,60 @@ from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
-class GrokApiException(Exception):
- """Grok API 业务异常"""
+# HTTP错误映射
+HTTP_ERROR_MAP = {
+ 400: ("invalid_request_error", "请求格式错误或缺少必填参数"),
+ 401: ("invalid_request_error", "令牌认证失败"),
+ 403: ("permission_error", "没有权限访问此资源"),
+ 404: ("invalid_request_error", "请求的资源不存在"),
+ 429: ("rate_limit_error", "请求频率超出限制，请稍后再试"),
+ 500: ("api_error", "内部服务器错误"),
+ 503: ("api_error", "服务暂时不可用"),
+}
- def __init__(self, message: str, error_code: str = None, details: dict = None):
+# Grok错误码映射
+GROK_STATUS_MAP = {
+ "NO_AUTH_TOKEN": status.HTTP_401_UNAUTHORIZED,
+ "INVALID_TOKEN": status.HTTP_401_UNAUTHORIZED,
+ "HTTP_ERROR": status.HTTP_502_BAD_GATEWAY,
+ "NETWORK_ERROR": status.HTTP_503_SERVICE_UNAVAILABLE,
+ "JSON_ERROR": status.HTTP_502_BAD_GATEWAY,
+ "API_ERROR": status.HTTP_502_BAD_GATEWAY,
+ "STREAM_ERROR": status.HTTP_502_BAD_GATEWAY,
+ "NO_RESPONSE": status.HTTP_502_BAD_GATEWAY,
+ "TOKEN_SAVE_ERROR": status.HTTP_500_INTERNAL_SERVER_ERROR,
+ "NO_AVAILABLE_TOKEN": status.HTTP_503_SERVICE_UNAVAILABLE,
+}
+
+GROK_TYPE_MAP = {
+ "NO_AUTH_TOKEN": "authentication_error",
+ "INVALID_TOKEN": "authentication_error",
+ "HTTP_ERROR": "api_error",
+ "NETWORK_ERROR": "api_error",
+ "JSON_ERROR": "api_error",
+ "API_ERROR": "api_error",
+ "STREAM_ERROR": "api_error",
+ "NO_RESPONSE": "api_error",
+ "TOKEN_SAVE_ERROR": "api_error",
+ "NO_AVAILABLE_TOKEN": "api_error",
+}
+
+
+class GrokApiException(Exception):
+ """Grok API业务异常"""
+
+ def __init__(self, message: str, error_code: str = None, details: dict = None, context: dict = None):
self.message = message
self.error_code = error_code
self.details = details or {}
+ self.context = context or {}
super().__init__(self.message)
def build_error_response(message: str, error_type: str, code: str = None, param: str = None) -> dict:
"""构建OpenAI兼容的错误响应"""
- error = {
- "message": message,
- "type": error_type,
- }
+ error = {"message": message, "type": error_type}
+
if code:
error["code"] = code
if param:
@@ -32,18 +70,8 @@ def build_error_response(message: str, error_type: str, code: str = None, param:
async def http_exception_handler(_: Request, exc: StarletteHTTPException) -> JSONResponse:
"""处理HTTP异常"""
- error_map = {
- 400: ("invalid_request_error", "请求格式错误或缺少必填参数。"),
- 401: ("invalid_request_error", "令牌认证失败。"),
- 403: ("permission_error", "没有权限访问此资源。"),
- 404: ("invalid_request_error", "请求的资源不存在。"),
- 429: ("rate_limit_error", "请求频率超出限制，请稍后再试。"),
- 500: ("api_error", "内部服务器错误。"),
- 503: ("api_error", "服务暂时不可用。"),
- }
-
- error_type, default_message = error_map.get(exc.status_code, ("api_error", str(exc.detail)))
- message = str(exc.detail) if exc.detail else default_message
+ error_type, default_msg = HTTP_ERROR_MAP.get(exc.status_code, ("api_error", str(exc.detail)))
+ message = str(exc.detail) if exc.detail else default_msg
return JSONResponse(
status_code=exc.status_code,
@@ -55,7 +83,7 @@ async def validation_exception_handler(_: Request, exc: RequestValidationError)
"""处理验证错误"""
errors = exc.errors()
param = errors[0]["loc"][-1] if errors and errors[0].get("loc") else None
- message = errors[0]["msg"] if errors and errors[0].get("msg") else "请求参数错误。"
+ message = errors[0]["msg"] if errors and errors[0].get("msg") else "请求参数错误"
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
@@ -64,36 +92,9 @@ async def validation_exception_handler(_: Request, exc: RequestValidationError)
async def grok_api_exception_handler(_: Request, exc: GrokApiException) -> JSONResponse:
- """处理Grok API业务异常"""
- # 根据错误码映射HTTP状态码
- status_code_map = {
- "NO_AUTH_TOKEN": status.HTTP_401_UNAUTHORIZED,
- "INVALID_TOKEN": status.HTTP_401_UNAUTHORIZED,
- "HTTP_ERROR": status.HTTP_502_BAD_GATEWAY,
- "NETWORK_ERROR": status.HTTP_503_SERVICE_UNAVAILABLE,
- "JSON_ERROR": status.HTTP_502_BAD_GATEWAY,
- "API_ERROR": status.HTTP_502_BAD_GATEWAY,
- "STREAM_ERROR": status.HTTP_502_BAD_GATEWAY,
- "NO_RESPONSE": status.HTTP_502_BAD_GATEWAY,
- "TOKEN_SAVE_ERROR": status.HTTP_500_INTERNAL_SERVER_ERROR,
- "NO_AVAILABLE_TOKEN": status.HTTP_503_SERVICE_UNAVAILABLE,
- }
-
- http_status = status_code_map.get(exc.error_code, status.HTTP_500_INTERNAL_SERVER_ERROR)
- error_type_map = {
- "NO_AUTH_TOKEN": "authentication_error",
- "INVALID_TOKEN": "authentication_error",
- "HTTP_ERROR": "api_error",
- "NETWORK_ERROR": "api_error",
- "JSON_ERROR": "api_error",
- "API_ERROR": "api_error",
- "STREAM_ERROR": "api_error",
- "NO_RESPONSE": "api_error",
- "TOKEN_SAVE_ERROR": "api_error",
- "NO_AVAILABLE_TOKEN": "api_error",
- }
-
- error_type = error_type_map.get(exc.error_code, "api_error")
+ """处理Grok API异常"""
+ http_status = GROK_STATUS_MAP.get(exc.error_code, status.HTTP_500_INTERNAL_SERVER_ERROR)
+ error_type = GROK_TYPE_MAP.get(exc.error_code, "api_error")
return JSONResponse(
status_code=http_status,
@@ -105,15 +106,12 @@ async def global_exception_handler(_: Request, exc: Exception) -> JSONResponse:
"""处理未捕获异常"""
return JSONResponse(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
- content=build_error_response(
- "服务器遇到意外错误，请重试。",
- "api_error"
- )
+ content=build_error_response("服务器遇到意外错误，请重试", "api_error")
)
def register_exception_handlers(app) -> None:
- """注册OpenAI兼容的异常处理器"""
+ """注册异常处理器"""
app.add_exception_handler(StarletteHTTPException, http_exception_handler)
app.add_exception_handler(RequestValidationError, validation_exception_handler)
app.add_exception_handler(GrokApiException, grok_api_exception_handler)
diff --git a/app/core/logger.py b/app/core/logger.py
index cdbbb6d..8129ff8 100644
--- a/app/core/logger.py
+++ b/app/core/logger.py
@@ -1,27 +1,32 @@
-"""全局日志模块"""
+"""全局日志模块 - 单例模式的日志管理器"""
import sys
import logging
from pathlib import Path
from logging.handlers import RotatingFileHandler
+
from app.core.config import setting
+# 过滤模式
+FILTER_PATTERNS = [
+ "chunk: b'", # SSE原始字节
+ "Got event:", # SSE事件
+ "Closing", # SSE关闭
+]
+
+
class MCPLogFilter(logging.Filter):
- """MCP 日志过滤器 - 过滤掉包含大量数据的 DEBUG 日志"""
+ """MCP日志过滤器 - 过滤大量数据的DEBUG日志"""
- def filter(self, record):
- # 过滤掉包含原始字节数据的 SSE 日志
- if record.name == "sse_starlette.sse" and "chunk: b'" in record.getMessage():
- return False
-
- # 过滤掉 SSE 的一些冗余日志
+ def filter(self, record: logging.LogRecord) -> bool:
+ """过滤日志"""
+ # 过滤SSE的DEBUG日志
if record.name == "sse_starlette.sse" and record.levelno == logging.DEBUG:
msg = record.getMessage()
- if any(x in msg for x in ["Got event:", "Closing", "chunk:"]):
- return False
+ return not any(p in msg for p in FILTER_PATTERNS)
- # 过滤掉 MCP streamable_http 的一些 DEBUG 日志
+ # 过滤MCP streamable_http的DEBUG日志
if "mcp.server.streamable_http" in record.name and record.levelno == logging.DEBUG:
return False
@@ -29,65 +34,76 @@ class MCPLogFilter(logging.Filter):
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_handler.addFilter(mcp_filter)
+ 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)
-
- # FastMCP 相关日志 - 关闭
- logging.getLogger("mcp").setLevel(logging.CRITICAL)
- logging.getLogger("fastmcp").setLevel(logging.CRITICAL)
+ # 配置第三方库
+ 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:
"""调试日志"""
@@ -110,5 +126,5 @@ class LoggerManager:
self.logger.critical(msg)
-# 全局日志器实例
+# 全局实例
logger = LoggerManager()
diff --git a/app/core/storage.py b/app/core/storage.py
index 942ee3e..db96fc3 100644
--- a/app/core/storage.py
+++ b/app/core/storage.py
@@ -1,7 +1,7 @@
"""存储抽象层 - 支持文件、MySQL和Redis存储"""
import os
-import json
+import orjson
import toml
import asyncio
import warnings
@@ -47,7 +47,7 @@ class BaseStorage(ABC):
class FileStorage(BaseStorage):
- """文件存储实现"""
+ """文件存储"""
def __init__(self, data_dir: Path):
self.data_dir = data_dir
@@ -60,89 +60,87 @@ class FileStorage(BaseStorage):
"""初始化文件存储"""
self.data_dir.mkdir(parents=True, exist_ok=True)
- # 初始化token文件
if not self.token_file.exists():
- await self._write_file(self.token_file, json.dumps({"sso": {}, "ssoSuper": {}}, indent=2, ensure_ascii=False))
- logger.info("[Storage] 创建新的token文件")
+ await self._write(self.token_file, orjson.dumps({"sso": {}, "ssoSuper": {}}, option=orjson.OPT_INDENT_2).decode())
+ logger.info("[Storage] 创建token文件")
- # 初始化配置文件
if not self.config_file.exists():
- default_config = {
+ default = {
"global": {"api_keys": [], "admin_username": "admin", "admin_password": "admin"},
"grok": {"proxy_url": "", "cf_clearance": "", "x_statsig_id": ""}
}
- await self._write_file(self.config_file, toml.dumps(default_config))
- logger.info("[Storage] 创建新的配置文件")
+ await self._write(self.config_file, toml.dumps(default))
+ logger.info("[Storage] 创建配置文件")
- async def _read_file(self, file_path: Path) -> str:
- """读取文件内容"""
- async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
+ async def _read(self, path: Path) -> str:
+ """读取文件"""
+ async with aiofiles.open(path, "r", encoding="utf-8") as f:
return await f.read()
- async def _write_file(self, file_path: Path, content: str) -> None:
- """写入文件内容"""
- async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
+ async def _write(self, path: Path, content: str) -> None:
+ """写入文件"""
+ async with aiofiles.open(path, "w", encoding="utf-8") as f:
await f.write(content)
- async def _load_json(self, file_path: Path, default: Dict[str, Any], lock: asyncio.Lock) -> Dict[str, Any]:
- """加载JSON文件"""
+ async def _load_json(self, path: Path, default: Dict, lock: asyncio.Lock) -> Dict[str, Any]:
+ """加载JSON"""
try:
async with lock:
- if not file_path.exists():
+ if not path.exists():
return default
- return json.loads(await self._read_file(file_path))
+ return orjson.loads(await self._read(path))
except Exception as e:
- logger.error(f"[Storage] 加载{file_path.name}失败: {e}")
+ logger.error(f"[Storage] 加载{path.name}失败: {e}")
return default
- async def _save_json(self, file_path: Path, data: Dict[str, Any], lock: asyncio.Lock) -> None:
- """保存JSON文件"""
+ async def _save_json(self, path: Path, data: Dict, lock: asyncio.Lock) -> None:
+ """保存JSON"""
try:
async with lock:
- await self._write_file(file_path, json.dumps(data, indent=2, ensure_ascii=False))
+ await self._write(path, orjson.dumps(data, option=orjson.OPT_INDENT_2).decode())
except Exception as e:
- logger.error(f"[Storage] 保存{file_path.name}失败: {e}")
+ logger.error(f"[Storage] 保存{path.name}失败: {e}")
raise
- async def _load_toml(self, file_path: Path, default: Dict[str, Any], lock: asyncio.Lock) -> Dict[str, Any]:
- """加载TOML文件"""
+ async def _load_toml(self, path: Path, default: Dict, lock: asyncio.Lock) -> Dict[str, Any]:
+ """加载TOML"""
try:
async with lock:
- if not file_path.exists():
+ if not path.exists():
return default
- return toml.loads(await self._read_file(file_path))
+ return toml.loads(await self._read(path))
except Exception as e:
- logger.error(f"[Storage] 加载{file_path.name}失败: {e}")
+ logger.error(f"[Storage] 加载{path.name}失败: {e}")
return default
- async def _save_toml(self, file_path: Path, data: Dict[str, Any], lock: asyncio.Lock) -> None:
- """保存TOML文件"""
+ async def _save_toml(self, path: Path, data: Dict, lock: asyncio.Lock) -> None:
+ """保存TOML"""
try:
async with lock:
- await self._write_file(file_path, toml.dumps(data))
+ await self._write(path, toml.dumps(data))
except Exception as e:
- logger.error(f"[Storage] 保存{file_path.name}失败: {e}")
+ logger.error(f"[Storage] 保存{path.name}失败: {e}")
raise
async def load_tokens(self) -> Dict[str, Any]:
- """加载token数据"""
+ """加载token"""
return await self._load_json(self.token_file, {"sso": {}, "ssoSuper": {}}, self._token_lock)
async def save_tokens(self, data: Dict[str, Any]) -> None:
- """保存token数据"""
+ """保存token"""
await self._save_json(self.token_file, data, self._token_lock)
async def load_config(self) -> Dict[str, Any]:
- """加载配置数据"""
+ """加载配置"""
return await self._load_toml(self.config_file, {"global": {}, "grok": {}}, self._config_lock)
async def save_config(self, data: Dict[str, Any]) -> None:
- """保存配置数据"""
+ """保存配置"""
await self._save_toml(self.config_file, data, self._config_lock)
class MysqlStorage(BaseStorage):
- """MySQL存储实现"""
+ """MySQL存储"""
def __init__(self, database_url: str, data_dir: Path):
self.database_url = database_url
@@ -155,22 +153,15 @@ class MysqlStorage(BaseStorage):
try:
import aiomysql
parsed = self._parse_url(self.database_url)
- logger.info(f"[Storage] 解析数据库连接: {parsed['user']}@{parsed['host']}:{parsed['port']}/{parsed['db']}")
+ logger.info(f"[Storage] MySQL: {parsed['user']}@{parsed['host']}:{parsed['port']}/{parsed['db']}")
- # 创建数据库
await self._create_db(parsed)
-
- # 创建连接池
self._pool = await aiomysql.create_pool(
host=parsed['host'], port=parsed['port'], user=parsed['user'],
password=parsed['password'], db=parsed['db'], charset="utf8mb4",
autocommit=True, maxsize=10
)
-
- # 创建表
await self._create_tables()
-
- # 初始化文件存储和同步数据
await self._file.init_db()
await self._sync_data()
@@ -181,26 +172,26 @@ class MysqlStorage(BaseStorage):
raise
def _parse_url(self, url: str) -> Dict[str, Any]:
- """解析数据库URL"""
- parsed = urlparse(url)
+ """解析URL"""
+ p = urlparse(url)
return {
- 'user': unquote(parsed.username) if parsed.username else "",
- 'password': unquote(parsed.password) if parsed.password else "",
- 'host': parsed.hostname,
- 'port': parsed.port or 3306,
- 'db': parsed.path[1:] if parsed.path else "grok2api"
+ 'user': unquote(p.username) if p.username else "",
+ 'password': unquote(p.password) if p.password else "",
+ 'host': p.hostname,
+ 'port': p.port or 3306,
+ 'db': p.path[1:] if p.path else "grok2api"
}
- async def _create_db(self, parsed: Dict[str, Any]) -> None:
+ async def _create_db(self, parsed: Dict) -> None:
"""创建数据库"""
import aiomysql
- temp_pool = await aiomysql.create_pool(
+ pool = await aiomysql.create_pool(
host=parsed['host'], port=parsed['port'], user=parsed['user'],
password=parsed['password'], charset="utf8mb4", autocommit=True, maxsize=1
)
try:
- async with temp_pool.acquire() as conn:
+ async with pool.acquire() as conn:
async with conn.cursor() as cursor:
with warnings.catch_warnings():
warnings.filterwarnings('ignore', message='.*database exists')
@@ -208,10 +199,10 @@ class MysqlStorage(BaseStorage):
f"CREATE DATABASE IF NOT EXISTS `{parsed['db']}` "
f"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
)
- logger.info(f"[Storage] MySQL数据库 '{parsed['db']}' 已就绪")
+ logger.info(f"[Storage] 数据库 '{parsed['db']}' 就绪")
finally:
- temp_pool.close()
- await temp_pool.wait_closed()
+ pool.close()
+ await pool.wait_closed()
async def _create_tables(self) -> None:
"""创建表"""
@@ -240,7 +231,7 @@ class MysqlStorage(BaseStorage):
warnings.filterwarnings('ignore', message='.*already exists')
for sql in tables.values():
await cursor.execute(sql)
- logger.info("[Storage] MySQL表创建/验证成功")
+ logger.info("[Storage] MySQL表就绪")
async def _sync_data(self) -> None:
"""同步数据"""
@@ -252,39 +243,33 @@ class MysqlStorage(BaseStorage):
await self._file.save_tokens(data)
else:
await self._file.save_config(data)
- logger.info(f"[Storage] {table.split('_')[1]}数据已从数据库同步到文件")
+ logger.info(f"[Storage] {table.split('_')[1]}数据已从DB同步")
else:
- if table == "grok_tokens":
- file_data = await self._file.load_tokens()
- if file_data.get(key) or file_data.get("ssoSuper"):
- await self._save_db(table, file_data)
- logger.info("[Storage] Token数据已从文件初始化到数据库")
- else:
- file_data = await self._file.load_config()
- if file_data.get(key) or file_data.get("grok"):
- await self._save_db(table, file_data)
- logger.info("[Storage] 配置数据已从文件初始化到数据库")
+ file_data = await (self._file.load_tokens() if table == "grok_tokens" else self._file.load_config())
+ if file_data.get(key) or (table == "grok_tokens" and file_data.get("ssoSuper")):
+ await self._save_db(table, file_data)
+ logger.info(f"[Storage] {table.split('_')[1]}数据已初始化到DB")
except Exception as e:
- logger.warning(f"[Storage] 数据同步失败: {e}")
+ logger.warning(f"[Storage] 同步失败: {e}")
- async def _load_db(self, table: str) -> Optional[Dict[str, Any]]:
- """从数据库加载数据"""
+ async def _load_db(self, table: str) -> Optional[Dict]:
+ """从DB加载"""
try:
async with self._pool.acquire() as conn:
async with conn.cursor() as cursor:
await cursor.execute(f"SELECT data FROM {table} ORDER BY id DESC LIMIT 1")
result = await cursor.fetchone()
- return json.loads(result[0]) if result else None
+ return orjson.loads(result[0]) if result else None
except Exception as e:
- logger.error(f"[Storage] 从数据库加载{table}失败: {e}")
+ logger.error(f"[Storage] 加载{table}失败: {e}")
return None
- async def _save_db(self, table: str, data: Dict[str, Any]) -> None:
- """保存数据到数据库"""
+ async def _save_db(self, table: str, data: Dict) -> None:
+ """保存到DB"""
try:
async with self._pool.acquire() as conn:
async with conn.cursor() as cursor:
- json_data = json.dumps(data, ensure_ascii=False)
+ json_data = orjson.dumps(data).decode()
await cursor.execute(f"SELECT id FROM {table} ORDER BY id DESC LIMIT 1")
result = await cursor.fetchone()
@@ -293,24 +278,24 @@ class MysqlStorage(BaseStorage):
else:
await cursor.execute(f"INSERT INTO {table} (data) VALUES (%s)", (json_data,))
except Exception as e:
- logger.error(f"[Storage] 保存数据到{table}失败: {e}")
+ logger.error(f"[Storage] 保存{table}失败: {e}")
raise
async def load_tokens(self) -> Dict[str, Any]:
- """加载token数据"""
+ """加载token"""
return await self._file.load_tokens()
async def save_tokens(self, data: Dict[str, Any]) -> None:
- """保存token数据"""
+ """保存token"""
await self._file.save_tokens(data)
await self._save_db("grok_tokens", data)
async def load_config(self) -> Dict[str, Any]:
- """加载配置数据"""
+ """加载配置"""
return await self._file.load_config()
async def save_config(self, data: Dict[str, Any]) -> None:
- """保存配置数据"""
+ """保存配置"""
await self._file.save_config(data)
await self._save_db("grok_settings", data)
@@ -319,11 +304,11 @@ class MysqlStorage(BaseStorage):
if self._pool:
self._pool.close()
await self._pool.wait_closed()
- logger.info("[Storage] MySQL连接池已关闭")
+ logger.info("[Storage] MySQL已关闭")
class RedisStorage(BaseStorage):
- """Redis存储实现"""
+ """Redis存储"""
def __init__(self, redis_url: str, data_dir: Path):
self.redis_url = redis_url
@@ -336,20 +321,17 @@ class RedisStorage(BaseStorage):
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}")
+ logger.info(f"[Storage] Redis: {parsed['host']}:{parsed['port']}/{parsed['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
)
- # 测试连接
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()
@@ -363,17 +345,16 @@ class RedisStorage(BaseStorage):
"""解析Redis URL"""
if url.startswith('redis://'):
url = url[8:]
- parsed = urlparse(f'//{url}')
+ p = 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
+ 'host': p.hostname or 'localhost',
+ 'port': p.port or 6379,
+ 'db': int(p.path.lstrip('/')) if p.path and p.path != '/' else 0,
+ 'username': unquote(p.username) if p.username else None,
+ 'password': unquote(p.password) if p.password else None
}
- # Redis 6+ 默认用户名
if result['password'] and not result['username']:
result['username'] = 'default'
@@ -388,44 +369,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 +413,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 +437,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 +461,5 @@ class StorageManager:
await self._storage.close()
-# 全局存储管理器实例
+# 全局实例
storage_manager = StorageManager()
diff --git a/app/models/grok_models.py b/app/models/grok_models.py
index 707ad09..6fab827 100644
--- a/app/models/grok_models.py
+++ b/app/models/grok_models.py
@@ -1,122 +1,127 @@
-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-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": "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": "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
}
+}
+
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"
@@ -127,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
@@ -137,15 +142,22 @@ class Models(Enum):
return model in _MODEL_CONFIG
@classmethod
- def to_grok(cls, model: str) -> tuple[str, str]:
- """转换为Grok内部模型名和模式类型"""
+ def to_grok(cls, model: str) -> Tuple[str, str]:
+ """转换为Grok内部模型名和模式
+
+ Returns:
+ (模型名, 模式类型) 元组
+ """
config = _MODEL_CONFIG.get(model)
- if config:
- return config["grok_model"]
- return model, "MODEL_MODE_FAST"
+ return config["grok_model"] if config else (model, "MODEL_MODE_FAST")
@classmethod
def to_rate_limit(cls, model: str) -> str:
- """转换为速率限制接口模型名"""
+ """转换为速率限制模型名"""
config = _MODEL_CONFIG.get(model)
- return config["rate_limit_model"] if config else model
\ No newline at end of file
+ return config["rate_limit_model"] if config else model
+
+ @classmethod
+ def get_all_model_names(cls) -> list[str]:
+ """获取所有模型名称"""
+ return list(_MODEL_CONFIG.keys())
\ No newline at end of file
diff --git a/app/models/openai_schema.py b/app/models/openai_schema.py
index 37b6eeb..540eed9 100644
--- a/app/models/openai_schema.py
+++ b/app/models/openai_schema.py
@@ -1,16 +1,18 @@
-"""OpenAI 请求-响应模型"""
+"""OpenAI 请求-响应模型定义"""
from fastapi import HTTPException
from typing import Optional, List, Union, Dict, Any
from pydantic import BaseModel, Field, field_validator
+from app.models.grok_models import Models
+
class OpenAIChatRequest(BaseModel):
- """OpenAI 聊天请求模型"""
+ """OpenAI聊天请求"""
model: str = Field(..., description="模型名称", min_length=1)
messages: List[Dict[str, Any]] = Field(..., description="消息列表", min_length=1)
- stream: bool = Field(False, description="启用流式响应")
+ stream: bool = Field(False, description="流式响应")
temperature: Optional[float] = Field(0.7, ge=0, le=2, description="采样温度")
max_tokens: Optional[int] = Field(None, ge=1, le=100000, description="最大Token数")
top_p: Optional[float] = Field(1.0, ge=0, le=1, description="采样参数")
@@ -20,31 +22,19 @@ class OpenAIChatRequest(BaseModel):
def validate_messages(cls, v):
"""验证消息格式"""
if not v:
- raise HTTPException(
- status_code=400,
- detail="消息列表不能为空"
- )
+ raise HTTPException(status_code=400, detail="消息列表不能为空")
for msg in v:
if not isinstance(msg, dict):
- raise HTTPException(
- status_code=400,
- detail="每个消息必须是一个字典"
- )
+ raise HTTPException(status_code=400, detail="每个消息必须是字典")
if 'role' not in msg:
- raise HTTPException(
- status_code=400,
- detail="消息缺少必填字段 'role'"
- )
+ raise HTTPException(status_code=400, detail="消息缺少 'role' 字段")
if 'content' not in msg:
- raise HTTPException(
- status_code=400,
- detail="消息缺少必填字段 'content'"
- )
+ raise HTTPException(status_code=400, detail="消息缺少 'content' 字段")
if msg['role'] not in ['system', 'user', 'assistant']:
raise HTTPException(
status_code=400,
- detail=f"无效的角色 '{msg['role']}', 必须是 'system', 'user' 或 'assistant'"
+ detail=f"无效角色 '{msg['role']}', 必须是 system/user/assistant"
)
return v
@@ -53,29 +43,27 @@ class OpenAIChatRequest(BaseModel):
@field_validator('model')
def validate_model(cls, v):
"""验证模型名称"""
- allowed_models = [
- 'grok-3-fast', 'grok-4-fast', 'grok-4-fast-expert',
- 'grok-4-expert', 'grok-4-heavy', 'grok-imagine-0.9'
- ]
- if v not in allowed_models:
+ if not Models.is_valid_model(v):
+ supported = Models.get_all_model_names()
raise HTTPException(
status_code=400,
- detail=f"不支持的模型 '{v}', 支持的模型: {', '.join(allowed_models)}"
+ detail=f"不支持的模型 '{v}', 支持: {', '.join(supported)}"
)
return v
+
class OpenAIChatCompletionMessage(BaseModel):
"""聊天完成消息"""
role: str = Field(..., description="角色")
- content: str = Field(..., description="消息内容")
+ content: str = Field(..., description="内容")
reference_id: Optional[str] = Field(default=None, description="参考ID")
annotations: Optional[List[str]] = Field(default=None, description="注释")
class OpenAIChatCompletionChoice(BaseModel):
"""聊天完成选项"""
- index: int = Field(..., description="选项索引")
- message: OpenAIChatCompletionMessage = Field(..., description="响应消息")
+ index: int = Field(..., description="索引")
+ message: OpenAIChatCompletionMessage = Field(..., description="消息")
logprobs: Optional[float] = Field(default=None, description="对数概率")
finish_reason: str = Field(default="stop", description="完成原因")
@@ -85,20 +73,20 @@ class OpenAIChatCompletionResponse(BaseModel):
id: str = Field(..., description="响应ID")
object: str = Field("chat.completion", description="对象类型")
created: int = Field(..., description="创建时间戳")
- model: str = Field(..., description="使用的模型")
- choices: List[OpenAIChatCompletionChoice] = Field(..., description="响应选项")
+ model: str = Field(..., description="模型")
+ choices: List[OpenAIChatCompletionChoice] = Field(..., description="选项")
usage: Optional[Dict[str, Any]] = Field(None, description="令牌使用")
class OpenAIChatCompletionChunkMessage(BaseModel):
- """流式响应消息片段"""
+ """流式消息片段"""
role: str = Field(..., description="角色")
- content: str = Field(..., description="消息内容")
+ content: str = Field(..., description="内容")
class OpenAIChatCompletionChunkChoice(BaseModel):
- """流式响应选项"""
- index: int = Field(..., description="选项索引")
+ """流式选项"""
+ index: int = Field(..., description="索引")
delta: Optional[Union[Dict[str, Any], OpenAIChatCompletionChunkMessage]] = Field(
None, description="Delta数据"
)
@@ -106,10 +94,10 @@ class OpenAIChatCompletionChunkChoice(BaseModel):
class OpenAIChatCompletionChunkResponse(BaseModel):
- """流式聊天完成响应"""
+ """流式聊天响应"""
id: str = Field(..., description="响应ID")
object: str = Field(default="chat.completion.chunk", description="对象类型")
created: int = Field(..., description="创建时间戳")
- model: str = Field(..., description="使用的模型")
+ model: str = Field(..., description="模型")
system_fingerprint: Optional[str] = Field(default=None, description="系统指纹")
- choices: List[OpenAIChatCompletionChunkChoice] = Field(..., description="响应选项")
\ No newline at end of file
+ choices: List[OpenAIChatCompletionChunkChoice] = Field(..., description="选项")
\ No newline at end of file
diff --git a/app/services/grok/cache.py b/app/services/grok/cache.py
index 21d8631..3380d87 100644
--- a/app/services/grok/cache.py
+++ b/app/services/grok/cache.py
@@ -1,9 +1,9 @@
-"""缓存服务模块"""
+"""缓存服务模块 - 提供图片和视频的下载、缓存和清理功能"""
import asyncio
import base64
from pathlib import Path
-from typing import Optional
+from typing import Optional, Tuple
from curl_cffi.requests import AsyncSession
from app.core.config import setting
@@ -11,170 +11,188 @@ 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] 文件缓存成功")
+ 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
- }
-
- # 使用缓存代理
- proxy_url = setting.get_cache_proxy()
- proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else {}
+ proxy = setting.get_proxy("cache")
+ proxies = {"http": proxy, "https": proxy} if proxy else {}
- if proxy_url:
- logger.debug(f"[{self.cache_type.upper()}Cache] 使用缓存代理: {proxy_url.split('@')[-1] if '@' in proxy_url else proxy_url}")
+ if proxy:
+ self._log("debug", f"使用代理: {proxy.split('@')[-1] if '@' in proxy else proxy}")
async with AsyncSession() as session:
- logger.debug(f"[{self.cache_type.upper()}Cache] 缓存生成视频文件: https://assets.grok.com{file_path}")
+ url = f"{ASSETS_URL}{file_path}"
+ self._log("debug", f"下载: {url}")
+
response = await session.get(
- f"https://assets.grok.com{file_path}",
- headers=headers,
+ url,
+ headers=self._build_headers(file_path, auth_token),
proxies=proxies,
- timeout=timeout,
+ timeout=timeout or self.timeout,
allow_redirects=True,
impersonate="chrome133a"
)
response.raise_for_status()
+
cache_path.write_bytes(response.content)
- logger.debug(f"[{self.cache_type.upper()}Cache] 文件缓存成功")
- asyncio.create_task(self.cleanup_cache())
+ self._log("debug", "缓存成功")
+
+ # 异步清理（带错误处理）
+ asyncio.create_task(self._safe_cleanup())
return cache_path
+
except Exception as e:
- logger.error(f"[{self.cache_type.upper()}Cache] 下载文件失败: {e}")
+ self._log("error", f"下载失败: {e}")
return None
def get_cached(self, file_path: str) -> Optional[Path]:
- """获取缓存的文件路径"""
- cache_path = self._cache_path(file_path)
- return cache_path if cache_path.exists() else None
+ """获取已缓存的文件"""
+ path = self._get_path(file_path)
+ return path if path.exists() else None
- async def cleanup_cache(self):
- """清理缓存目录，确保不超过配置的大小限制"""
+ async def _safe_cleanup(self):
+ """安全清理（捕获异常）"""
try:
- max_size_mb = setting.global_config.get(f"{self.cache_type}_cache_max_size_mb", 500)
- max_size_bytes = max_size_mb * 1024 * 1024
-
- files = [(fp, (stat := fp.stat()).st_size, stat.st_mtime)
- for fp in self.cache_dir.glob("*") if fp.is_file()]
- total_size = sum(size for _, size, _ in files)
-
- if total_size <= max_size_bytes:
- logger.debug(f"[{self.cache_type.upper()}Cache] 缓存大小统计: {total_size / 1024 / 1024:.2f}MB，未达到限制")
- return
-
- logger.info(f"[{self.cache_type.upper()}Cache] 清理缓存，当前大小 {total_size / 1024 / 1024:.2f}MB 超过限制 {max_size_mb}MB")
- files.sort(key=lambda x: x[2])
-
- for file_path, size, _ in files:
- if total_size <= max_size_bytes:
- break
- file_path.unlink()
- total_size -= size
- logger.debug(f"[{self.cache_type.upper()}Cache] 清理缓存文件: {file_path}")
-
- logger.info(f"[{self.cache_type.upper()}Cache] 清理缓存完成，当前大小 {total_size / 1024 / 1024:.2f}MB")
+ await self.cleanup()
except Exception as e:
- logger.error(f"[{self.cache_type.upper()}Cache] 清理缓存失败: {e}")
+ self._log("error", f"后台清理失败: {e}")
+
+ async def cleanup(self):
+ """清理超限缓存"""
+ if self._cleanup_lock.locked():
+ return
+
+ async with self._cleanup_lock:
+ try:
+ max_mb = setting.global_config.get(f"{self.cache_type}_cache_max_size_mb", 500)
+ max_bytes = max_mb * 1024 * 1024
+
+ # 获取文件信息 (path, size, mtime)
+ files = [(f, (s := f.stat()).st_size, s.st_mtime)
+ for f in self.cache_dir.glob("*") if f.is_file()]
+ total = sum(size for _, size, _ in files)
+
+ if total <= max_bytes:
+ return
+
+ self._log("info", f"清理缓存 {total/1024/1024:.1f}MB -> {max_mb}MB")
+
+ # 删除最旧的文件
+ for path, size, _ in sorted(files, key=lambda x: x[2]):
+ if total <= max_bytes:
+ break
+ path.unlink()
+ total -= size
+
+ self._log("info", f"清理完成: {total/1024/1024:.1f}MB")
+ except Exception as e:
+ self._log("error", f"清理失败: {e}")
-class ImageCacheService(CacheService):
+class ImageCache(CacheService):
"""图片缓存服务"""
def __init__(self):
- super().__init__("image")
+ super().__init__("image", timeout=30.0)
- async def download_image(self, image_path: str, auth_token: str) -> Optional[Path]:
- """下载并缓存图片"""
- return await self.download_file(image_path, auth_token, timeout=30.0)
-
- def get_cached(self, image_path: str) -> Optional[Path]:
- """获取缓存的图片路径"""
- return super().get_cached(image_path)
+ async def download_image(self, path: str, token: str) -> Optional[Path]:
+ """下载图片"""
+ return await self.download(path, token)
@staticmethod
def to_base64(image_path: Path) -> Optional[str]:
- """将图片转换为 base64 编码"""
+ """图片转base64"""
try:
if not image_path.exists():
- logger.error(f"[ImageCache] 图片文件不存在: {image_path}")
+ logger.error(f"[ImageCache] 文件不存在: {image_path}")
return None
- with open(image_path, "rb") as f:
- base64_data = base64.b64encode(f.read()).decode('utf-8')
-
- mime_type = {'.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png',
- '.gif': 'image/gif', '.webp': 'image/webp'}.get(image_path.suffix.lower(), 'image/jpeg')
-
- return f"data:{mime_type};base64,{base64_data}"
+ data = base64.b64encode(image_path.read_bytes()).decode()
+ mime = MIME_TYPES.get(image_path.suffix.lower(), DEFAULT_MIME)
+ return f"data:{mime};base64,{data}"
except Exception as e:
- logger.error(f"[ImageCache] 图片转 base64 失败: {e}")
+ logger.error(f"[ImageCache] 转换失败: {e}")
return None
- async def download_base64(self, image_path: str, auth_token: str) -> Optional[str]:
- """下载图片并转换为 base64 编码（转换后立即删除缓存文件）"""
+ async def download_base64(self, path: str, token: str) -> Optional[str]:
+ """下载并转为base64（自动删除临时文件）"""
try:
- cache_path = await self.download_file(image_path, auth_token, timeout=30.0)
+ cache_path = await self.download(path, token)
if not cache_path:
return None
- base64_str = self.to_base64(cache_path)
-
+ result = self.to_base64(cache_path)
+
+ # 清理临时文件
try:
cache_path.unlink()
- logger.debug(f"[ImageCache] 已删除临时文件: {cache_path}")
except Exception as e:
logger.warning(f"[ImageCache] 删除临时文件失败: {e}")
- return base64_str
+ return result
except Exception as e:
- logger.error(f"[ImageCache] 下载并转换 base64 失败: {e}")
+ logger.error(f"[ImageCache] 下载base64失败: {e}")
return None
-class VideoCacheService(CacheService):
+class VideoCache(CacheService):
"""视频缓存服务"""
def __init__(self):
- super().__init__("video")
+ super().__init__("video", timeout=60.0)
- async def download_video(self, video_path: str, auth_token: str) -> Optional[Path]:
- """下载并缓存视频"""
- return await self.download_file(video_path, auth_token, timeout=60.0)
-
- def get_cached(self, video_path: str) -> Optional[Path]:
- """获取缓存的视频路径"""
- return super().get_cached(video_path)
+ async def download_video(self, path: str, token: str) -> Optional[Path]:
+ """下载视频"""
+ return await self.download(path, token)
# 全局实例
-image_cache_service = ImageCacheService()
-video_cache_service = VideoCacheService()
+image_cache_service = ImageCache()
+video_cache_service = VideoCache()
diff --git a/app/services/grok/client.py b/app/services/grok/client.py
index 34727ab..5b8968e 100644
--- a/app/services/grok/client.py
+++ b/app/services/grok/client.py
@@ -1,9 +1,8 @@
-"""Grok API 客户端模块"""
+"""Grok API 客户端 - 处理OpenAI到Grok的请求转换和响应处理"""
import asyncio
-import json
-from typing import Dict, List, Tuple, Any
-
+import orjson
+from typing import Dict, List, Tuple, Any, Optional
from curl_cffi import requests as curl_requests
from app.core.config import setting
@@ -16,136 +15,148 @@ 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"
-MAX_RETRY = 3 # 最大重试次数
+
+# 常量
+API_ENDPOINT = "https://grok.com/rest/app-chat/conversations/new"
+TIMEOUT = 120
+BROWSER = "chrome133a"
+MAX_RETRY = 3
+MAX_UPLOADS = 5
class GrokClient:
"""Grok API 客户端"""
+
+ _upload_sem = asyncio.Semaphore(MAX_UPLOADS)
@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)
-
- # 提取消息内容和图片URL
- content, image_urls = GrokClient._extract_content(messages)
- model_name, model_mode = Models.to_grok(model)
- is_video_model = Models.get_model_info(model).get("is_video_model", False)
+ 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:
- if len(image_urls) > 1:
- logger.warning(f"[Client] 视频模型只允许一张图片，当前有{len(image_urls)}张，只使用第一张")
- image_urls = image_urls[:1]
-
- # 重试逻辑
- return await GrokClient._try(model, content, image_urls, model_name, model_mode, is_video_model, stream)
+ # 获取模型信息
+ info = Models.get_model_info(model)
+ grok_model, mode = Models.to_grok(model)
+ is_video = info.get("is_video_model", False)
+
+ # 视频模型限制
+ 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)
@staticmethod
- async def _try(model: str, content: str, image_urls: List[str], model_name: str, model_mode: str, is_video: bool, stream: bool):
- """带重试的请求执行"""
+ async def _retry(model: str, content: str, images: List[str], grok_model: str, mode: str, is_video: bool, stream: bool):
+ """重试请求"""
last_err = None
for i in range(MAX_RETRY):
try:
- # 获取token
- auth_token = token_manager.get_token(model)
+ token = token_manager.get_token(model)
+ img_ids, img_uris = await GrokClient._upload(images, token)
- # 上传图片
- imgs, uris = await GrokClient._upload_imgs(image_urls, auth_token)
-
- # 视频模型 - 创建会话
+ # 视频模型创建会话
post_id = None
- if is_video and imgs and uris:
- try:
- create_result = await PostCreateManager.create(imgs[0], uris[0], auth_token)
- if create_result and create_result.get("success"):
- post_id = create_result.get("post_id")
- else:
- logger.warning(f"[Client] 创建会话失败，继续使用原有流程")
- except Exception as e:
- logger.warning(f"[Client] 创建会话异常: {e}，继续使用原有流程")
+ 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, model_name, model_mode, imgs, uris, is_video, post_id)
- return await GrokClient._send_request(payload, auth_token, model, stream, post_id)
+ 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
- # 401/429 可重试，其他错误直接抛出
+ # 仅401/429可重试
if e.error_code not in ["HTTP_ERROR", "NO_AVAILABLE_TOKEN"]:
raise
- # 检查是否为可重试的状态码
status = e.context.get("status") if e.context else None
if status not in [401, 429]:
raise
if i < MAX_RETRY - 1:
- logger.warning(f"[Client] 请求失败(状态码:{status}), 重试 {i+1}/{MAX_RETRY}")
- await asyncio.sleep(0.5) # 短暂延迟
- else:
- logger.error(f"[Client] 重试{MAX_RETRY}次后仍失败")
+ logger.warning(f"[Client] 失败(状态:{status}), 重试 {i+1}/{MAX_RETRY}")
+ await asyncio.sleep(0.5)
- raise last_err if last_err else GrokApiException("请求失败", "REQUEST_ERROR")
+ 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) -> Tuple[List[str], List[str]]:
- """上传图片并返回附件ID列表"""
- image_attachments = []
- image_uris = []
- # 并发上传所有图片
- tasks = [ImageUploadManager.upload(url, auth_token) for url in image_urls]
- results = await asyncio.gather(*tasks, return_exceptions=True)
-
- for url, (file_id, file_uri) in zip(image_urls, results):
- if isinstance(file_id, Exception):
- logger.warning(f"[Client] 图片上传失败: {url}, 错误: {file_id}")
- elif file_id:
- image_attachments.append(file_id)
- image_uris.append(file_uri)
-
- return image_attachments, image_uris
+ 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._upload_sem:
+ 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 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], image_uris: List[str], is_video_model: bool = False, post_id: str = None) -> 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,
@@ -160,144 +171,96 @@ 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 and image_uris:
- image_url = image_uris[0]
-
- # 构建 URL 消息
- if post_id:
- image_message = f"https://grok.com/imagine/{post_id} {content} --mode=custom"
- else:
- image_message = f"https://assets.grok.com/post/{image_url} {content} --mode=custom"
-
- payload = {
- "temporary": True,
- "modelName": "grok-3",
- "message": image_message,
- "fileAttachments": image_attachments,
- "toolOverrides": {"videoGen": True}
- }
-
- return payload
@staticmethod
- async def _send_request(payload: dict, auth_token: str, model: str, stream: bool, post_id: str = None):
- """发送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)
+ # 构建请求
+ headers = GrokClient._build_headers(token)
if model == "grok-imagine-0.9":
- # 传入会话ID
- file_attachments = payload.get("fileAttachments", [])
- referer_id = post_id if post_id else (file_attachments[0] if file_attachments else "")
- if referer_id:
- headers["Referer"] = f"https://grok.com/imagine/{referer_id}"
+ ref_id = post_id or payload.get("fileAttachments", [""])[0]
+ if ref_id:
+ headers["Referer"] = f"https://grok.com/imagine/{ref_id}"
- # 使用服务代理
- proxy_url = setting.get_service_proxy()
- proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else None
+ proxy = setting.get_proxy("service")
+ proxies = {"http": proxy, "https": proxy} if proxy else None
- # 构建请求参数
- request_kwargs = {
- "headers": headers,
- "data": json.dumps(payload),
- "impersonate": IMPERSONATE_BROWSER,
- "timeout": REQUEST_TIMEOUT,
- "stream": True,
- "proxies": proxies
- }
-
- # 在线程池中执行同步HTTP请求，避免阻塞事件循环
+ # 执行请求
response = await asyncio.to_thread(
curl_requests.post,
- GROK_API_ENDPOINT,
- **request_kwargs
+ API_ENDPOINT,
+ headers=headers,
+ data=orjson.dumps(payload),
+ impersonate=BROWSER,
+ timeout=TIMEOUT,
+ stream=True,
+ proxies=proxies
)
-
- # 处理非成功响应
+
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)
-
+ GrokClient._handle_error(response, token)
+
+ # 成功 - 重置失败计数
+ asyncio.create_task(token_manager.reset_failure(token))
+
+ # 处理响应
+ 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}")
+ 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
+ logger.error(f"[Client] 请求错误: {e}")
+ raise GrokApiException(f"请求错误: {e}", "REQUEST_ERROR") from e
@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 _handle_error(response, auth_token: str):
- """处理错误响应"""
- # 处理 403 错误
+ def _handle_error(response, token: str):
+ """处理错误"""
if response.status_code == 403:
- error_message = "服务器IP被Block，请尝试 1. 更换服务器IP 2. 使用代理IP 3. 服务器登陆Grok.com，过盾后F12找到CF值填入后台设置"
- error_data = {"cf_blocked": True, "status": 403}
- logger.warning(f"[Client] {error_message}")
+ msg = "您的IP被拦截，请尝试以下方法之一: 1.更换IP 2.使用代理 3.配置CF值"
+ data = {"cf_blocked": True, "status": 403}
+ logger.warning(f"[Client] {msg}")
else:
- # 其他错误尝试解析 JSON
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 str(e)
-
- # 记录Token失败
- asyncio.create_task(token_manager.record_failure(auth_token, response.status_code, error_message))
-
+ data = response.json()
+ msg = str(data)
+ except:
+ data = response.text
+ msg = data[:200] if data else "未知错误"
+
+ asyncio.create_task(token_manager.record_failure(token, response.status_code, msg))
raise GrokApiException(
- f"请求失败: {response.status_code} - {error_message}",
+ f"请求失败: {response.status_code} - {msg}",
"HTTP_ERROR",
- {"status": response.status_code, "data": error_data}
+ {"status": response.status_code, "data": data}
)
@staticmethod
- async def _process_response(response, auth_token: str, model: str, stream: bool):
- """处理API响应"""
- if stream:
- result = GrokResponseProcessor.process_stream(response, auth_token)
- asyncio.create_task(GrokClient._update_rate_limits(auth_token, model))
- else:
- result = await GrokResponseProcessor.process_normal(response, auth_token, model)
- asyncio.create_task(GrokClient._update_rate_limits(auth_token, model))
-
- return result
-
- @staticmethod
- async def _update_rate_limits(auth_token: str, model: str):
- """异步更新速率限制信息"""
+ async def _update_limits(token: str, model: str):
+ """更新速率限制"""
try:
- await token_manager.check_limits(auth_token, model)
+ await token_manager.check_limits(token, model)
except Exception as e:
- logger.error(f"[Client] 更新速率限制失败: {e}")
\ No newline at end of file
+ logger.error(f"[Client] 更新限制失败: {e}")
\ No newline at end of file
diff --git a/app/services/grok/create.py b/app/services/grok/create.py
index 29621d4..72be355 100644
--- a/app/services/grok/create.py
+++ b/app/services/grok/create.py
@@ -1,6 +1,6 @@
-"""Post创建管理器"""
+"""Post创建管理器 - 用于视频生成前的会话创建"""
-import json
+import orjson
from typing import Dict, Any, Optional
from curl_cffi.requests import AsyncSession
@@ -9,75 +9,65 @@ from app.core.exception import GrokApiException
from app.core.config import setting
from app.core.logger import logger
-# 常量定义
-CREATE_ENDPOINT = "https://grok.com/rest/media/post/create"
-REQUEST_TIMEOUT = 30
-IMPERSONATE_BROWSER = "chrome133a"
+
+# 常量
+ENDPOINT = "https://grok.com/rest/media/post/create"
+TIMEOUT = 30
+BROWSER = "chrome133a"
class PostCreateManager:
- """
- Grok 会话创建管理器
-
- 提供图片会话创建功能，用于视频生成前的准备工作
- """
+ """会话创建管理器"""
@staticmethod
async def create(file_id: str, file_uri: str, auth_token: str) -> Optional[Dict[str, Any]]:
- """
- 创建会话记录
+ """创建会话记录
Args:
- file_id: 上传后的文件ID
- file_uri: 上传后的文件URI
+ file_id: 文件ID
+ file_uri: 文件URI
auth_token: 认证令牌
Returns:
- 创建的会话信息，包含会话ID等
+ 会话信息字典，包含post_id等
"""
- try:
- # 验证参数
- if not file_id or not file_uri:
- raise GrokApiException("会话ID或URI缺失", "INVALID_PARAMS")
-
- if not auth_token:
- raise GrokApiException("认证令牌缺失", "NO_AUTH_TOKEN")
+ # 参数验证
+ if not file_id or not file_uri:
+ raise GrokApiException("文件ID或URI缺失", "INVALID_PARAMS")
+ if not auth_token:
+ raise GrokApiException("认证令牌缺失", "NO_AUTH_TOKEN")
- # 构建创建数据
- media_url = f"https://assets.grok.com/{file_uri}"
-
- create_data = {
- "media_url": media_url,
+ try:
+ # 构建请求
+ data = {
+ "media_url": f"https://assets.grok.com/{file_uri}",
"media_type": "MEDIA_POST_TYPE_IMAGE"
}
-
- # 获取认证令牌和cookie
- 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
+ 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
+ }
+
+ proxy = setting.grok_config.get("proxy_url", "")
+ proxies = {"http": proxy, "https": proxy} if proxy else None
- # 发送异步请求
+ # 发送请求
async with AsyncSession() as session:
response = await session.post(
- CREATE_ENDPOINT,
- headers={
- **get_dynamic_headers("/rest/media/post/create"),
- "Cookie": cookie,
- },
- json=create_data,
- impersonate=IMPERSONATE_BROWSER,
- timeout=REQUEST_TIMEOUT,
- proxies=proxies,
+ ENDPOINT,
+ headers=headers,
+ json=data,
+ impersonate=BROWSER,
+ timeout=TIMEOUT,
+ proxies=proxies
)
- # 检查响应
if response.status_code == 200:
result = response.json()
post_id = result.get("post", {}).get("id", "")
- logger.debug(f"[PostCreate] 创建会话成功，会话ID: {post_id}")
+ logger.debug(f"[PostCreate] 成功，会话ID: {post_id}")
return {
"post_id": post_id,
"file_id": file_id,
@@ -85,19 +75,19 @@ class PostCreateManager:
"success": True,
"data": result
}
- else:
- error_msg = f"状态码: {response.status_code}"
- try:
- error_data = response.json()
- error_msg = f"{error_msg}, 详情: {error_data}"
- except:
- error_msg = f"{error_msg}, 详情: {response.text[:200]}"
-
- logger.error(f"[PostCreate] 创建会话失败: {error_msg}")
- raise GrokApiException(f"创建会话失败: {error_msg}", "CREATE_ERROR")
+
+ # 错误处理
+ 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
+ logger.error(f"[PostCreate] 异常: {e}")
+ raise GrokApiException(f"创建异常: {e}", "CREATE_ERROR") from e
diff --git a/app/services/grok/processer.py b/app/services/grok/processer.py
index aa2d5be..4eeb5b9 100644
--- a/app/services/grok/processer.py
+++ b/app/services/grok/processer.py
@@ -1,10 +1,10 @@
-"""Grok API 响应处理器模块"""
+"""Grok API 响应处理器 - 处理流式和非流式响应"""
-import json
+import orjson
import uuid
import time
import asyncio
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Tuple
from app.core.config import setting
from app.core.exception import GrokApiException
@@ -21,61 +21,43 @@ from app.services.grok.cache import image_cache_service, video_cache_service
class StreamTimeoutManager:
- """流式响应超时管理器"""
+ """流式响应超时管理"""
- def __init__(self, chunk_timeout: int = 120, first_response_timeout: int = 30, total_timeout: int = 600):
- """初始化超时管理器
-
- Args:
- chunk_timeout: 数据块间隔超时（秒）
- first_response_timeout: 首次响应超时（秒）
- total_timeout: 总超时限制（秒，0表示不限制）
- """
+ def __init__(self, chunk_timeout: int = 120, first_timeout: int = 30, total_timeout: int = 600):
self.chunk_timeout = chunk_timeout
- self.first_response_timeout = first_response_timeout
+ self.first_timeout = first_timeout
self.total_timeout = total_timeout
-
self.start_time = asyncio.get_event_loop().time()
self.last_chunk_time = self.start_time
- self.first_chunk_received = False
+ self.first_received = False
- def check_timeout(self) -> tuple[bool, str]:
- """检查是否超时
+ def check_timeout(self) -> Tuple[bool, str]:
+ """检查超时"""
+ now = asyncio.get_event_loop().time()
- Returns:
- (is_timeout, timeout_message): 是否超时及超时信息
- """
- current_time = asyncio.get_event_loop().time()
+ if not self.first_received and now - self.start_time > self.first_timeout:
+ return True, f"首次响应超时({self.first_timeout}秒)"
- # 检查首次响应超时
- if not self.first_chunk_received:
- if current_time - self.start_time > self.first_response_timeout:
- return True, f"首次响应超时 ({self.first_response_timeout}秒未收到首个数据块)"
+ if self.total_timeout > 0 and now - self.start_time > self.total_timeout:
+ return True, f"总超时({self.total_timeout}秒)"
- # 检查总超时
- if self.total_timeout > 0:
- if current_time - self.start_time > self.total_timeout:
- return True, f"流式响应总超时 ({self.total_timeout}秒)"
-
- # 检查数据块间隔超时
- if self.first_chunk_received:
- if current_time - self.last_chunk_time > self.chunk_timeout:
- return True, f"数据块间隔超时 ({self.chunk_timeout}秒无新数据)"
+ if self.first_received and now - self.last_chunk_time > self.chunk_timeout:
+ return True, f"数据块超时({self.chunk_timeout}秒)"
return False, ""
- def mark_chunk_received(self):
- """标记收到数据块"""
+ def mark_received(self):
+ """标记收到数据"""
self.last_chunk_time = asyncio.get_event_loop().time()
- self.first_chunk_received = True
+ self.first_received = True
- def get_total_duration(self) -> float:
- """获取总耗时（秒）"""
+ def duration(self) -> float:
+ """获取总耗时"""
return asyncio.get_event_loop().time() - self.start_time
class GrokResponseProcessor:
- """Grok API 响应处理器"""
+ """Grok响应处理器"""
@staticmethod
async def process_normal(response, auth_token: str, model: str = None) -> OpenAIChatCompletionResponse:
@@ -86,7 +68,7 @@ class GrokResponseProcessor:
if not chunk:
continue
- data = json.loads(chunk.decode("utf-8"))
+ data = orjson.loads(chunk)
# 错误检查
if error := data.get("error"):
@@ -96,133 +78,57 @@ class GrokResponseProcessor:
{"code": error.get("code")}
)
- # 提取响应数据
grok_resp = data.get("result", {}).get("response", {})
- # 提取视频数据
+ # 视频响应
if video_resp := grok_resp.get("streamingVideoGenerationResponse"):
if video_url := video_resp.get("videoUrl"):
- logger.debug(f"[Processor] 检测到视频生成: {video_url}")
- full_video_url = f"https://assets.grok.com/{video_url}"
-
- # 下载并缓存视频
- try:
- cache_path = await video_cache_service.download_video(f"/{video_url}", auth_token)
- if cache_path:
- video_path = video_url.replace('/', '-')
- base_url = setting.global_config.get("base_url", "")
- local_video_url = f"{base_url}/images/{video_path}" if base_url else f"/images/{video_path}"
- content = f'\n'
- else:
- content = f'\n'
- except Exception as e:
- logger.warning(f"[Processor] 缓存视频失败: {e}")
- content = f'\n'
-
- # 返回视频响应
- result = OpenAIChatCompletionResponse(
- id=f"chatcmpl-{uuid.uuid4()}",
- object="chat.completion",
- created=int(time.time()),
- model=model or "grok-imagine-0.9",
- choices=[OpenAIChatCompletionChoice(
- index=0,
- message=OpenAIChatCompletionMessage(
- role="assistant",
- content=content
- ),
- finish_reason="stop"
- )],
- usage=None
- )
+ content = await GrokResponseProcessor._build_video_content(video_url, auth_token)
+ result = GrokResponseProcessor._build_response(content, model or "grok-imagine-0.9")
response_closed = True
response.close()
return result
- # 提取模型响应
+ # 模型响应
model_response = grok_resp.get("modelResponse")
if not model_response:
continue
- # 检查 modelResponse 中的错误
if error_msg := model_response.get("error"):
- raise GrokApiException(
- f"模型响应错误: {error_msg}",
- "MODEL_ERROR"
- )
+ raise GrokApiException(f"模型错误: {error_msg}", "MODEL_ERROR")
- # 构建响应内容
- model_name = model_response.get("model")
+ # 构建内容
content = model_response.get("message", "")
+ model_name = model_response.get("model")
- # 提取图片数据
+ # 处理图片
if images := model_response.get("generatedImageUrls"):
- # 获取图片返回模式
- image_mode = setting.global_config.get("image_mode", "url")
+ content = await GrokResponseProcessor._append_images(content, images, auth_token)
- for img in images:
- try:
- if image_mode == "base64":
- # base64 模式：下载并转换为 base64
- base64_str = await image_cache_service.download_base64(f"/{img}", auth_token)
- if base64_str:
- content += f"\n"
- 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
@@ -234,15 +140,15 @@ class GrokResponseProcessor:
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()),
@@ -251,275 +157,278 @@ class GrokResponseProcessor:
index=chunk_index,
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)} bytes")
+ logger.debug(f"[Processor] 解析响应: {len(grok_resp)} bytes")
if not grok_resp:
continue
- # 更新超时计时器
- timeout_manager.mark_chunk_received()
+ timeout_mgr.mark_received()
- # 更新模型名称
+ # 更新模型
if user_resp := grok_resp.get("userResponse"):
if m := user_resp.get("model"):
model = m
- # 提取视频数据
+ # 视频处理
if video_resp := grok_resp.get("streamingVideoGenerationResponse"):
progress = video_resp.get("progress", 0)
v_url = video_resp.get("videoUrl")
- # 处理进度更新
+ # 进度更新
if progress > last_video_progress:
last_video_progress = progress
-
- # 思考状态
if show_thinking:
- # 添加 标签
if not video_progress_started:
- content = f"视频已生成{progress}%\n"
+ content = f"视频已生成{progress}%\\n"
video_progress_started = True
elif progress < 100:
- content = f"视频已生成{progress}%\n"
+ content = f"视频已生成{progress}%\\n"
else:
- # 关闭 标签
- content = f"视频已生成{progress}%\n"
-
+ content = f"视频已生成{progress}%\\n"
yield make_chunk(content)
chunk_index += 1
- # 处理视频URL
+ # 视频URL
if v_url:
- logger.debug(f"[Processor] 视频生成完成")
- 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}"
- video_content = f'\n'
- else:
- video_content = f'\n'
- except Exception as e:
- logger.warning(f"[Processor] 缓存视频失败: {e}")
- video_content = f'\n'
-
+ logger.debug("[Processor] 视频生成完成")
+ video_content = await GrokResponseProcessor._build_video_content(v_url, auth_token)
yield make_chunk(video_content)
chunk_index += 1
continue
- # 检查生成模式
+ # 图片模式
if grok_resp.get("imageAttachmentInfo"):
is_image = True
- # 获取token
token = grok_resp.get("token", "")
- # 提取图片数据
+ # 图片处理
if is_image:
if model_resp := grok_resp.get("modelResponse"):
- # 获取图片返回模式
image_mode = setting.global_config.get("image_mode", "url")
-
- # 初始化内容变量
content = ""
- # 生成图片链接并缓存
for img in model_resp.get("generatedImageUrls", []):
try:
if image_mode == "base64":
- # base64 模式：下载并转换为 base64
+ # Base64模式 - 分块发送
base64_str = await image_cache_service.download_base64(f"/{img}", auth_token)
if base64_str:
- # 分块发送 base64 数据，每 8KB 一个 chunk
- markdown_prefix = "![Generated Image](data:"
- markdown_suffix = ")\n"
-
- # 提取 data URL 的 mime 和 base64 部分
+ # 分块发送大数据
if 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)
+ yield make_chunk(f"![Generated Image](data:{parts[0]},")
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)
+ # 8KB分块
+ for i in range(0, len(parts[1]), 8192):
+ yield make_chunk(parts[1][i:i+8192])
chunk_index += 1
-
- # 发送后缀
- yield make_chunk(markdown_suffix)
+ yield make_chunk(")\\n")
chunk_index += 1
else:
- yield make_chunk(f"\n")
+ yield make_chunk(f"\\n")
chunk_index += 1
else:
- yield make_chunk(f"\n")
+ yield make_chunk(f"\\n")
chunk_index += 1
else:
- yield make_chunk(f"\n")
+ yield make_chunk(f"\\n")
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}"
- content += f"\n"
+ content += f"\\n"
except Exception as e:
logger.warning(f"[Processor] 处理图片失败: {e}")
- content += f"\n"
+ content += f"\\n"
- # 发送内容
yield make_chunk(content.strip(), "stop")
return
elif token:
yield make_chunk(token)
chunk_index += 1
- # 提取对话数据
+ # 对话处理
else:
- # 过滤 list 格式的 token
if isinstance(token, list):
continue
- # 过滤特定标签
if any(tag in token for tag in filtered_tags if token):
continue
- # 获取当前状态
current_is_thinking = grok_resp.get("isThinking", False)
message_tag = grok_resp.get("messageTag")
- # 跳过后续的 标签
if thinking_finished and current_is_thinking:
continue
- # 检查 toolUsageCardId
+ # 搜索结果处理
if grok_resp.get("toolUsageCardId"):
if web_search := grok_resp.get("webSearchResults"):
if current_is_thinking:
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"
+ preview_clean = preview.replace("\\n", "") if isinstance(preview, str) else ""
+ token += f'\\n- [{title}]({url} "{preview_clean}")'
+ token += "\\n"
else:
- # show_thinking=false 时跳过 thinking 状态下的搜索结果
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:
- # 进入 thinking 状态
if show_thinking:
- content = f"\n{content}"
+ content = f"\\n{content}"
else:
should_skip = True
elif is_thinking and not current_is_thinking:
- # 退出 thinking 状态
if show_thinking:
- content = f"\n\n{content}"
+ content = f"\\n\\n{content}"
thinking_finished = True
elif current_is_thinking:
- # 处于 thinking 状态中
if not show_thinking:
should_skip = True
- # 只在不需要跳过时才发送
if not should_skip:
yield make_chunk(content)
chunk_index += 1
is_thinking = current_is_thinking
- except (json.JSONDecodeError, UnicodeDecodeError) as e:
- logger.warning(f"[Processor] 解析chunk失败: {e}")
+ except (orjson.JSONDecodeError, UnicodeDecodeError) as e:
+ logger.warning(f"[Processor] 解析失败: {e}")
continue
except Exception as e:
- logger.warning(f"[Processor] 处理chunk出错: {e}")
+ logger.warning(f"[Processor] 处理出错: {e}")
continue
- # 发送结束块
yield make_chunk("", "stop")
-
- # 发送流结束标记
yield "data: [DONE]\n\n"
-
- # 记录流式响应统计
- logger.info(f"[Processor] 流式响应完成，总耗时: {timeout_manager.get_total_duration():.2f}秒")
+ logger.info(f"[Processor] 流式完成，耗时: {timeout_mgr.duration():.2f}秒")
except Exception as e:
- logger.error(f"[Processor] 流式处理发生严重错误: {e}")
+ logger.error(f"[Processor] 严重错误: {e}")
yield make_chunk(f"处理错误: {e}", "error")
- # 发送流结束标记
yield "data: [DONE]\n\n"
finally:
- # 确保响应对象被关闭
if not response_closed and hasattr(response, 'close'):
try:
response.close()
- logger.debug("[Processor] 流式响应对象已关闭")
+ logger.debug("[Processor] 响应已关闭")
except Exception as e:
- logger.warning(f"[Processor] 关闭流式响应对象时出错: {e}")
\ No newline at end of file
+ logger.warning(f"[Processor] 关闭失败: {e}")
+
+ @staticmethod
+ async def _build_video_content(video_url: str, auth_token: str) -> str:
+ """构建视频内容"""
+ logger.debug(f"[Processor] 检测到视频: {video_url}")
+ full_url = f"https://assets.grok.com/{video_url}"
+
+ try:
+ cache_path = await video_cache_service.download_video(f"/{video_url}", auth_token)
+ if cache_path:
+ video_path = video_url.replace('/', '-')
+ base_url = setting.global_config.get("base_url", "")
+ local_url = f"{base_url}/images/{video_path}" if base_url else f"/images/{video_path}"
+ return f'\\n'
+ except Exception as e:
+ logger.warning(f"[Processor] 缓存视频失败: {e}")
+
+ return f'\\n'
+
+ @staticmethod
+ async def _append_images(content: str, images: list, auth_token: str) -> str:
+ """追加图片到内容"""
+ image_mode = setting.global_config.get("image_mode", "url")
+
+ for img in images:
+ try:
+ if image_mode == "base64":
+ base64_str = await image_cache_service.download_base64(f"/{img}", auth_token)
+ if base64_str:
+ content += f"\\n"
+ 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
+ )
\ No newline at end of file
diff --git a/app/services/grok/statsig.py b/app/services/grok/statsig.py
index 5e6e240..a3d7017 100644
--- a/app/services/grok/statsig.py
+++ b/app/services/grok/statsig.py
@@ -1,4 +1,4 @@
-"""Grok 请求头管理模块"""
+"""Grok 请求头管理 - 生成动态请求头和Statsig ID"""
import base64
import random
@@ -10,95 +10,71 @@ from app.core.logger import logger
from app.core.config import setting
-def _generate_random_string(length: int, use_letters: bool = True) -> str:
- """生成随机字符串
-
- Args:
- length: 字符串长度
- use_letters: 是否使用字母（True）或数字+字母（False）
-
- Returns:
- 随机字符串
- """
- if use_letters:
- # 生成随机字母（小写）
- return ''.join(random.choices(string.ascii_lowercase, k=length))
- else:
- # 生成随机数字和字母组合（小写）
- return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))
+# 基础请求头
+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
+ """生成x-statsig-id
- 随机选择两种格式之一：
+ 随机选择两种格式：
1. e:TypeError: Cannot read properties of null (reading 'children['xxxxx']')
- 其中 xxxxx 是5位随机字符串
2. e:TypeError: Cannot read properties of undefined (reading 'xxxxxxxxxx')
- 其中 xxxxxxxxxx 是10位随机字母
-
- Returns:
- base64 编码后的字符串
"""
- # 随机选择一种格式
- format_type = random.choice([1, 2])
-
- if format_type == 1:
- # 格式1: children['xxxxx']
- random_str = _generate_random_string(5, use_letters=False)
- error_msg = f"e:TypeError: Cannot read properties of null (reading 'children['{random_str}']')"
+ 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:
- # 格式2: 'xxxxxxxxxx'
- random_str = _generate_random_string(10, use_letters=True)
- error_msg = f"e:TypeError: Cannot read properties of undefined (reading '{random_str}')"
+ rand = _random_string(10)
+ msg = f"e:TypeError: Cannot read properties of undefined (reading '{rand}')"
- # base64 编码
- encoded = base64.b64encode(error_msg.encode('utf-8')).decode('utf-8')
- return encoded
+ return base64.b64encode(msg.encode()).decode()
def get_dynamic_headers(pathname: str = "/rest/app-chat/conversations/new") -> Dict[str, str]:
"""获取请求头
-
+
Args:
pathname: 请求路径
-
+
Returns:
- 请求头字典
+ 完整的请求头字典
"""
- # 检查是否启用动态生成
- dynamic_statsig = setting.grok_config.get("dynamic_statsig", False)
-
- if dynamic_statsig:
- # 动态生成 x-statsig-id
+ # 获取或生成statsig-id
+ if setting.grok_config.get("dynamic_statsig", False):
statsig_id = _generate_statsig_id()
- logger.debug(f"[Statsig] 动态生成值 {statsig_id}")
+ logger.debug(f"[Statsig] 动态生成: {statsig_id}")
else:
- # 使用配置文件中的固定值
statsig_id = setting.grok_config.get("x_statsig_id")
- logger.debug(f"[Statsig] 使用固定值 {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-environment=production,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
- "x-statsig-id": statsig_id,
- "x-xai-request-id": str(uuid.uuid4())
- }
+ # 构建请求头
+ headers = BASE_HEADERS.copy()
+ headers["x-statsig-id"] = statsig_id
+ headers["x-xai-request-id"] = str(uuid.uuid4())
+ headers["Content-Type"] = "text/plain;charset=UTF-8" if "upload-file" in pathname else "application/json"
return headers
\ No newline at end of file
diff --git a/app/services/grok/token.py b/app/services/grok/token.py
index 5cc5159..1e58008 100644
--- a/app/services/grok/token.py
+++ b/app/services/grok/token.py
@@ -1,6 +1,6 @@
-"""Grok Token 管理器模块"""
+"""Grok Token 管理器 - 单例模式的Token负载均衡和状态管理"""
-import json
+import orjson
import time
import asyncio
import aiofiles
@@ -14,37 +14,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 +43,65 @@ class GrokTokenManager:
self._file_lock = asyncio.Lock()
self.token_file.parent.mkdir(parents=True, exist_ok=True)
self._storage = None
-
- # 同步加载初始数据
self._load_data()
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: {}
- }
-
+ """加载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)
+ self.token_data = orjson.loads(f.read())
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))
+ await f.write(orjson.dumps(self.token_data, option=orjson.OPT_INDENT_2).decode())
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)})
@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] = {
@@ -138,310 +115,233 @@ class GrokTokenManager:
"tags": [],
"note": ""
}
- added_count += 1
+ count += 1
await self._save_data()
- logger.info(f"[Token] 成功添加 {added_count} 个 {token_type.value} Token")
+ 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")
+ 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的标签"""
+ """更新Token标签"""
if token not in self.token_data[token_type.value]:
- logger.warning(f"[Token] Token不存在: {token[:10]}...")
- raise GrokApiException(
- "Token不存在",
- "TOKEN_NOT_FOUND",
- {"token": token[:10]}
- )
-
- # 确保tags是列表且不包含空字符串
- cleaned_tags = [tag.strip() for tag in tags if tag and tag.strip()]
- self.token_data[token_type.value][token]["tags"] = cleaned_tags
+ 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
await self._save_data()
- logger.info(f"[Token] 成功更新Token {token[:10]}... 的标签: {cleaned_tags}")
+ logger.info(f"[Token] 更新标签: {token[:10]}... -> {cleaned}")
async def update_token_note(self, token: str, token_type: TokenType, note: str) -> None:
- """更新Token的备注"""
+ """更新Token备注"""
if token not in self.token_data[token_type.value]:
- logger.warning(f"[Token] Token不存在: {token[:10]}...")
- raise GrokApiException(
- "Token不存在",
- "TOKEN_NOT_FOUND",
- {"token": token[:10]}
- )
+ raise GrokApiException("Token不存在", "TOKEN_NOT_FOUND", {"token": token[:10]})
self.token_data[token_type.value][token]["note"] = note.strip()
-
await self._save_data()
- logger.info(f"[Token] 成功更新Token {token[:10]}... 的备注")
+ logger.info(f"[Token] 更新备注: {token[:10]}...")
def get_tokens(self) -> Dict[str, Any]:
- """获取所有Token数据"""
+ """获取所有Token"""
return self.token_data.copy()
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"""
+ 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():
+ if data.get("status") == "expired":
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)
-
- # 准备请求
- 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
-
- headers = get_dynamic_headers("/rest/rate-limits")
- headers["Cookie"] = cookie
-
- # 获取代理配置
- proxy_url = setting.grok_config.get("proxy_url", "")
- proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else None
+ 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"] = f"{auth_token};{cf}" if cf else auth_token
+
+ proxy = setting.grok_config.get("proxy_url", "")
+ proxies = {"http": proxy, "https": proxy} if proxy else None
- # 发送异步请求
async with AsyncSession() as session:
response = await session.post(
- RATE_LIMIT_ENDPOINT,
+ RATE_LIMIT_API,
headers=headers,
json=payload,
- impersonate=IMPERSONATE_BROWSER,
- timeout=REQUEST_TIMEOUT,
+ impersonate=BROWSER,
+ timeout=TIMEOUT,
proxies=proxies
)
if response.status_code == 200:
- rate_limit_data = response.json()
-
- # 保存速率限制信息
- sso_value = self._extract_sso(auth_token)
- if sso_value:
+ data = response.json()
+ sso = self._extract_sso(auth_token)
+
+ if sso:
if model == "grok-4-heavy":
- await self.update_limits(sso_value, normal=None, heavy=rate_limit_data.get("remainingQueries", -1))
- logger.info(f"[Token] 更新 Token 限制: sso={sso_value[:10]}..., heavy={rate_limit_data.get('remainingQueries', -1)}")
+ 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_value, normal=rate_limit_data.get("remainingTokens", -1), heavy=None)
- logger.info(f"[Token] 更新 Token 限制: sso={sso_value[:10]}..., basic={rate_limit_data.get('remainingTokens', -1)}")
-
- return rate_limit_data
+ 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:
- logger.warning(f"[Token] 获取速率限制失败，状态码: {response.status_code}")
-
- # 根据HTTP状态码处理不同的错误
- sso_value = self._extract_sso(auth_token)
- if sso_value:
+ logger.warning(f"[Token] 获取限制失败: {response.status_code}")
+ sso = self._extract_sso(auth_token)
+ if sso:
if response.status_code == 401:
- # Token失效
- await self.record_failure(auth_token, 401, "Token authentication failed")
+ await self.record_failure(auth_token, 401, "Token失效")
elif response.status_code == 403:
- # 服务器被block，不改变token状态，但记录失败信息
- await self.record_failure(auth_token, 403, "Server blocked (Cloudflare)")
+ await self.record_failure(auth_token, 403, "服务器被Block")
else:
- # 其他错误，设置为限流状态
- await self.record_failure(auth_token, response.status_code, f"Rate limit or other error: {response.status_code}")
-
+ 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
-
+ self.token_data[token_type][sso]["heavyremainingQueries"] = heavy
await self._save_data()
- logger.info(f"[Token] 已更新Token {sso_value[:10]}... 的限制信息")
+ 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: 服务器IP被Block，不影响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错误是服务器IP被Block，不是Token问题
- if status_code == STATSIG_INVALID_CODE:
- logger.warning(
- f"[Token] 服务器IP被Block (403)，请 1. 更换服务器IP 2. 使用代理IP "
- f"3. 服务器登陆Grok.com，过盾后F12找到CF值填入后台设置"
- )
+ 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 status == TOKEN_INVALID and data["failedCount"] >= MAX_FAILURES:
+ data["status"] = "expired"
+ logger.error(f"[Token] 标记失效: {sso[:10]}... (连续401错误{data['failedCount']}次)")
await self._save_data()
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
-
+ if data.get("failedCount", 0) > 0:
+ data["failedCount"] = 0
+ data["lastFailureTime"] = None
+ data["lastFailureReason"] = None
await self._save_data()
- logger.info(f"[Token] Token {sso_value[:10]}... 失败计数已重置")
+ logger.info(f"[Token] 重置失败计数: {sso[:10]}...")
except Exception as e:
- logger.error(f"[Token] 重置Token失败计数时发生错误: {str(e)}")
+ logger.error(f"[Token] 重置失败错误: {e}")
-# 全局Token管理器实例
+# 全局实例
token_manager = GrokTokenManager()
diff --git a/app/services/grok/upload.py b/app/services/grok/upload.py
index 8738b2e..9bea85a 100644
--- a/app/services/grok/upload.py
+++ b/app/services/grok/upload.py
@@ -1,10 +1,9 @@
-"""图片上传管理器"""
+"""图片上传管理器 - 支持Base64和URL图片上传"""
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,136 +11,134 @@ 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
+ # 请求配置
+ 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,
+ }
- proxy_url = setting.grok_config.get("proxy_url", "")
- proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else None
+ proxy = setting.grok_config.get("proxy_url", "")
+ 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,
+ 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", "")
file_uri = result.get("fileUri", "")
- logger.debug(f"[Upload] 图片上传成功，文件ID: {file_id}")
+ logger.debug(f"[Upload] 成功，ID: {file_id}")
return file_id, file_uri
return "", ""
except Exception as e:
- logger.warning(f"[Upload] 上传图片失败: {e}")
- return ""
+ logger.warning(f"[Upload] 失败: {e}")
+ return "", ""
@staticmethod
- def _is_url(image_input: str) -> bool:
- """检查输入是否为有效的URL"""
+ def _is_url(input_str: str) -> bool:
+ """检查是否为URL"""
try:
- result = urlparse(image_input)
+ result = urlparse(input_str)
return all([result.scheme, result.netloc]) and result.scheme in ['http', 'https']
- except Exception as e:
- logger.warning(f"[Upload] URL解析失败: {e}")
+ except:
return False
@staticmethod
async def _download(url: str) -> Tuple[str, str]:
- """下载图片并转换为Base64"""
+ """下载图片并转Base64
+
+ Returns:
+ (base64_string, mime_type) 元组
+ """
try:
async with AsyncSession() as session:
response = await session.get(url, timeout=5)
response.raise_for_status()
- # 获取内容类型
- content_type = response.headers.get('content-type', DEFAULT_MIME_TYPE)
+ content_type = response.headers.get('content-type', DEFAULT_MIME)
if not content_type.startswith('image/'):
- content_type = DEFAULT_MIME_TYPE
+ content_type = DEFAULT_MIME
- # 转换为 Base64
- image_base64 = base64.b64encode(response.content).decode('utf-8')
- return image_base64, content_type
+ b64 = base64.b64encode(response.content).decode()
+ return b64, content_type
except Exception as e:
- logger.warning(f"[Upload] 下载图片失败: {e}")
+ logger.warning(f"[Upload] 下载失败: {e}")
return "", ""
@staticmethod
def _get_info(image_data: str, mime_type: Optional[str] = None) -> Tuple[str, str]:
- """获取图片文件名和MIME类型"""
- # mime_type 有值，直接使用
+ """获取文件名和MIME类型
+
+ Returns:
+ (file_name, mime_type) 元组
+ """
+ # 已提供MIME类型
if mime_type:
- extension = mime_type.split("/")[1] if "/" in mime_type else DEFAULT_EXTENSION
- file_name = f"image.{extension}"
- return file_name, mime_type
+ ext = mime_type.split("/")[1] if "/" in mime_type else DEFAULT_EXT
+ return f"image.{ext}", mime_type
- # mime_type 没有值，使用默认值
- mime_type = DEFAULT_MIME_TYPE
- extension = DEFAULT_EXTENSION
+ # 从Base64提取
+ mime = DEFAULT_MIME
+ ext = DEFAULT_EXT
- # 从 Base64 数据中提取 MIME 类型
if "data:image" in image_data:
- match = re.search(r"data:([a-zA-Z0-9]+/[a-zA-Z0-9-.+]+);base64,", image_data)
- if match:
- mime_type = match.group(1)
- extension = mime_type.split("/")[1]
+ if match := re.search(r"data:([a-zA-Z0-9]+/[a-zA-Z0-9-.+]+);base64,", image_data):
+ mime = match.group(1)
+ ext = mime.split("/")[1]
- file_name = f"image.{extension}"
- return file_name, mime_type
\ No newline at end of file
+ return f"image.{ext}", mime
\ No newline at end of file
diff --git a/data/temp/video/users-8522ce45-679b-4e0e-a0f7-bb18f434eb6b-generated-15f7113f-5d16-4ff1-bdaa-a2eabd66671c-generated_video.mp4 b/data/temp/video/users-8522ce45-679b-4e0e-a0f7-bb18f434eb6b-generated-15f7113f-5d16-4ff1-bdaa-a2eabd66671c-generated_video.mp4
new file mode 100644
index 0000000..195951e
Binary files /dev/null and b/data/temp/video/users-8522ce45-679b-4e0e-a0f7-bb18f434eb6b-generated-15f7113f-5d16-4ff1-bdaa-a2eabd66671c-generated_video.mp4 differ