From d168a685c15c4b2a349f00591e860cc1ac6f3c15 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Wed, 11 Sep 2024 19:12:32 +0800
Subject: [PATCH] fix: cohere SafetyMode

---
 relay/channel/cohere/dto.go          | 2 +-
 relay/channel/cohere/relay-cohere.go | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/relay/channel/cohere/dto.go b/relay/channel/cohere/dto.go
index 7f50a15..e7452fd 100644
--- a/relay/channel/cohere/dto.go
+++ b/relay/channel/cohere/dto.go
@@ -8,7 +8,7 @@ type CohereRequest struct {
 	Message     string        `json:"message"`
 	Stream      bool          `json:"stream"`
 	MaxTokens   int           `json:"max_tokens"`
-	SafetyMode  string        `json:"safety_mode"`
+	SafetyMode  string        `json:"safety_mode,omitempty"`
 }
 
 type ChatHistory struct {
diff --git a/relay/channel/cohere/relay-cohere.go b/relay/channel/cohere/relay-cohere.go
index adec316..132039b 100644
--- a/relay/channel/cohere/relay-cohere.go
+++ b/relay/channel/cohere/relay-cohere.go
@@ -22,7 +22,9 @@ func requestOpenAI2Cohere(textRequest dto.GeneralOpenAIRequest) *CohereRequest {
 		Message:     "",
 		Stream:      textRequest.Stream,
 		MaxTokens:   textRequest.GetMaxTokens(),
-		SafetyMode:  common.CohereSafetySetting,
+	}
+	if common.CohereSafetySetting != "NONE" {
+		cohereReq.SafetyMode = common.CohereSafetySetting
 	}
 	if cohereReq.MaxTokens == 0 {
 		cohereReq.MaxTokens = 4000