diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index c27103a..c2dc480 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -864,9 +864,7 @@ async def refresh_tokens_api(data: dict): batch_size = get_config("performance.usage_batch_size") async def _refresh_one(t): - return await mgr.sync_usage( - t, "grok-3", consume_on_fail=False, is_usage=False - ) + return await mgr.sync_usage(t, consume_on_fail=False, is_usage=False) raw_results = await run_in_batches( unique_tokens, @@ -918,9 +916,7 @@ async def refresh_tokens_api_async(data: dict): try: async def _refresh_one(t: str): - return await mgr.sync_usage( - t, "grok-3", consume_on_fail=False, is_usage=False - ) + return await mgr.sync_usage(t, consume_on_fail=False, is_usage=False) async def _on_item(item: str, res: dict): task.record(bool(res.get("ok"))) diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index 7550c82..e273433 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -9,11 +9,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 UpstreamException -from app.services.grok.utils.headers import apply_statsig, build_sso_cookie -from app.services.grok.utils.retry import retry_on_status - -LIMITS_API = "https://grok.com/rest/rate-limits" +from app.services.reverse import RateLimitsReverse _USAGE_SEMAPHORE = asyncio.Semaphore(25) _USAGE_SEM_VALUE = 25 @@ -22,52 +18,12 @@ _USAGE_SEM_VALUE = 25 class UsageService: """用量查询服务""" - def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("network.base_proxy_url") - self.timeout = get_config("network.timeout") - - def _build_headers(self, token: str) -> dict: - """构建请求头""" - user_agent = get_config("security.user_agent") - headers = { - "Accept": "*/*", - "Accept-Encoding": "gzip, deflate, br, zstd", - "Accept-Language": "zh-CN,zh;q=0.9", - "Baggage": "sentry-environment=production,sentry-release=d6add6fb0460641fd482d767a335ef72b9b6abb8,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c", - "Cache-Control": "no-cache", - "Content-Type": "application/json", - "Origin": "https://grok.com", - "Pragma": "no-cache", - "Priority": "u=1, i", - "Referer": "https://grok.com/", - "Sec-Ch-Ua": '"Google Chrome";v="136", "Chromium";v="136", "Not(A:Brand";v="24"', - "Sec-Ch-Ua-Arch": "arm", - "Sec-Ch-Ua-Bitness": "64", - "Sec-Ch-Ua-Mobile": "?0", - "Sec-Ch-Ua-Model": "", - "Sec-Ch-Ua-Platform": '"macOS"', - "Sec-Fetch-Dest": "empty", - "Sec-Fetch-Mode": "cors", - "Sec-Fetch-Site": "same-origin", - "User-Agent": user_agent, - } - - apply_statsig(headers) - headers["Cookie"] = build_sso_cookie(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: + async def get(self, token: str) -> Dict: """ 获取速率限制信息 Args: token: 认证 Token - model_name: 模型名称 Returns: 响应数据 @@ -86,61 +42,15 @@ class UsageService: _USAGE_SEM_VALUE = value _USAGE_SEMAPHORE = asyncio.Semaphore(value) async with _USAGE_SEMAPHORE: - # 定义状态码提取器 - def extract_status(e: Exception) -> int | None: - 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} - browser = get_config("security.browser") - - async with AsyncSession() as session: - response = await session.post( - LIMITS_API, - headers=headers, - json=payload, - impersonate=browser, - timeout=self.timeout, - proxies=self._build_proxies(), - ) - - if response.status_code == 200: - data = response.json() - remaining = data.get("remainingTokens", 0) - logger.info( - f"Usage sync success: remaining={remaining}, token={token[:10]}..." - ) - return data - - logger.error( - f"Usage sync failed: status={response.status_code}, token={token[:10]}..." - ) - - raise UpstreamException( - message=f"Failed to get usage stats: {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)}, - ) - - # 带重试的执行 try: - result = await retry_on_status( - do_request, extract_status=extract_status + async with AsyncSession() as session: + response = await RateLimitsReverse.request(session, token) + data = response.json() + remaining = data.get("remainingTokens", 0) + logger.info( + f"Usage sync success: remaining={remaining}, token={token[:10]}..." ) - return result + return data except Exception: # 最后一次失败已经被记录 diff --git a/app/services/reverse/__init__.py b/app/services/reverse/__init__.py index 158e51d..cb4e88a 100644 --- a/app/services/reverse/__init__.py +++ b/app/services/reverse/__init__.py @@ -6,6 +6,7 @@ from .assets_download import AssetsDownloadReverse from .assets_list import AssetsListReverse from .assets_upload import AssetsUploadReverse from .media_post import MediaPostReverse +from .rate_limits import RateLimitsReverse from .utils.headers import build_headers from .utils.statsig import StatsigGenerator @@ -16,6 +17,7 @@ __all__ = [ "AssetsListReverse", "AssetsUploadReverse", "MediaPostReverse", + "RateLimitsReverse", "StatsigGenerator", "build_headers", ] diff --git a/app/services/reverse/rate_limits.py b/app/services/reverse/rate_limits.py new file mode 100644 index 0000000..aa85260 --- /dev/null +++ b/app/services/reverse/rate_limits.py @@ -0,0 +1,108 @@ +""" +Reverse interface: rate limits. +""" + +import orjson +from typing import Any +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 +from app.services.token.service import TokenService +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status + +RATE_LIMITS_API = "https://grok.com/rest/rate-limits" + + +class RateLimitsReverse: + """/rest/rate-limits reverse interface.""" + + @staticmethod + async def request(session: AsyncSession, token: str) -> Any: + """Fetch rate limits from Grok. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + + Returns: + Any: The response from the request. + """ + try: + # Get proxies + base_proxy = get_config("network.base_proxy_url") + proxies = {"http": base_proxy, "https": base_proxy} if base_proxy else None + + # Build headers + headers = build_headers( + cookie_token=token, + content_type="application/json", + origin="https://grok.com", + referer="https://grok.com/", + ) + + # Build payload + payload = { + "requestKind": "DEFAULT", + "modelName": "grok-4-1-thinking-1129", + } + + # Curl Config + timeout = get_config("network.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.post( + RATE_LIMITS_API, + headers=headers, + data=orjson.dumps(payload), + timeout=timeout, + proxies=proxies, + impersonate=browser, + ) + + if response.status_code != 200: + logger.error( + f"RateLimitsReverse: Request failed, {response.status_code}", + extra={"error_type": "UpstreamException"}, + ) + raise UpstreamException( + message=f"RateLimitsReverse: Request failed, {response.status_code}", + details={"status": response.status_code}, + ) + + return response + + return await retry_on_status(_do_request) + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + if status == 401: + try: + await TokenService.record_fail( + token, status, "rate_limits_auth_failed" + ) + except Exception: + pass + raise + + # Handle other non-upstream exceptions + logger.error( + f"RateLimitsReverse: Request failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"RateLimitsReverse: Request failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + + +__all__ = ["RateLimitsReverse"] diff --git a/app/services/token/manager.py b/app/services/token/manager.py index b5e28f1..e8e6d41 100644 --- a/app/services/token/manager.py +++ b/app/services/token/manager.py @@ -329,7 +329,6 @@ class TokenManager: async def sync_usage( self, token_str: str, - model_name: str, fallback_effort: EffortType = EffortType.LOW, consume_on_fail: bool = True, is_usage: bool = True, @@ -341,7 +340,6 @@ class TokenManager: Args: token_str: Token 字符串（可带 sso= 前缀） - model_name: 模型名称（用于 API 查询） fallback_effort: 降级时的消耗力度 consume_on_fail: 失败时是否降级扣费 is_usage: 是否记录为一次使用（影响 use_count） @@ -367,7 +365,7 @@ class TokenManager: from app.services.grok.services.usage import UsageService usage_service = UsageService() - result = await usage_service.get(token_str, model_name=model_name) + result = await usage_service.get(token_str) if result and "remainingTokens" in result: old_quota = target_token.quota diff --git a/app/services/token/service.py b/app/services/token/service.py index 75b23a6..63e635c 100644 --- a/app/services/token/service.py +++ b/app/services/token/service.py @@ -43,22 +43,19 @@ class TokenService: return await manager.consume(token, effort) @staticmethod - async def sync_usage( - token: str, model: str, effort: EffortType = EffortType.LOW - ) -> bool: + async def sync_usage(token: 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) + return await manager.sync_usage(token, effort) @staticmethod async def record_fail(token: str, status_code: int = 401, reason: str = "") -> bool: