xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-18 16:36:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into latest

Browse Source
This commit is contained in:
CaIon 2023-11-12 20:45:10 +08:00
commit 6b4d30d76d
6 changed files with 40 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
.github/workflows/docker-image.yml vendored
View File

@ -1,32 +0,0 @@
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 }}

4
controller/relay-audio.go
View File

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

2
controller/relay-image.go
View File

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

2
controller/relay-mj.go
View File

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

6
controller/relay-text.go
View File

@ -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))
}
if consumeQuota && preConsumedQuota > 0 {
err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
if err != nil {
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 {
go func(ctx context.Context) {
// return pre-consumed quota
err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0)
if err != nil {
common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
}
@ -434,7 +434,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
quota = 0
}
quotaDelta := quota - preConsumedQuota
err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
if err != nil {
common.LogError(ctx, "error consuming token remain quota: "+err.Error())
}

45
model/token.go
View File

@ -178,26 +178,40 @@ func decreaseTokenQuota(id int, quota int) (err error) {
return err
}
func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
func PreConsumeTokenQuota(tokenId int, quota int) (userQuota int, err error) {
if quota < 0 {
return errors.New("quota 不能为负数！")
return 0, errors.New("quota 不能为负数！")
}
token, err := GetTokenById(tokenId)
if err != nil {
return err
return 0, err
}
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 {
return err
return 0, err
}
if userQuota < quota {
return errors.New("用户额度不足")
return userQuota, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
}
quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-quota < common.QuotaRemindThreshold
noMoreQuota := userQuota-quota <= 0
if !token.UnlimitedQuota {
err = DecreaseTokenQuota(tokenId, quota)
if err != nil {
return userQuota, err
}
}
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 {
go func() {
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)
} else {
err = IncreaseUserQuota(token.UserId, -quota)