feat: 添加Web验证页面并优化密钥管理功能

app/api/gemini_routes.py

@@ -6,7 +6,7 @@ from app.core.logger import get_gemini_logger
 from app.core.security import SecurityService
 from app.schemas.gemini_models import GeminiRequest
 from app.services.gemini_chat_service import GeminiChatService
-from app.services.key_manager import KeyManager
+from app.services.key_manager import KeyManager, get_key_manager_instance
 from app.services.model_service import ModelService
 from app.services.chat.retry_handler import RetryHandler
 
@@ -16,13 +16,20 @@ logger = get_gemini_logger()
 
 # 初始化服务
 security_service = SecurityService(settings.ALLOWED_TOKENS, settings.AUTH_TOKEN)
-key_manager = KeyManager(settings.API_KEYS)
+
+async def get_key_manager():
+    return await get_key_manager_instance()
+
+async def get_next_working_key_wrapper(key_manager: KeyManager = Depends(get_key_manager)):
+    return await key_manager.get_next_working_key()
+
 model_service = ModelService(settings.MODEL_SEARCH)
 
 
 @router.get("/models")
 @router_v1beta.get("/models")
-async def list_models(_=Depends(security_service.verify_key)):
+async def list_models(_=Depends(security_service.verify_key),
+    key_manager: KeyManager = Depends(get_key_manager)):
     """获取可用的Gemini模型列表"""
     logger.info("-" * 50 + "list_gemini_models" + "-" * 50)
     logger.info("Handling Gemini models list request")
@@ -40,12 +47,13 @@ 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=key_manager, key_arg="api_key")
+@RetryHandler(max_retries=3, key_manager=Depends(get_key_manager), key_arg="api_key")
 async def generate_content(
         model_name: str,
         request: GeminiRequest,
         _=Depends(security_service.verify_goog_api_key),
-        api_key: str = Depends(key_manager.get_next_working_key),
+        api_key: str = Depends(get_next_working_key_wrapper),
+        key_manager: KeyManager = Depends(get_key_manager)
 ):
     chat_service = GeminiChatService(settings.BASE_URL, key_manager)
     """非流式生成内容"""
