feat: 合并main

This commit is contained in:
suziheng
2025-04-21 16:36:54 +08:00
138 changed files with 8404 additions and 3879 deletions

View File

@@ -2,10 +2,13 @@ package model
import (
"context"
"github.com/songquanpeng/one-api/common"
"gorm.io/gorm"
"sort"
"strings"
"gorm.io/gorm"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/utils"
)
type Ability struct {
@@ -49,6 +52,7 @@ func GetRandomSatisfiedChannel(group string, model string, ignoreFirstPriority b
func (channel *Channel) AddAbilities() error {
models_ := strings.Split(channel.Models, ",")
models_ = utils.DeDuplication(models_)
groups_ := strings.Split(channel.Group, ",")
abilities := make([]Ability, 0, len(models_))
for _, model := range models_ {

View File

@@ -4,26 +4,31 @@ import (
"context"
"fmt"
"gorm.io/gorm"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger"
"gorm.io/gorm"
)
type Log struct {
Id int `json:"id"`
UserId int `json:"user_id" gorm:"index"`
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_type"`
Type int `json:"type" gorm:"index:idx_created_at_type"`
Content string `json:"content"`
Username string `json:"username" gorm:"index:index_username_model_name,priority:2;default:''"`
TokenName string `json:"token_name" gorm:"index;default:''"`
ModelName string `json:"model_name" gorm:"index;index:index_username_model_name,priority:1;default:''"`
Quota int `json:"quota" gorm:"default:0"`
PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
ChannelId int `json:"channel" gorm:"index"`
Id int `json:"id"`
UserId int `json:"user_id" gorm:"index"`
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_type"`
Type int `json:"type" gorm:"index:idx_created_at_type"`
Content string `json:"content"`
Username string `json:"username" gorm:"index:index_username_model_name,priority:2;default:''"`
TokenName string `json:"token_name" gorm:"index;default:''"`
ModelName string `json:"model_name" gorm:"index;index:index_username_model_name,priority:1;default:''"`
Quota int `json:"quota" gorm:"default:0"`
PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
ChannelId int `json:"channel" gorm:"index"`
RequestId string `json:"request_id" gorm:"default:''"`
ElapsedTime int64 `json:"elapsed_time" gorm:"default:0"` // unit is ms
IsStream bool `json:"is_stream" gorm:"default:false"`
SystemPromptReset bool `json:"system_prompt_reset" gorm:"default:false"`
}
const (
@@ -32,9 +37,21 @@ const (
LogTypeConsume
LogTypeManage
LogTypeSystem
LogTypeTest
)
func RecordLog(userId int, logType int, content string) {
func recordLogHelper(ctx context.Context, log *Log) {
requestId := helper.GetRequestID(ctx)
log.RequestId = requestId
err := LOG_DB.Create(log).Error
if err != nil {
logger.Error(ctx, "failed to record log: "+err.Error())
return
}
logger.Infof(ctx, "record log: %+v", log)
}
func RecordLog(ctx context.Context, userId int, logType int, content string) {
if logType == LogTypeConsume && !config.LogConsumeEnabled {
return
}
@@ -45,13 +62,10 @@ func RecordLog(userId int, logType int, content string) {
Type: logType,
Content: content,
}
err := LOG_DB.Create(log).Error
if err != nil {
logger.SysError("failed to record log: " + err.Error())
}
recordLogHelper(ctx, log)
}
func RecordTopupLog(userId int, content string, quota int) {
func RecordTopupLog(ctx context.Context, userId int, content string, quota int) {
log := &Log{
UserId: userId,
Username: GetUsernameById(userId),
@@ -60,34 +74,23 @@ func RecordTopupLog(userId int, content string, quota int) {
Content: content,
Quota: quota,
}
err := LOG_DB.Create(log).Error
if err != nil {
logger.SysError("failed to record log: " + err.Error())
}
recordLogHelper(ctx, log)
}
func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int64, content string) {
logger.Info(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
func RecordConsumeLog(ctx context.Context, log *Log) {
if !config.LogConsumeEnabled {
return
}
log := &Log{
UserId: userId,
Username: GetUsernameById(userId),
CreatedAt: helper.GetTimestamp(),
Type: LogTypeConsume,
Content: content,
PromptTokens: promptTokens,
CompletionTokens: completionTokens,
TokenName: tokenName,
ModelName: modelName,
Quota: int(quota),
ChannelId: channelId,
}
err := LOG_DB.Create(log).Error
if err != nil {
logger.Error(ctx, "failed to record log: "+err.Error())
}
log.Username = GetUsernameById(log.UserId)
log.CreatedAt = helper.GetTimestamp()
log.Type = LogTypeConsume
recordLogHelper(ctx, log)
}
func RecordTestLog(ctx context.Context, log *Log) {
log.CreatedAt = helper.GetTimestamp()
log.Type = LogTypeTest
recordLogHelper(ctx, log)
}
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int) (logs []*Log, err error) {

View File

@@ -1,11 +1,14 @@
package model
import (
"context"
"errors"
"fmt"
"gorm.io/gorm"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/helper"
"gorm.io/gorm"
)
const (
@@ -48,7 +51,7 @@ func GetRedemptionById(id int) (*Redemption, error) {
return &redemption, err
}
func Redeem(key string, userId int) (quota int64, err error) {
func Redeem(ctx context.Context, key string, userId int) (quota int64, err error) {
if key == "" {
return 0, errors.New("未提供兑换码")
}
@@ -82,7 +85,7 @@ func Redeem(key string, userId int) (quota int64, err error) {
if err != nil {
return 0, errors.New("兑换失败," + err.Error())
}
RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s", common.LogQuota(redemption.Quota)))
RecordLog(ctx, userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s", common.LogQuota(redemption.Quota)))
return redemption.Quota, nil
}

View File

@@ -3,12 +3,14 @@ package model
import (
"errors"
"fmt"
"gorm.io/gorm"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/message"
"gorm.io/gorm"
)
const (
@@ -238,16 +240,31 @@ func PreConsumeTokenQuota(tokenId int, quota int64) (err error) {
if err != nil {
logger.SysError("failed to fetch user email: " + err.Error())
}
prompt := "您的额度即将用尽"
prompt := "额度提醒"
var contentText string
if noMoreQuota {
prompt = "您的额度已用尽"
contentText = "您的额度已用尽"
} else {
contentText = "您的额度即将用尽"
}
if email != "" {
topUpLink := fmt.Sprintf("%s/topup", config.ServerAddress)
err = message.SendEmail(prompt, email,
fmt.Sprintf("%s当前剩余额度为 %d为了不影响您的使用请及时充值。<br/>充值链接:<a href='%s'>%s</a>", prompt, userQuota, topUpLink, topUpLink))
content := message.EmailTemplate(
prompt,
fmt.Sprintf(`
<p>您好!</p>
<p>%s当前剩余额度为 <strong>%d</strong>。</p>
<p>为了不影响您的使用,请及时充值。</p>
<p style="text-align: center; margin: 30px 0;">
<a href="%s" style="background-color: #007bff; color: white; padding: 12px 24px; text-decoration: none; border-radius: 4px; display: inline-block;">立即充值</a>
</p>
<p style="color: #666;">如果按钮无法点击,请复制以下链接到浏览器中打开:</p>
<p style="background-color: #f8f8f8; padding: 10px; border-radius: 4px; word-break: break-all;">%s</p>
`, contentText, userQuota, topUpLink, topUpLink),
)
err = message.SendEmail(prompt, email, content)
if err != nil {
logger.SysError("failed to send email" + err.Error())
logger.SysError("failed to send email: " + err.Error())
}
}
}()

View File

@@ -1,16 +1,19 @@
package model
import (
"context"
"errors"
"fmt"
"strings"
"gorm.io/gorm"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/blacklist"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/random"
"gorm.io/gorm"
"strings"
)
const (
@@ -92,7 +95,7 @@ func GetUserById(id int, selectAll bool) (*User, error) {
if selectAll {
err = DB.First(&user, "id = ?", id).Error
} else {
err = DB.Omit("password").First(&user, "id = ?", id).Error
err = DB.Omit("password", "access_token").First(&user, "id = ?", id).Error
}
return &user, err
}
@@ -114,7 +117,7 @@ func DeleteUserById(id int) (err error) {
return user.Delete()
}
func (user *User) Insert(inviterId int) error {
func (user *User) Insert(ctx context.Context, inviterId int) error {
var err error
if user.Password != "" {
user.Password, err = common.Password2Hash(user.Password)
@@ -130,16 +133,16 @@ func (user *User) Insert(inviterId int) error {
return result.Error
}
if config.QuotaForNewUser > 0 {
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(config.QuotaForNewUser)))
RecordLog(ctx, user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(config.QuotaForNewUser)))
}
if inviterId != 0 {
if config.QuotaForInvitee > 0 {
_ = IncreaseUserQuota(user.Id, config.QuotaForInvitee)
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(config.QuotaForInvitee)))
RecordLog(ctx, user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(config.QuotaForInvitee)))
}
if config.QuotaForInviter > 0 {
_ = IncreaseUserQuota(inviterId, config.QuotaForInviter)
RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(config.QuotaForInviter)))
RecordLog(ctx, inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(config.QuotaForInviter)))
}
}
// create default token