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
This commit is contained in:
@@ -3,8 +3,10 @@ package service
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
"math"
|
||||
"one-api/common"
|
||||
constant2 "one-api/constant"
|
||||
"one-api/dto"
|
||||
"one-api/model"
|
||||
relaycommon "one-api/relay/common"
|
||||
@@ -99,7 +101,7 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
|
||||
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 {
|
||||
return err
|
||||
}
|
||||
@@ -222,7 +224,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
|
||||
} else {
|
||||
quotaDelta := quota - preConsumedQuota
|
||||
if quotaDelta != 0 {
|
||||
err := model.PostConsumeQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true)
|
||||
err := PostConsumeQuota(relayInfo, quotaDelta, preConsumedQuota, true)
|
||||
if err != nil {
|
||||
common.LogError(ctx, "error consuming token remain quota: "+err.Error())
|
||||
}
|
||||
@@ -239,3 +241,88 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
|
||||
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)
|
||||
}
|
||||
|
||||
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()))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -3,7 +3,10 @@ package service
|
||||
import (
|
||||
"fmt"
|
||||
"one-api/common"
|
||||
"one-api/constant"
|
||||
"one-api/dto"
|
||||
"one-api/model"
|
||||
"strings"
|
||||
)
|
||||
|
||||
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()))
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user