Commit 995b3147 by JustSong

perf: use a goroutine to handle quota post consumption (#364)

parent c563e1ed
...@@ -305,51 +305,54 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode { ...@@ -305,51 +305,54 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
} }
var textResponse TextResponse var textResponse TextResponse
tokenName := c.GetString("token_name")
channelId := c.GetInt("channel_id")
defer func() { defer func() {
c.Writer.Flush() // c.Writer.Flush()
if consumeQuota { go func() {
quota := 0 if consumeQuota {
completionRatio := 1.0 quota := 0
if strings.HasPrefix(textRequest.Model, "gpt-3.5") { completionRatio := 1.0
completionRatio = 1.333333 if strings.HasPrefix(textRequest.Model, "gpt-3.5") {
} completionRatio = 1.333333
if strings.HasPrefix(textRequest.Model, "gpt-4") { }
completionRatio = 2 if strings.HasPrefix(textRequest.Model, "gpt-4") {
} completionRatio = 2
}
promptTokens = textResponse.Usage.PromptTokens promptTokens = textResponse.Usage.PromptTokens
completionTokens = textResponse.Usage.CompletionTokens completionTokens = textResponse.Usage.CompletionTokens
quota = promptTokens + int(float64(completionTokens)*completionRatio) quota = promptTokens + int(float64(completionTokens)*completionRatio)
quota = int(float64(quota) * ratio) quota = int(float64(quota) * ratio)
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 quotaDelta := quota - preConsumedQuota
err := model.PostConsumeTokenQuota(tokenId, quotaDelta) err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }
err = model.CacheUpdateUserQuota(userId) err = model.CacheUpdateUserQuota(userId)
if err != nil { if err != nil {
common.SysError("error update user quota cache: " + err.Error()) common.SysError("error update user quota cache: " + err.Error())
} }
if quota != 0 { if quota != 0 {
tokenName := c.GetString("token_name") logContent := fmt.Sprintf("模型倍率 %.2f,分组倍率 %.2f", modelRatio, groupRatio)
logContent := fmt.Sprintf("模型倍率 %.2f,分组倍率 %.2f", modelRatio, groupRatio) model.RecordConsumeLog(userId, promptTokens, completionTokens, textRequest.Model, tokenName, quota, logContent)
model.RecordConsumeLog(userId, promptTokens, completionTokens, textRequest.Model, tokenName, quota, logContent) model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
channelId := c.GetInt("channel_id") model.UpdateChannelUsedQuota(channelId, quota)
model.UpdateChannelUsedQuota(channelId, quota) }
} }
} }()
}() }()
switch apiType { switch apiType {
case APITypeOpenAI: case APITypeOpenAI:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment