Skip to content

phsl / new-api

Go to a project
Commit 0fcd243f by 1808837298@qq.com
Options

feat: Add user notification settings with quota warning and multiple notification methods 

- Implement user notification settings with email and webhook options
- Add new user settings for quota warning threshold and notification preferences
- Create backend API and database support for user notification configuration
- Enhance frontend personal settings with notification configuration UI
- Support custom notification email and webhook URL
- Add service layer for sending user notifications
parent a353ea3e
Showing with 488 additions and 240 deletions
common/logger.go
...@@ -100,6 +100,14 @@ func LogQuota(quota int) string { ...@@ -100,6 +100,14 @@ func LogQuota(quota int) string {
} }
} }
func FormatQuota(quota int) string {
if DisplayInCurrencyEnabled {
return fmt.Sprintf("＄%.6f", float64(quota)/QuotaPerUnit)
} else {
return fmt.Sprintf("%d", quota)
}
}
// LogJson 仅供测试使用 only for test // LogJson 仅供测试使用 only for test
func LogJson(ctx context.Context, msg string, obj any) { func LogJson(ctx context.Context, msg string, obj any) {
jsonStr, err := json.Marshal(obj) jsonStr, err := json.Marshal(obj)
......
common/model-ratio.go
...@@ -477,20 +477,6 @@ func GetAudioCompletionRatio(name string) float64 { ...@@ -477,20 +477,6 @@ func GetAudioCompletionRatio(name string) float64 {
return 2 return 2
} }
//func GetAudioPricePerMinute(name string) float64 {
// if strings.HasPrefix(name, "gpt-4o-realtime") {
// return 0.06
// }
// return 0.06
//}
//
//func GetAudioCompletionPricePerMinute(name string) float64 {
// if strings.HasPrefix(name, "gpt-4o-realtime") {
// return 0.24
// }
// return 0.24
//}
func GetCompletionRatioMap() map[string]float64 { func GetCompletionRatioMap() map[string]float64 {
if CompletionRatio == nil { if CompletionRatio == nil {
CompletionRatio = defaultCompletionRatio CompletionRatio = defaultCompletionRatio
......
constant/env.go
...@@ -44,5 +44,5 @@ func InitEnv() { ...@@ -44,5 +44,5 @@ func InitEnv() {
} }
} }
// 是否生成初始令牌，默认关闭。 // GenerateDefaultToken 是否生成初始令牌，默认关闭。
var GenerateDefaultToken = common.GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false) var GenerateDefaultToken = common.GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false)
constant/user_setting.go 0 → 100644
package constant
var (
UserSettingNotifyType = "notify_type" // QuotaWarningType 额度预警类型
UserSettingQuotaWarningThreshold = "quota_warning_threshold" // QuotaWarningThreshold 额度预警阈值
UserSettingWebhookUrl = "webhook_url" // WebhookUrl webhook地址
UserSettingWebhookSecret = "webhook_secret" // WebhookSecret webhook密钥
UserSettingNotificationEmail = "notification_email" // NotificationEmail 通知邮箱地址
)
var (
NotifyTypeEmail = "email" // Email 邮件
NotifyTypeWebhook = "webhook" // Webhook
)
controller/user.go
...@@ -4,6 +4,7 @@ import ( ...@@ -4,6 +4,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"net/http" "net/http"
"net/url"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
"one-api/setting" "one-api/setting"
...@@ -913,3 +914,115 @@ func TopUp(c *gin.Context) { ...@@ -913,3 +914,115 @@ func TopUp(c *gin.Context) {
}) })
return return
} }
type UpdateUserSettingRequest struct {
QuotaWarningType string `json:"notify_type"`
QuotaWarningThreshold int `json:"quota_warning_threshold"`
WebhookUrl string `json:"webhook_url,omitempty"`
WebhookSecret string `json:"webhook_secret,omitempty"`
NotificationEmail string `json:"notification_email,omitempty"`
}
func UpdateUserSetting(c *gin.Context) {
var req UpdateUserSettingRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "无效的参数",
})
return
}
// 验证预警类型
if req.QuotaWarningType != constant.NotifyTypeEmail && req.QuotaWarningType != constant.NotifyTypeWebhook {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "无效的预警类型",
})
return
}
// 验证预警阈值
if req.QuotaWarningThreshold <= 0 {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "预警阈值必须大于0",
})
return
}
// 如果是webhook类型,验证webhook地址
if req.QuotaWarningType == constant.NotifyTypeWebhook {
if req.WebhookUrl == "" {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "Webhook地址不能为空",
})
return
}
// 验证URL格式
if _, err := url.ParseRequestURI(req.WebhookUrl); err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "无效的Webhook地址",
})
return
}
}
// 如果是邮件类型，验证邮箱地址
if req.QuotaWarningType == constant.NotifyTypeEmail && req.NotificationEmail != "" {
// 验证邮箱格式
if !strings.Contains(req.NotificationEmail, "@") {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "无效的邮箱地址",
})
return
}
}
userId := c.GetInt("id")
user, err := model.GetUserById(userId, true)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
// 构建设置
settings := map[string]interface{}{
constant.UserSettingNotifyType: req.QuotaWarningType,
constant.UserSettingQuotaWarningThreshold: req.QuotaWarningThreshold,
}
// 如果是webhook类型,添加webhook相关设置
if req.QuotaWarningType == constant.NotifyTypeWebhook {
settings[constant.UserSettingWebhookUrl] = req.WebhookUrl
if req.WebhookSecret != "" {
settings[constant.UserSettingWebhookSecret] = req.WebhookSecret
}
}
// 如果提供了通知邮箱，添加到设置中
if req.QuotaWarningType == constant.NotifyTypeEmail && req.NotificationEmail != "" {
settings[constant.UserSettingNotificationEmail] = req.NotificationEmail
}
// 更新用户设置
user.SetSetting(settings)
if err := user.Update(false); err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "更新设置失败: " + err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "设置已更新",
})
}
docker-compose.yml
...@@ -24,7 +24,7 @@ services: ...@@ -24,7 +24,7 @@ services:
- redis - redis
- mysql - mysql
healthcheck: healthcheck:
test: [ "CMD-SHELL", "wget -q -O - http://localhost:3000/api/status | grep -o '\"success\":\\s*true' | awk -F: '{print $2}'" ] test: ["CMD-SHELL", "wget -q -O - http://localhost:3000/api/status | grep -o '\"success\":\\s*true' | awk -F: '{print $$2}'"]
interval: 30s interval: 30s
timeout: 10s timeout: 10s
retries: 3 retries: 3
......
dto/notify.go 0 → 100644
package dto
type Notify struct {
Type string `json:"type"`
Title string `json:"title"`
Content string `json:"content"`
Values []interface{} `json:"values"`
}
const ContentValueParam = "{{value}}"
const (
NotifyTypeQuotaExceed = "quota_exceed"
NotifyTypeChannelUpdate = "channel_update"
)
func NewNotify(t string, title string, content string, values []interface{}) Notify {
return Notify{
Type: t,
Title: title,
Content: content,
Values: values,
}
}
dto/openai_request.go
...@@ -18,6 +18,7 @@ type GeneralOpenAIRequest struct { ...@@ -18,6 +18,7 @@ type GeneralOpenAIRequest struct {
Model string `json:"model,omitempty"` Model string `json:"model,omitempty"`
Messages []Message `json:"messages,omitempty"` Messages []Message `json:"messages,omitempty"`
Prompt any `json:"prompt,omitempty"` Prompt any `json:"prompt,omitempty"`
Prefix any `json:"prefix,omitempty"`
Suffix any `json:"suffix,omitempty"` Suffix any `json:"suffix,omitempty"`
Stream bool `json:"stream,omitempty"` Stream bool `json:"stream,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"` StreamOptions *StreamOptions `json:"stream_options,omitempty"`
......
model/token.go
...@@ -3,13 +3,11 @@ package model ...@@ -3,13 +3,11 @@ package model
import ( import (
"errors" "errors"
"fmt" "fmt"
"github.com/bytedance/gopkg/util/gopool"
"gorm.io/gorm"
"one-api/common" "one-api/common"
relaycommon "one-api/relay/common"
"one-api/setting"
"strconv"
"strings" "strings"
"github.com/bytedance/gopkg/util/gopool"
"gorm.io/gorm"
) )
type Token struct { type Token struct {
...@@ -322,80 +320,3 @@ func decreaseTokenQuota(id int, quota int) (err error) { ...@@ -322,80 +320,3 @@ func decreaseTokenQuota(id int, quota int) (err error) {
).Error ).Error
return err return err
} }
func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
if quota < 0 {
return errors.New("quota 不能为负数！")
}
if relayInfo.IsPlayground {
return nil
}
//if relayInfo.TokenUnlimited {
// return nil
//}
token, err := GetTokenById(relayInfo.TokenId)
if err != nil {
return err
}
if !relayInfo.TokenUnlimited && token.RemainQuota < quota {
return errors.New("令牌额度不足")
}
err = DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
if err != nil {
return err
}
return nil
}
func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
if quota > 0 {
err = DecreaseUserQuota(relayInfo.UserId, quota)
} else {
err = IncreaseUserQuota(relayInfo.UserId, -quota)
}
if err != nil {
return err
}
if !relayInfo.IsPlayground {
if quota > 0 {
err = DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
} else {
err = IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
}
if err != nil {
return err
}
}
if sendEmail {
if (quota + preConsumedQuota) != 0 {
quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-(quota+preConsumedQuota) < common.QuotaRemindThreshold
noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0
if quotaTooLow || noMoreQuota {
go func() {
email, err := GetUserEmail(relayInfo.UserId)
if err != nil {
common.SysError("failed to fetch user email: " + err.Error())
}
prompt := "您的额度即将用尽"
if noMoreQuota {
prompt = "您的额度已用尽"
}
if email != "" {
topUpLink := fmt.Sprintf("%s/topup", setting.ServerAddress)
err = common.SendEmail(prompt, email,
fmt.Sprintf("%s，当前剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota, topUpLink, topUpLink))
if err != nil {
common.SysError("failed to send email" + err.Error())
}
common.SysLog("user quota is low, consumed quota: " + strconv.Itoa(quota) + ", user quota: " + strconv.Itoa(userQuota))
}
}()
}
}
}
return nil
}
model/user.go
package model package model
import ( import (
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"one-api/common" "one-api/common"
...@@ -38,6 +39,7 @@ type User struct { ...@@ -38,6 +39,7 @@ type User struct {
InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"` InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
DeletedAt gorm.DeletedAt `gorm:"index"` DeletedAt gorm.DeletedAt `gorm:"index"`
LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"` LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"`
Setting string `json:"setting" gorm:"type:text;column:setting"`
} }
func (user *User) GetAccessToken() string { func (user *User) GetAccessToken() string {
...@@ -51,6 +53,22 @@ func (user *User) SetAccessToken(token string) { ...@@ -51,6 +53,22 @@ func (user *User) SetAccessToken(token string) {
user.AccessToken = &token user.AccessToken = &token
} }
func (user *User) GetSetting() map[string]interface{} {
if user.Setting == "" {
return nil
}
return common.StrToMap(user.Setting)
}
func (user *User) SetSetting(setting map[string]interface{}) {
settingBytes, err := json.Marshal(setting)
if err != nil {
common.SysError("failed to marshal setting: " + err.Error())
return
}
user.Setting = string(settingBytes)
}
// CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil // CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
func CheckUserExistOrDeleted(username string, email string) (bool, error) { func CheckUserExistOrDeleted(username string, email string) (bool, error) {
var user User var user User
...@@ -315,8 +333,8 @@ func (user *User) Update(updatePassword bool) error { ...@@ -315,8 +333,8 @@ func (user *User) Update(updatePassword bool) error {
return err return err
} }
// 更新缓存 // Update cache
return updateUserCache(user.Id, user.Username, user.Group, user.Quota, user.Status) return updateUserCache(*user)
} }
func (user *User) Edit(updatePassword bool) error { func (user *User) Edit(updatePassword bool) error {
...@@ -344,8 +362,8 @@ func (user *User) Edit(updatePassword bool) error { ...@@ -344,8 +362,8 @@ func (user *User) Edit(updatePassword bool) error {
return err return err
} }
// 更新缓存 // Update cache
return updateUserCache(user.Id, user.Username, user.Group, user.Quota, user.Status) return updateUserCache(*user)
} }
func (user *User) Delete() error { func (user *User) Delete() error {
...@@ -371,8 +389,8 @@ func (user *User) HardDelete() error { ...@@ -371,8 +389,8 @@ func (user *User) HardDelete() error {
// ValidateAndFill check password & user status // ValidateAndFill check password & user status
func (user *User) ValidateAndFill() (err error) { func (user *User) ValidateAndFill() (err error) {
// When querying with struct, GORM will only query with non-zero fields, // When querying with struct, GORM will only query with non-zero fields,
// that means if your field’s value is 0, '', false or other zero values, // that means if your field's value is 0, '', false or other zero values,
// it won’t be used to build query conditions // it won't be used to build query conditions
password := user.Password password := user.Password
username := strings.TrimSpace(user.Username) username := strings.TrimSpace(user.Username)
if username == "" || password == "" { if username == "" || password == "" {
...@@ -531,7 +549,6 @@ func GetUserQuota(id int, fromDB bool) (quota int, err error) { ...@@ -531,7 +549,6 @@ func GetUserQuota(id int, fromDB bool) (quota int, err error) {
return quota, nil return quota, nil
} }
// Don't return error - fall through to DB // Don't return error - fall through to DB
//common.SysError("failed to get user quota from cache: " + err.Error())
} }
fromDB = true fromDB = true
err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
...@@ -580,6 +597,35 @@ func GetUserGroup(id int, fromDB bool) (group string, err error) { ...@@ -580,6 +597,35 @@ func GetUserGroup(id int, fromDB bool) (group string, err error) {
return group, nil return group, nil
} }
// GetUserSetting gets setting from Redis first, falls back to DB if needed
func GetUserSetting(id int, fromDB bool) (settingMap map[string]interface{}, err error) {
var setting string
defer func() {
// Update Redis cache asynchronously on successful DB read
if shouldUpdateRedis(fromDB, err) {
gopool.Go(func() {
if err := updateUserSettingCache(id, setting); err != nil {
common.SysError("failed to update user setting cache: " + err.Error())
}
})
}
}()
if !fromDB && common.RedisEnabled {
setting, err := getUserSettingCache(id)
if err == nil {
return setting, nil
}
// Don't return error - fall through to DB
}
fromDB = true
err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&setting).Error
if err != nil {
return map[string]interface{}{}, err
}
return common.StrToMap(setting), nil
}
func IncreaseUserQuota(id int, quota int) (err error) { func IncreaseUserQuota(id int, quota int) (err error) {
if quota < 0 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
...@@ -725,10 +771,10 @@ func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool { ...@@ -725,10 +771,10 @@ func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
return !errors.Is(err, gorm.ErrRecordNotFound) return !errors.Is(err, gorm.ErrRecordNotFound)
} }
func (u *User) FillUserByLinuxDOId() error { func (user *User) FillUserByLinuxDOId() error {
if u.LinuxDOId == "" { if user.LinuxDOId == "" {
return errors.New("linux do id is empty") return errors.New("linux do id is empty")
} }
err := DB.Where("linux_do_id = ?", u.LinuxDOId).First(u).Error err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error
return err return err
} }
model/user_cache.go
package model package model
import ( import (
"encoding/json"
"fmt" "fmt"
"one-api/common" "one-api/common"
"one-api/constant" "one-api/constant"
"strconv"
"time" "time"
"github.com/bytedance/gopkg/util/gopool"
) )
// Change UserCache struct to userCache // UserCache struct remains the same as it represents the cached data structure
type userCache struct { type UserCache struct {
Id int `json:"id"` Id int `json:"id"`
Group string `json:"group"` Group string `json:"group"`
Email string `json:"email"`
Quota int `json:"quota"` Quota int `json:"quota"`
Status int `json:"status"` Status int `json:"status"`
Role int `json:"role"`
Username string `json:"username"` Username string `json:"username"`
Setting string `json:"setting"`
} }
// Rename all exported functions to private ones func (user *UserCache) GetSetting() map[string]interface{} {
// invalidateUserCache clears all user related cache if user.Setting == "" {
func invalidateUserCache(userId int) error {
if !common.RedisEnabled {
return nil return nil
} }
return common.StrToMap(user.Setting)
}
keys := []string{ func (user *UserCache) SetSetting(setting map[string]interface{}) {
fmt.Sprintf(constant.UserGroupKeyFmt, userId), settingBytes, err := json.Marshal(setting)
fmt.Sprintf(constant.UserQuotaKeyFmt, userId), if err != nil {
fmt.Sprintf(constant.UserEnabledKeyFmt, userId), common.SysError("failed to marshal setting: " + err.Error())
fmt.Sprintf(constant.UserUsernameKeyFmt, userId), return
}
for _, key := range keys {
if err := common.RedisDel(key); err != nil {
return fmt.Errorf("failed to delete cache key %s: %w", key, err)
}
} }
return nil user.Setting = string(settingBytes)
} }
// updateUserGroupCache updates user group cache // getUserCacheKey returns the key for user cache
func updateUserGroupCache(userId int, group string) error { func getUserCacheKey(userId int) string {
if !common.RedisEnabled { return fmt.Sprintf("user:%d", userId)
return nil
}
return common.RedisSet(
fmt.Sprintf(constant.UserGroupKeyFmt, userId),
group,
time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second,
)
} }
// updateUserQuotaCache updates user quota cache // invalidateUserCache clears user cache
func updateUserQuotaCache(userId int, quota int) error { func invalidateUserCache(userId int) error {
if !common.RedisEnabled { if !common.RedisEnabled {
return nil return nil
} }
return common.RedisSet( return common.RedisHDelObj(getUserCacheKey(userId))
fmt.Sprintf(constant.UserQuotaKeyFmt, userId),
fmt.Sprintf("%d", quota),
time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second,
)
} }
// updateUserStatusCache updates user status cache // updateUserCache updates all user cache fields using hash
func updateUserStatusCache(userId int, userEnabled bool) error { func updateUserCache(user User) error {
if !common.RedisEnabled { if !common.RedisEnabled {
return nil return nil
} }
enabled := "0"
if userEnabled {
enabled = "1"
}
return common.RedisSet(
fmt.Sprintf(constant.UserEnabledKeyFmt, userId),
enabled,
time.Duration(constant.UserId2StatusCacheSeconds)*time.Second,
)
}
// updateUserNameCache updates username cache cache := &UserCache{
func updateUserNameCache(userId int, username string) error { Id: user.Id,
if !common.RedisEnabled { Group: user.Group,
return nil Quota: user.Quota,
Status: user.Status,
Username: user.Username,
Setting: user.Setting,
Email: user.Email,
} }
return common.RedisSet(
fmt.Sprintf(constant.UserUsernameKeyFmt, userId), return common.RedisHSetObj(
username, getUserCacheKey(user.Id),
cache,
time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second, time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second,
) )
} }
// updateUserCache updates all user cache fields // GetUserCache gets complete user cache from hash
func updateUserCache(userId int, username string, userGroup string, quota int, status int) error { func GetUserCache(userId int) (userCache *UserCache, err error) {
if !common.RedisEnabled { var user *User
return nil var fromDB bool
defer func() {
// Update Redis cache asynchronously on successful DB read
if shouldUpdateRedis(fromDB, err) && user != nil {
gopool.Go(func() {
if err := updateUserCache(*user); err != nil {
common.SysError("failed to update user status cache: " + err.Error())
} }
})
if err := updateUserGroupCache(userId, userGroup); err != nil {
return fmt.Errorf("update group cache: %w", err)
} }
}()
if err := updateUserQuotaCache(userId, quota); err != nil { // Try getting from Redis first
return fmt.Errorf("update quota cache: %w", err) err = common.RedisHGetObj(getUserCacheKey(userId), &userCache)
if err == nil {
return userCache, nil
} }
if err := updateUserStatusCache(userId, status == common.UserStatusEnabled); err != nil { // If Redis fails, get from DB
return fmt.Errorf("update status cache: %w", err) fromDB = true
user, err = GetUserById(userId, false)
if err != nil {
return nil, err // Return nil and error if DB lookup fails
} }
if err := updateUserNameCache(userId, username); err != nil { // Create cache object from user data
return fmt.Errorf("update username cache: %w", err) userCache = &UserCache{
Id: user.Id,
Group: user.Group,
Quota: user.Quota,
Status: user.Status,
Username: user.Username,
Setting: user.Setting,
Email: user.Email,
} }
return userCache, nil
}
// Add atomic quota operations using hash fields
func cacheIncrUserQuota(userId int, delta int64) error {
if !common.RedisEnabled {
return nil return nil
}
return common.RedisHIncrBy(getUserCacheKey(userId), "Quota", delta)
} }
// getUserGroupCache gets user group from cache func cacheDecrUserQuota(userId int, delta int64) error {
return cacheIncrUserQuota(userId, -delta)
}
// Helper functions to get individual fields if needed
func getUserGroupCache(userId int) (string, error) { func getUserGroupCache(userId int) (string, error) {
if !common.RedisEnabled { cache, err := GetUserCache(userId)
return "", nil if err != nil {
return "", err
} }
return common.RedisGet(fmt.Sprintf(constant.UserGroupKeyFmt, userId)) return cache.Group, nil
} }
// getUserQuotaCache gets user quota from cache
func getUserQuotaCache(userId int) (int, error) { func getUserQuotaCache(userId int) (int, error) {
if !common.RedisEnabled { cache, err := GetUserCache(userId)
return 0, nil
}
quotaStr, err := common.RedisGet(fmt.Sprintf(constant.UserQuotaKeyFmt, userId))
if err != nil { if err != nil {
return 0, err return 0, err
} }
return strconv.Atoi(quotaStr) return cache.Quota, nil
} }
// getUserStatusCache gets user status from cache
func getUserStatusCache(userId int) (int, error) { func getUserStatusCache(userId int) (int, error) {
if !common.RedisEnabled { cache, err := GetUserCache(userId)
return 0, nil
}
statusStr, err := common.RedisGet(fmt.Sprintf(constant.UserEnabledKeyFmt, userId))
if err != nil { if err != nil {
return 0, err return 0, err
} }
return strconv.Atoi(statusStr) return cache.Status, nil
} }
// getUserNameCache gets username from cache
func getUserNameCache(userId int) (string, error) { func getUserNameCache(userId int) (string, error) {
if !common.RedisEnabled { cache, err := GetUserCache(userId)
return "", nil if err != nil {
return "", err
} }
return common.RedisGet(fmt.Sprintf(constant.UserUsernameKeyFmt, userId)) return cache.Username, nil
} }
// getUserCache gets complete user cache func getUserSettingCache(userId int) (map[string]interface{}, error) {
func getUserCache(userId int) (*userCache, error) { setting := make(map[string]interface{})
if !common.RedisEnabled { cache, err := GetUserCache(userId)
return nil, nil
}
group, err := getUserGroupCache(userId)
if err != nil { if err != nil {
return nil, fmt.Errorf("get group cache: %w", err) return setting, err
} }
return cache.GetSetting(), nil
}
quota, err := getUserQuotaCache(userId) // New functions for individual field updates
if err != nil { func updateUserStatusCache(userId int, status bool) error {
return nil, fmt.Errorf("get quota cache: %w", err) if !common.RedisEnabled {
return nil
} }
statusInt := common.UserStatusEnabled
status, err := getUserStatusCache(userId) if !status {
if err != nil { statusInt = common.UserStatusDisabled
return nil, fmt.Errorf("get status cache: %w", err)
} }
return common.RedisHSetField(getUserCacheKey(userId), "Status", fmt.Sprintf("%d", statusInt))
}
username, err := getUserNameCache(userId) func updateUserQuotaCache(userId int, quota int) error {
if err != nil { if !common.RedisEnabled {
return nil, fmt.Errorf("get username cache: %w", err) return nil
} }
return common.RedisHSetField(getUserCacheKey(userId), "Quota", fmt.Sprintf("%d", quota))
}
return &userCache{ func updateUserGroupCache(userId int, group string) error {
Id: userId, if !common.RedisEnabled {
Group: group, return nil
Quota: quota, }
Status: status, return common.RedisHSetField(getUserCacheKey(userId), "Group", group)
Username: username,
}, nil
} }
// Add atomic quota operations func updateUserNameCache(userId int, username string) error {
func cacheIncrUserQuota(userId int, delta int64) error {
if !common.RedisEnabled { if !common.RedisEnabled {
return nil return nil
} }
key := fmt.Sprintf(constant.UserQuotaKeyFmt, userId) return common.RedisHSetField(getUserCacheKey(userId), "Username", username)
return common.RedisIncr(key, delta)
} }
func cacheDecrUserQuota(userId int, delta int64) error { func updateUserSettingCache(userId int, setting string) error {
return cacheIncrUserQuota(userId, -delta) if !common.RedisEnabled {
return nil
}
return common.RedisHSetField(getUserCacheKey(userId), "Setting", setting)
} }
relay/relay-mj.go
...@@ -194,7 +194,7 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse { ...@@ -194,7 +194,7 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
} }
defer func(ctx context.Context) { defer func(ctx context.Context) {
if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 { if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 {
err := model.PostConsumeQuota(relayInfo, userQuota, quota, 0, true) err := service.PostConsumeQuota(relayInfo, quota, 0, true)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }
...@@ -500,7 +500,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons ...@@ -500,7 +500,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons
defer func(ctx context.Context) { defer func(ctx context.Context) {
if consumeQuota && midjResponseWithStatus.StatusCode == 200 { if consumeQuota && midjResponseWithStatus.StatusCode == 200 {
err := model.PostConsumeQuota(relayInfo, userQuota, quota, 0, true) err := service.PostConsumeQuota(relayInfo, quota, 0, true)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }
......
relay/relay-text.go
...@@ -272,7 +272,7 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo ...@@ -272,7 +272,7 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo
return 0, 0, service.OpenAIErrorWrapperLocal(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) return 0, 0, service.OpenAIErrorWrapperLocal(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
} }
if userQuota-preConsumedQuota < 0 { if userQuota-preConsumedQuota < 0 {
return 0, 0, service.OpenAIErrorWrapperLocal(fmt.Errorf("chat pre-consumed quota failed, user quota: %d, need quota: %d", userQuota, preConsumedQuota), "insufficient_user_quota", http.StatusBadRequest) return 0, 0, service.OpenAIErrorWrapperLocal(fmt.Errorf("chat pre-consumed quota failed, user quota: %s, need quota: %d", common.FormatQuota(userQuota), preConsumedQuota), "insufficient_user_quota", http.StatusBadRequest)
} }
if userQuota > 100*preConsumedQuota { if userQuota > 100*preConsumedQuota {
// 用户额度充足，判断令牌额度是否充足 // 用户额度充足，判断令牌额度是否充足
...@@ -282,18 +282,18 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo ...@@ -282,18 +282,18 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo
if tokenQuota > 100*preConsumedQuota { if tokenQuota > 100*preConsumedQuota {
// 令牌额度充足，信任令牌 // 令牌额度充足，信任令牌
preConsumedQuota = 0 preConsumedQuota = 0
common.LogInfo(c, fmt.Sprintf("user %d quota %d and token %d quota %d are enough, trusted and no need to pre-consume", relayInfo.UserId, userQuota, relayInfo.TokenId, tokenQuota)) common.LogInfo(c, fmt.Sprintf("user %d quota %s and token %d quota %d are enough, trusted and no need to pre-consume", relayInfo.UserId, common.FormatQuota(userQuota), relayInfo.TokenId, tokenQuota))
} }
} else { } else {
// in this case, we do not pre-consume quota // in this case, we do not pre-consume quota
// because the user has enough quota // because the user has enough quota
preConsumedQuota = 0 preConsumedQuota = 0
common.LogInfo(c, fmt.Sprintf("user %d with unlimited token has enough quota %d, trusted and no need to pre-consume", relayInfo.UserId, userQuota)) common.LogInfo(c, fmt.Sprintf("user %d with unlimited token has enough quota %s, trusted and no need to pre-consume", relayInfo.UserId, common.FormatQuota(userQuota)))
} }
} }
if preConsumedQuota > 0 { if preConsumedQuota > 0 {
err = model.PreConsumeTokenQuota(relayInfo, preConsumedQuota) err = service.PreConsumeTokenQuota(relayInfo, preConsumedQuota)
if err != nil { if err != nil {
return 0, 0, service.OpenAIErrorWrapperLocal(err, "pre_consume_token_quota_failed", http.StatusForbidden) return 0, 0, service.OpenAIErrorWrapperLocal(err, "pre_consume_token_quota_failed", http.StatusForbidden)
} }
...@@ -310,7 +310,7 @@ func returnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, us ...@@ -310,7 +310,7 @@ func returnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, us
go func() { go func() {
relayInfoCopy := *relayInfo relayInfoCopy := *relayInfo
err := model.PostConsumeQuota(&relayInfoCopy, userQuota, -preConsumedQuota, 0, false) err := service.PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false)
if err != nil { if err != nil {
common.SysError("error return pre-consumed quota: " + err.Error()) common.SysError("error return pre-consumed quota: " + err.Error())
} }
...@@ -368,7 +368,7 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelN ...@@ -368,7 +368,7 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelN
//} //}
quotaDelta := quota - preConsumedQuota quotaDelta := quota - preConsumedQuota
if quotaDelta != 0 { if quotaDelta != 0 {
err := model.PostConsumeQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true) err := service.PostConsumeQuota(relayInfo, quotaDelta, preConsumedQuota, true)
if err != nil { if err != nil {
common.LogError(ctx, "error consuming token remain quota: "+err.Error()) common.LogError(ctx, "error consuming token remain quota: "+err.Error())
} }
......
relay/relay_task.go
...@@ -113,7 +113,7 @@ func RelayTaskSubmit(c *gin.Context, relayMode int) (taskErr *dto.TaskError) { ...@@ -113,7 +113,7 @@ func RelayTaskSubmit(c *gin.Context, relayMode int) (taskErr *dto.TaskError) {
// release quota // release quota
if relayInfo.ConsumeQuota && taskErr == nil { if relayInfo.ConsumeQuota && taskErr == nil {
err := model.PostConsumeQuota(relayInfo.ToRelayInfo(), userQuota, quota, 0, true) err := service.PostConsumeQuota(relayInfo.ToRelayInfo(), quota, 0, true)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }
......
router/api-router.go
...@@ -56,6 +56,7 @@ func SetApiRouter(router *gin.Engine) { ...@@ -56,6 +56,7 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.POST("/pay", controller.RequestEpay) selfRoute.POST("/pay", controller.RequestEpay)
selfRoute.POST("/amount", controller.RequestAmount) selfRoute.POST("/amount", controller.RequestAmount)
selfRoute.POST("/aff_transfer", controller.TransferAffQuota) selfRoute.POST("/aff_transfer", controller.TransferAffQuota)
selfRoute.PUT("/setting", controller.UpdateUserSetting)
} }
adminRoute := userRoute.Group("/") adminRoute := userRoute.Group("/")
......
service/quota.go
...@@ -3,8 +3,10 @@ package service ...@@ -3,8 +3,10 @@ package service
import ( import (
"errors" "errors"
"fmt" "fmt"
"github.com/bytedance/gopkg/util/gopool"
"math" "math"
"one-api/common" "one-api/common"
constant2 "one-api/constant"
"one-api/dto" "one-api/dto"
"one-api/model" "one-api/model"
relaycommon "one-api/relay/common" relaycommon "one-api/relay/common"
...@@ -99,7 +101,7 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag ...@@ -99,7 +101,7 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
return errors.New(fmt.Sprintf("令牌额度不足，剩余额度为 %d", token.RemainQuota)) return errors.New(fmt.Sprintf("令牌额度不足，剩余额度为 %d", token.RemainQuota))
} }
err = model.PostConsumeQuota(relayInfo, 0, quota, 0, false) err = PostConsumeQuota(relayInfo, quota, 0, false)
if err != nil { if err != nil {
return err return err
} }
...@@ -222,7 +224,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, ...@@ -222,7 +224,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
} else { } else {
quotaDelta := quota - preConsumedQuota quotaDelta := quota - preConsumedQuota
if quotaDelta != 0 { if quotaDelta != 0 {
err := model.PostConsumeQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true) err := PostConsumeQuota(relayInfo, quotaDelta, preConsumedQuota, true)
if err != nil { if err != nil {
common.LogError(ctx, "error consuming token remain quota: "+err.Error()) common.LogError(ctx, "error consuming token remain quota: "+err.Error())
} }
...@@ -239,3 +241,88 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, ...@@ -239,3 +241,88 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.PromptTokens, usage.CompletionTokens, logModel, model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.PromptTokens, usage.CompletionTokens, logModel,
tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other) tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other)
} }
func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
if quota < 0 {
return errors.New("quota 不能为负数！")
}
if relayInfo.IsPlayground {
return nil
}
//if relayInfo.TokenUnlimited {
// return nil
//}
token, err := model.GetTokenById(relayInfo.TokenId)
if err != nil {
return err
}
if !relayInfo.TokenUnlimited && token.RemainQuota < quota {
return errors.New("令牌额度不足")
}
err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
if err != nil {
return err
}
return nil
}
func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (err error) {
if quota > 0 {
err = model.DecreaseUserQuota(relayInfo.UserId, quota)
} else {
err = model.IncreaseUserQuota(relayInfo.UserId, -quota)
}
if err != nil {
return err
}
if !relayInfo.IsPlayground {
if quota > 0 {
err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
} else {
err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
}
if err != nil {
return err
}
}
if sendEmail {
if (quota + preConsumedQuota) != 0 {
checkAndSendQuotaNotify(relayInfo.UserId, quota, preConsumedQuota)
}
}
return nil
}
func checkAndSendQuotaNotify(userId int, quota int, preConsumedQuota int) {
gopool.Go(func() {
userCache, err := model.GetUserCache(userId)
if err != nil {
common.SysError("failed to get user cache: " + err.Error())
}
userSetting := userCache.GetSetting()
threshold := common.QuotaRemindThreshold
if userCustomThreshold, ok := userSetting[constant2.UserSettingQuotaWarningThreshold]; ok {
threshold = int(userCustomThreshold.(float64))
}
//noMoreQuota := userCache.Quota-(quota+preConsumedQuota) <= 0
quotaTooLow := false
consumeQuota := quota + preConsumedQuota
if userCache.Quota-consumeQuota < threshold {
quotaTooLow = true
}
if quotaTooLow {
prompt := "您的额度即将用尽"
topUpLink := fmt.Sprintf("%s/topup", setting.ServerAddress)
content := "{{value}}，当前剩余额度为 {{value}}，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='{{value}}'>{{value}}</a>"
err = NotifyUser(userCache, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, []interface{}{prompt, common.FormatQuota(userCache.Quota), topUpLink, topUpLink}))
if err != nil {
common.SysError(fmt.Sprintf("failed to send quota notify to user %d: %s", userId, err.Error()))
}
}
})
}
service/user_notify.go
...@@ -3,7 +3,10 @@ package service ...@@ -3,7 +3,10 @@ package service
import ( import (
"fmt" "fmt"
"one-api/common" "one-api/common"
"one-api/constant"
"one-api/dto"
"one-api/model" "one-api/model"
"strings"
) )
func notifyRootUser(subject string, content string) { func notifyRootUser(subject string, content string) {
...@@ -15,3 +18,43 @@ func notifyRootUser(subject string, content string) { ...@@ -15,3 +18,43 @@ func notifyRootUser(subject string, content string) {
common.SysError(fmt.Sprintf("failed to send email: %s", err.Error())) common.SysError(fmt.Sprintf("failed to send email: %s", err.Error()))
} }
} }
func NotifyUser(user *model.UserCache, data dto.Notify) error {
userSetting := user.GetSetting()
notifyType, ok := userSetting[constant.UserSettingNotifyType]
if !ok {
notifyType = constant.NotifyTypeEmail
}
switch notifyType {
case constant.NotifyTypeEmail:
userEmail := user.Email
// check setting email
if settingEmail, ok := userSetting[constant.UserSettingNotificationEmail]; ok {
userEmail = settingEmail.(string)
}
if userEmail == "" {
common.SysLog(fmt.Sprintf("user %d has no email, skip sending email", user.Id))
return nil
}
return sendEmailNotify(userEmail, data)
case constant.NotifyTypeWebhook:
webhookURL, ok := userSetting[constant.UserSettingWebhookUrl]
if !ok {
common.SysError(fmt.Sprintf("user %d has no webhook url, skip sending webhook", user.Id))
return nil
}
// TODO: 实现webhook通知
_ = webhookURL // 临时处理未使用警告，等待webhook实现
}
return nil // 添加缺失的return
}
func sendEmailNotify(userEmail string, data dto.Notify) error {
// make email content
content := data.Content
// 处理占位符
for _, value := range data.Values {
content = strings.Replace(content, dto.ContentValueParam, fmt.Sprintf("%v", value), 1)
}
return common.SendEmail(data.Title, userEmail, content)
}
web/src/helpers/render.js
...@@ -386,7 +386,7 @@ export function renderQuotaWithPrompt(quota, digits) { ...@@ -386,7 +386,7 @@ export function renderQuotaWithPrompt(quota, digits) {
let displayInCurrency = localStorage.getItem('display_in_currency'); let displayInCurrency = localStorage.getItem('display_in_currency');
displayInCurrency = displayInCurrency === 'true'; displayInCurrency = displayInCurrency === 'true';
if (displayInCurrency) { if (displayInCurrency) {
return '|' + i18next.t('等价金额') + ': ' + renderQuota(quota, digits) + ''; return ' | ' + i18next.t('等价金额') + ': ' + renderQuota(quota, digits) + '';
} }
return ''; return '';
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment