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

feat: able to fetch every request's cost

Browse Source
This commit is contained in:
Laisky.Cai
2024-04-23 00:58:25 +00:00
parent 7047d9605e
commit 84a6817314
14 changed files with 119 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
relay/controller/helper.go
View File

@@ -160,12 +160,11 @@ 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) {
func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64) (quota int64) {
if usage == nil {
logger.Error(ctx, "usage is nil, which is unexpected")
return
}
var quota int64
completionRatio := billingratio.GetCompletionRatio(textRequest.Model)
promptTokens := usage.PromptTokens
completionTokens := usage.CompletionTokens
@@ -193,6 +192,8 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.M
model.RecordConsumeLog(ctx, meta.UserId, meta.ChannelId, promptTokens, completionTokens, textRequest.Model, meta.TokenName, quota, logContent)
model.UpdateUserUsedQuotaAndRequestCount(meta.UserId, quota)
model.UpdateChannelUsedQuota(meta.ChannelId, quota)
return
}
func getMappedModelName(modelName string, mapping map[string]string) (string, bool) {

20
relay/controller/text.go
View File

@@ -8,7 +8,9 @@ import (
"strings"
"github.com/Laisky/errors/v2"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/common/logger"
"github.com/Laisky/one-api/model"
"github.com/Laisky/one-api/relay"
"github.com/Laisky/one-api/relay/adaptor/openai"
"github.com/Laisky/one-api/relay/apitype"
@@ -16,11 +18,11 @@ import (
billingratio "github.com/Laisky/one-api/relay/billing/ratio"
"github.com/Laisky/one-api/relay/channeltype"
"github.com/Laisky/one-api/relay/meta"
"github.com/Laisky/one-api/relay/model"
relaymodel "github.com/Laisky/one-api/relay/model"
"github.com/gin-gonic/gin"
)
func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode {
func RelayTextHelper(c *gin.Context) *relaymodel.ErrorWithStatusCode {
ctx := c.Request.Context()
meta := meta.GetByContext(c)
// get & validate textRequest
@@ -108,7 +110,19 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode {
billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
return respErr
}
// post-consume quota
go postConsumeQuota(ctx, usage, meta, textRequest, ratio, preConsumedQuota, modelRatio, groupRatio)
go func() {
quota := postConsumeQuota(c, usage, meta, textRequest, ratio, preConsumedQuota, modelRatio, groupRatio)
docu := model.NewUserRequestCost(
c.GetInt(ctxkey.Id),
c.GetString(ctxkey.RequestId),
quota,
)
if err = docu.Insert(); err != nil {
logger.Errorf(c, "insert user request cost failed: %+v", err)
}
}()
return nil
}