Update Cohere Safety Setting

Qianxia666 · Sep 1, 2024 · 0b4317c · 0b4317c
1 parent 4b1e83c
commit 0b4317c
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -66,7 +66,7 @@
 - `GET_MEDIA_TOKEN_NOT_STREAM`：是否在非流（`stream=false`）情况下统计图片token，默认为 `true`。
 - `UPDATE_TASK`：是否更新异步任务（Midjourney、Suno），默认为 `true`，关闭后将不会更新任务进度。
 - `GEMINI_MODEL_MAP`：Gemini模型指定版本(v1/v1beta)，使用“模型:版本”指定，","分隔，例如：-e GEMINI_MODEL_MAP="gemini-1.5-pro-latest:v1beta,gemini-1.5-pro-001:v1beta"，为空则使用默认配置
-
+- `COHERE_SAFETY_SETTING`：Cohere模型[安全设置](https://docs.cohere.com/docs/safety-modes#overview)，可选值为 `NONE`, `CONTEXTUAL`，`STRICT`，默认为 `NONE`。
 ## 部署
 ### 部署要求
 - 本地数据库（默认）：SQLite（Docker 部署默认使用 SQLite，必须挂载 `/data` 目录到宿主机）

diff --git a/common/constants.go b/common/constants.go
@@ -112,6 +112,9 @@ var RelayTimeout = GetEnvOrDefault("RELAY_TIMEOUT", 0) // unit is second
 
 var GeminiSafetySetting = GetEnvOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE")
 
+// https://docs.cohere.com/docs/safety-modes Type; NONE/CONTEXTUAL/STRICT
+var CohereSafetySetting = GetEnvOrDefaultString("COHERE_SAFETY_SETTING", "NONE")
+
 const (
 	RequestIdKey = "X-Oneapi-Request-Id"
 )

diff --git a/relay/channel/cohere/dto.go b/relay/channel/cohere/dto.go
@@ -8,6 +8,7 @@ type CohereRequest struct {
 	Message     string        `json:"message"`
 	Stream      bool          `json:"stream"`
 	MaxTokens   int           `json:"max_tokens"`
+	SafetyMode  string        `json:"safety_mode"`
 }
 
 type ChatHistory struct {

diff --git a/relay/channel/cohere/relay-cohere.go b/relay/channel/cohere/relay-cohere.go
@@ -22,6 +22,7 @@ func requestOpenAI2Cohere(textRequest dto.GeneralOpenAIRequest) *CohereRequest {
 		Message:     "",
 		Stream:      textRequest.Stream,
 		MaxTokens:   textRequest.GetMaxTokens(),
+		SafetyMode:  common.CohereSafetySetting,
 	}
 	if cohereReq.MaxTokens == 0 {
 		cohereReq.MaxTokens = 4000
@@ -44,6 +45,7 @@ func requestOpenAI2Cohere(textRequest dto.GeneralOpenAIRequest) *CohereRequest {
 			})
 		}
 	}
+
 	return &cohereReq
 }