From 35cfebee121a38b78f3dafc9793f85b068a75f5a Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Sat, 15 Jul 2023 19:06:51 +0800
Subject: [PATCH] feat: retry on failed (close #112)

---
 common/constants.go                    |  1 +
 controller/relay.go                    | 20 +++++++++++++++-----
 model/option.go                        |  3 +++
 web/src/components/OperationSetting.js | 17 ++++++++++++++++-
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/common/constants.go b/common/constants.go
index a97fda0e..4402b9d9 100644
--- a/common/constants.go
+++ b/common/constants.go
@@ -68,6 +68,7 @@ var AutomaticDisableChannelEnabled = false
 var QuotaRemindThreshold = 1000
 var PreConsumedQuota = 500
 var ApproximateTokenEnabled = false
+var RetryTimes = 0
 
 var RootUserEmail = ""
 
diff --git a/controller/relay.go b/controller/relay.go
index 013e56d5..45054945 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"one-api/common"
+	"strconv"
 	"strings"
 
 	"github.com/gin-gonic/gin"
@@ -132,12 +133,21 @@ func Relay(c *gin.Context) {
 		err = relayTextHelper(c, relayMode)
 	}
 	if err != nil {
-		if err.StatusCode == http.StatusTooManyRequests {
-			err.OpenAIError.Message = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
+		retryTimesStr := c.Query("retry")
+		retryTimes, _ := strconv.Atoi(retryTimesStr)
+		if retryTimesStr == "" {
+			retryTimes = common.RetryTimes
+		}
+		if retryTimes > 0 {
+			c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s?retry=%d", c.Request.URL.Path, retryTimes-1))
+		} else {
+			if err.StatusCode == http.StatusTooManyRequests {
+				err.OpenAIError.Message = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
+			}
+			c.JSON(err.StatusCode, gin.H{
+				"error": err.OpenAIError,
+			})
 		}
-		c.JSON(err.StatusCode, gin.H{
-			"error": err.OpenAIError,
-		})
 		channelId := c.GetInt("channel_id")
 		common.SysError(fmt.Sprintf("relay error (channel #%d): %s", channelId, err.Message))
 		// https://platform.openai.com/docs/guides/error-codes/api-errors
diff --git a/model/option.go b/model/option.go
index 35aeec4c..e7bc6806 100644
--- a/model/option.go
+++ b/model/option.go
@@ -68,6 +68,7 @@ func InitOptionMap() {
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMap["ChatLink"] = common.ChatLink
 	common.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(common.QuotaPerUnit, 'f', -1, 64)
+	common.OptionMap["RetryTimes"] = strconv.Itoa(common.RetryTimes)
 	common.OptionMapRWMutex.Unlock()
 	loadOptionsFromDatabase()
 }
@@ -196,6 +197,8 @@ func updateOptionMap(key string, value string) (err error) {
 		common.QuotaRemindThreshold, _ = strconv.Atoi(value)
 	case "PreConsumedQuota":
 		common.PreConsumedQuota, _ = strconv.Atoi(value)
+	case "RetryTimes":
+		common.RetryTimes, _ = strconv.Atoi(value)
 	case "ModelRatio":
 		err = common.UpdateModelRatioByJSONString(value)
 	case "GroupRatio":
diff --git a/web/src/components/OperationSetting.js b/web/src/components/OperationSetting.js
index 69100c85..2adc7fa4 100644
--- a/web/src/components/OperationSetting.js
+++ b/web/src/components/OperationSetting.js
@@ -20,6 +20,7 @@ const OperationSetting = () => {
     DisplayInCurrencyEnabled: '',
     DisplayTokenStatEnabled: '',
     ApproximateTokenEnabled: '',
+    RetryTimes: 0,
   });
   const [originInputs, setOriginInputs] = useState({});
   let [loading, setLoading] = useState(false);
@@ -122,6 +123,9 @@ const OperationSetting = () => {
         if (originInputs['QuotaPerUnit'] !== inputs.QuotaPerUnit) {
           await updateOption('QuotaPerUnit', inputs.QuotaPerUnit);
         }
+        if (originInputs['RetryTimes'] !== inputs.RetryTimes) {
+          await updateOption('RetryTimes', inputs.RetryTimes);
+        }
         break;
     }
   };
@@ -133,7 +137,7 @@ const OperationSetting = () => {
           <Header as='h3'>
             通用设置
           </Header>
-          <Form.Group widths={3}>
+          <Form.Group widths={4}>
             <Form.Input
               label='充值链接'
               name='TopUpLink'
@@ -162,6 +166,17 @@ const OperationSetting = () => {
               step='0.01'
               placeholder='一单位货币能兑换的额度'
             />
+            <Form.Input
+              label='失败重试次数'
+              name='RetryTimes'
+              type={'number'}
+              step='1'
+              min='0'
+              onChange={handleInputChange}
+              autoComplete='new-password'
+              value={inputs.RetryTimes}
+              placeholder='失败重试次数'
+            />
           </Form.Group>
           <Form.Group inline>
             <Form.Checkbox