From 5abb0a9c4eae2565df2ad326228a54f36666f960 Mon Sep 17 00:00:00 2001
From: "1808837298@qq.com" <1808837298@qq.com>
Date: Tue, 18 Jun 2024 11:03:24 +0800
Subject: [PATCH] fix: auto ban

---
 controller/channel-test.go |  7 ++++++-
 controller/relay.go        |  2 +-
 service/channel.go         | 19 +++++++++++--------
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/controller/channel-test.go b/controller/channel-test.go
index db03e75..469e3c8 100644
--- a/controller/channel-test.go
+++ b/controller/channel-test.go
@@ -219,7 +219,12 @@ func testAllChannels(notify bool) error {
 			if channel.AutoBan != nil && *channel.AutoBan == 0 {
 				ban = false
 			}
-			if isChannelEnabled && service.ShouldDisableChannel(openaiErr, -1) && ban {
+			openAiErrWithStatus := dto.OpenAIErrorWithStatusCode{
+				StatusCode: -1,
+				Error:      *openaiErr,
+				LocalError: false,
+			}
+			if isChannelEnabled && service.ShouldDisableChannel(&openAiErrWithStatus) && ban {
 				service.DisableChannel(channel.Id, channel.Name, err.Error())
 			}
 			if !isChannelEnabled && service.ShouldEnableChannel(err, openaiErr, channel.Status) {
diff --git a/controller/relay.go b/controller/relay.go
index a066e5d..e10aa46 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -128,7 +128,7 @@ func shouldRetry(c *gin.Context, channelId int, openaiErr *dto.OpenAIErrorWithSt
 func processChannelError(c *gin.Context, channelId int, err *dto.OpenAIErrorWithStatusCode) {
 	autoBan := c.GetBool("auto_ban")
 	common.LogError(c.Request.Context(), fmt.Sprintf("relay error (channel #%d, status code: %d): %s", channelId, err.StatusCode, err.Error.Message))
-	if service.ShouldDisableChannel(&err.Error, err.StatusCode) && autoBan {
+	if service.ShouldDisableChannel(err) && autoBan {
 		channelName := c.GetString("channel_name")
 		service.DisableChannel(channelId, channelName, err.Error.Message)
 	}
diff --git a/service/channel.go b/service/channel.go
index 7ee12b2..628605b 100644
--- a/service/channel.go
+++ b/service/channel.go
@@ -24,17 +24,20 @@ func EnableChannel(channelId int, channelName string) {
 	notifyRootUser(subject, content)
 }
 
-func ShouldDisableChannel(err *relaymodel.OpenAIError, statusCode int) bool {
+func ShouldDisableChannel(err *relaymodel.OpenAIErrorWithStatusCode) bool {
 	if !common.AutomaticDisableChannelEnabled {
 		return false
 	}
 	if err == nil {
 		return false
 	}
-	if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden {
+	if err.LocalError {
+		return false
+	}
+	if err.StatusCode == http.StatusUnauthorized || err.StatusCode == http.StatusForbidden {
 		return true
 	}
-	switch err.Code {
+	switch err.Error.Code {
 	case "invalid_api_key":
 		return true
 	case "account_deactivated":
@@ -42,7 +45,7 @@ func ShouldDisableChannel(err *relaymodel.OpenAIError, statusCode int) bool {
 	case "billing_not_active":
 		return true
 	}
-	switch err.Type {
+	switch err.Error.Type {
 	case "insufficient_quota":
 		return true
 	// https://docs.anthropic.com/claude/reference/errors
@@ -53,13 +56,13 @@ func ShouldDisableChannel(err *relaymodel.OpenAIError, statusCode int) bool {
 	case "forbidden":
 		return true
 	}
-	if strings.HasPrefix(err.Message, "Your credit balance is too low") { // anthropic
+	if strings.HasPrefix(err.Error.Message, "Your credit balance is too low") { // anthropic
 		return true
-	} else if strings.HasPrefix(err.Message, "This organization has been disabled.") {
+	} else if strings.HasPrefix(err.Error.Message, "This organization has been disabled.") {
 		return true
-	} else if strings.HasPrefix(err.Message, "You exceeded your current quota") {
+	} else if strings.HasPrefix(err.Error.Message, "You exceeded your current quota") {
 		return true
-	} else if strings.HasPrefix(err.Message, "Permission denied") {
+	} else if strings.HasPrefix(err.Error.Message, "Permission denied") {
 		return true
 	}
 	return false