Browse Source

fix: 完善令牌预扣费逻辑

CaIon 2 years ago
parent
commit
64e9e9cc20
2 changed files with 4 additions and 4 deletions
  1. 3 2
      controller/relay-text.go
  2. 1 2
      middleware/auth.go

+ 3 - 2
controller/relay-text.go

@@ -52,7 +52,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	tokenId := c.GetInt("token_id")
 	userId := c.GetInt("id")
 	group := c.GetString("group")
-	tokenQuota := c.GetInt("token_quota")
+	tokenUnlimited := c.GetBool("token_unlimited_quota")
 	startTime := time.Now()
 	var textRequest GeneralOpenAIRequest
 
@@ -263,8 +263,9 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	}
 	if userQuota > 100*preConsumedQuota {
 		// 用户额度充足,判断令牌额度是否充足
-		if tokenQuota != -1 {
+		if !tokenUnlimited {
 			// 非无限令牌,判断令牌额度是否充足
+			tokenQuota := c.GetInt("token_quota")
 			if tokenQuota > 100*preConsumedQuota {
 				// 令牌额度充足,信任令牌
 				preConsumedQuota = 0

+ 1 - 2
middleware/auth.go

@@ -115,10 +115,9 @@ func TokenAuth() func(c *gin.Context) {
 		c.Set("id", token.UserId)
 		c.Set("token_id", token.Id)
 		c.Set("token_name", token.Name)
+		c.Set("token_unlimited_quota", token.UnlimitedQuota)
 		if !token.UnlimitedQuota {
 			c.Set("token_quota", token.RemainQuota)
-		} else {
-			c.Set("token_quota", -1)
 		}
 		if token.ModelLimitsEnabled {
 			c.Set("token_model_limit_enabled", true)