From 42908e883400f2c379a5c9e3b278a91d5ef55ffe Mon Sep 17 00:00:00 2001
From: magicdmer <magicdmer@163.com>
Date: Wed, 15 Jan 2025 16:49:30 +0800
Subject: [PATCH] refactor: (GeminiProvider) optimize safety settings handling

- Extract safety threshold logic into getModelSafetySetting method
- gemini-exp-* models not support 'OFF', must use 'BLOCK_NONE'
---
 src/renderer/src/providers/GeminiProvider.ts | 27 ++++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/renderer/src/providers/GeminiProvider.ts b/src/renderer/src/providers/GeminiProvider.ts
index 347a6d8d..31dad53c 100644
--- a/src/renderer/src/providers/GeminiProvider.ts
+++ b/src/renderer/src/providers/GeminiProvider.ts
@@ -112,10 +112,15 @@ export default class GeminiProvider extends BaseProvider {
     }
   }
 
+  private getModelSafetySetting(modelId: string): HarmBlockThreshold {
+    return modelId.includes('gemini-exp-') ? HarmBlockThreshold.BLOCK_NONE : "OFF" as HarmBlockThreshold
+  }
+  
   public async completions({ messages, assistant, onChunk, onFilterMessages }: CompletionsParams) {
     const defaultModel = getDefaultModel()
     const model = assistant.model || defaultModel
     const { contextCount, maxTokens, streamOutput } = getAssistantSettings(assistant)
+    const safetyThreshold = this.getModelSafetySetting(model.id)
 
     const userMessages = filterContextMessages(takeRight(messages, contextCount + 2))
     onFilterMessages(userMessages)
@@ -145,14 +150,26 @@ export default class GeminiProvider extends BaseProvider {
           ...this.getCustomParameters(assistant)
         },
         safetySettings: [
-          { category: HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold: 'OFF' as HarmBlockThreshold },
+          { 
+            category: HarmCategory.HARM_CATEGORY_HATE_SPEECH, 
+            threshold: safetyThreshold
+          },
           {
             category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
-            threshold: 'OFF' as HarmBlockThreshold
+            threshold: safetyThreshold
           },
-          { category: HarmCategory.HARM_CATEGORY_HARASSMENT, threshold: 'OFF' as HarmBlockThreshold },
-          { category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold: 'OFF' as HarmBlockThreshold },
-          { category: 'HARM_CATEGORY_CIVIC_INTEGRITY' as HarmCategory, threshold: 'BLOCK_NONE' as HarmBlockThreshold }
+          { 
+            category: HarmCategory.HARM_CATEGORY_HARASSMENT, 
+            threshold: safetyThreshold
+          },
+          { 
+            category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, 
+            threshold: safetyThreshold
+          },
+          { 
+            category: 'HARM_CATEGORY_CIVIC_INTEGRITY' as HarmCategory, 
+            threshold: safetyThreshold
+          }
         ]
       },
       this.requestOptions