Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

55 строки
1.6 KiB

  1. package model
  2. import (
  3. "github.com/QuantumNous/new-api/common"
  4. "gorm.io/gorm"
  5. )
  6. type UserModelRateLimit struct {
  7. Id int `json:"id" gorm:"primaryKey"`
  8. UserId int `json:"user_id" gorm:"uniqueIndex:idx_user_model_rate_limit"`
  9. Model string `json:"model" gorm:"size:255;uniqueIndex:idx_user_model_rate_limit"`
  10. Rpm int `json:"rpm"`
  11. CreatedAt int64 `json:"created_at"`
  12. UpdatedAt int64 `json:"updated_at"`
  13. }
  14. // GetUserModelRateLimits 查询指定用户的所有 RPM 配置
  15. func GetUserModelRateLimits(userId int) ([]UserModelRateLimit, error) {
  16. var list []UserModelRateLimit
  17. err := DB.Where("user_id = ?", userId).Find(&list).Error
  18. return list, err
  19. }
  20. // SetUserModelRateLimits 覆盖式写入:事务中先硬删再批量插入。
  21. // 空 slice 时仅删除该用户所有配置。
  22. func SetUserModelRateLimits(userId int, items []UserModelRateLimit) error {
  23. return DB.Transaction(func(tx *gorm.DB) error {
  24. // 先硬删该用户所有配置
  25. if err := tx.Where("user_id = ?", userId).Delete(&UserModelRateLimit{}).Error; err != nil {
  26. return err
  27. }
  28. if len(items) == 0 {
  29. return nil
  30. }
  31. now := common.GetTimestamp()
  32. for i := range items {
  33. items[i].UserId = userId
  34. items[i].Id = 0 // 让数据库自增
  35. items[i].CreatedAt = now
  36. items[i].UpdatedAt = now
  37. }
  38. return tx.Create(&items).Error
  39. })
  40. }
  41. // GetUserModelRpm 查单个用户+模型的 RPM,不命中返回 (0, false)
  42. func GetUserModelRpm(userId int, model string) (int, bool) {
  43. var item UserModelRateLimit
  44. err := DB.Where("user_id = ? AND model = ?", userId, model).First(&item).Error
  45. if err != nil {
  46. return 0, false
  47. }
  48. return item.Rpm, true
  49. }