From 0b4317ce282b5a761373207e54219ff5213b676a Mon Sep 17 00:00:00 2001
From: HynoR <20227709+HynoR@users.noreply.github.com>
Date: Mon, 2 Sep 2024 06:47:49 +0700
Subject: [PATCH] Update Cohere Safety Setting

---
 README.md                            | 2 +-
 common/constants.go                  | 3 +++
 relay/channel/cohere/dto.go          | 1 +
 relay/channel/cohere/relay-cohere.go | 2 ++
 4 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index b35edb7..8f43a74 100644
--- 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
index 51144fd..3700fce 100644
--- 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
index b2c2739..7f50a15 100644
--- 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
index 6c9c238..adec316 100644
--- 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
 }