@@ -69,12 +77,13 @@ async def generate_content(
 
 @router.post("/models/{model_name}:streamGenerateContent")
 @router_v1beta.post("/models/{model_name}:streamGenerateContent")
-@RetryHandler(max_retries=3, key_manager=key_manager, key_arg="api_key")
+@RetryHandler(max_retries=3, key_manager=Depends(get_key_manager), key_arg="api_key")
 async def stream_generate_content(
         model_name: str,
         request: GeminiRequest,
         _=Depends(security_service.verify_goog_api_key),
-        api_key: str = Depends(key_manager.get_next_working_key),
+        api_key: str = Depends(get_next_working_key_wrapper),
+        key_manager: KeyManager = Depends(get_key_manager)
 ):
     chat_service = GeminiChatService(settings.BASE_URL, key_manager)
     """流式生成内容"""

app/api/openai_routes.py

@@ -8,7 +8,7 @@ from app.schemas.openai_models import ChatRequest, EmbeddingRequest, ImageGenera
 from app.services.chat.retry_handler import RetryHandler
 from app.services.embedding_service import EmbeddingService
 from app.services.image_create_service import ImageCreateService
-from app.services.key_manager import KeyManager
+from app.services.key_manager import KeyManager, get_key_manager_instance
 from app.services.model_service import ModelService
 from app.services.openai_chat_service import OpenAIChatService
 
@@ -17,15 +17,22 @@ logger = get_openai_logger()
 
 # 初始化服务
 security_service = SecurityService(settings.ALLOWED_TOKENS, settings.AUTH_TOKEN)
-key_manager = KeyManager(settings.API_KEYS)
 model_service = ModelService(settings.MODEL_SEARCH)
 embedding_service = EmbeddingService(settings.BASE_URL)
 image_create_service = ImageCreateService()
 
+async def get_key_manager():
+    return await get_key_manager_instance()
+
+async def get_next_working_key_wrapper(key_manager: KeyManager = Depends(get_key_manager)):
+    return await key_manager.get_next_working_key()
 
 @router.get("/v1/models")
 @router.get("/hf/v1/models")
-async def list_models(_=Depends(security_service.verify_authorization)):
+async def list_models(
+    _=Depends(security_service.verify_authorization),
+    key_manager: KeyManager = Depends(get_key_manager)
+):
     logger.info("-" * 50 + "list_models" + "-" * 50)
     logger.info("Handling models list request")
     api_key = await key_manager.get_next_working_key()
@@ -39,11 +46,12 @@ async def list_models(_=Depends(security_service.verify_authorization)):
 
 @router.post("/v1/chat/completions")
 @router.post("/hf/v1/chat/completions")
-@RetryHandler(max_retries=3, key_manager=key_manager, key_arg="api_key")
+@RetryHandler(max_retries=3, key_manager=Depends(get_key_manager), key_arg="api_key")
 async def chat_completion(
-        request: ChatRequest,
-        _=Depends(security_service.verify_authorization),
-        api_key: str = Depends(key_manager.get_next_working_key),
+    request: ChatRequest,
+    _=Depends(security_service.verify_authorization),
+    api_key: str = Depends(get_next_working_key_wrapper),
+    key_manager: KeyManager = Depends(get_key_manager)
 ):
     # 如果model是imagen3,使用paid_key
     if request.model == f"{settings.CREATE_IMAGE_MODEL}-chat":
@@ -58,23 +66,21 @@ async def chat_completion(
         if request.model == f"{settings.CREATE_IMAGE_MODEL}-chat":
             response = await chat_service.create_image_chat_completion(request=request)
         else:
-            response = await chat_service.create_chat_completion(request,api_key)
+            response = await chat_service.create_chat_completion(request, api_key)
         # 处理流式响应
         if request.stream:
             return StreamingResponse(response, media_type="text/event-stream")
         logger.info("Chat completion request successful")
         return response
-
     except Exception as e:
         logger.error(f"Chat completion failed after retries: {str(e)}")
         raise HTTPException(status_code=500, detail="Chat completion failed") from e
 
-
 @router.post("/v1/images/generations")
 @router.post("/hf/v1/images/generations")
 async def generate_image(
-        request: ImageGenerationRequest,
-        _=Depends(security_service.verify_authorization),
+    request: ImageGenerationRequest,
+    _=Depends(security_service.verify_authorization),
 ):
     logger.info("-" * 50 + "generate_image" + "-" * 50)
     logger.info(f"Handling image generation request for prompt: {request.prompt}")
@@ -83,17 +89,16 @@ async def generate_image(
         response = image_create_service.generate_images(request)
         logger.info("Image generation request successful")
         return response
-
     except Exception as e:
         logger.error(f"Image generation request failed: {str(e)}")
         raise HTTPException(status_code=500, detail="Image generation request failed") from e
 
-
 @router.post("/v1/embeddings")
 @router.post("/hf/v1/embeddings")
 async def embedding(
-        request: EmbeddingRequest,
-        _=Depends(security_service.verify_authorization),
+    request: EmbeddingRequest,
+    _=Depends(security_service.verify_authorization),
+    key_manager: KeyManager = Depends(get_key_manager)
 ):
     logger.info("-" * 50 + "embedding" + "-" * 50)
     logger.info(f"Handling embedding request for model: {request.model}")
@@ -109,11 +114,11 @@ async def embedding(
         logger.error(f"Embedding request failed: {str(e)}")
         raise HTTPException(status_code=500, detail="Embedding request failed") from e
 
-
 @router.get("/v1/keys/list")
 @router.get("/hf/v1/keys/list")
 async def get_keys_list(
-        _=Depends(security_service.verify_auth_token),
+    _=Depends(security_service.verify_auth_token),
+    key_manager: KeyManager = Depends(get_key_manager)
 ):
     """获取有效和无效的API key列表"""
     logger.info("-" * 50 + "get_keys_list" + "-" * 50)