Files
opencatd-open/team/model/usage.go
2025-01-01 23:10:01 +08:00

73 lines
2.1 KiB
Go

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