From f8cb3d99d8d4aab269b5dc9557693e537345acff Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Sun, 1 Oct 2023 11:41:46 +0800
Subject: [PATCH] update sparkdesk prize

---
 auth/payment.go    | 2 ++
 utils/tokenizer.go | 7 ++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/auth/payment.go b/auth/payment.go
index b32b2a39..44f6fe5d 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 6f2d7b8d..d808c174 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: