feat: custom gemini safety settings issue #1

Author: soulteary

models/gemini/define.go

@@ -6,6 +6,7 @@ const (
 	ENV_GEMINI_MODEL_ALIAS = "GEMINI_MODEL_ALIAS"
 	ENV_GEMINI_API_KEY     = "GEMINI_API_KEY"
 	ENV_GEMINI_MODEL       = "GEMINI_MODEL"
+	ENV_GEMINI_SAFETY      = "GEMINI_SAFETY"
 
 	ENV_GEMINI_HTTP_PROXY  = "GEMINI_HTTP_PROXY"
 	ENV_GEMINI_SOCKS_PROXY = "GEMINI_SOCKS_PROXY"
@@ -17,6 +18,14 @@ const (
 	DEFAULT_REST_API_ENTRYPOINT   = "https://generativelanguage.googleapis.com"
 )
 
+const (
+	DEFAULT_SAFETY_THRESHOLD_NONE   = "BLOCK_NONE"
+	DEFAULT_SAFETY_THRESHOLD_LESS   = "BLOCK_ONLY_HIGH"
+	DEFAULT_SAFETY_THRESHOLD_MEDIUM = "BLOCK_MEDIUM_AND_ABOVE"
+	DEFAULT_SAFETY_THRESHOLD_HIGH   = "BLOCK_LOW_AND_ABOVE"
+	DEFAULT_SAFETY_THRESHOLD_UNSET  = "HARM_BLOCK_THRESHOLD_UNSPECIFIED"
+)
+
 const (
 	DEFAULT_GEMINI_API_VER = DEFAULT_REST_API_VERSION
 	DEFAULT_GEMINI_MODEL   = "gemini-pro"

models/gemini/proxy.go

@@ -94,9 +94,13 @@ func getDirector(req *http.Request, body []byte, c *gin.Context, requestConverte
 
 		// set default safety settings
 		var safetySettings []GeminiSafetySettings
+		safetyThreshold := fn.GetStringOrDefaultFromEnv(ENV_GEMINI_SAFETY, DEFAULT_SAFETY_THRESHOLD_UNSET)
+		if safetyThreshold != DEFAULT_SAFETY_THRESHOLD_NONE && safetyThreshold != DEFAULT_SAFETY_THRESHOLD_UNSET && safetyThreshold != DEFAULT_SAFETY_THRESHOLD_LESS && safetyThreshold != DEFAULT_SAFETY_THRESHOLD_MEDIUM && safetyThreshold != DEFAULT_SAFETY_THRESHOLD_HIGH {
+			safetyThreshold = DEFAULT_SAFETY_THRESHOLD_UNSET
+		}
 		safetySettings = append(safetySettings, GeminiSafetySettings{
 			Category:  "HARM_CATEGORY_DANGEROUS_CONTENT",
-			Threshold: "BLOCK_ONLY_HIGH",
+			Threshold: safetyThreshold,
 		})
 		payload.SafetySettings = safetySettings