fix: monitor

labring · Dec 22, 2024 · 4db4fa6 · 4db4fa6
1 parent c872d96
commit 4db4fa6
Showing 1 changed file with 30 additions and 7 deletions.
diff --git a/service/aiproxy/controller/relay.go b/service/aiproxy/controller/relay.go
@@ -41,12 +41,31 @@ func relayHelper(meta *meta.Meta, c *gin.Context) *model.ErrorWithStatusCode {
 	}
 }
 
-func RelayHelper(meta *meta.Meta, c *gin.Context) *model.ErrorWithStatusCode {
+func RelayHelper(meta *meta.Meta, c *gin.Context) (*model.ErrorWithStatusCode, bool) {
 	err := relayHelper(meta, c)
-	if err := monitor.AddRequest(c.Request.Context(), meta.OriginModelName, int64(meta.Channel.ID), err != nil); err != nil {
-		log.Errorf("add request failed: %+v", err)
+	if err == nil {
+		if err := monitor.AddRequest(
+			c.Request.Context(),
+			meta.OriginModelName,
+			int64(meta.Channel.ID),
+			false,
+		); err != nil {
+			log.Errorf("add request failed: %+v", err)
+		}
+		return nil, false
+	}
+	if shouldRetry(c, err.StatusCode) {
+		if err := monitor.AddRequest(
+			c.Request.Context(),
+			meta.OriginModelName,
+			int64(meta.Channel.ID),
+			true,
+		); err != nil {
+			log.Errorf("add request failed: %+v", err)
+		}
+		return err, true
 	}
-	return err
+	return nil, false
 }
 
 func getChannelWithFallback(model string, failedChannelIDs ...int) (*dbmodel.Channel, error) {
@@ -88,14 +107,14 @@ func Relay(c *gin.Context) {
 	}
 
 	meta := middleware.NewMetaByContext(c, channel)
-	bizErr := RelayHelper(meta, c)
+	bizErr, retry := RelayHelper(meta, c)
 	if bizErr == nil {
 		return
 	}
 	failedChannelIDs = append(failedChannelIDs, channel.ID)
 	requestID := c.GetString(ctxkey.RequestID)
 	var retryTimes int64
-	if shouldRetry(c, bizErr.StatusCode) {
+	if retry {
 		retryTimes = config.GetRetryTimes()
 	}
 	for i := retryTimes; i > 0; i-- {
@@ -117,10 +136,13 @@ func Relay(c *gin.Context) {
 		}
 		c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
 		meta.Reset(newChannel)
-		bizErr = RelayHelper(meta, c)
+		bizErr, retry = RelayHelper(meta, c)
 		if bizErr == nil {
 			return
 		}
+		if !retry {
+			break
+		}
 		failedChannelIDs = append(failedChannelIDs, newChannel.ID)
 	}
 	if bizErr != nil {
@@ -139,6 +161,7 @@ func Relay(c *gin.Context) {
 	}
 }
 
+// 仅当是channel错误时，才需要重试，用户请求参数错误时，不需要重试
 func shouldRetry(_ *gin.Context, statusCode int) bool {
 	if statusCode == http.StatusTooManyRequests ||
 		statusCode == http.StatusGatewayTimeout ||