fix #3 safty setting

eatgrass · Mar 5, 2024 · d5f0b6c · d5f0b6c
1 parent f53371c
commit d5f0b6c
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 19 deletions.
diff --git a/src/GeminiService.ts b/src/GeminiService.ts
@@ -1,7 +1,7 @@
 import GeminiAssistantPlugin from 'main'
 import { GoogleGenerativeAI } from 'generative-ai.mjs'
-import type { GeminiPrompt } from 'Settings'
 
+import type { GeminiPrompt } from 'Settings'
 export type Model = 'gemini-pro' | 'gemini-pro-vision'
 
 export default class Gemini {
@@ -33,7 +33,9 @@ export default class Gemini {
         if (!option.prompt || option.prompt.length == 0) {
             return
         }
+
         const model = this.genAI.getGenerativeModel({ model: option.model })
+        const threshold = this.plugin.getSettings().saftyThreshold
 
         model.generationConfig = {
             stopSequences: option.config.stopSequences,
@@ -42,6 +44,30 @@ export default class Gemini {
             topP: option.config.topP,
             topK: option.config.topK,
         }
+
+        model.safetySettings = [
+            {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                threshold,
+            },
+            {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                threshold,
+            },
+            {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                threshold,
+            },
+            {
+                category: 'HARM_CATEGORY_HATE_SPEECH',
+                threshold,
+            },
+        ]
+
+        // model.safetySettings = {
+        // 	(HarmCategory as any).HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+        // 	(HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+        // }
         return await model.generateContentStream(option.prompt)
     }
 
@@ -61,8 +87,30 @@ export class GeminiChat {
     private session
 
     constructor(plugin: GeminiAssistantPlugin, model: any) {
+        const threshold = plugin.getSettings().saftyThreshold
         this.model = model
-        this.session = this.model.startChat()
+        const safetySettings = [
+            {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                threshold,
+            },
+            {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                threshold,
+            },
+            {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                threshold,
+            },
+            {
+                category: 'HARM_CATEGORY_HATE_SPEECH',
+                threshold,
+            },
+        ]
+        this.session = this.model.startChat({
+            history: [],
+            safetySettings,
+        })
     }
 
     public async send(msg: string) {

diff --git a/src/Settings.ts b/src/Settings.ts
@@ -35,6 +35,7 @@ export interface Settings {
     model: Model
     prompts: GeminiPrompt[]
     chat: GeminiChat
+    saftyThreshold: string
 }
 
 export const DEFAULT_GEMINI_CONFIGS: Record<Model, GeminiConfig> = {
@@ -82,6 +83,7 @@ export const DEFAULT_SETTINGS: Settings = {
         config: DEFAULT_GEMINI_CONFIGS['gemini-pro'],
         type: 'chat',
     },
+    saftyThreshold: 'HARM_BLOCK_THRESHOLD_UNSPECIFIED',
 }
 
 export default class GeminiSettings extends PluginSettingTab {
@@ -109,17 +111,16 @@ export default class GeminiSettings extends PluginSettingTab {
         const { containerEl } = this
         containerEl.empty()
 
-        new Setting(containerEl).setName('API key').addText((text) => {
+        new Setting(containerEl).setName('API Key').addText((text) => {
             text.setValue(this.settings.apiKey)
             text.onChange((apiKey) => {
                 this.updateSettings({ apiKey })
-				this.plugin.updateApiKey(apiKey)
-
+                this.plugin.updateApiKey(apiKey)
             })
         })
 
         new Setting(containerEl)
-            .setName('Chat setting')
+            .setName('Chat Setting')
             .addExtraButton((button) => {
                 button.setIcon('settings')
                 button.setTooltip('Edit')
@@ -134,6 +135,25 @@ export default class GeminiSettings extends PluginSettingTab {
                 })
             })
 
+        new Setting(containerEl).setHeading().setName('Model')
+
+        new Setting(containerEl)
+            .setName('Harm Content Block Threshold')
+            .addDropdown((dropdown) => {
+                dropdown.addOptions({
+                    HARM_BLOCK_THRESHOLD_UNSPECIFIED: 'Default',
+                    BLOCK_LOW_AND_ABOVE: 'Low',
+                    BLOCK_MEDIUM_AND_ABOVE: 'Medium',
+                    BLOCK_ONLY_HIGH: 'High',
+                    BLOCK_NONE: 'None',
+                })
+                dropdown.setValue(this.settings.saftyThreshold)
+
+                dropdown.onChange((value) => {
+                    this.updateSettings({ saftyThreshold: value })
+                })
+            })
+
         new Setting(containerEl)
             .setHeading()
             .setName('Prompts')
@@ -142,7 +162,7 @@ export default class GeminiSettings extends PluginSettingTab {
                 button.setTooltip('Add prompt')
                 button.onClick(() => {
                     let prompt: GeminiPrompt = {
-                        display: 'Custom prompt',
+                        display: 'Custom Prompt',
                         scope: Scope.SELECTION,
                         model: 'gemini-pro',
                         config: DEFAULT_GEMINI_CONFIGS['gemini-pro'],