diff --git a/app/handler/retry_handler.py b/app/handler/retry_handler.py index 37fec3c..0161e27 100644 --- a/app/handler/retry_handler.py +++ b/app/handler/retry_handler.py @@ -2,7 +2,7 @@ from functools import wraps from typing import Callable, TypeVar -from app.core.constants import MAX_RETRIES +from app.config.config import settings from app.log.logger import get_retry_logger T = TypeVar("T") @@ -12,8 +12,7 @@ logger = get_retry_logger() class RetryHandler: """重试处理装饰器""" - def __init__(self, max_retries: int = MAX_RETRIES, key_arg: str = "api_key"): - self.max_retries = max_retries + def __init__(self, key_arg: str = "api_key"): self.key_arg = key_arg def __call__(self, func: Callable[..., T]) -> Callable[..., T]: @@ -21,14 +20,14 @@ class RetryHandler: async def wrapper(*args, **kwargs) -> T: last_exception = None - for attempt in range(self.max_retries): + 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 {self.max_retries}" + f"API call failed with error: {str(e)}. Attempt {retries} of {settings.MAX_RETRIES}" ) # 从函数参数中获取 key_manager diff --git a/app/router/gemini_routes.py b/app/router/gemini_routes.py index 5512a7f..c47f68d 100644 --- a/app/router/gemini_routes.py +++ b/app/router/gemini_routes.py @@ -99,7 +99,7 @@ async def list_models( @router.post("/models/{model_name}:generateContent") @router_v1beta.post("/models/{model_name}:generateContent") -@RetryHandler(max_retries=settings.MAX_RETRIES, key_arg="api_key") +@RetryHandler(key_arg="api_key") async def generate_content( model_name: str, request: GeminiRequest, @@ -128,7 +128,7 @@ async def generate_content( @router.post("/models/{model_name}:streamGenerateContent") @router_v1beta.post("/models/{model_name}:streamGenerateContent") -@RetryHandler(max_retries=settings.MAX_RETRIES, key_arg="api_key") +@RetryHandler(key_arg="api_key") async def stream_generate_content( model_name: str, request: GeminiRequest, @@ -156,6 +156,7 @@ async def stream_generate_content( # 注意:流本身的错误需要在服务层或流迭代中处理,这里只返回流响应 return StreamingResponse(response_stream, media_type="text/event-stream") + @router.post("/reset-all-fail-counts") async def reset_all_key_fail_counts(key_type: str = None, key_manager: KeyManager = Depends(get_key_manager)): """批量重置Gemini API密钥的失败计数,可选择性地仅重置有效或无效密钥""" @@ -250,7 +251,6 @@ async def reset_selected_key_fail_counts( return JSONResponse({"success": False, "message": f"批量重置处理失败: {str(e)}"}, status_code=500) - @router.post("/reset-fail-count/{api_key}") async def reset_key_fail_count(api_key: str, key_manager: KeyManager = Depends(get_key_manager)): """重置指定Gemini API密钥的失败计数""" @@ -266,6 +266,7 @@ async def reset_key_fail_count(api_key: str, key_manager: KeyManager = Depends(g logger.error(f"Failed to reset key failure count: {str(e)}") return JSONResponse({"success": False, "message": f"重置失败: {str(e)}"}, status_code=500) + @router.post("/verify-key/{api_key}") async def verify_key(api_key: str, chat_service: GeminiChatService = Depends(get_chat_service), key_manager: KeyManager = Depends(get_key_manager)): """验证Gemini API密钥的有效性""" diff --git a/app/router/openai_compatiable_routes.py b/app/router/openai_compatiable_routes.py index d03377e..f4d6ebc 100644 --- a/app/router/openai_compatiable_routes.py +++ b/app/router/openai_compatiable_routes.py @@ -52,7 +52,7 @@ async def list_models( @router.post("/openai/v1/chat/completions") -@RetryHandler(max_retries=settings.MAX_RETRIES, key_arg="api_key") +@RetryHandler(key_arg="api_key") async def chat_completion( request: ChatRequest, _=Depends(security_service.verify_authorization), diff --git a/app/router/openai_routes.py b/app/router/openai_routes.py index b1b993e..68ff476 100644 --- a/app/router/openai_routes.py +++ b/app/router/openai_routes.py @@ -59,7 +59,7 @@ async def list_models( @router.post("/v1/chat/completions") @router.post("/hf/v1/chat/completions") -@RetryHandler(max_retries=settings.MAX_RETRIES, key_arg="api_key") +@RetryHandler(key_arg="api_key") async def chat_completion( request: ChatRequest, _=Depends(security_service.verify_authorization),