mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-05 23:56:37 +08:00
refactor: 优化 RetryHandler 装饰器以支持动态 KeyManager 注入
This commit is contained in:
@@ -12,9 +12,8 @@ logger = get_retry_logger()
|
||||
class RetryHandler:
|
||||
"""重试处理装饰器"""
|
||||
|
||||
def __init__(self, max_retries: int = 3, key_manager: KeyManager = None, key_arg: str = "api_key"):
|
||||
def __init__(self, max_retries: int = 3, key_arg: str = "api_key"):
|
||||
self.max_retries = max_retries
|
||||
self.key_manager = key_manager
|
||||
self.key_arg = key_arg
|
||||
|
||||
def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
|
||||
@@ -29,9 +28,11 @@ class RetryHandler:
|
||||
last_exception = e
|
||||
logger.warning(f"API call failed with error: {str(e)}. Attempt {attempt + 1} of {self.max_retries}")
|
||||
|
||||
if self.key_manager:
|
||||
# 从函数参数中获取 key_manager
|
||||
key_manager = kwargs.get('key_manager')
|
||||
if key_manager:
|
||||
old_key = kwargs.get(self.key_arg)
|
||||
new_key = await self.key_manager.handle_api_failure(old_key)
|
||||
new_key = await key_manager.handle_api_failure(old_key)
|
||||
kwargs[self.key_arg] = new_key
|
||||
logger.info(f"Switched to new API key: {new_key}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user