支持统计用户的使用量

store/db.go

@@ -11,6 +11,8 @@ import (
 
 var db *gorm.DB
 
+var usage *gorm.DB
+
 func init() {
 	if _, err := os.Stat("db"); os.IsNotExist(err) {
 		errDir := os.MkdirAll("db", 0755)
@@ -31,4 +33,13 @@ func init() {
 	}
 	LoadKeysCache()
 	LoadAuthCache()
+
+	usage, err = gorm.Open(sqlite.Open("./db/usage.db"), &gorm.Config{})
+	if err != nil {
+		panic(err)
+	}
+	err = usage.AutoMigrate(&DailyUsage{}, &Usage{})
+	if err != nil {
+		panic(err)
+	}
 }

store/usage.go

@@ -0,0 +1,149 @@
+package store
+
+import (
+	"errors"
+	"time"
+
+	"github.com/Sakurasan/to"
+	"gorm.io/gorm"
+)
+
+type DailyUsage struct {
+	ID              int       `gorm:"column:id"`
+	UserID          int       `gorm:"column:user_id";primaryKey`
+	Date            time.Time `gorm:"column:date"`
+	SKU             string    `gorm:"column:sku"`
+	PromptUnits     int       `gorm:"column:prompt_units"`
+	CompletionUnits int       `gorm:"column:completion_units"`
+	TotalUnit       int       `gorm:"column:total_unit"`
+	Cost            string    `gorm:"column:cost"`
+}
+
+func (DailyUsage) TableName() string {
+	return "daily_usages"
+}
+
+type Usage struct {
+	ID              int       `gorm:"column:id"`
+	PromptHash      string    `gorm:"column:prompt_hash"`
+	UserID          int       `gorm:"column:user_id"`
+	SKU             string    `gorm:"column:sku"`
+	PromptUnits     int       `gorm:"column:prompt_units"`
+	CompletionUnits int       `gorm:"column:completion_units"`
+	TotalUnit       int       `gorm:"column:total_unit"`
+	Cost            string    `gorm:"column:cost"`
+	Date            time.Time `gorm:"column:date"`
+}
+
+func (Usage) TableName() string {
+	return "usages"
+}
+
+type Summary struct {
+	UserId             int     `gorm:"column:user_id"`
+	SumPromptUnits     int     `gorm:"column:sum_prompt_units"`
+	SumCompletionUnits int     `gorm:"column:sum_completion_units"`
+	SumTotalUnit       int     `gorm:"column:sum_total_unit"`
+	SumCost            float64 `gorm:"column:sum_cost"`
+}
+type CalcUsage struct {
+	UserID    int    `json:"userId,omitempty"`
+	TotalUnit int    `json:"totalUnit,omitempty"`
+	Cost      string `json:"cost,omitempty"`
+}
+
+func QueryUsage(from, to string) ([]CalcUsage, error) {
+	var results = []CalcUsage{}
+	err := usage.Model(&DailyUsage{}).Select(`user_id, 
+	--SUM(prompt_units) AS prompt_units,
+	-- SUM(completion_units) AS completion_units,
+	SUM(total_unit) AS total_unit,
+	printf('%.6f', SUM(cost)) AS cost`).
+		Group("user_id").
+		Where("date >= ? AND date < ?", from, to).
+		Find(&results).Error
+	if err != nil {
+		return nil, err
+	}
+	return results, nil
+}
+
+type Tokens struct {
+	UserID          int
+	PromptCount     int
+	CompletionCount int
+	TotalTokens     int
+	Cost            string
+	Model           string
+	PromptHash      string
+}
+
+func Record(chatlog *Tokens) (err error) {
+	u := &Usage{
+		UserID:          chatlog.UserID,
+		SKU:             chatlog.Model,
+		PromptHash:      chatlog.PromptHash,
+		PromptUnits:     chatlog.PromptCount,
+		CompletionUnits: chatlog.CompletionCount,
+		TotalUnit:       chatlog.TotalTokens,
+		Cost:            to.String(chatlog.Cost),
+		Date:            time.Now(),
+	}
+	err = usage.Create(u).Error
+	return
+
+}
+
+func SumDaily(userid int) error {
+	var count int64
+	err := usage.Model(&DailyUsage{}).Where("user_id = ? and date = ?", userid, time.Date(time.Now().Year(), time.Now().Month(), time.Now().Day(), 0, 0, 0, 0, time.UTC)).Count(&count).Error
+	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return err
+	}
+	if count == 0 {
+		if err := insertSumDaily(userid); err != nil {
+			return err
+		}
+	} else {
+		if err := updateSumDaily(userid, time.Date(time.Now().Year(), time.Now().Month(), time.Now().Day(), 0, 0, 0, 0, time.UTC)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func insertSumDaily(uid int) error {
+	nowstr := time.Date(time.Now().Year(), time.Now().Month(), time.Now().Day(), 0, 0, 0, 0, time.UTC)
+	err := usage.Exec(`INSERT INTO daily_usages
+		(user_id, date, sku, prompt_units, completion_units, total_unit, cost)
+		SELECT 
+		user_id,
+		?,
+		sku,
+		SUM(prompt_units) AS sum_prompt_units,
+		SUM(completion_units) AS sum_completion_units,
+		SUM(total_unit) AS sum_total_unit,
+		SUM(cost) AS sum_cost
+		FROM usages 
+		WHERE date >= ?
+		AND user_id = ?`, nowstr, nowstr, uid).Error
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func updateSumDaily(uid int, date time.Time) error {
+	// var u = Summary{}
+	err := usage.Model(&Usage{}).Exec(`UPDATE daily_usages
+	SET 
+	prompt_units = (SELECT SUM(prompt_units) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date),
+	completion_units = (SELECT SUM(completion_units) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date),
+	total_unit = (SELECT SUM(total_unit) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date),
+	cost = (SELECT SUM(cost) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date)
+	WHERE user_id = ? AND date >= ?`, uid, date).Error
+	if err != nil {
+		return err
+	}
+	return nil
+}

store/userdb.go

@@ -73,6 +73,15 @@ func GetUserByName(name string) (*User, error) {
 	return &user, nil
 }
 
+func GetUserID(authkey string) (int, error) {
+	var user User
+	result := db.Where(&User{Token: authkey}).First(&user)
+	if result.Error != nil {
+		return 0, result.Error
+	}
+	return int(user.ID), nil
+}
+
 func GetAllUsers() ([]*User, error) {
 	var users []*User
 	result := db.Find(&users)