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