This commit is contained in:
CaIon
2023-08-14 22:16:32 +08:00
parent c134604cee
commit 8f2119e410
33 changed files with 3224 additions and 1138 deletions

View File

@@ -3,6 +3,7 @@ package model
import (
"gorm.io/gorm"
"one-api/common"
"strings"
)
type Log struct {
@@ -17,6 +18,7 @@ type Log struct {
Quota int `json:"quota" gorm:"default:0"`
PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
TokenId int `json:"token_id" gorm:"default:0;index"`
}
const (
@@ -27,6 +29,11 @@ const (
LogTypeSystem
)
func GetLogByKey(key string) (logs []*Log, err error) {
err = DB.Joins("left join tokens on tokens.id = logs.token_id").Where("tokens.key = ?", strings.Split(key, "-")[1]).Find(&logs).Error
return logs, err
}
func RecordLog(userId int, logType int, content string) {
if logType == LogTypeConsume && !common.LogConsumeEnabled {
return
@@ -44,7 +51,7 @@ func RecordLog(userId int, logType int, content string) {
}
}
func RecordConsumeLog(userId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string) {
func RecordConsumeLog(userId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string, tokenId int) {
if !common.LogConsumeEnabled {
return
}
@@ -59,6 +66,7 @@ func RecordConsumeLog(userId int, promptTokens int, completionTokens int, modelN
TokenName: tokenName,
ModelName: modelName,
Quota: quota,
TokenId: tokenId,
}
err := DB.Create(log).Error
if err != nil {

View File

@@ -88,6 +88,14 @@ func InitDB() (err error) {
if err != nil {
return err
}
err = db.AutoMigrate(&Midjourney{})
if err != nil {
return err
}
err = db.AutoMigrate(&TopUp{})
if err != nil {
return err
}
common.SysLog("database migrated")
err = createRootAccountIfNeed()
return err

87
model/midjourney.go Normal file
View File

@@ -0,0 +1,87 @@
package model
type Midjourney struct {
Id int `json:"id"`
Code int `json:"code"`
UserId int `json:"user_id" gorm:"index"`
Action string `json:"action"`
MjId string `json:"mj_id" gorm:"index"`
Prompt string `json:"prompt"`
PromptEn string `json:"prompt_en"`
Description string `json:"description"`
State string `json:"state"`
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
FinishTime int64 `json:"finish_time"`
ImageUrl string `json:"image_url"`
Status string `json:"status"`
Progress string `json:"progress"`
FailReason string `json:"fail_reason"`
}
func GetAllUserTask(userId int, startIdx int, num int) []*Midjourney {
var tasks []*Midjourney
var err error
err = DB.Where("user_id = ?", userId).Order("id desc").Limit(num).Offset(startIdx).Find(&tasks).Error
if err != nil {
return nil
}
return tasks
}
func GetAllTasks(startIdx int, num int) []*Midjourney {
var tasks []*Midjourney
var err error
err = DB.Order("id desc").Limit(num).Offset(startIdx).Find(&tasks).Error
if err != nil {
return nil
}
return tasks
}
func GetAllUnFinishTasks() []*Midjourney {
var tasks []*Midjourney
var err error
// get all tasks progress is not 100%
err = DB.Where("progress != ?", "100%").Find(&tasks).Error
if err != nil {
return nil
}
return tasks
}
func GetByMJId(mjId string) *Midjourney {
var mj *Midjourney
var err error
err = DB.Where("mj_id = ?", mjId).First(&mj).Error
if err != nil {
return nil
}
return mj
}
func GetMjByuId(id int) *Midjourney {
var mj *Midjourney
var err error
err = DB.Where("id = ?", id).First(&mj).Error
if err != nil {
return nil
}
return mj
}
func UpdateProgress(id int, progress string) error {
return DB.Model(&Midjourney{}).Where("id = ?", id).Update("progress", progress).Error
}
func (midjourney *Midjourney) Insert() error {
var err error
err = DB.Create(midjourney).Error
return err
}
func (midjourney *Midjourney) Update() error {
var err error
err = DB.Save(midjourney).Error
return err
}

View File

@@ -69,6 +69,10 @@ func InitOptionMap() {
common.OptionMap["ChatLink"] = common.ChatLink
common.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(common.QuotaPerUnit, 'f', -1, 64)
common.OptionMap["RetryTimes"] = strconv.Itoa(common.RetryTimes)
common.OptionMap["NormalPrice"] = strconv.FormatFloat(common.NormalPrice, 'f', -1, 64)
common.OptionMap["StablePrice"] = strconv.FormatFloat(common.StablePrice, 'f', -1, 64)
common.OptionMap["BasePrice"] = strconv.FormatFloat(common.BasePrice, 'f', -1, 64)
common.OptionMapRWMutex.Unlock()
loadOptionsFromDatabase()
}
@@ -207,6 +211,12 @@ func updateOptionMap(key string, value string) (err error) {
common.TopUpLink = value
case "ChatLink":
common.ChatLink = value
case "NormalPrice":
common.NormalPrice, _ = strconv.ParseFloat(value, 64)
case "BasePrice":
common.BasePrice, _ = strconv.ParseFloat(value, 64)
case "StablePrice":
common.StablePrice, _ = strconv.ParseFloat(value, 64)
case "ChannelDisableThreshold":
common.ChannelDisableThreshold, _ = strconv.ParseFloat(value, 64)
case "QuotaPerUnit":

43
model/topup.go Normal file
View File

@@ -0,0 +1,43 @@
package model
type TopUp struct {
Id int `json:"id"`
UserId int `json:"user_id" gorm:"index"`
Amount int `json:"amount"`
Money int `json:"money"`
TradeNo string `json:"trade_no"`
CreateTime int64 `json:"create_time"`
Status string `json:"status"`
}
func (topUp *TopUp) Insert() error {
var err error
err = DB.Create(topUp).Error
return err
}
func (topUp *TopUp) Update() error {
var err error
err = DB.Save(topUp).Error
return err
}
func GetTopUpById(id int) *TopUp {
var topUp *TopUp
var err error
err = DB.Where("id = ?", id).First(&topUp).Error
if err != nil {
return nil
}
return topUp
}
func GetTopUpByTradeNo(tradeNo string) *TopUp {
var topUp *TopUp
var err error
err = DB.Where("trade_no = ?", tradeNo).First(&topUp).Error
if err != nil {
return nil
}
return topUp
}

View File

@@ -28,6 +28,8 @@ type User struct {
Group string `json:"group" gorm:"type:varchar(32);default:'default'"`
AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
StableMode bool `json:"stable_mode" gorm:"type:tinyint;default:0;column:stable_mode"`
MaxPrice string `json:"max_price" gorm:"type:varchar(32);default:'7'"`
}
func GetMaxUserId() int {
@@ -116,7 +118,14 @@ func (user *User) Update(updatePassword bool) error {
return err
}
}
err = DB.Model(user).Updates(user).Error
newUser := *user
err = DB.Model(user).UpdateColumns(map[string]interface{}{
"stable_mode": user.StableMode,
"max_price": user.MaxPrice,
}).Error
DB.First(&user, user.Id)
err = DB.Model(user).Updates(newUser).Error
return err
}