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

fix: refactor cost calculation logic for web-search tools and improve quota handling

Browse Source
This commit is contained in:
Laisky.Cai 2025-03-13 09:33:03 +00:00
parent 82225818e3
commit 2f34798aa9
2 changed files with 51 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
relay/adaptor/openai/adaptor.go
View File

@ -208,6 +208,7 @@ func (a *Adaptor) DoResponse(c *gin.Context,
// ------------------------------------- // -------------------------------------
// calculate web-search tool cost // calculate web-search tool cost
// ------------------------------------- // -------------------------------------
if usage != nil {
searchContextSize := "medium" searchContextSize := "medium"
var req *model.GeneralOpenAIRequest var req *model.GeneralOpenAIRequest
if vi, ok := c.Get(ctxkey.ConvertedRequest); ok { if vi, ok := c.Get(ctxkey.ConvertedRequest); ok {
@ -250,6 +251,7 @@ func (a *Adaptor) DoResponse(c *gin.Context,
} }
} }
} }
}
return return
} }

15
relay/controller/helper.go
View File

@ -96,7 +96,15 @@ func preConsumeQuota(c *gin.Context, textRequest *relaymodel.GeneralOpenAIReques
return preConsumedQuota, nil 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) (quota int64) { 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 { if usage == nil {
logger.Error(ctx, "usage is nil, which is unexpected") logger.Error(ctx, "usage is nil, which is unexpected")
return return
@ -108,17 +116,18 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.M
// but the behavior of third-party providers may differ, so for now we do not add them manually. // but the behavior of third-party providers may differ, so for now we do not add them manually.
// completionTokens := usage.CompletionTokens + usage.CompletionTokensDetails.ReasoningTokens // completionTokens := usage.CompletionTokens + usage.CompletionTokensDetails.ReasoningTokens
completionTokens := usage.CompletionTokens completionTokens := usage.CompletionTokens
quota = int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio)) quota = int64(math.Ceil((float64(promptTokens)+float64(completionTokens)*completionRatio)*ratio)) + usage.ToolsCost
if ratio != 0 && quota <= 0 { if ratio != 0 && quota <= 0 {
quota = 1 quota = 1
} }
totalTokens := promptTokens + completionTokens totalTokens := promptTokens + completionTokens
if totalTokens == 0 { if totalTokens == 0 {
// in this case, must be some error happened // in this case, must be some error happened
// we cannot just return, because we may have to return the pre-consumed quota // we cannot just return, because we may have to return the pre-consumed quota
quota = 0 quota = 0
} }
quotaDelta := quota - preConsumedQuota + usage.ToolsCost quotaDelta := quota - preConsumedQuota
err := model.PostConsumeTokenQuota(meta.TokenId, quotaDelta) err := model.PostConsumeTokenQuota(meta.TokenId, quotaDelta)
if err != nil { if err != nil {
logger.Error(ctx, "error consuming token remain quota: "+err.Error()) logger.Error(ctx, "error consuming token remain quota: "+err.Error())