diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index ea88209..9e04bb6 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -288,7 +288,11 @@ class GrokChatService: # 重试机制 def extract_status(e: Exception) -> int | None: if isinstance(e, UpstreamException) and e.details: - return e.details.get("status") + status = e.details.get("status") + # 429 不在内层重试,由外层跨 token 重试处理 + if status == 429: + return None + return status return None session = None @@ -387,22 +391,41 @@ class ChatService: token_mgr = await get_token_manager() await token_mgr.reload_if_stale() - token = None - for pool_name in ModelService.pool_candidates_for_model(model): - token = token_mgr.get_token(pool_name) - if token: - break + # 解析参数(只需解析一次) + think = {"enabled": True, "disabled": False}.get(thinking) + is_stream = stream if stream is not None else get_config("chat.stream") - if not token: - logger.info("No available tokens, attempting to refresh cooling tokens...") - result = await token_mgr.refresh_cooling_tokens() - if result.get("recovered", 0) > 0: - for pool_name in ModelService.pool_candidates_for_model(model): - token = token_mgr.get_token(pool_name) - if token: - break + # 构造请求(只需构造一次) + chat_request = ChatRequest( + model=model, messages=messages, stream=is_stream, think=think + ) + + # 跨 Token 重试循环 + tried_tokens = set() + max_token_retries = 3 + last_error = None + + for attempt in range(max_token_retries): + # 选择 token(排除已失败的) + token = None + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name, exclude=tried_tokens) + if token: + break + + if not token and not tried_tokens: + # 首次就无 token,尝试刷新 + logger.info("No available tokens, attempting to refresh cooling tokens...") + result = await token_mgr.refresh_cooling_tokens() + if result.get("recovered", 0) > 0: + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name) + if token: + break if not token: + if last_error: + raise last_error raise AppException( message="No available tokens. Please try again later.", error_type=ErrorType.RATE_LIMIT.value, @@ -410,43 +433,63 @@ class ChatService: status_code=429, ) - # 解析参数 - think = {"enabled": True, "disabled": False}.get(thinking) - is_stream = stream if stream is not None else get_config("chat.stream") + tried_tokens.add(token) - # 构造请求 - chat_request = ChatRequest( - model=model, messages=messages, stream=is_stream, think=think + try: + # 请求 Grok + service = GrokChatService() + response, _, model_name = await service.chat_openai(token, chat_request) + + # 处理响应 + if is_stream: + logger.debug(f"Processing stream response: model={model}") + processor = StreamProcessor(model_name, token, think) + return wrap_stream_with_usage( + processor.process(response), token_mgr, token, model + ) + + # 非流式 + logger.debug(f"Processing non-stream response: model={model}") + result = await CollectProcessor(model_name, token).process(response) + 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.info(f"Chat completed: model={model}, effort={effort.value}") + except Exception as e: + logger.warning(f"Failed to record usage: {e}") + return result + + except UpstreamException as e: + status_code = e.details.get("status") if e.details else None + last_error = e + + if status_code == 429: + # 配额不足,标记 token 为 cooling 并换 token 重试 + await token_mgr.mark_rate_limited(token) + logger.warning( + f"Token {token[:10]}... rate limited (429), " + f"trying next token (attempt {attempt + 1}/{max_token_retries})" + ) + continue + + # 非 429 错误,不换 token,直接抛出 + raise + + # 所有 token 都 429,抛出最后的错误 + if last_error: + raise last_error + raise AppException( + message="No available tokens. Please try again later.", + error_type=ErrorType.RATE_LIMIT.value, + code="rate_limit_exceeded", + status_code=429, ) - # 请求 Grok - service = GrokChatService() - response, _, model_name = await service.chat_openai(token, chat_request) - - # 处理响应 - if is_stream: - logger.debug(f"Processing stream response: model={model}") - processor = StreamProcessor(model_name, token, think) - return wrap_stream_with_usage( - processor.process(response), token_mgr, token, model - ) - - # 非流式 - logger.debug(f"Processing non-stream response: model={model}") - result = await CollectProcessor(model_name, token).process(response) - 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.info(f"Chat completed: model={model}, effort={effort.value}") - except Exception as e: - logger.warning(f"Failed to record usage: {e}") - return result - __all__ = [ "GrokChatService", diff --git a/app/services/token/manager.py b/app/services/token/manager.py index d72cbd4..39423ed 100644 --- a/app/services/token/manager.py +++ b/app/services/token/manager.py @@ -185,12 +185,13 @@ class TokenManager: if self._dirty: self._schedule_save() - def get_token(self, pool_name: str = "ssoBasic") -> Optional[str]: + def get_token(self, pool_name: str = "ssoBasic", exclude: set = None) -> Optional[str]: """ 获取可用 Token Args: pool_name: Token 池名称 + exclude: 需要排除的 token 字符串集合 Returns: Token 字符串或 None @@ -200,7 +201,7 @@ class TokenManager: logger.warning(f"Pool '{pool_name}' not found") return None - token_info = pool.select() + token_info = pool.select(exclude=exclude) if not token_info: logger.warning(f"No available token in pool '{pool_name}'") return None @@ -435,6 +436,37 @@ class TokenManager: logger.warning(f"Token {raw_token[:10]}...: not found for failure record") return False + async def mark_rate_limited(self, token_str: str) -> bool: + """ + 将 Token 标记为配额耗尽(COOLING) + + 当 Grok API 返回 429 时调用,将 quota 设为 0 并标记 COOLING, + 使该 Token 不再被选中,等待下次 Scheduler 刷新恢复。 + + Args: + token_str: Token 字符串 + + Returns: + 是否成功 + """ + raw_token = token_str.removeprefix("sso=") + + for pool in self.pools.values(): + token = pool.get(raw_token) + if token: + old_quota = token.quota + token.quota = 0 + token.status = TokenStatus.COOLING + logger.warning( + f"Token {raw_token[:10]}...: marked as rate limited " + f"(quota {old_quota} -> 0, status -> cooling)" + ) + self._schedule_save() + return True + + logger.warning(f"Token {raw_token[:10]}...: not found for rate limit marking") + return False + # ========== 管理功能 ========== async def add(self, token: str, pool_name: str = "ssoBasic") -> bool: diff --git a/app/services/token/pool.py b/app/services/token/pool.py index 405a84b..55a9103 100644 --- a/app/services/token/pool.py +++ b/app/services/token/pool.py @@ -28,7 +28,7 @@ class TokenPool: """获取 Token""" return self._tokens.get(token_str) - def select(self) -> Optional[TokenInfo]: + def select(self, exclude: set = None) -> Optional[TokenInfo]: """ 选择一个可用 Token 策略: @@ -41,6 +41,7 @@ class TokenPool: t for t in self._tokens.values() if t.status == TokenStatus.ACTIVE and t.quota > 0 + and (not exclude or t.token not in exclude) ] if not available: