mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-04 23:19:17 +08:00
refactor: 重构Gemini和OpenAI聊天服务以支持工具和安全设置
- 将 `_build_payload`、`_build_tools`、`_get_safety_settings` 和 `_has_image_parts` 函数从 `OpenAIChatService` 和 `GeminiChatService` 类中提取为独立的函数。 - 将 `_handle_stream_response` 和 `_handle_normal_response` 函数从 `GeminiResponseHandler` 和 `OpenAIResponseHandler` 类中提取为独立的函数。 - 将 `_extract_text` 函数从 `OpenAIResponseHandler` 类中提取为独立的函数, 并在 `GeminiResponseHandler` 中复用。 - 将 `_convert_image` 函数从 `OpenAIMessageConverter` 类中提取为独立的函数。 - 优化 `OpenAIChatService` 和 `GeminiChatService` 中的代码结构, 使其更清晰。 - 优化 `app/api/openai_routes.py` 和 `app/api/gemini_routes.py` 中的路由函数, 移除不必要的参数。
This commit is contained in:
+13
-10
@@ -1,5 +1,4 @@
|
||||
from http.client import HTTPException
|
||||
from fastapi import APIRouter, Depends
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from app.core.config import settings
|
||||
@@ -10,6 +9,7 @@ from app.services.gemini_chat_service import GeminiChatService
|
||||
from app.services.key_manager import KeyManager
|
||||
from app.services.model_service import ModelService
|
||||
from app.services.chat.retry_handler import RetryHandler
|
||||
|
||||
router = APIRouter(prefix="/gemini/v1beta")
|
||||
router_v1beta = APIRouter(prefix="/v1beta")
|
||||
logger = get_gemini_logger()
|
||||
@@ -22,26 +22,29 @@ model_service = ModelService(settings.MODEL_SEARCH)
|
||||
|
||||
@router.get("/models")
|
||||
@router_v1beta.get("/models")
|
||||
async def list_models(
|
||||
key: str = None,
|
||||
token: str = Depends(security_service.verify_key),
|
||||
):
|
||||
async def list_models(_=Depends(security_service.verify_key)):
|
||||
"""获取可用的Gemini模型列表"""
|
||||
logger.info("-" * 50 + "list_gemini_models" + "-" * 50)
|
||||
logger.info("Handling Gemini models list request")
|
||||
api_key = await key_manager.get_next_working_key()
|
||||
logger.info(f"Using API key: {api_key}")
|
||||
models_json = model_service.get_gemini_models(api_key)
|
||||
models_json["models"].append({"name": "models/gemini-2.0-flash-exp-search", "version": "2.0", "displayName": "Gemini 2.0 Flash Search Experimental", "description": "Gemini 2.0 Flash Search Experimental", "inputTokenLimit": 32767, "outputTokenLimit": 8192, "supportedGenerationMethods": ["generateContent", "countTokens"], "temperature": 1, "topP": 0.95, "topK": 64, "maxTemperature": 2})
|
||||
models_json["models"].append({"name": "models/gemini-2.0-flash-exp-search", "version": "2.0",
|
||||
"displayName": "Gemini 2.0 Flash Search Experimental",
|
||||
"description": "Gemini 2.0 Flash Search Experimental", "inputTokenLimit": 32767,
|
||||
"outputTokenLimit": 8192,
|
||||
"supportedGenerationMethods": ["generateContent", "countTokens"], "temperature": 1,
|
||||
"topP": 0.95, "topK": 64, "maxTemperature": 2})
|
||||
return models_json
|
||||
|
||||
|
||||
@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")
|
||||
async def generate_content(
|
||||
model_name: str,
|
||||
request: GeminiRequest,
|
||||
x_goog_api_key: str = Depends(security_service.verify_goog_api_key),
|
||||
_=Depends(security_service.verify_goog_api_key),
|
||||
api_key: str = Depends(key_manager.get_next_working_key),
|
||||
):
|
||||
chat_service = GeminiChatService(settings.BASE_URL, key_manager)
|
||||
@@ -70,7 +73,7 @@ async def generate_content(
|
||||
async def stream_generate_content(
|
||||
model_name: str,
|
||||
request: GeminiRequest,
|
||||
x_goog_api_key: str = Depends(security_service.verify_goog_api_key),
|
||||
_=Depends(security_service.verify_goog_api_key),
|
||||
api_key: str = Depends(key_manager.get_next_working_key),
|
||||
):
|
||||
chat_service = GeminiChatService(settings.BASE_URL, key_manager)
|
||||
@@ -81,7 +84,7 @@ async def stream_generate_content(
|
||||
logger.info(f"Using API key: {api_key}")
|
||||
|
||||
try:
|
||||
response_stream =chat_service.stream_generate_content(
|
||||
response_stream = chat_service.stream_generate_content(
|
||||
model=model_name,
|
||||
request=request,
|
||||
api_key=api_key
|
||||
|
||||
+12
-20
@@ -1,16 +1,15 @@
|
||||
from http.client import HTTPException
|
||||
from fastapi import APIRouter, Depends, Header
|
||||
from fastapi import HTTPException, APIRouter, Depends
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.logger import get_openai_logger
|
||||
from app.core.security import SecurityService
|
||||
from app.schemas.openai_models import ChatRequest, EmbeddingRequest
|
||||
from app.services.chat.retry_handler import RetryHandler
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
from app.services.key_manager import KeyManager
|
||||
from app.services.model_service import ModelService
|
||||
from app.services.openai_chat_service import OpenAIChatService
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
from app.schemas.openai_models import ChatRequest, EmbeddingRequest
|
||||
from app.core.config import settings
|
||||
from app.core.logger import get_openai_logger
|
||||
|
||||
router = APIRouter()
|
||||
logger = get_openai_logger()
|
||||
@@ -24,10 +23,7 @@ embedding_service = EmbeddingService(settings.BASE_URL)
|
||||
|
||||
@router.get("/v1/models")
|
||||
@router.get("/hf/v1/models")
|
||||
async def list_models(
|
||||
authorization: str = Header(None),
|
||||
token: str = Depends(security_service.verify_authorization),
|
||||
):
|
||||
async def list_models(_=Depends(security_service.verify_authorization)):
|
||||
logger.info("-" * 50 + "list_models" + "-" * 50)
|
||||
logger.info("Handling models list request")
|
||||
api_key = await key_manager.get_next_working_key()
|
||||
@@ -43,10 +39,9 @@ async def list_models(
|
||||
@router.post("/hf/v1/chat/completions")
|
||||
@RetryHandler(max_retries=3, key_manager=key_manager, key_arg="api_key")
|
||||
async def chat_completion(
|
||||
request: ChatRequest,
|
||||
authorization: str = Header(None),
|
||||
token: str = 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(key_manager.get_next_working_key),
|
||||
):
|
||||
chat_service = OpenAIChatService(settings.BASE_URL, key_manager)
|
||||
logger.info("-" * 50 + "chat_completion" + "-" * 50)
|
||||
@@ -67,15 +62,13 @@ async def chat_completion(
|
||||
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/embeddings")
|
||||
@router.post("/hf/v1/embeddings")
|
||||
async def embedding(
|
||||
request: EmbeddingRequest,
|
||||
authorization: str = Header(None),
|
||||
token: str = Depends(security_service.verify_authorization),
|
||||
request: EmbeddingRequest,
|
||||
_=Depends(security_service.verify_authorization),
|
||||
):
|
||||
logger.info("-" * 50 + "embedding" + "-" * 50)
|
||||
logger.info(f"Handling embedding request for model: {request.model}")
|
||||
@@ -95,8 +88,7 @@ async def embedding(
|
||||
@router.get("/v1/keys/list")
|
||||
@router.get("/hf/v1/keys/list")
|
||||
async def get_keys_list(
|
||||
authorization: str = Header(None),
|
||||
token: str = Depends(security_service.verify_auth_token),
|
||||
_=Depends(security_service.verify_auth_token),
|
||||
):
|
||||
"""获取有效和无效的API key列表"""
|
||||
logger.info("-" * 50 + "get_keys_list" + "-" * 50)
|
||||
|
||||
Reference in New Issue
Block a user