erik/mobibot
erik/mobibot
1
1
Fork 0
Code Issues Releases Activity

Added safety settings

Browse source
This commit is contained in:
Erik C. Thauvin 2024-04-17 16:22:37 -07:00
parent 6fca32451f
commit aa15e98d29
Signed by: erik
GPG key ID: 776702A6A2DA330E
1 changed files with 22 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

22
src/main/kotlin/net/thauvin/erik/mobibot/modules/Gemini.kt
View file

@ -33,6 +33,8 @@ package net.thauvin.erik.mobibot.modules
import com.google.cloud.vertexai.VertexAI import com.google.cloud.vertexai.VertexAI
import com.google.cloud.vertexai.api.GenerationConfig import com.google.cloud.vertexai.api.GenerationConfig
import com.google.cloud.vertexai.api.HarmCategory
import com.google.cloud.vertexai.api.SafetySetting
import com.google.cloud.vertexai.generativeai.ChatSession import com.google.cloud.vertexai.generativeai.ChatSession
import com.google.cloud.vertexai.generativeai.GenerativeModel import com.google.cloud.vertexai.generativeai.GenerativeModel
import com.google.cloud.vertexai.generativeai.ResponseHandler import com.google.cloud.vertexai.generativeai.ResponseHandler
@ -41,6 +43,7 @@ import net.thauvin.erik.mobibot.Utils.sendMessage
import org.pircbotx.hooks.types.GenericMessageEvent import org.pircbotx.hooks.types.GenericMessageEvent
import org.slf4j.Logger import org.slf4j.Logger
import org.slf4j.LoggerFactory import org.slf4j.LoggerFactory
import java.util.*
class Gemini : AbstractModule() { class Gemini : AbstractModule() {
@ -109,9 +112,28 @@ class Gemini : AbstractModule() {
try { try {
VertexAI(projectId, location).use { vertexAI -> VertexAI(projectId, location).use { vertexAI ->
val generationConfig = GenerationConfig.newBuilder().setMaxOutputTokens(maxToken).build() val generationConfig = GenerationConfig.newBuilder().setMaxOutputTokens(maxToken).build()
val safetySettings = Arrays.asList(
SafetySetting.newBuilder()
.setCategory(HarmCategory.HARM_CATEGORY_HATE_SPEECH)
.setThreshold(SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE)
.build(),
SafetySetting.newBuilder()
.setCategory(HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT)
.setThreshold(SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE)
.build(),
SafetySetting.newBuilder()
.setCategory(HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT)
.setThreshold(SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE)
.build(),
SafetySetting.newBuilder()
.setCategory(HarmCategory.HARM_CATEGORY_HARASSMENT)
.setThreshold(SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE)
.build()
)
val model = GenerativeModel.Builder().setModelName("gemini-pro-vision") val model = GenerativeModel.Builder().setModelName("gemini-pro-vision")
.setGenerationConfig(generationConfig) .setGenerationConfig(generationConfig)
.setVertexAi(vertexAI).build() .setVertexAi(vertexAI).build()
.withSafetySettings(safetySettings)
val session = ChatSession(model) val session = ChatSession(model)
val response = session.sendMessage(query) val response = session.sendMessage(query)