diff --git a/auth/payment.go b/auth/payment.go index b32b2a3..44f6fe5 100644 --- a/auth/payment.go +++ b/auth/payment.go @@ -78,6 +78,8 @@ func CanEnableModel(db *sql.DB, user *User, model string) bool { return user != nil && user.GetQuota(db) >= 5 case globals.GPT432k, globals.GPT432k0613, globals.GPT432k0314: return user != nil && user.GetQuota(db) >= 50 + case globals.SparkDesk: + return user != nil && user.GetQuota(db) >= 1 default: return true } diff --git a/utils/tokenizer.go b/utils/tokenizer.go index 6f2d7b8..d808c17 100644 --- a/utils/tokenizer.go +++ b/utils/tokenizer.go @@ -49,7 +49,8 @@ func GetWeightByModel(model string) int { globals.GPT4, globals.GPT40314, - globals.GPT40613: + globals.GPT40613, + globals.SparkDesk: return 3 case globals.GPT3Turbo0301, globals.GPT3Turbo16k0301: return 4 // every message follows <|start|>{role/name}\n{content}<|end|>\n @@ -103,6 +104,8 @@ func CountInputToken(model string, v []globals.Message) float32 { return float32(CountTokenPrice(v, model)) / 1000 * 2.1 case globals.GPT432k: return float32(CountTokenPrice(v, model)) / 1000 * 4.2 + case globals.SparkDesk: + return float32(CountTokenPrice(v, model)) / 1000 * 0.36 case globals.Claude2, globals.Claude2100k: return 0 default: @@ -120,6 +123,8 @@ func CountOutputToken(model string, t int) float32 { return float32(t*GetWeightByModel(model)) / 1000 * 4.3 case globals.GPT432k: return float32(t*GetWeightByModel(model)) / 1000 * 8.6 + case globals.SparkDesk: + return float32(t*GetWeightByModel(model)) / 1000 * 0.36 case globals.Claude2, globals.Claude2100k: return 0 default: