update sparkdesk prize

coaidev · Oct 1, 2023 · f8cb3d9 · f8cb3d9
1 parent 91d650a
commit f8cb3d9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git 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
@@ -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: