GithubBackup/gemini-balance
1
0
Fork 0
mirror of https://github.com/snailyp/gemini-balance.git synced 2026-05-31 05:09:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 优化 RetryHandler 装饰器以支持动态 KeyManager 注入

Browse Source
This commit is contained in:
yinpeng
2025-02-12 17:10:02 +08:00
parent 44f956e4e4
commit b3851441f1
3 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
app/api/gemini_routes.py
View File

@@ -47,7 +47,7 @@ async def list_models(_=Depends(security_service.verify_key),
@router.post("/models/{model_name}:generateContent")
@router_v1beta.post("/models/{model_name}:generateContent")
@RetryHandler(max_retries=3, key_manager=Depends(get_key_manager), key_arg="api_key")
@RetryHandler(max_retries=3, key_arg="api_key")
async def generate_content(
model_name: str,
request: GeminiRequest,
@@ -77,7 +77,7 @@ async def generate_content(
@router.post("/models/{model_name}:streamGenerateContent")
@router_v1beta.post("/models/{model_name}:streamGenerateContent")
@RetryHandler(max_retries=3, key_manager=Depends(get_key_manager), key_arg="api_key")
@RetryHandler(max_retries=3, key_arg="api_key")
async def stream_generate_content(
model_name: str,
request: GeminiRequest,

2
app/api/openai_routes.py
View File

@@ -46,7 +46,7 @@ async def list_models(
@router.post("/v1/chat/completions")
@router.post("/hf/v1/chat/completions")
@RetryHandler(max_retries=3, key_manager=Depends(get_key_manager), key_arg="api_key")
@RetryHandler(max_retries=3, key_arg="api_key")
async def chat_completion(
request: ChatRequest,
_=Depends(security_service.verify_authorization),

9
app/services/chat/retry_handler.py
View File

@@ -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}")