refactor: 修改安全设置逻辑以匹配特定模型

app/services/gemini_chat_service.py

@@ -81,7 +81,7 @@ class GeminiChatService:
         
     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:
+        if model == "gemini-2.0-flash-exp":
             return [
                 {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
                 {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},

app/services/openai_chat_service.py

@@ -11,6 +11,8 @@ from app.core.config import settings
 from app.services.key_manager import KeyManager
 
 logger = get_openai_logger()
+
+
 class OpenAIChatService:
     """聊天服务"""
 
@@ -37,59 +39,55 @@ class OpenAIChatService:
         return self._handle_normal_completion(request.model, payload, api_key)
 
     def _handle_normal_completion(
-        self,
-        model: str,
-        payload: Dict[str, Any],
-        api_key: str
+        self, model: str, payload: Dict[str, Any], api_key: str
     ) -> Dict[str, Any]:
         """处理普通聊天完成"""
         response = self.api_client.generate_content(payload, model, api_key)
         return self.response_handler.handle_response(
-            response,
-            model,
-            stream=False,
-            finish_reason="stop"
+            response, model, stream=False, finish_reason="stop"
         )
 
     async def _handle_stream_completion(
-        self,
-        model: str,
-        payload: Dict[str, Any],
-        api_key: str
+        self, model: str, payload: Dict[str, Any], api_key: str
     ) -> AsyncGenerator[str, None]:
         """处理流式聊天完成，添加重试逻辑"""
         retries = 0
         max_retries = 3
         while retries < max_retries:
             try:
-                async for line in self.api_client.stream_generate_content(payload, model, api_key):
+                async for line in self.api_client.stream_generate_content(
+                    payload, model, api_key
+                ):
                     # print(line)
                     if line.startswith("data:"):
                         chunk = json.loads(line[6:])
                         openai_chunk = self.response_handler.handle_response(
-                            chunk,
-                            model,
-                            stream=True,
-                            finish_reason=None
+                            chunk, model, stream=True, finish_reason=None
                         )
                         if openai_chunk:
                             yield f"data: {json.dumps(openai_chunk)}\n\n"
                 yield f"data: {json.dumps(self.response_handler.handle_response({}, model, stream=True, finish_reason='stop'))}\n\n"
                 yield "data: [DONE]\n\n"
                 logger.info("Streaming completed successfully")
-                break # 成功后退出循环
+                break  # 成功后退出循环
             except Exception as e:
                 retries += 1
-                logger.warning(f"Streaming API call failed with error: {str(e)}. Attempt {retries} of {max_retries}")
+                logger.warning(
+                    f"Streaming API call failed with error: {str(e)}. Attempt {retries} of {max_retries}"
+                )
                 api_key = await self.key_manager.handle_api_failure(api_key)
                 logger.info(f"Switched to new API key: {api_key}")
                 if retries >= max_retries:
-                    logger.error(f"Max retries ({max_retries}) reached for streaming. Raising error")
+                    logger.error(
+                        f"Max retries ({max_retries}) reached for streaming. Raising error"
+                    )
                     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]:
+    def _build_payload(
+        self, request: ChatRequest, messages: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
         """构建请求payload"""
         return {
             "contents": messages,
@@ -98,20 +96,24 @@ class OpenAIChatService:
                 "maxOutputTokens": request.max_tokens,
                 "stopSequences": request.stop,
                 "topP": request.top_p,
-                "topK": request.top_k
+                "topK": request.top_k,
             },
             "tools": self._build_tools(request, messages),
-            "safetySettings": self._get_safety_settings(request.model)
+            "safetySettings": self._get_safety_settings(request.model),
         }
 
-    def _build_tools(self, request: ChatRequest, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    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):
+        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": {}})
@@ -128,18 +130,23 @@ class OpenAIChatService:
 
     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:
+        # 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"}
+                {"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"}
+            {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"},
         ]