Просмотр исходного кода

feat: able to change gemini safety setting

CaIon 2 лет назад
Родитель
Сommit
aca8d25372
4 измененных файлов с 28 добавлено и 19 удалено
  1. 2 0
      common/constants.go
  2. 7 0
      common/utils.go
  3. 18 18
      controller/relay-gemini.go
  4. 1 1
      controller/relay.go

+ 2 - 0
common/constants.go

@@ -104,6 +104,8 @@ var BatchUpdateInterval = GetOrDefault("BATCH_UPDATE_INTERVAL", 5)
 
 var RelayTimeout = GetOrDefault("RELAY_TIMEOUT", 0) // unit is second
 
+var GeminiSafetySetting = GetOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE")
+
 const (
 	RequestIdKey = "X-Oneapi-Request-Id"
 )

+ 7 - 0
common/utils.go

@@ -202,6 +202,13 @@ func GetOrDefault(env string, defaultValue int) int {
 	return num
 }
 
+func GetOrDefaultString(env string, defaultValue string) string {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	return os.Getenv(env)
+}
+
 func MessageWithRequestId(message string, id string) string {
 	return fmt.Sprintf("%s (request id: %s)", message, id)
 }

+ 18 - 18
controller/relay-gemini.go

@@ -60,24 +60,24 @@ type GeminiChatGenerationConfig struct {
 func requestOpenAI2Gemini(textRequest GeneralOpenAIRequest) *GeminiChatRequest {
 	geminiRequest := GeminiChatRequest{
 		Contents: make([]GeminiChatContent, 0, len(textRequest.Messages)),
-		//SafetySettings: []GeminiChatSafetySettings{
-		//	{
-		//		Category:  "HARM_CATEGORY_HARASSMENT",
-		//		Threshold: "BLOCK_ONLY_HIGH",
-		//	},
-		//	{
-		//		Category:  "HARM_CATEGORY_HATE_SPEECH",
-		//		Threshold: "BLOCK_ONLY_HIGH",
-		//	},
-		//	{
-		//		Category:  "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-		//		Threshold: "BLOCK_ONLY_HIGH",
-		//	},
-		//	{
-		//		Category:  "HARM_CATEGORY_DANGEROUS_CONTENT",
-		//		Threshold: "BLOCK_ONLY_HIGH",
-		//	},
-		//},
+		SafetySettings: []GeminiChatSafetySettings{
+			{
+				Category:  "HARM_CATEGORY_HARASSMENT",
+				Threshold: common.GeminiSafetySetting,
+			},
+			{
+				Category:  "HARM_CATEGORY_HATE_SPEECH",
+				Threshold: common.GeminiSafetySetting,
+			},
+			{
+				Category:  "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+				Threshold: common.GeminiSafetySetting,
+			},
+			{
+				Category:  "HARM_CATEGORY_DANGEROUS_CONTENT",
+				Threshold: common.GeminiSafetySetting,
+			},
+		},
 		GenerationConfig: GeminiChatGenerationConfig{
 			Temperature:     textRequest.Temperature,
 			TopP:            textRequest.TopP,

+ 1 - 1
controller/relay.go

@@ -332,7 +332,7 @@ func Relay(c *gin.Context) {
 			retryTimes = common.RetryTimes
 		}
 		if retryTimes > 0 {
-			c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s?retry=%d", c.Request.URL.Path, retryTimes-1))
+			c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s?retry=%d&error=%s", c.Request.URL.Path, retryTimes-1, err.Message))
 		} else {
 			if err.StatusCode == http.StatusTooManyRequests {
 				//err.OpenAIError.Message = "当前分组上游负载已饱和,请稍后再试"