mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-07 16:56:36 +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,6 +13,76 @@ from app.services.key_manager import KeyManager
|
||||
logger = get_openai_logger()
|
||||
|
||||
|
||||
def _has_image_parts(contents: List[Dict[str, Any]]) -> bool:
|
||||
"""判断消息是否包含图片部分"""
|
||||
for content in contents:
|
||||
if "parts" in content:
|
||||
for part in content["parts"]:
|
||||
if "image_url" in part or "inline_data" in part:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _build_tools(
|
||||
request: ChatRequest, messages: List[Dict[str, Any]]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""构建工具"""
|
||||
tools = []
|
||||
model = request.model
|
||||
|
||||
if (
|
||||
settings.TOOLS_CODE_EXECUTION_ENABLED
|
||||
and not (model.endswith("-search") or "-thinking" in model)
|
||||
and not _has_image_parts(messages)
|
||||
):
|
||||
tools.append({"code_execution": {}})
|
||||
if model.endswith("-search"):
|
||||
tools.append({"googleSearch": {}})
|
||||
return tools
|
||||
|
||||
|
||||
def _get_safety_settings(model: str) -> List[Dict[str, str]]:
|
||||
"""获取安全设置"""
|
||||
# if (
|
||||
# "2.0" in model
|
||||
# and "gemini-2.0-flash-thinking-exp" not in model
|
||||
# and "gemini-2.0-pro-exp" not in model
|
||||
# ):
|
||||
if model == "gemini-2.0-flash-exp":
|
||||
return [
|
||||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"},
|
||||
]
|
||||
return [
|
||||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"},
|
||||
]
|
||||
|
||||
|
||||
def _build_payload(
|
||||
request: ChatRequest, messages: List[Dict[str, Any]]
|
||||
) -> Dict[str, Any]:
|
||||
"""构建请求payload"""
|
||||
return {
|
||||
"contents": messages,
|
||||
"generationConfig": {
|
||||
"temperature": request.temperature,
|
||||
"maxOutputTokens": request.max_tokens,
|
||||
"stopSequences": request.stop,
|
||||
"topP": request.top_p,
|
||||
"topK": request.top_k,
|
||||
},
|
||||
"tools": _build_tools(request, messages),
|
||||
"safetySettings": _get_safety_settings(request.model),
|
||||
}
|
||||
|
||||
|
||||
class OpenAIChatService:
|
||||
"""聊天服务"""
|
||||
|
||||
@@ -32,7 +102,7 @@ class OpenAIChatService:
|
||||
messages = self.message_converter.convert(request.messages)
|
||||
|
||||
# 构建请求payload
|
||||
payload = self._build_payload(request, messages)
|
||||
payload = _build_payload(request, messages)
|
||||
|
||||
if request.stream:
|
||||
return self._handle_stream_completion(request.model, payload, api_key)
|
||||
@@ -84,69 +154,3 @@ class OpenAIChatService:
|
||||
yield f"data: {json.dumps({'error': 'Streaming failed after retries'})}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
break
|
||||
|
||||
def _build_payload(
|
||||
self, request: ChatRequest, messages: List[Dict[str, Any]]
|
||||
) -> Dict[str, Any]:
|
||||
"""构建请求payload"""
|
||||
return {
|
||||
"contents": messages,
|
||||
"generationConfig": {
|
||||
"temperature": request.temperature,
|
||||
"maxOutputTokens": request.max_tokens,
|
||||
"stopSequences": request.stop,
|
||||
"topP": request.top_p,
|
||||
"topK": request.top_k,
|
||||
},
|
||||
"tools": self._build_tools(request, messages),
|
||||
"safetySettings": self._get_safety_settings(request.model),
|
||||
}
|
||||
|
||||
def _build_tools(
|
||||
self, request: ChatRequest, messages: List[Dict[str, Any]]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""构建工具"""
|
||||
tools = []
|
||||
model = request.model
|
||||
|
||||
if (
|
||||
settings.TOOLS_CODE_EXECUTION_ENABLED
|
||||
and not (model.endswith("-search") or "-thinking" in model)
|
||||
and not self._has_image_parts(messages)
|
||||
):
|
||||
tools.append({"code_execution": {}})
|
||||
if model.endswith("-search"):
|
||||
tools.append({"googleSearch": {}})
|
||||
return tools
|
||||
|
||||
def _has_image_parts(self, contents: List[Dict[str, Any]]) -> bool:
|
||||
"""判断消息是否包含图片部分"""
|
||||
for content in contents:
|
||||
if "parts" in content:
|
||||
for part in content["parts"]:
|
||||
if "image_url" in part or "inline_data" in part:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _get_safety_settings(self, model: str) -> List[Dict[str, str]]:
|
||||
"""获取安全设置"""
|
||||
# if (
|
||||
# "2.0" in model
|
||||
# and "gemini-2.0-flash-thinking-exp" not in model
|
||||
# and "gemini-2.0-pro-exp" not in model
|
||||
# ):
|
||||
if model == "gemini-2.0-flash-exp":
|
||||
return [
|
||||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
|
||||
{"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"},
|
||||
]
|
||||
return [
|
||||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"},
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user