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