mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-05-07 06:03:11 +08:00
a6a1e7fb52
将 API 调用的最大重试次数 (`MAX_RETRIES`) 的配置移至 `app.config.config.settings`。 - 修改 `RetryHandler` 以直接从全局设置读取 `MAX_RETRIES`。 - 更新使用 `RetryHandler` 的路由装饰器，移除冗余的 `max_retries` 参数传递。 这使得重试次数的配置更加集中和易于管理。
51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
|
|
from functools import wraps
|
|
from typing import Callable, TypeVar
|
|
|
|
from app.config.config import settings
|
|
from app.log.logger import get_retry_logger
|
|
|
|
T = TypeVar("T")
|
|
logger = get_retry_logger()
|
|
|
|
|
|
class RetryHandler:
|
|
"""重试处理装饰器"""
|
|
|
|
def __init__(self, key_arg: str = "api_key"):
|
|
self.key_arg = key_arg
|
|
|
|
def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
|
|
@wraps(func)
|
|
async def wrapper(*args, **kwargs) -> T:
|
|
last_exception = None
|
|
|
|
for attempt in range(settings.MAX_RETRIES):
|
|
retries = attempt + 1
|
|
try:
|
|
return await func(*args, **kwargs)
|
|
except Exception as e:
|
|
last_exception = e
|
|
logger.warning(
|
|
f"API call failed with error: {str(e)}. Attempt {retries} of {settings.MAX_RETRIES}"
|
|
)
|
|
|
|
# 从函数参数中获取 key_manager
|
|
key_manager = kwargs.get("key_manager")
|
|
if key_manager:
|
|
old_key = kwargs.get(self.key_arg)
|
|
new_key = await key_manager.handle_api_failure(old_key, retries)
|
|
if new_key:
|
|
kwargs[self.key_arg] = new_key
|
|
logger.info(f"Switched to new API key: {new_key}")
|
|
else:
|
|
logger.error(f"No valid API key available after {retries} retries.")
|
|
break
|
|
|
|
logger.error(
|
|
f"All retry attempts failed, raising final exception: {str(last_exception)}"
|
|
)
|
|
raise last_exception
|
|
|
|
return wrapper
|