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

fix: 修复高并发下，高额度用户使用低额度令牌没有预扣费而导致令牌大额欠费

Browse Source
This commit is contained in:
CaIon 2024-01-11 14:53:55 +08:00
parent a1e3c9a710
commit f0ffe14437
2 changed files with 19 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
controller/relay-text.go
View File

@ -53,6 +53,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
tokenId := c.GetInt("token_id") tokenId := c.GetInt("token_id")
userId := c.GetInt("id") userId := c.GetInt("id")
group := c.GetString("group") group := c.GetString("group")
tokenUnlimited := c.GetBool("token_unlimited_quota")
var textRequest GeneralOpenAIRequest var textRequest GeneralOpenAIRequest
err := common.UnmarshalBodyReusable(c, &textRequest) err := common.UnmarshalBodyReusable(c, &textRequest)
if err != nil { if err != nil {
@ -240,15 +241,20 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError)
} }
if userQuota > 100*preConsumedQuota { if userQuota > 100*preConsumedQuota {
// in this case, we do not pre-consume quota // 用户额度充足，判断令牌额度是否充足
// because the user has enough quota if !tokenUnlimited {
preConsumedQuota = 0 // 非无限令牌，判断令牌额度是否充足
common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota)) tokenQuota := c.GetInt("token_quota")
} if tokenQuota > 100*preConsumedQuota {
if preConsumedQuota > 0 { // 令牌额度充足，信任令牌
err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota) preConsumedQuota = 0
if err != nil { common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d quota %d and token %d quota %d are enough, trusted and no need to pre-consume", userId, userQuota, tokenId, tokenQuota))
return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden) }
} else {
// in this case, we do not pre-consume quota
// because the user has enough quota
preConsumedQuota = 0
common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d with unlimited token has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
} }
} }
var requestBody io.Reader var requestBody io.Reader

4
middleware/auth.go
View File

@ -106,6 +106,10 @@ func TokenAuth() func(c *gin.Context) {
c.Set("id", token.UserId) c.Set("id", token.UserId)
c.Set("token_id", token.Id) c.Set("token_id", token.Id)
c.Set("token_name", token.Name) c.Set("token_name", token.Name)
c.Set("token_unlimited_quota", token.UnlimitedQuota)
if !token.UnlimitedQuota {
c.Set("token_quota", token.RemainQuota)
}
if len(parts) > 1 { if len(parts) > 1 {
if model.IsAdmin(token.UserId) { if model.IsAdmin(token.UserId) {
c.Set("channelId", parts[1]) c.Set("channelId", parts[1])