package model import ( "github.com/QuantumNous/new-api/common" "gorm.io/gorm" ) type UserModelRateLimit struct { Id int `json:"id" gorm:"primaryKey"` UserId int `json:"user_id" gorm:"uniqueIndex:idx_user_model_rate_limit"` Model string `json:"model" gorm:"size:255;uniqueIndex:idx_user_model_rate_limit"` Rpm int `json:"rpm"` CreatedAt int64 `json:"created_at"` UpdatedAt int64 `json:"updated_at"` } // GetUserModelRateLimits 查询指定用户的所有 RPM 配置 func GetUserModelRateLimits(userId int) ([]UserModelRateLimit, error) { var list []UserModelRateLimit err := DB.Where("user_id = ?", userId).Find(&list).Error return list, err } // SetUserModelRateLimits 覆盖式写入:事务中先硬删再批量插入。 // 空 slice 时仅删除该用户所有配置。 func SetUserModelRateLimits(userId int, items []UserModelRateLimit) error { return DB.Transaction(func(tx *gorm.DB) error { // 先硬删该用户所有配置 if err := tx.Where("user_id = ?", userId).Delete(&UserModelRateLimit{}).Error; err != nil { return err } if len(items) == 0 { return nil } now := common.GetTimestamp() for i := range items { items[i].UserId = userId items[i].Id = 0 // 让数据库自增 items[i].CreatedAt = now items[i].UpdatedAt = now } return tx.Create(&items).Error }) } // GetUserModelRpm 查单个用户+模型的 RPM,不命中返回 (0, false) func GetUserModelRpm(userId int, model string) (int, bool) { var item UserModelRateLimit err := DB.Where("user_id = ? AND model = ?", userId, model).First(&item).Error if err != nil { return 0, false } return item.Rpm, true }