fix: recore req to redis

labring · Jan 6, 2025 · 07c4cb5 · 07c4cb5
1 parent e560a0f
commit 07c4cb5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/service/aiproxy/common/rpmlimit/rate-limit.go b/service/aiproxy/common/rpmlimit/rate-limit.go
@@ -137,8 +137,15 @@ func GetRPM(ctx context.Context, group, model string) (int64, error) {
 }
 
 func redisRateLimitRequest(ctx context.Context, group, model string, maxRequestNum int64, duration time.Duration) (bool, error) {
-	rdb := common.RDB
-	result, err := rdb.Eval(
+	result, err := PushRequest(ctx, group, model, duration)
+	if err != nil {
+		return false, err
+	}
+	return result <= maxRequestNum, nil
+}
+
+func PushRequest(ctx context.Context, group, model string, duration time.Duration) (int64, error) {
+	result, err := common.RDB.Eval(
 		ctx,
 		pushRequestScript,
 		[]string{
@@ -148,9 +155,9 @@ func redisRateLimitRequest(ctx context.Context, group, model string, maxRequestN
 		time.Now().UnixMilli(),
 	).Int64()
 	if err != nil {
-		return false, err
+		return 0, err
 	}
-	return result <= maxRequestNum, nil
+	return result, nil
 }
 
 func RateLimit(ctx context.Context, group, model string, maxRequestNum int64, duration time.Duration) (bool, error) {

diff --git a/service/aiproxy/middleware/distributor.go b/service/aiproxy/middleware/distributor.go
@@ -88,6 +88,11 @@ func checkGroupModelRPMAndTPM(c *gin.Context, group *model.GroupCache, mc *model
 		if !ok {
 			return fmt.Errorf("group (%s) is requesting too frequently", group.ID)
 		}
+	} else if common.RedisEnabled {
+		_, err := rpmlimit.PushRequest(c.Request.Context(), group.ID, mc.Model, time.Minute)
+		if err != nil {
+			log.Errorf("push request error: %s", err.Error())
+		}
 	}
 
 	if adjustedModelConfig.TPM > 0 {