diff --git a/addition/generation/api.go b/addition/generation/api.go index c75fe89..b56b92c 100644 --- a/addition/generation/api.go +++ b/addition/generation/api.go @@ -62,9 +62,9 @@ func GenerateAPI(c *gin.Context) { } check, plan := auth.CanEnableModelWithSubscription(db, cache, user, form.Model) - if !check { + if check != nil { conn.Send(globals.GenerationSegmentResponse{ - Message: "You don't have enough quota to use this model.", + Message: check.Error(), Quota: 0, End: true, }) diff --git a/app/src/components/PopupDialog.tsx b/app/src/components/PopupDialog.tsx index 250df0e..d2261d6 100644 --- a/app/src/components/PopupDialog.tsx +++ b/app/src/components/PopupDialog.tsx @@ -37,6 +37,8 @@ export type PopupDialogProps = { cancelLabel?: string; confirmLabel?: string; + + componentProps?: any; }; type PopupFieldProps = PopupDialogProps & { @@ -50,6 +52,7 @@ function PopupField({ onValueChange, value, placeholder, + componentProps, }: PopupFieldProps) { switch (type) { case popupTypes.Text: @@ -62,6 +65,7 @@ function PopupField({ }} value={value} placeholder={placeholder} + {...componentProps} /> ); @@ -72,6 +76,7 @@ function PopupField({ value={parseFloat(value)} onValueChange={(v) => setValue(v.toString())} placeholder={placeholder} + {...componentProps} /> ); @@ -82,6 +87,7 @@ function PopupField({ onCheckedChange={(state: boolean) => { setValue(state.toString()); }} + {...componentProps} /> ); diff --git a/app/src/components/admin/UserTable.tsx b/app/src/components/admin/UserTable.tsx index 858e38d..6642458 100644 --- a/app/src/components/admin/UserTable.tsx +++ b/app/src/components/admin/UserTable.tsx @@ -137,6 +137,7 @@ function OperationMenu({ user, onRefresh }: OperationMenuProps) { onValueChange={getNumber} open={quotaOpen} setOpen={setQuotaOpen} + componentProps={{ acceptNegative: true }} onSubmit={async (value) => { const quota = parseNumber(value); const resp = await quotaOperation(user.id, quota); @@ -155,6 +156,7 @@ function OperationMenu({ user, onRefresh }: OperationMenuProps) { onValueChange={getNumber} open={quotaSetOpen} setOpen={setQuotaSetOpen} + componentProps={{ acceptNegative: true }} onSubmit={async (value) => { const quota = parseNumber(value); const resp = await quotaOperation(user.id, quota, true); diff --git a/app/src/components/ui/number-input.tsx b/app/src/components/ui/number-input.tsx index 0367985..0bd3d9a 100644 --- a/app/src/components/ui/number-input.tsx +++ b/app/src/components/ui/number-input.tsx @@ -40,6 +40,8 @@ const NumberInput = React.forwardRef( return v.match(exp)?.join("") || ""; } + if (v === "-" && props.acceptNegative) return v; + // replace -0124.5 to -124.5, 0043 to 43, 2.000 to 2.000 const exp = /^[-+]?0+(?=[0-9]+(\.[0-9]+)?$)/; v = v.replace(exp, ""); diff --git a/auth/rule.go b/auth/rule.go index d1f0ddb..bb00b9d 100644 --- a/auth/rule.go +++ b/auth/rule.go @@ -3,27 +3,52 @@ package auth import ( "chat/channel" "database/sql" + "fmt" "github.com/go-redis/redis/v8" ) +const ( + ErrNotAuthenticated = "not authenticated error (model: %s)" + ErrNotSetPrice = "the price of the model is not set error (model: %s)" + 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) -func CanEnableModel(db *sql.DB, user *User, model string) bool { +func CanEnableModel(db *sql.DB, user *User, model string) error { isAuth := user != nil charge := channel.ChargeInstance.GetCharge(model) if !charge.IsBilling() { // return if is the user is authenticated or anonymous is allowed for this model - return charge.SupportAnonymous() || isAuth + if charge.SupportAnonymous() || isAuth { + return nil + } + + return fmt.Errorf(ErrNotAuthenticated, model) + } + + if !isAuth { + return fmt.Errorf(ErrNotAuthenticated, model) } // return if the user is authenticated and has enough quota - return isAuth && user.GetQuota(db) >= charge.GetLimit() + limit := charge.GetLimit() + if limit == -1 { + return fmt.Errorf(ErrNotSetPrice, model) + } + + quota := user.GetQuota(db) + if quota < limit { + return fmt.Errorf(ErrNotEnoughQuota, model, limit, quota) + } + + return nil } -func CanEnableModelWithSubscription(db *sql.DB, cache *redis.Client, user *User, model string) (canEnable bool, usePlan bool) { +func CanEnableModelWithSubscription(db *sql.DB, cache *redis.Client, user *User, model string) (canEnable error, usePlan bool) { // use subscription quota first if user != nil && HandleSubscriptionUsage(db, cache, user, model) { - return true, true + return nil, true } return CanEnableModel(db, user, model), false } diff --git a/channel/charge.go b/channel/charge.go index 1d20f30..460d807 100644 --- a/channel/charge.go +++ b/channel/charge.go @@ -283,7 +283,7 @@ func (c *Charge) GetLimit() float32 { // 1k input tokens + 1k output tokens return c.GetInput() + c.GetOutput() default: - return 0 + return -1 } } diff --git a/manager/chat.go b/manager/chat.go index 24ddaad..42ea977 100644 --- a/manager/chat.go +++ b/manager/chat.go @@ -15,7 +15,6 @@ import ( ) const defaultMessage = "empty response" -const defaultQuotaMessage = "You don't have enough quota or you don't have permission to use this model. please [buy](/buy) or [subscribe](/subscribe) to get more." func CollectQuota(c *gin.Context, user *auth.User, buffer *utils.Buffer, uncountable bool, err error) { db := utils.GetDBFromContext(c) @@ -73,13 +72,14 @@ func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conve Conversation: instance.GetId(), }) - if !check { + if check != nil { + message := check.Error() conn.Send(globals.ChatSegmentResponse{ - Message: defaultQuotaMessage, + Message: message, Quota: 0, End: true, }) - return defaultQuotaMessage + return message } if form := ExtractCacheData(conn.GetCtx(), &CacheProps{ diff --git a/manager/chat_completions.go b/manager/chat_completions.go index 918fd78..41cf7e9 100644 --- a/manager/chat_completions.go +++ b/manager/chat_completions.go @@ -55,8 +55,8 @@ func ChatRelayAPI(c *gin.Context) { } check := auth.CanEnableModel(db, user, form.Model) - if !check { - sendErrorResponse(c, fmt.Errorf("quota exceeded"), "quota_exceeded_error") + if check != nil { + sendErrorResponse(c, check, "quota_exceeded_error") return } diff --git a/manager/completions.go b/manager/completions.go index 543f8b2..b3dba51 100644 --- a/manager/completions.go +++ b/manager/completions.go @@ -27,8 +27,8 @@ func NativeChatHandler(c *gin.Context, user *auth.User, model string, message [] cache := utils.GetCacheFromContext(c) check, plan := auth.CanEnableModelWithSubscription(db, cache, user, model) - if !check { - return defaultQuotaMessage, 0 + if check != nil { + return check.Error(), 0 } if form := ExtractCacheData(c, &CacheProps{ diff --git a/manager/images.go b/manager/images.go index b925e73..5cff199 100644 --- a/manager/images.go +++ b/manager/images.go @@ -49,8 +49,8 @@ func ImagesRelayAPI(c *gin.Context) { } check := auth.CanEnableModel(db, user, form.Model) - if !check { - sendErrorResponse(c, fmt.Errorf("quota exceeded"), "quota_exceeded_error") + if check != nil { + sendErrorResponse(c, check, "quota_exceeded_error") return }