feat: 添加成功时自动启用通道功能, close #27

Calcium-Ion · Feb 1, 2024 · 5c8f8b4 · 5c8f8b4
1 parent affe511
commit 5c8f8b4
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 9 deletions.
diff --git a/common/constants.go b/common/constants.go
@@ -87,6 +87,7 @@ var QuotaForInviter = 0
 var QuotaForInvitee = 0
 var ChannelDisableThreshold = 5.0
 var AutomaticDisableChannelEnabled = false
+var AutomaticEnableChannelEnabled = false
 var QuotaRemindThreshold = 1000
 var PreConsumedQuota = 500
 

diff --git a/controller/channel-test.go b/controller/channel-test.go
@@ -78,6 +78,9 @@ func testChannel(channel *model.Channel, request ChatRequest) (err error, openai
 		return err, nil
 	}
 	if response.Usage.CompletionTokens == 0 {
+		if response.Error.Message == "" {
+			response.Error.Message = "补全 tokens 非预期返回 0"
+		}
 		return errors.New(fmt.Sprintf("type %s, code %v, message %s", response.Error.Type, response.Error.Code, response.Error.Message)), &response.Error
 	}
 	return nil, nil
@@ -146,12 +149,23 @@ var testAllChannelsRunning bool = false
 
 // disable & notify
 func disableChannel(channelId int, channelName string, reason string) {
-	if common.RootUserEmail == "" {
-		common.RootUserEmail = model.GetRootUserEmail()
-	}
 	model.UpdateChannelStatusById(channelId, common.ChannelStatusAutoDisabled)
 	subject := fmt.Sprintf("通道「%s」（#%d）已被禁用", channelName, channelId)
 	content := fmt.Sprintf("通道「%s」（#%d）已被禁用，原因：%s", channelName, channelId, reason)
+	notifyRootUser(subject, content)
+}
+
+func enableChannel(channelId int, channelName string) {
+	model.UpdateChannelStatusById(channelId, common.ChannelStatusEnabled)
+	subject := fmt.Sprintf("通道「%s」（#%d）已被启用", channelName, channelId)
+	content := fmt.Sprintf("通道「%s」（#%d）已被启用", channelName, channelId)
+	notifyRootUser(subject, content)
+}
+
+func notifyRootUser(subject string, content string) {
+	if common.RootUserEmail == "" {
+		common.RootUserEmail = model.GetRootUserEmail()
+	}
 	err := common.SendEmail(subject, common.RootUserEmail, content)
 	if err != nil {
 		common.SysError(fmt.Sprintf("failed to send email: %s", err.Error()))
@@ -180,9 +194,7 @@ func testAllChannels(notify bool) error {
 	}
 	go func() {
 		for _, channel := range channels {
-			if channel.Status != common.ChannelStatusEnabled {
-				continue
-			}
+			isChannelEnabled := channel.Status == common.ChannelStatusEnabled
 			tik := time.Now()
 			err, openaiErr := testChannel(channel, *testRequest)
 			tok := time.Now()
@@ -201,9 +213,12 @@ func testAllChannels(notify bool) error {
 			if channel.AutoBan != nil && *channel.AutoBan == 0 {
 				ban = false
 			}
-			if shouldDisableChannel(openaiErr, -1) && ban {
+			if isChannelEnabled && shouldDisableChannel(openaiErr, -1) && ban {
 				disableChannel(channel.Id, channel.Name, err.Error())
 			}
+			if !isChannelEnabled && shouldEnableChannel(err, openaiErr) {
+				enableChannel(channel.Id, channel.Name)
+			}
 			channel.UpdateResponseTime(milliseconds)
 			time.Sleep(common.RequestInterval)
 		}

diff --git a/controller/relay-utils.go b/controller/relay-utils.go
@@ -258,6 +258,19 @@ func shouldDisableChannel(err *OpenAIError, statusCode int) bool {
 	return false
 }
 
+func shouldEnableChannel(err error, openAIErr *OpenAIError) bool {
+	if !common.AutomaticEnableChannelEnabled {
+		return false
+	}
+	if err != nil {
+		return false
+	}
+	if openAIErr != nil {
+		return false
+	}
+	return true
+}
+
 func setEventStreamHeaders(c *gin.Context) {
 	c.Writer.Header().Set("Content-Type", "text/event-stream")
 	c.Writer.Header().Set("Cache-Control", "no-cache")

diff --git a/model/option.go b/model/option.go
@@ -34,6 +34,7 @@ func InitOptionMap() {
 	common.OptionMap["TurnstileCheckEnabled"] = strconv.FormatBool(common.TurnstileCheckEnabled)
 	common.OptionMap["RegisterEnabled"] = strconv.FormatBool(common.RegisterEnabled)
 	common.OptionMap["AutomaticDisableChannelEnabled"] = strconv.FormatBool(common.AutomaticDisableChannelEnabled)
+	common.OptionMap["AutomaticEnableChannelEnabled"] = strconv.FormatBool(common.AutomaticEnableChannelEnabled)
 	common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled)
 	common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled)
 	common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled)
@@ -158,6 +159,8 @@ func updateOptionMap(key string, value string) (err error) {
 			common.EmailDomainRestrictionEnabled = boolValue
 		case "AutomaticDisableChannelEnabled":
 			common.AutomaticDisableChannelEnabled = boolValue
+		case "AutomaticEnableChannelEnabled":
+			common.AutomaticEnableChannelEnabled = boolValue
 		case "LogConsumeEnabled":
 			common.LogConsumeEnabled = boolValue
 		case "DisplayInCurrencyEnabled":

diff --git a/web/src/components/ChannelsTable.js b/web/src/components/ChannelsTable.js
@@ -478,7 +478,7 @@ const ChannelsTable = () => {
         const res = await API.get(`/api/channel/test`);
         const {success, message} = res.data;
         if (success) {
-            showInfo('已成功开始测试所有已启用通道，请刷新页面查看结果。');
+            showInfo('已成功开始测试所有通道，请刷新页面查看结果。');
         } else {
             showError(message);
         }
@@ -702,7 +702,7 @@ const ChannelsTable = () => {
                         onConfirm={testAllChannels}
                         position={isMobile()?'top':'top'}
                     >
-                        <Button theme='light' type='warning' style={{marginRight: 8}}>测试所有已启用通道</Button>
+                        <Button theme='light' type='warning' style={{marginRight: 8}}>测试所有通道</Button>
                     </Popconfirm>
                     <Popconfirm
                         title="确定？"

diff --git a/web/src/components/OperationSetting.js b/web/src/components/OperationSetting.js
@@ -18,6 +18,7 @@ const OperationSetting = () => {
         ChatLink2: '', // 添加的新状态变量
         QuotaPerUnit: 0,
         AutomaticDisableChannelEnabled: '',
+        AutomaticEnableChannelEnabled: '',
         ChannelDisableThreshold: 0,
         LogConsumeEnabled: '',
         DisplayInCurrencyEnabled: '',
@@ -332,6 +333,12 @@ const OperationSetting = () => {
                             name='AutomaticDisableChannelEnabled'
                             onChange={handleInputChange}
                         />
+                        <Form.Checkbox
+                            checked={inputs.AutomaticEnableChannelEnabled === 'true'}
+                            label='成功时自动启用通道'
+                            name='AutomaticEnableChannelEnabled'
+                            onChange={handleInputChange}
+                        />
                     </Form.Group>
                     <Form.Button onClick={() => {
                         submitConfig('monitor').then();