diff --git a/controller/relay-text.go b/controller/relay-text.go index 9ce11601..35b1c543 100644 --- a/controller/relay-text.go +++ b/controller/relay-text.go @@ -201,6 +201,10 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode { if err != nil { common.SysError("error consuming token remain quota: " + err.Error()) } + err = model.CacheUpdateUserQuota(userId) + if err != nil { + common.SysError("error update user quota cache: " + err.Error()) + } if quota != 0 { tokenName := c.GetString("token_name") logContent := fmt.Sprintf("模型倍率 %.2f,分组倍率 %.2f", modelRatio, groupRatio) diff --git a/model/cache.go b/model/cache.go index e25b64f4..284e29ad 100644 --- a/model/cache.go +++ b/model/cache.go @@ -83,6 +83,18 @@ func CacheGetUserQuota(id int) (quota int, err error) { return quota, err } +func CacheUpdateUserQuota(id int) error { + if !common.RedisEnabled { + return nil + } + quota, err := GetUserQuota(id) + if err != nil { + return err + } + err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), UserId2QuotaCacheSeconds*time.Second) + return err +} + func CacheIsUserEnabled(userId int) bool { if !common.RedisEnabled { return IsUserEnabled(userId)