feat: support unset price blocked for non-admin users

This commit is contained in:
Zhang Minghan 2024-02-19 21:14:29 +08:00
parent afa19d75f8
commit 8708bca503
3 changed files with 14 additions and 5 deletions

View File

@ -9,15 +9,21 @@ import (
const ( const (
ErrNotAuthenticated = "not authenticated error (model: %s)" ErrNotAuthenticated = "not authenticated error (model: %s)"
ErrNotSetPrice = "the price of the model is not set error (model: %s)" ErrNotSetPrice = "the price of the model is not set (model: %s)"
ErrNotEnoughQuota = "user quota is not enough error (model: %s, minimum quota: %0.2f, your quota: %0.2f)" ErrNotEnoughQuota = "user quota is not enough error (model: %s, minimum quota: %0.2f, your quota: %0.2f)"
) )
// CanEnableModel returns whether the model can be enabled (without subscription) // CanEnableModel returns whether the model can be enabled (without subscription)
func CanEnableModel(db *sql.DB, user *User, model string) error { func CanEnableModel(db *sql.DB, user *User, model string) error {
isAuth := user != nil isAuth := user != nil
isAdmin := isAuth && user.IsAdmin(db)
charge := channel.ChargeInstance.GetCharge(model) charge := channel.ChargeInstance.GetCharge(model)
if charge.IsUnsetType() && !isAdmin {
return fmt.Errorf(ErrNotSetPrice, model)
}
if !charge.IsBilling() { if !charge.IsBilling() {
// return if is the user is authenticated or anonymous is allowed for this model // return if is the user is authenticated or anonymous is allowed for this model
if charge.SupportAnonymous() || isAuth { if charge.SupportAnonymous() || isAuth {
@ -33,9 +39,6 @@ func CanEnableModel(db *sql.DB, user *User, model string) error {
// return if the user is authenticated and has enough quota // return if the user is authenticated and has enough quota
limit := charge.GetLimit() limit := charge.GetLimit()
if limit == -1 {
return fmt.Errorf(ErrNotSetPrice, model)
}
quota := user.GetQuota(db) quota := user.GetQuota(db)
if quota < limit { if quota < limit {

View File

@ -74,6 +74,7 @@ func (m *ChargeManager) GetCharge(model string) *Charge {
return &Charge{ return &Charge{
Type: globals.NonBilling, Type: globals.NonBilling,
Anonymous: false, Anonymous: false,
Unset: true,
} }
} }
@ -236,6 +237,10 @@ func (m *ChargeManager) GetRuleByModel(model string) *Charge {
return nil return nil
} }
func (c *Charge) IsUnsetType() bool {
return c.Unset
}
func (c *Charge) GetType() string { func (c *Charge) GetType() string {
if c.Type == "" { if c.Type == "" {
return globals.NonBilling return globals.NonBilling
@ -283,7 +288,7 @@ func (c *Charge) GetLimit() float32 {
// 1k input tokens + 1k output tokens // 1k input tokens + 1k output tokens
return c.GetInput() + c.GetOutput() return c.GetInput() + c.GetOutput()
default: default:
return -1 return 0
} }
} }

View File

@ -39,6 +39,7 @@ type Charge struct {
Input float32 `json:"input" mapstructure:"input"` Input float32 `json:"input" mapstructure:"input"`
Output float32 `json:"output" mapstructure:"output"` Output float32 `json:"output" mapstructure:"output"`
Anonymous bool `json:"anonymous" mapstructure:"anonymous"` Anonymous bool `json:"anonymous" mapstructure:"anonymous"`
Unset bool `json:"-" mapstructure:"-"`
} }
type ChargeSequence []*Charge type ChargeSequence []*Charge