ai-gateway/controller/billing.go
2024-04-21 19:43:23 +08:00

96 lines
2.1 KiB
Go

package controller
import (
"github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/ctxkey"
"github.com/songquanpeng/one-api/model"
relaymodel "github.com/songquanpeng/one-api/relay/model"
)
func GetSubscription(c *gin.Context) {
var remainQuota int64
var usedQuota int64
var err error
var token *model.Token
var expiredTime int64
if config.DisplayTokenStatEnabled {
tokenId := c.GetInt(ctxkey.TokenId)
token, err = model.GetTokenById(tokenId)
expiredTime = token.ExpiredTime
remainQuota = token.RemainQuota
usedQuota = token.UsedQuota
} else {
userId := c.GetInt(ctxkey.Id)
remainQuota, err = model.GetUserQuota(userId)
if err != nil {
usedQuota, err = model.GetUserUsedQuota(userId)
}
}
if expiredTime <= 0 {
expiredTime = 0
}
if err != nil {
Error := relaymodel.Error{
Message: err.Error(),
Type: "upstream_error",
}
c.JSON(200, gin.H{
"error": Error,
})
return
}
quota := remainQuota + usedQuota
amount := float64(quota)
if config.DisplayInCurrencyEnabled {
amount /= config.QuotaPerUnit
}
if token != nil && token.UnlimitedQuota {
amount = 100000000
}
subscription := OpenAISubscriptionResponse{
Object: "billing_subscription",
HasPaymentMethod: true,
SoftLimitUSD: amount,
HardLimitUSD: amount,
SystemHardLimitUSD: amount,
AccessUntil: expiredTime,
}
c.JSON(200, subscription)
return
}
func GetUsage(c *gin.Context) {
var quota int64
var err error
var token *model.Token
if config.DisplayTokenStatEnabled {
tokenId := c.GetInt(ctxkey.TokenId)
token, err = model.GetTokenById(tokenId)
quota = token.UsedQuota
} else {
userId := c.GetInt(ctxkey.Id)
quota, err = model.GetUserUsedQuota(userId)
}
if err != nil {
Error := relaymodel.Error{
Message: err.Error(),
Type: "one_api_error",
}
c.JSON(200, gin.H{
"error": Error,
})
return
}
amount := float64(quota)
if config.DisplayInCurrencyEnabled {
amount /= config.QuotaPerUnit
}
usage := OpenAIUsageResponse{
Object: "list",
TotalUsage: amount * 100,
}
c.JSON(200, usage)
return
}