xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-11 10:53:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: refactor postConsumeQuota function to return quota and update user request cost handling

Browse Source
This commit is contained in:
Laisky.Cai
2024-11-22 03:12:09 +00:00
parent f322f8029d
commit 1f103202a6
2 changed files with 20 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
relay/controller/helper.go
View File

@@ -91,7 +91,7 @@ func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIR
return preConsumedQuota, nil
}
func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64, systemPromptReset bool) {
func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64, systemPromptReset bool) (quota int64) {
if usage == nil {
logger.Error(ctx, "usage is nil, which is unexpected")
return
@@ -100,7 +100,7 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.M
completionRatio := billingratio.GetCompletionRatio(textRequest.Model, meta.ChannelType)
promptTokens := usage.PromptTokens
completionTokens := usage.CompletionTokens
quota := int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio))
quota = int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio))
if ratio != 0 && quota <= 0 {
quota = 1
}
@@ -128,7 +128,7 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.M
model.UpdateUserUsedQuotaAndRequestCount(meta.UserId, quota)
model.UpdateChannelUsedQuota(meta.ChannelId, quota)
return
return quota
}
func getMappedModelName(modelName string, mapping map[string]string) (string, bool) {