Merge branch 'main' into latest

.github/workflows/docker-image.yml

-name: Docker Image CI
-
-on:
-  push:
-    branches: [ "main" ]
-  pull_request:
-    branches: [ "main" ]
-
-jobs:
-
-  build:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v3
-    - uses: docker/login-action@v3.0.0
-      with:
-        username: ${{ secrets.DOCKERHUB_USERNAME }}
-        password: ${{ secrets.DOCKERHUB_TOKEN }}
-    - name: Extract metadata (tags, labels) for Docker
-      id: meta
-      uses: docker/metadata-action@v3
-      with:
-        images: calciumion/neko-api
-    - name: Build the Docker image
-      uses: docker/build-push-action@v5.0.0
-      with:
-          context: .
-          push: true
-          tags: ${{ steps.meta.outputs.tags }}
-          labels: ${{ steps.meta.outputs.labels }}

controller/relay-audio.go

@@ -44,7 +44,7 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		preConsumedQuota = 0
 		preConsumedQuota = 0
 	}
 	}
 	if preConsumedQuota > 0 {
 	if preConsumedQuota > 0 {
-		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
+		userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
 		if err != nil {
 		if err != nil {
 			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
 		}
@@ -99,7 +99,7 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		go func() {
 		go func() {
 			quota := countTokenText(audioResponse.Text, audioModel)
 			quota := countTokenText(audioResponse.Text, audioModel)
 			quotaDelta := quota - preConsumedQuota
 			quotaDelta := quota - preConsumedQuota
-			err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
+			err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
 			if err != nil {
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
 			}

controller/relay-image.go

@@ -147,7 +147,7 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	var textResponse ImageResponse
 	var textResponse ImageResponse
 	defer func(ctx context.Context) {
 	defer func(ctx context.Context) {
 		if consumeQuota {
 		if consumeQuota {
-			err := model.PostConsumeTokenQuota(tokenId, quota)
+			err := model.PostConsumeTokenQuota(tokenId, userId, quota, 0)
 			if err != nil {
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
 			}

controller/relay-mj.go

@@ -359,7 +359,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 
 
 	defer func(ctx context.Context) {
 	defer func(ctx context.Context) {
 		if consumeQuota {
 		if consumeQuota {
-			err := model.PostConsumeTokenQuota(tokenId, quota)
+			err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0)
 			if err != nil {
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
 			}

controller/relay-text.go

@@ -233,7 +233,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 		//common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
 		//common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
 	}
 	}
 	if consumeQuota && preConsumedQuota > 0 {
 	if consumeQuota && preConsumedQuota > 0 {
-		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
+		userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
 		if err != nil {
 		if err != nil {
 			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
 		}
@@ -400,7 +400,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 			if preConsumedQuota != 0 {
 			if preConsumedQuota != 0 {
 				go func(ctx context.Context) {
 				go func(ctx context.Context) {
 					// return pre-consumed quota
 					// return pre-consumed quota
-					err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
+					err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0)
 					if err != nil {
 					if err != nil {
 						common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
 						common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
 					}
 					}
@@ -434,7 +434,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 					quota = 0
 					quota = 0
 				}
 				}
 				quotaDelta := quota - preConsumedQuota
 				quotaDelta := quota - preConsumedQuota
-				err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
+				err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
 				if err != nil {
 				if err != nil {
 					common.LogError(ctx, "error consuming token remain quota: "+err.Error())
 					common.LogError(ctx, "error consuming token remain quota: "+err.Error())
 				}
 				}

model/token.go

@@ -178,26 +178,40 @@ func decreaseTokenQuota(id int, quota int) (err error) {
 	return err
 	return err
 }
 }
 
 
-func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
+func PreConsumeTokenQuota(tokenId int, quota int) (userQuota int, err error) {
 	if quota < 0 {
 	if quota < 0 {
-		return errors.New("quota 不能为负数！")
+		return 0, errors.New("quota 不能为负数！")
 	}
 	}
 	token, err := GetTokenById(tokenId)
 	token, err := GetTokenById(tokenId)
 	if err != nil {
 	if err != nil {
-		return err
+		return 0, err
 	}
 	}
 	if !token.UnlimitedQuota && token.RemainQuota < quota {
 	if !token.UnlimitedQuota && token.RemainQuota < quota {
-		return errors.New("令牌额度不足")
+		return 0, errors.New("令牌额度不足")
 	}
 	}
-	userQuota, err := GetUserQuota(token.UserId)
+	userQuota, err = GetUserQuota(token.UserId)
 	if err != nil {
 	if err != nil {
-		return err
+		return 0, err
 	}
 	}
 	if userQuota < quota {
 	if userQuota < quota {
-		return errors.New("用户额度不足")
+		return userQuota, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
+	}
+	if !token.UnlimitedQuota {
+		err = DecreaseTokenQuota(tokenId, quota)
+		if err != nil {
+			return userQuota, err
 		}
 		}
-	quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-quota < common.QuotaRemindThreshold
-	noMoreQuota := userQuota-quota <= 0
+	}
+	err = DecreaseUserQuota(token.UserId, quota)
+	return userQuota, err
+}
+
+func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int) (err error) {
+	token, err := GetTokenById(tokenId)
+
+	if quota > 0 {
+		quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-(quota+preConsumedQuota) < common.QuotaRemindThreshold
+		noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0
 		if quotaTooLow || noMoreQuota {
 		if quotaTooLow || noMoreQuota {
 			go func() {
 			go func() {
 				email, err := GetUserEmail(token.UserId)
 				email, err := GetUserEmail(token.UserId)
@@ -218,19 +232,6 @@ func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
 				}
 				}
 			}()
 			}()
 		}
 		}
-	if !token.UnlimitedQuota {
-		err = DecreaseTokenQuota(tokenId, quota)
-		if err != nil {
-			return err
-		}
-	}
-	err = DecreaseUserQuota(token.UserId, quota)
-	return err
-}
-
-func PostConsumeTokenQuota(tokenId int, quota int) (err error) {
-	token, err := GetTokenById(tokenId)
-	if quota > 0 {
 		err = DecreaseUserQuota(token.UserId, quota)
 		err = DecreaseUserQuota(token.UserId, quota)
 	} else {
 	} else {
 		err = IncreaseUserQuota(token.UserId, -quota)
 		err = IncreaseUserQuota(token.UserId, -quota)