Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: update cached user quota after post-consuming (close #204)

Browse Source
This commit is contained in:
JustSong 2023-06-27 19:22:58 +08:00
parent 0941e294bf
commit 737672fb0b
2 changed files with 16 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
controller/relay-text.go
View File

@ -201,6 +201,10 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) 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 { if quota != 0 {
tokenName := c.GetString("token_name") tokenName := c.GetString("token_name")
logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio) logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)

12
model/cache.go
View File

@ -83,6 +83,18 @@ func CacheGetUserQuota(id int) (quota int, err error) {
return quota, err 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 { func CacheIsUserEnabled(userId int) bool {
if !common.RedisEnabled { if !common.RedisEnabled {
return IsUserEnabled(userId) return IsUserEnabled(userId)