diff --git a/auth/payment.go b/auth/payment.go index 35b2fde..b32b2a3 100644 --- a/auth/payment.go +++ b/auth/payment.go @@ -1,7 +1,9 @@ package auth import ( + "chat/globals" "chat/utils" + "database/sql" "encoding/json" "github.com/spf13/viper" ) @@ -62,3 +64,39 @@ func Pay(username string, amount float32) bool { resp, _ := utils.Unmarshal[PaymentResponse](converter) return resp.Type } + +func ReduceDalle(db *sql.DB, user *User) bool { + if user.GetQuota(db) < 1 { + return false + } + return user.UseQuota(db, 1) +} + +func CanEnableModel(db *sql.DB, user *User, model string) bool { + switch model { + case globals.GPT4, globals.GPT40613, globals.GPT40314: + return user != nil && user.GetQuota(db) >= 5 + case globals.GPT432k, globals.GPT432k0613, globals.GPT432k0314: + return user != nil && user.GetQuota(db) >= 50 + default: + return true + } +} + +func CanEnableModelWithSubscription(db *sql.DB, user *User, model string, useReverse bool) bool { + if utils.Contains(model, globals.GPT4Array) { + if useReverse { + return true + } + } + return CanEnableModel(db, user, model) +} + +func BuyQuota(db *sql.DB, user *User, quota int) bool { + money := float32(quota) * 0.1 + if Pay(user.Username, money) { + user.IncreaseQuota(db, float32(quota)) + return true + } + return false +} diff --git a/auth/usage.go b/auth/usage.go deleted file mode 100644 index ac2fbb3..0000000 --- a/auth/usage.go +++ /dev/null @@ -1,98 +0,0 @@ -package auth - -import ( - "chat/globals" - "chat/utils" - "database/sql" -) - -// Price Calculation -// 10 nio points = ¥1 -// from 2023-9-6, 1 USD = 7.3124 CNY -// -// GPT-4 price (8k-context) -// Input Output -// $0.03 / 1K tokens $0.06 / 1K tokens -// ¥0.21 / 1K tokens ¥0.43 / 1K tokens -// 2.1 nio / 1K tokens 4.3 nio / 1K tokens -// -// GPT-4 price (32k-context) -// Input Output -// $0.06 / 1K tokens $0.12 / 1K tokens -// ¥0.43 / 1K tokens ¥0.86 / 1K tokens -// 4.3 nio / 1K tokens 8.6 nio / 1K tokens - -// Dalle price (512x512) -// $0.018 / per image -// ¥0.13 / per image -// 1 nio / per image - -func CountInputToken(model string, v []globals.Message) float32 { - switch model { - case globals.GPT3Turbo: - return 0 - case globals.GPT3Turbo16k: - return 0 - case globals.GPT4: - return float32(utils.CountTokenPrice(v, model)) / 1000 * 2.1 - case globals.GPT432k: - return float32(utils.CountTokenPrice(v, model)) / 1000 * 4.2 - case globals.Claude2, globals.Claude2100k: - return 0 - default: - return 0 - } -} - -func CountOutputToken(model string, t int) float32 { - switch model { - case globals.GPT3Turbo: - return 0 - case globals.GPT3Turbo16k: - return 0 - case globals.GPT4: - return float32(t*utils.GetWeightByModel(model)) / 1000 * 4.3 - case globals.GPT432k: - return float32(t*utils.GetWeightByModel(model)) / 1000 * 8.6 - case globals.Claude2, globals.Claude2100k: - return 0 - default: - return 0 - } -} - -func ReduceDalle(db *sql.DB, user *User) bool { - if user.GetQuota(db) < 1 { - return false - } - return user.UseQuota(db, 1) -} - -func CanEnableModel(db *sql.DB, user *User, model string) bool { - switch model { - case globals.GPT4, globals.GPT40613, globals.GPT40314: - return user != nil && user.GetQuota(db) >= 5 - case globals.GPT432k, globals.GPT432k0613, globals.GPT432k0314: - return user != nil && user.GetQuota(db) >= 50 - default: - return true - } -} - -func CanEnableModelWithSubscription(db *sql.DB, user *User, model string, useReverse bool) bool { - if utils.Contains(model, globals.GPT4Array) { - if useReverse { - return true - } - } - return CanEnableModel(db, user, model) -} - -func BuyQuota(db *sql.DB, user *User, quota int) bool { - money := float32(quota) * 0.1 - if Pay(user.Username, money) { - user.IncreaseQuota(db, float32(quota)) - return true - } - return false -} diff --git a/main.go b/main.go index 961a14e..8aa7505 100644 --- a/main.go +++ b/main.go @@ -20,7 +20,7 @@ func main() { app := gin.Default() middleware.RegisterMiddleware(app) - + { auth.Register(app) manager.Register(app) diff --git a/utils/buffer.go b/utils/buffer.go index 8ff071a..822f9ea 100644 --- a/utils/buffer.go +++ b/utils/buffer.go @@ -1,7 +1,6 @@ package utils import ( - "chat/auth" "chat/globals" ) @@ -19,7 +18,7 @@ func NewBuffer(model string, history []globals.Message) *Buffer { Cursor: 0, Times: 0, Model: model, - Quota: auth.CountInputToken(model, history), + Quota: CountInputToken(model, history), } } @@ -28,7 +27,7 @@ func (b *Buffer) GetCursor() int { } func (b *Buffer) GetQuota() float32 { - return b.Quota + auth.CountOutputToken(b.Model, b.ReadTimes()) + return b.Quota + CountOutputToken(b.Model, b.ReadTimes()) } func (b *Buffer) Write(data string) string { diff --git a/utils/tokenizer.go b/utils/tokenizer.go index e0f0e68..6f2d7b8 100644 --- a/utils/tokenizer.go +++ b/utils/tokenizer.go @@ -7,9 +7,30 @@ import ( "strings" ) -// Using https://github.com/pkoukk/tiktoken-go -// To count number of tokens of openai chat messages -// OpenAI Cookbook: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb +// Using https://github.com/pkoukk/tiktoken-go +// To count number of tokens of openai chat messages +// OpenAI Cookbook: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb + +// Price Calculation +// 10 nio points = ¥1 +// from 2023-9-6, 1 USD = 7.3124 CNY +// +// GPT-4 price (8k-context) +// Input Output +// $0.03 / 1K tokens $0.06 / 1K tokens +// ¥0.21 / 1K tokens ¥0.43 / 1K tokens +// 2.1 nio / 1K tokens 4.3 nio / 1K tokens +// +// GPT-4 price (32k-context) +// Input Output +// $0.06 / 1K tokens $0.12 / 1K tokens +// ¥0.43 / 1K tokens ¥0.86 / 1K tokens +// 4.3 nio / 1K tokens 8.6 nio / 1K tokens + +// Dalle price (512x512) +// $0.018 / per image +// ¥0.13 / per image +// 1 nio / per image func GetWeightByModel(model string) int { switch model { @@ -71,3 +92,37 @@ func NumTokensFromMessages(messages []globals.Message, model string) (tokens int func CountTokenPrice(messages []globals.Message, model string) int { return NumTokensFromMessages(messages, model) } + +func CountInputToken(model string, v []globals.Message) float32 { + switch model { + case globals.GPT3Turbo: + return 0 + case globals.GPT3Turbo16k: + return 0 + case globals.GPT4: + return float32(CountTokenPrice(v, model)) / 1000 * 2.1 + case globals.GPT432k: + return float32(CountTokenPrice(v, model)) / 1000 * 4.2 + case globals.Claude2, globals.Claude2100k: + return 0 + default: + return 0 + } +} + +func CountOutputToken(model string, t int) float32 { + switch model { + case globals.GPT3Turbo: + return 0 + case globals.GPT3Turbo16k: + return 0 + case globals.GPT4: + return float32(t*GetWeightByModel(model)) / 1000 * 4.3 + case globals.GPT432k: + return float32(t*GetWeightByModel(model)) / 1000 * 8.6 + case globals.Claude2, globals.Claude2100k: + return 0 + default: + return 0 + } +}