mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-05 15:36:45 +08:00
refactor(retry): 统一管理 API 最大重试次数配置
将 API 调用的最大重试次数 (`MAX_RETRIES`) 的配置移至 `app.config.config.settings`。 - 修改 `RetryHandler` 以直接从全局设置读取 `MAX_RETRIES`。 - 更新使用 `RetryHandler` 的路由装饰器,移除冗余的 `max_retries` 参数传递。 这使得重试次数的配置更加集中和易于管理。
This commit is contained in:
@@ -2,7 +2,7 @@
|
|||||||
from functools import wraps
|
from functools import wraps
|
||||||
from typing import Callable, TypeVar
|
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
|
from app.log.logger import get_retry_logger
|
||||||
|
|
||||||
T = TypeVar("T")
|
T = TypeVar("T")
|
||||||
@@ -12,8 +12,7 @@ logger = get_retry_logger()
|
|||||||
class RetryHandler:
|
class RetryHandler:
|
||||||
"""重试处理装饰器"""
|
"""重试处理装饰器"""
|
||||||
|
|
||||||
def __init__(self, max_retries: int = MAX_RETRIES, key_arg: str = "api_key"):
|
def __init__(self, key_arg: str = "api_key"):
|
||||||
self.max_retries = max_retries
|
|
||||||
self.key_arg = key_arg
|
self.key_arg = key_arg
|
||||||
|
|
||||||
def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
|
def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
|
||||||
@@ -21,14 +20,14 @@ class RetryHandler:
|
|||||||
async def wrapper(*args, **kwargs) -> T:
|
async def wrapper(*args, **kwargs) -> T:
|
||||||
last_exception = None
|
last_exception = None
|
||||||
|
|
||||||
for attempt in range(self.max_retries):
|
for attempt in range(settings.MAX_RETRIES):
|
||||||
retries = attempt + 1
|
retries = attempt + 1
|
||||||
try:
|
try:
|
||||||
return await func(*args, **kwargs)
|
return await func(*args, **kwargs)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
last_exception = e
|
last_exception = e
|
||||||
logger.warning(
|
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
|
# 从函数参数中获取 key_manager
|
||||||
|
|||||||
@@ -99,7 +99,7 @@ async def list_models(
|
|||||||
|
|
||||||
@router.post("/models/{model_name}:generateContent")
|
@router.post("/models/{model_name}:generateContent")
|
||||||
@router_v1beta.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(
|
async def generate_content(
|
||||||
model_name: str,
|
model_name: str,
|
||||||
request: GeminiRequest,
|
request: GeminiRequest,
|
||||||
@@ -128,7 +128,7 @@ async def generate_content(
|
|||||||
|
|
||||||
@router.post("/models/{model_name}:streamGenerateContent")
|
@router.post("/models/{model_name}:streamGenerateContent")
|
||||||
@router_v1beta.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(
|
async def stream_generate_content(
|
||||||
model_name: str,
|
model_name: str,
|
||||||
request: GeminiRequest,
|
request: GeminiRequest,
|
||||||
@@ -156,6 +156,7 @@ async def stream_generate_content(
|
|||||||
# 注意:流本身的错误需要在服务层或流迭代中处理,这里只返回流响应
|
# 注意:流本身的错误需要在服务层或流迭代中处理,这里只返回流响应
|
||||||
return StreamingResponse(response_stream, media_type="text/event-stream")
|
return StreamingResponse(response_stream, media_type="text/event-stream")
|
||||||
|
|
||||||
|
|
||||||
@router.post("/reset-all-fail-counts")
|
@router.post("/reset-all-fail-counts")
|
||||||
async def reset_all_key_fail_counts(key_type: str = None, key_manager: KeyManager = Depends(get_key_manager)):
|
async def reset_all_key_fail_counts(key_type: str = None, key_manager: KeyManager = Depends(get_key_manager)):
|
||||||
"""批量重置Gemini API密钥的失败计数,可选择性地仅重置有效或无效密钥"""
|
"""批量重置Gemini API密钥的失败计数,可选择性地仅重置有效或无效密钥"""
|
||||||
@@ -250,7 +251,6 @@ async def reset_selected_key_fail_counts(
|
|||||||
return JSONResponse({"success": False, "message": f"批量重置处理失败: {str(e)}"}, status_code=500)
|
return JSONResponse({"success": False, "message": f"批量重置处理失败: {str(e)}"}, status_code=500)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/reset-fail-count/{api_key}")
|
@router.post("/reset-fail-count/{api_key}")
|
||||||
async def reset_key_fail_count(api_key: str, key_manager: KeyManager = Depends(get_key_manager)):
|
async def reset_key_fail_count(api_key: str, key_manager: KeyManager = Depends(get_key_manager)):
|
||||||
"""重置指定Gemini API密钥的失败计数"""
|
"""重置指定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)}")
|
logger.error(f"Failed to reset key failure count: {str(e)}")
|
||||||
return JSONResponse({"success": False, "message": f"重置失败: {str(e)}"}, status_code=500)
|
return JSONResponse({"success": False, "message": f"重置失败: {str(e)}"}, status_code=500)
|
||||||
|
|
||||||
|
|
||||||
@router.post("/verify-key/{api_key}")
|
@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)):
|
async def verify_key(api_key: str, chat_service: GeminiChatService = Depends(get_chat_service), key_manager: KeyManager = Depends(get_key_manager)):
|
||||||
"""验证Gemini API密钥的有效性"""
|
"""验证Gemini API密钥的有效性"""
|
||||||
|
|||||||
@@ -52,7 +52,7 @@ async def list_models(
|
|||||||
|
|
||||||
|
|
||||||
@router.post("/openai/v1/chat/completions")
|
@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(
|
async def chat_completion(
|
||||||
request: ChatRequest,
|
request: ChatRequest,
|
||||||
_=Depends(security_service.verify_authorization),
|
_=Depends(security_service.verify_authorization),
|
||||||
|
|||||||
@@ -59,7 +59,7 @@ async def list_models(
|
|||||||
|
|
||||||
@router.post("/v1/chat/completions")
|
@router.post("/v1/chat/completions")
|
||||||
@router.post("/hf/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(
|
async def chat_completion(
|
||||||
request: ChatRequest,
|
request: ChatRequest,
|
||||||
_=Depends(security_service.verify_authorization),
|
_=Depends(security_service.verify_authorization),
|
||||||
|
|||||||
Reference in New Issue
Block a user