fix import cycle

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
+}

main.go

@@ -20,7 +20,7 @@ func main() {
 
 	app := gin.Default()
 	middleware.RegisterMiddleware(app)
-	
+
 	{
 		auth.Register(app)
 		manager.Register(app)

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 {

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
+	}
+}