From 93710c1e78b474295a8b496fa887d9cd440e8486 Mon Sep 17 00:00:00 2001 From: kangfenmao Date: Wed, 25 Sep 2024 09:49:19 +0800 Subject: [PATCH] fix: gemini safety settings #110 --- src/renderer/src/providers/GeminiProvider.ts | 22 ++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/renderer/src/providers/GeminiProvider.ts b/src/renderer/src/providers/GeminiProvider.ts index c28b395c..fa383c9b 100644 --- a/src/renderer/src/providers/GeminiProvider.ts +++ b/src/renderer/src/providers/GeminiProvider.ts @@ -1,4 +1,12 @@ -import { Content, GoogleGenerativeAI, InlineDataPart, Part, TextPart } from '@google/generative-ai' +import { + Content, + GoogleGenerativeAI, + HarmBlockThreshold, + HarmCategory, + InlineDataPart, + Part, + TextPart +} from '@google/generative-ai' import { getAssistantSettings, getDefaultModel, getTopNamingModel } from '@renderer/services/assistant' import { EVENT_NAMES } from '@renderer/services/event' import { filterContextMessages, filterMessages } from '@renderer/services/messages' @@ -68,10 +76,20 @@ export default class GeminiProvider extends BaseProvider { generationConfig: { maxOutputTokens: maxTokens, temperature: assistant?.settings?.temperature - } + }, + safetySettings: [ + { category: HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH }, + { + category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, + threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH + }, + { category: HarmCategory.HARM_CATEGORY_HARASSMENT, threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH }, + { category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH } + ] }) const chat = geminiModel.startChat({ history }) + const messageContents = await this.getMessageContents(userLastMessage!) const userMessagesStream = await chat.sendMessageStream(messageContents.parts)