From ad77ad8f2b719cbcd2076e9265bfbdc79cd22c9e Mon Sep 17 00:00:00 2001
From: KurolVeko <128911459+Kurolviko@users.noreply.github.com>
Date: Fri, 30 Aug 2024 17:00:51 +0800
Subject: [PATCH] Lower Gemini's safety thresholds

Gemini's default safety thresholds are set too high, resulting in frequent censorship of generated text. I have lowered the thresholds for all four safety categories according to Google's documentation.
---
 bot/gemini/google_gemini_bot.py | 44 ++++++++++++++++++++++++++-------
 1 file changed, 35 insertions(+), 9 deletions(-)

diff --git a/bot/gemini/google_gemini_bot.py b/bot/gemini/google_gemini_bot.py
index 8a4100a..86eb124 100644
--- a/bot/gemini/google_gemini_bot.py
+++ b/bot/gemini/google_gemini_bot.py
@@ -14,6 +14,7 @@ from bridge.reply import Reply, ReplyType
 from common.log import logger
 from config import conf
 from bot.baidu.baidu_wenxin_session import BaiduWenxinSession
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
 
 
 # OpenAI对话模型API (可用)
@@ -38,16 +39,41 @@ class GoogleGeminiBot(Bot):
             gemini_messages = self._convert_to_gemini_messages(self.filter_messages(session.messages))
             genai.configure(api_key=self.api_key)
             model = genai.GenerativeModel(self.model)
-            response = model.generate_content(gemini_messages)
-            reply_text = response.text
-            self.sessions.session_reply(reply_text, session_id)
-            logger.info(f"[Gemini] reply={reply_text}")
-            return Reply(ReplyType.TEXT, reply_text)
+            
+            # 添加安全设置
+            safety_settings = {
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+            }
+            
+            # 生成回复，包含安全设置
+            response = model.generate_content(
+                gemini_messages,
+                safety_settings=safety_settings
+            )
+            if response.candidates and response.candidates[0].content:
+                reply_text = response.candidates[0].content.parts[0].text
+                logger.info(f"[Gemini] reply={reply_text}")
+                self.sessions.session_reply(reply_text, session_id)
+                return Reply(ReplyType.TEXT, reply_text)
+            else:
+                # 没有有效响应内容，可能内容被屏蔽，输出安全评分
+                logger.warning("[Gemini] No valid response generated. Checking safety ratings.")
+                if hasattr(response, 'candidates') and response.candidates:
+                    for rating in response.candidates[0].safety_ratings:
+                        logger.warning(f"Safety rating: {rating.category} - {rating.probability}")
+                error_message = "No valid response generated due to safety constraints."
+                self.sessions.session_reply(error_message, session_id)
+                return Reply(ReplyType.ERROR, error_message)
+                    
         except Exception as e:
-            logger.error("[Gemini] fetch reply error, may contain unsafe content")
-            logger.error(e)
-            return Reply(ReplyType.ERROR, "invoke [Gemini] api failed!")
-
+            logger.error(f"[Gemini] Error generating response: {str(e)}", exc_info=True)
+            error_message = "Failed to invoke [Gemini] api!"
+            self.sessions.session_reply(error_message, session_id)
+            return Reply(ReplyType.ERROR, error_message)
+            
     def _convert_to_gemini_messages(self, messages: list):
         res = []
         for msg in messages: