refact:openteam

team/model/usage.go

@@ -0,0 +1,72 @@
+package model
+
+import (
+	"net/http"
+	"opencatd-open/store"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+type DailyUsage struct {
+	ID               int64     `gorm:"column:id;primaryKey;autoIncrement"`
+	UserID           int64     `gorm:"column:user_id;index:idx_user_id"`
+	TokenId          int64     `gorm:"column:token_id;index:idx_token_id"`
+	Capability       string    `gorm:"column:capability;index:idx_capability;comment:模型能力"`
+	Date             time.Time `gorm:"column:date;autoCreateTime;index:idx_date"`
+	Model            string    `gorm:"column:model"`
+	Stream           bool      `gorm:"column:stream"`
+	PromptTokens     int       `gorm:"column:prompt_tokens"`
+	CompletionTokens int       `gorm:"column:completion_tokens"`
+	TotalTokens      int       `gorm:"column:total_tokens"`
+	Cost             string    `gorm:"column:cost"`
+	CreatedAt        time.Time `gorm:"column:created_at;autoCreateTime"`
+}
+
+func (DailyUsage) TableName() string {
+	return "daily_usages"
+}
+
+type Usage struct {
+	ID              int       `gorm:"column:id"`
+	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"
+}
+
+func HandleUsage(c *gin.Context) {
+	fromStr := c.Query("from")
+	toStr := c.Query("to")
+	getMonthStartAndEnd := func() (start, end string) {
+		loc, _ := time.LoadLocation("Local")
+		now := time.Now().In(loc)
+
+		year, month, _ := now.Date()
+
+		startOfMonth := time.Date(year, month, 1, 0, 0, 0, 0, loc)
+		endOfMonth := startOfMonth.AddDate(0, 1, 0)
+
+		start = startOfMonth.Format("2006-01-02")
+		end = endOfMonth.Format("2006-01-02")
+		return
+	}
+	if fromStr == "" || toStr == "" {
+		fromStr, toStr = getMonthStartAndEnd()
+	}
+
+	usage, err := store.QueryUsage(fromStr, toStr)
+	if err != nil {
+		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(200, usage)
+}