ai-gateway/controller/billing.go
Laisky.Cai d548a01c59
feat: Handle errors, validate model names, and calculate quota usage (#978)
- Improved error handling in various modules for better stability and responsiveness.
- Optimized code in several files for improved efficiency and readability.
- Enhanced user experience by providing more detailed error responses in the controller.
- Strengthened security by ignoring sensitive files in `.gitignore`.
2024-02-12 21:35:40 +08:00

95 lines
2.0 KiB
Go

package controller
import (
"github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/model"
"github.com/songquanpeng/one-api/relay/channel/openai"
)
func GetSubscription(c *gin.Context) {
var remainQuota int
var usedQuota int
var err error
var token *model.Token
var expiredTime int64
if config.DisplayTokenStatEnabled {
tokenId := c.GetInt("token_id")
token, err = model.GetTokenById(tokenId)
expiredTime = token.ExpiredTime
remainQuota = token.RemainQuota
usedQuota = token.UsedQuota
} else {
userId := c.GetInt("id")
remainQuota, err = model.GetUserQuota(userId)
if err != nil {
usedQuota, err = model.GetUserUsedQuota(userId)
}
}
if expiredTime <= 0 {
expiredTime = 0
}
if err != nil {
Error := openai.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 int
var err error
var token *model.Token
if config.DisplayTokenStatEnabled {
tokenId := c.GetInt("token_id")
token, err = model.GetTokenById(tokenId)
quota = token.UsedQuota
} else {
userId := c.GetInt("id")
quota, err = model.GetUserUsedQuota(userId)
}
if err != nil {
Error := openai.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
}