Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🐛 fix: preQuotaConsumption

Browse Source 
Signed-off-by: K8sCat <k8scat@gmail.com>
This commit is contained in:
K8sCat 2024-06-07 12:48:23 +08:00
parent 507ba44910
commit 7ec974f56f
No known key found for this signature in database
GPG Key ID: CFB6A9364A3D8FF5
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
relay/relay_util/quota.go
View File

@ -43,10 +43,12 @@ func NewQuota(c *gin.Context, modelName string, promptTokens int) (*Quota, *type
quota.groupRatio = common.GetGroupRatio(c.GetString("group"))
quota.inputRatio = quota.price.GetInput() * quota.groupRatio
if quota.price.Type == model.TimesPriceType {
quota.preConsumedQuota = int(1000 * quota.inputRatio)
} else {
quota.preConsumedQuota = int(float64(quota.promptTokens)*quota.inputRatio) + config.PreConsumedQuota
if quota.price.Input != 0 || quota.price.Output != 0 {
if quota.price.Type == model.TimesPriceType {
quota.preConsumedQuota = int(1000 * quota.inputRatio)
} else {
quota.preConsumedQuota = int(float64(quota.promptTokens)*quota.inputRatio) + config.PreConsumedQuota
}
}
errWithCode := quota.preQuotaConsumption()
@ -58,6 +60,10 @@ func NewQuota(c *gin.Context, modelName string, promptTokens int) (*Quota, *type
}
func (q *Quota) preQuotaConsumption() *types.OpenAIErrorWithStatusCode {
if q.preConsumedQuota == 0 {
return nil
}
userQuota, err := model.CacheGetUserQuota(q.userId)
if err != nil {
return common.ErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)