From 0a9de391431fab3dab2cd974280661373645b9d0 Mon Sep 17 00:00:00 2001
From: yinpeng <2291314224@qq.com>
Date: Sat, 18 Jan 2025 23:22:39 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=A0=B9=E6=8D=AE=E6=A8=A1=E5=9E=8B?=
 =?UTF-8?q?=E7=89=88=E6=9C=AC=E8=B0=83=E6=95=B4=E5=AE=89=E5=85=A8=E8=AE=BE?=
 =?UTF-8?q?=E7=BD=AE=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/services/gemini_chat_service.py | 14 +++++++++++---
 app/services/openai_chat_service.py | 14 +++++++++++---
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/app/services/gemini_chat_service.py b/app/services/gemini_chat_service.py
index 0ee8772..58df814 100644
--- a/app/services/gemini_chat_service.py
+++ b/app/services/gemini_chat_service.py
@@ -53,7 +53,7 @@ class GeminiChatService:
         return {
             "contents": payload.get("contents", []),
             "tools": self._build_tools(model, payload),
-            "safetySettings": self._get_safety_settings(),
+            "safetySettings": self._get_safety_settings(model),
             "generationConfig": payload.get("generationConfig", {}),
             "systemInstruction": payload.get("systemInstruction", [])
         }
@@ -78,12 +78,20 @@ class GeminiChatService:
                         return True
         return False
         
-    def _get_safety_settings(self) -> List[Dict[str, str]]:
+    def _get_safety_settings(self, model: str) -> List[Dict[str, str]]:
         """获取安全设置"""
+        if "2.0" in model and model != "gemini-2.0-flash-thinking-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": "BLOCK_ONLY_HIGH"}
+            ]
         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_ONLY_HIGH"}
         ]
\ No newline at end of file
diff --git a/app/services/openai_chat_service.py b/app/services/openai_chat_service.py
index ca67bae..a1a6cb5 100644
--- a/app/services/openai_chat_service.py
+++ b/app/services/openai_chat_service.py
@@ -100,7 +100,7 @@ class OpenAIChatService:
                 "topK": request.top_k
             },
             "tools": self._build_tools(request, messages),
-            "safetySettings": self._get_safety_settings()
+            "safetySettings": self._get_safety_settings(request.model)
         }
     
     def _build_tools(self, request: ChatRequest, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
@@ -125,12 +125,20 @@ class OpenAIChatService:
                         return True
         return False
         
-    def _get_safety_settings(self) -> List[Dict[str, str]]:
+    def _get_safety_settings(self, model: str) -> List[Dict[str, str]]:
         """获取安全设置"""
+        if "2.0" in model and model != "gemini-2.0-flash-thinking-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": "BLOCK_ONLY_HIGH"}
+            ]
         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_ONLY_HIGH"}
         ]
\ No newline at end of file