xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-09-17 08:46:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add sync lock for sub or add user's power

Browse Source
This commit is contained in:
RockYang 2024-09-03 12:09:36 +08:00
parent aa46fdb113
commit 13bf73e6cf
16 changed files with 341 additions and 423 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@ -9,6 +9,7 @@
* Bug修复：手机端角色和模型选择不生效
* Bug修复：用户登录过期之后聊天页面出现大量报错，需要刷新页面才能正常
* 功能优化：优化聊天页面 Websocket 断线重连代码，提高用户体验
* 功能优化：给算力增减服务全部加上数据库事务和同步锁
* 功能新增：支持 Luma 文生视频功能
## v4.1.2

29
api/handler/chatimpl/chat_handler.go
View File

@ -46,9 +46,10 @@ type ChatHandler struct {
licenseService *service.LicenseService
ReqCancelFunc *types.LMap[string, context.CancelFunc] // HttpClient 请求取消 handle function
ChatContexts *types.LMap[string, []types.Message] // 聊天上下文 Map [chatId] => []Message
userService *service.UserService
}
func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService) *ChatHandler {
func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService, userService *service.UserService) *ChatHandler {
return &ChatHandler{
BaseHandler: handler.BaseHandler{App: app, DB: db},
redis: redis,
@ -56,6 +57,7 @@ func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manag
licenseService: licenseService,
ReqCancelFunc: types.NewLMap[string, context.CancelFunc](),
ChatContexts: types.NewLMap[string, []types.Message](),
userService: userService,
}
}
@ -482,24 +484,15 @@ func (h *ChatHandler) subUserPower(userVo vo.User, session *types.ChatSession, p
if session.Model.Power > 0 {
power = session.Model.Power
}
res := h.DB.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("power", gorm.Expr("power - ?", power))
if res.Error == nil {
// 记录算力消费日志
var u model.User
h.DB.Where("id", userVo.Id).First(&u)
h.DB.Create(&model.PowerLog{
UserId: userVo.Id,
Username: userVo.Username,
Type: types.PowerConsume,
Amount: power,
Mark: types.PowerSub,
Balance: u.Power,
Model: session.Model.Value,
Remark: fmt.Sprintf("模型名称：%s, 提问长度：%d，回复长度：%d", session.Model.Name, promptTokens, replyTokens),
CreatedAt: time.Now(),
})
}
err := h.userService.DecreasePower(int(userVo.Id), power, model.PowerLog{
Type: types.PowerConsume,
Model: session.Model.Value,
Remark: fmt.Sprintf("模型名称：%s, 提问长度：%d，回复长度：%d", session.Model.Name, promptTokens, replyTokens),
})
if err != nil {
logger.Error(err)
}
}
func (h *ChatHandler) saveChatHistory(

54
api/handler/dalle_handler.go
View File

@ -11,32 +11,33 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/service/dalle"
"geekai/service/oss"
"geekai/store/model"
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/gorilla/websocket"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/gorilla/websocket"
"gorm.io/gorm"
"net/http"
)
type DallJobHandler struct {
BaseHandler
redis *redis.Client
service *dalle.Service
uploader *oss.UploaderManager
redis *redis.Client
dallService *dalle.Service
uploader *oss.UploaderManager
userService *service.UserService
}
func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service, manager *oss.UploaderManager) *DallJobHandler {
func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service, manager *oss.UploaderManager, userService *service.UserService) *DallJobHandler {
return &DallJobHandler{
service: service,
uploader: manager,
dallService: service,
uploader: manager,
userService: userService,
BaseHandler: BaseHandler{
App: app,
DB: db,
@ -61,14 +62,14 @@ func (h *DallJobHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.service.Clients.Put(uint(userId), client)
h.dallService.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
go func() {
for {
_, msg, err := client.Receive()
if err != nil {
client.Close()
h.service.Clients.Delete(uint(userId))
h.dallService.Clients.Delete(uint(userId))
return
}
@ -127,7 +128,7 @@ func (h *DallJobHandler) Image(c *gin.Context) {
return
}
h.service.PushTask(types.DallTask{
h.dallService.PushTask(types.DallTask{
JobId: job.Id,
UserId: uint(userId),
Prompt: data.Prompt,
@ -137,7 +138,7 @@ func (h *DallJobHandler) Image(c *gin.Context) {
Power: job.Power,
})
client := h.service.Clients.Get(job.UserId)
client := h.dallService.Clients.Get(job.UserId)
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@ -233,26 +234,11 @@ func (h *DallJobHandler) Remove(c *gin.Context) {
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
tx.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerRefund,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
err := h.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())

30
api/handler/markmap_handler.go
View File

@ -15,6 +15,7 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/store/model"
"geekai/utils"
"github.com/gin-gonic/gin"
@ -30,13 +31,15 @@ import (
// MarkMapHandler 生成思维导图
type MarkMapHandler struct {
BaseHandler
clients *types.LMap[int, *types.WsClient]
clients *types.LMap[int, *types.WsClient]
userService *service.UserService
}
func NewMarkMapHandler(app *core.AppServer, db *gorm.DB) *MarkMapHandler {
func NewMarkMapHandler(app *core.AppServer, db *gorm.DB, userService *service.UserService) *MarkMapHandler {
return &MarkMapHandler{
BaseHandler: BaseHandler{App: app, DB: db},
clients: types.NewLMap[int, *types.WsClient](),
userService: userService,
}
}
@ -185,22 +188,13 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
// 扣减算力
if chatModel.Power > 0 {
res = h.DB.Model(&model.User{}).Where("id", userId).UpdateColumn("power", gorm.Expr("power - ?", chatModel.Power))
if res.Error == nil {
// 记录算力消费日志
var u model.User
h.DB.Where("id", userId).First(&u)
h.DB.Create(&model.PowerLog{
UserId: u.Id,
Username: u.Username,
Type: types.PowerConsume,
Amount: chatModel.Power,
Mark: types.PowerSub,
Balance: u.Power,
Model: chatModel.Value,
Remark: fmt.Sprintf("AI绘制思维导图，模型名称：%s, ", chatModel.Value),
CreatedAt: time.Now(),
})
err = h.userService.DecreasePower(userId, chatModel.Power, model.PowerLog{
Type: types.PowerConsume,
Model: chatModel.Value,
Remark: fmt.Sprintf("AI绘制思维导图，模型名称：%s, ", chatModel.Value),
})
if err != nil {
return err
}
}

130
api/handler/mj_handler.go
View File

@ -30,16 +30,18 @@ import (
type MidJourneyHandler struct {
BaseHandler
service *mj.Service
snowflake *service.Snowflake
uploader *oss.UploaderManager
mjService *mj.Service
snowflake *service.Snowflake
uploader *oss.UploaderManager
userService *service.UserService
}
func NewMidJourneyHandler(app *core.AppServer, db *gorm.DB, snowflake *service.Snowflake, service *mj.Service, manager *oss.UploaderManager) *MidJourneyHandler {
func NewMidJourneyHandler(app *core.AppServer, db *gorm.DB, snowflake *service.Snowflake, service *mj.Service, manager *oss.UploaderManager, userService *service.UserService) *MidJourneyHandler {
return &MidJourneyHandler{
snowflake: snowflake,
service: service,
uploader: manager,
snowflake: snowflake,
mjService: service,
uploader: manager,
userService: userService,
BaseHandler: BaseHandler{
App: app,
DB: db,
@ -80,7 +82,7 @@ func (h *MidJourneyHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.service.Clients.Put(uint(userId), client)
h.mjService.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@ -196,7 +198,7 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
return
}
h.service.PushTask(types.MjTask{
h.mjService.PushTask(types.MjTask{
Id: job.Id,
TaskId: taskId,
Type: types.TaskType(data.TaskType),
@ -208,28 +210,22 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
Mode: h.App.SysConfig.MjMode,
})
client := h.service.Clients.Get(uint(job.UserId))
client := h.mjService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
// update user's power
tx := h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
user, _ := h.GetLoginUser(c)
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power - job.Power,
Mark: types.PowerSub,
Model: "mid-journey",
Remark: fmt.Sprintf("%s操作，任务ID：%s", opt, job.TaskId),
CreatedAt: time.Now(),
})
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "mid-journey",
Remark: fmt.Sprintf("%s操作，任务ID：%s", opt, job.TaskId),
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
}
@ -269,7 +265,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
return
}
h.service.PushTask(types.MjTask{
h.mjService.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskUpscale,
UserId: userId,
@ -280,27 +276,22 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
Mode: h.App.SysConfig.MjMode,
})
client := h.service.Clients.Get(uint(job.UserId))
client := h.mjService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
// update user's power
tx := h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
user, _ := h.GetLoginUser(c)
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power - job.Power,
Mark: types.PowerSub,
Model: "mid-journey",
Remark: fmt.Sprintf("Upscale 操作，任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
})
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "mid-journey",
Remark: fmt.Sprintf("Upscale 操作，任务ID：%s", job.TaskId),
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
}
@ -334,7 +325,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
return
}
h.service.PushTask(types.MjTask{
h.mjService.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskVariation,
UserId: userId,
@ -345,28 +336,21 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
Mode: h.App.SysConfig.MjMode,
})
client := h.service.Clients.Get(uint(job.UserId))
client := h.mjService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
// update user's power
tx := h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
user, _ := h.GetLoginUser(c)
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power - job.Power,
Mark: types.PowerSub,
Model: "mid-journey",
Remark: fmt.Sprintf("Variation 操作，任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
})
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "mid-journey",
Remark: fmt.Sprintf("Variation 操作，任务ID：%s", job.TaskId),
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
}
@ -465,25 +449,11 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
tx.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerRefund,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "mid-journey",
Remark: fmt.Sprintf("绘画任务失败，退回算力。任务ID：%s，Err: %s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
err := h.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "mid-journey",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
@ -498,7 +468,7 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
logger.Error("remove image failed: ", err)
}
client := h.service.Clients.Get(uint(job.UserId))
client := h.mjService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}

34
api/handler/redeem_handler.go
View File

@ -11,6 +11,7 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/store/model"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
@ -21,11 +22,12 @@ import (
type RedeemHandler struct {
BaseHandler
lock sync.Mutex
lock sync.Mutex
userService *service.UserService
}
func NewRedeemHandler(app *core.AppServer, db *gorm.DB) *RedeemHandler {
return &RedeemHandler{BaseHandler: BaseHandler{App: app, DB: db}}
func NewRedeemHandler(app *core.AppServer, db *gorm.DB, userService *service.UserService) *RedeemHandler {
return &RedeemHandler{BaseHandler: BaseHandler{App: app, DB: db}, userService: userService}
}
func (h *RedeemHandler) Verify(c *gin.Context) {
@ -59,7 +61,11 @@ func (h *RedeemHandler) Verify(c *gin.Context) {
}
tx := h.DB.Begin()
err := tx.Model(&model.User{}).Where("id", userId).UpdateColumn("power", gorm.Expr("power + ?", item.Power)).Error
err := h.userService.IncreasePower(int(userId), item.Power, model.PowerLog{
Type: types.PowerRedeem,
Model: "兑换码",
Remark: fmt.Sprintf("兑换码核销，算力：%d，兑换码：%s...", item.Power, item.Code[:10]),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
@ -76,26 +82,6 @@ func (h *RedeemHandler) Verify(c *gin.Context) {
return
}
// 记录算力充值日志
var user model.User
err = tx.Where("id", userId).First(&user).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
h.DB.Create(&model.PowerLog{
UserId: userId,
Username: user.Username,
Type: types.PowerRedeem,
Amount: item.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "兑换码",
Remark: fmt.Sprintf("兑换码核销，算力：%d，兑换码：%s...", item.Power, item.Code[:10]),
CreatedAt: time.Now(),
})
tx.Commit()
resp.SUCCESS(c)

81
api/handler/sd_handler.go
View File

@ -31,19 +31,27 @@ import (
type SdJobHandler struct {
BaseHandler
redis *redis.Client
service *sd.Service
uploader *oss.UploaderManager
snowflake *service.Snowflake
leveldb *store.LevelDB
redis *redis.Client
sdService *sd.Service
uploader *oss.UploaderManager
snowflake *service.Snowflake
leveldb *store.LevelDB
userService *service.UserService
}
func NewSdJobHandler(app *core.AppServer, db *gorm.DB, service *sd.Service, manager *oss.UploaderManager, snowflake *service.Snowflake, levelDB *store.LevelDB) *SdJobHandler {
func NewSdJobHandler(app *core.AppServer,
db *gorm.DB,
service *sd.Service,
manager *oss.UploaderManager,
snowflake *service.Snowflake,
userService *service.UserService,
levelDB *store.LevelDB) *SdJobHandler {
return &SdJobHandler{
service: service,
uploader: manager,
snowflake: snowflake,
leveldb: levelDB,
sdService: service,
uploader: manager,
snowflake: snowflake,
leveldb: levelDB,
userService: userService,
BaseHandler: BaseHandler{
App: app,
DB: db,
@ -68,7 +76,7 @@ func (h *SdJobHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.service.Clients.Put(uint(userId), client)
h.sdService.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@ -159,34 +167,27 @@ func (h *SdJobHandler) Image(c *gin.Context) {
return
}
h.service.PushTask(types.SdTask{
h.sdService.PushTask(types.SdTask{
Id: int(job.Id),
Type: types.TaskImage,
Params: params,
UserId: userId,
})
client := h.service.Clients.Get(uint(job.UserId))
client := h.sdService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
// update user's power
tx := h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
user, _ := h.GetLoginUser(c)
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power - job.Power,
Mark: types.PowerSub,
Model: "stable-diffusion",
Remark: fmt.Sprintf("绘图操作，任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
})
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "stable-diffusion",
Remark: fmt.Sprintf("绘图操作，任务ID：%s", job.TaskId),
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@ -290,25 +291,11 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
tx.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerRefund,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%s， Err: %s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%s， Err: %s", job.TaskId, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())

66
api/handler/suno_handler.go
View File

@ -11,6 +11,7 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/service/oss"
"geekai/service/suno"
"geekai/store/model"
@ -26,18 +27,20 @@ import (
type SunoHandler struct {
BaseHandler
service *suno.Service
uploader *oss.UploaderManager
sunoService *suno.Service
uploader *oss.UploaderManager
userService *service.UserService
}
func NewSunoHandler(app *core.AppServer, db *gorm.DB, service *suno.Service, uploader *oss.UploaderManager) *SunoHandler {
func NewSunoHandler(app *core.AppServer, db *gorm.DB, service *suno.Service, uploader *oss.UploaderManager, userService *service.UserService) *SunoHandler {
return &SunoHandler{
BaseHandler: BaseHandler{
App: app,
DB: db,
},
service: service,
uploader: uploader,
sunoService: service,
uploader: uploader,
userService: userService,
}
}
@ -58,7 +61,7 @@ func (h *SunoHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.service.Clients.Put(uint(userId), client)
h.sunoService.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@ -123,7 +126,7 @@ func (h *SunoHandler) Create(c *gin.Context) {
}
// 创建任务
h.service.PushTask(types.SunoTask{
h.sunoService.PushTask(types.SunoTask{
Id: job.Id,
UserId: job.UserId,
Type: job.Type,
@ -140,24 +143,17 @@ func (h *SunoHandler) Create(c *gin.Context) {
})
// update user's power
tx = h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
user, _ := h.GetLoginUser(c)
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power - job.Power,
Mark: types.PowerSub,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 文生歌曲，%s", job.ModelName),
CreatedAt: time.Now(),
})
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Remark: fmt.Sprintf("Suno 文生歌曲，%s", job.ModelName),
CreatedAt: time.Now(),
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
client := h.service.Clients.Get(uint(job.UserId))
client := h.sunoService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@ -239,25 +235,11 @@ func (h *SunoHandler) Remove(c *gin.Context) {
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
tx.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerRefund,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())

44
api/handler/user_handler.go
View File

@ -34,6 +34,7 @@ type UserHandler struct {
redis *redis.Client
licenseService *service.LicenseService
captcha *service.CaptchaService
userService *service.UserService
}
func NewUserHandler(
@ -42,6 +43,7 @@ func NewUserHandler(
searcher *xdb.Searcher,
client *redis.Client,
captcha *service.CaptchaService,
userService *service.UserService,
licenseService *service.LicenseService) *UserHandler {
return &UserHandler{
BaseHandler: BaseHandler{DB: db, App: app},
@ -49,6 +51,7 @@ func NewUserHandler(
redis: client,
captcha: captcha,
licenseService: licenseService,
userService: userService,
}
}
@ -155,10 +158,9 @@ func (h *UserHandler) Register(c *gin.Context) {
user.Nickname = fmt.Sprintf("极客学长@%d", utils.RandomNumber(6))
}
res := h.DB.Create(&user)
if res.Error != nil {
resp.ERROR(c, "保存数据失败")
logger.Error(res.Error)
tx := h.DB.Begin()
if err := tx.Create(&user).Error; err != nil {
resp.ERROR(c, err.Error())
return
}
@ -167,35 +169,35 @@ func (h *UserHandler) Register(c *gin.Context) {
// 增加邀请数量
h.DB.Model(&model.InviteCode{}).Where("code = ?", data.InviteCode).UpdateColumn("reg_num", gorm.Expr("reg_num + ?", 1))
if h.App.SysConfig.InvitePower > 0 {
h.DB.Model(&model.User{}).Where("id = ?", inviteCode.UserId).UpdateColumn("power", gorm.Expr("power + ?", h.App.SysConfig.InvitePower))
// 记录邀请算力充值日志
var inviter model.User
h.DB.Where("id", inviteCode.UserId).First(&inviter)
h.DB.Create(&model.PowerLog{
UserId: inviter.Id,
Username: inviter.Username,
Type: types.PowerInvite,
Amount: h.App.SysConfig.InvitePower,
Balance: inviter.Power,
Mark: types.PowerAdd,
Model: "",
Remark: fmt.Sprintf("邀请用户注册奖励，金额：%d，邀请码：%s，新用户：%s", h.App.SysConfig.InvitePower, inviteCode.Code, user.Username),
CreatedAt: time.Now(),
err := h.userService.IncreasePower(int(inviteCode.UserId), h.App.SysConfig.InvitePower, model.PowerLog{
Type: types.PowerInvite,
Model: "",
Remark: fmt.Sprintf("邀请用户注册奖励，金额：%d，邀请码：%s，新用户：%s", h.App.SysConfig.InvitePower, inviteCode.Code, user.Username),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
// 添加邀请记录
h.DB.Create(&model.InviteLog{
err := tx.Create(&model.InviteLog{
InviterId: inviteCode.UserId,
UserId: user.Id,
Username: user.Username,
InviteCode: inviteCode.Code,
Remark: fmt.Sprintf("奖励 %d 算力", h.App.SysConfig.InvitePower),
})
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
_ = h.redis.Del(c, key) // 注册成功，删除短信验证码
// 自动登录创建 token
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
"user_id": user.Id,

76
api/handler/video_handler.go
View File

@ -11,6 +11,7 @@ import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/service/oss"
"geekai/service/video"
"geekai/store/model"
@ -21,23 +22,24 @@ import (
"github.com/gorilla/websocket"
"gorm.io/gorm"
"net/http"
"time"
)
type VideoHandler struct {
BaseHandler
service *video.Service
uploader *oss.UploaderManager
videoService *video.Service
uploader *oss.UploaderManager
userService *service.UserService
}
func NewVideoHandler(app *core.AppServer, db *gorm.DB, service *video.Service, uploader *oss.UploaderManager) *VideoHandler {
func NewVideoHandler(app *core.AppServer, db *gorm.DB, service *video.Service, uploader *oss.UploaderManager, userService *service.UserService) *VideoHandler {
return &VideoHandler{
BaseHandler: BaseHandler{
App: app,
DB: db,
},
service: service,
uploader: uploader,
videoService: service,
uploader: uploader,
userService: userService,
}
}
@ -58,7 +60,7 @@ func (h *VideoHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.service.Clients.Put(uint(userId), client)
h.videoService.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@ -102,7 +104,7 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) {
}
// 创建任务
h.service.PushTask(types.VideoTask{
h.videoService.PushTask(types.VideoTask{
Id: job.Id,
UserId: userId,
Type: types.VideoLuma,
@ -111,24 +113,17 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) {
})
// update user's power
tx = h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
user, _ := h.GetLoginUser(c)
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power - job.Power,
Mark: types.PowerSub,
Model: "luma",
Remark: fmt.Sprintf("Luma 文生视频，任务ID：%d", job.Id),
CreatedAt: time.Now(),
})
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "luma",
Remark: fmt.Sprintf("Luma 文生视频，任务ID：%d", job.Id),
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
client := h.service.Clients.Get(uint(job.UserId))
client := h.videoService.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@ -194,25 +189,11 @@ func (h *VideoHandler) Remove(c *gin.Context) {
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
tx.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerRefund,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "luma",
Remark: fmt.Sprintf("Luma 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "luma",
Remark: fmt.Sprintf("Luma 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
@ -230,7 +211,14 @@ func (h *VideoHandler) Publish(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetLoginUserId(c)
publish := h.GetBool(c, "publish")
err := h.DB.Model(&model.VideoJob{}).Where("id", id).Where("user_id", userId).UpdateColumn("publish", publish).Error
var job model.VideoJob
err := h.DB.Where("id = ?", id).Where("user_id", userId).First(&job).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
err = h.DB.Model(&job).UpdateColumn("publish", publish).Error
if err != nil {
resp.ERROR(c, err.Error())
return

2
api/main.go
View File

@ -209,7 +209,7 @@ func main() {
s.CheckTaskNotify()
s.DownloadFiles()
}),
fx.Provide(service.NewUserService),
fx.Provide(payment.NewAlipayService),
fx.Provide(payment.NewHuPiPay),
fx.Provide(payment.NewJPayService),

45
api/service/dalle/service.go
View File

@ -35,9 +35,10 @@ type Service struct {
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
userService *service.UserService
}
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client) *Service {
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService) *Service {
return &Service{
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
@ -45,6 +46,7 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
notifyQueue: store.NewRedisQueue("DallE_Notify_Queue", redisCli),
Clients: types.NewLMap[uint, *types.WsClient](),
uploadManager: manager,
userService: userService,
}
}
@ -122,32 +124,23 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
return "", errors.New("insufficient of power")
}
// 更新用户算力
tx := s.db.Model(&model.User{}).Where("id", user.Id).UpdateColumn("power", gorm.Expr("power - ?", task.Power))
// 记录算力变化日志
if tx.Error == nil && tx.RowsAffected > 0 {
var u model.User
s.db.Where("id", user.Id).First(&u)
s.db.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: task.Power,
Balance: u.Power,
Mark: types.PowerSub,
Model: "dall-e-3",
Remark: fmt.Sprintf("绘画提示词：%s", utils.CutWords(task.Prompt, 10)),
CreatedAt: time.Now(),
})
// 扣减算力
err := s.userService.DecreasePower(int(user.Id), task.Power, model.PowerLog{
Type: types.PowerConsume,
Model: "dall-e-3",
Remark: fmt.Sprintf("绘画提示词：%s", utils.CutWords(task.Prompt, 10)),
})
if err != nil {
return "", fmt.Errorf("error with decrease power: %v", err)
}
// get image generation API KEY
var apiKey model.ApiKey
tx = s.db.Where("type", "dalle").
err = s.db.Where("type", "dalle").
Where("enabled", true).
Order("last_used_at ASC").First(&apiKey)
if tx.Error != nil {
return "", fmt.Errorf("no available DALL-E api key: %v", tx.Error)
Order("last_used_at ASC").First(&apiKey).Error
if err != nil {
return "", fmt.Errorf("no available DALL-E api key: %v", err)
}
var res imgRes
@ -181,13 +174,13 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
// update the api key last use time
s.db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
// update task progress
tx = s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
err = s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
"progress": 100,
"org_url": res.Data[0].Url,
"prompt": prompt,
})
if tx.Error != nil {
return "", fmt.Errorf("err with update database: %v", tx.Error)
}).Error
if err != nil {
return "", fmt.Errorf("err with update database: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: service.TaskStatusFailed})

83
api/service/user_service.go Normal file
View File

@ -0,0 +1,83 @@
package service
import (
"fmt"
"geekai/core/types"
"geekai/store/model"
"gorm.io/gorm"
"sync"
"time"
)
type UserService struct {
db *gorm.DB
lock sync.Mutex
}
func NewUserService(db *gorm.DB) *UserService {
return &UserService{db: db, lock: sync.Mutex{}}
}
// IncreasePower 增加用户算力
func (s *UserService) IncreasePower(userId int, power int, log model.PowerLog) error {
s.lock.Lock()
defer s.lock.Unlock()
tx := s.db.Begin()
err := tx.Model(&model.User{}).Where("id", userId).UpdateColumn("power", gorm.Expr("power + ?", power)).Error
if err != nil {
tx.Rollback()
return err
}
var user model.User
tx.Where("id", userId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: log.Type,
Amount: power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: log.Model,
Remark: log.Remark,
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
return err
}
tx.Commit()
return nil
}
// DecreasePower 减少用户算力
func (s *UserService) DecreasePower(userId int, power int, log model.PowerLog) error {
s.lock.Lock()
defer s.lock.Unlock()
tx := s.db.Begin()
err := tx.Model(&model.User{}).Where("id", userId).UpdateColumn("power", gorm.Expr("power - ?", power)).Error
if err != nil {
tx.Rollback()
return fmt.Errorf("扣减算力失败：%v", err)
}
var user model.User
tx.Where("id", userId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: log.Type,
Amount: power,
Balance: user.Power,
Mark: types.PowerSub,
Model: log.Model,
Remark: log.Remark,
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
return fmt.Errorf("记录算力日志失败：%v", err)
}
tx.Commit()
return nil
}

48
api/service/xxl_job_service.go
View File

@ -81,54 +81,6 @@ func (e *XXLJobExecutor) ClearOrders(cxt context.Context, param *xxl.RunReq) (ms
// 自动将 VIP 会员的算力补充到每月赠送的最大值
func (e *XXLJobExecutor) ResetVipPower(cxt context.Context, param *xxl.RunReq) (msg string) {
logger.Info("开始进行月底账号盘点...")
var users []model.User
res := e.db.Where("vip", 1).Where("status", 1).Find(&users)
if res.Error != nil {
return "No vip users found"
}
var sysConfig model.Config
res = e.db.Where("marker", "system").First(&sysConfig)
if res.Error != nil {
return "error with get system config: " + res.Error.Error()
}
var config types.SystemConfig
err := utils.JsonDecode(sysConfig.Config, &config)
if err != nil {
return "error with decode system config: " + err.Error()
}
for _, u := range users {
// 处理过期的 VIP
if u.ExpiredTime > 0 && u.ExpiredTime <= time.Now().Unix() {
u.Vip = false
e.db.Model(&model.User{}).Where("id", u.Id).UpdateColumn("vip", false)
continue
}
if u.Power < config.VipMonthPower {
power := config.VipMonthPower - u.Power
// update user
tx := e.db.Model(&model.User{}).Where("id", u.Id).UpdateColumn("power", gorm.Expr("power + ?", power))
// 记录算力变动日志
if tx.Error == nil {
var user model.User
e.db.Where("id", u.Id).First(&user)
e.db.Create(&model.PowerLog{
UserId: u.Id,
Username: u.Username,
Type: types.PowerRecharge,
Amount: power,
Mark: types.PowerAdd,
Balance: user.Power,
Model: "系统盘点",
Remark: fmt.Sprintf("VIP会员每月算力派发，：%d", config.VipMonthPower),
CreatedAt: time.Now(),
})
}
}
}
logger.Info("月底盘点完成！")
return "success"
}

38
deploy/docker-compose.yaml
View File

@ -27,17 +27,17 @@ services:
ports:
- "6380:6379"
xxl-job-admin:
container_name: geekai-xxl-job-admin
image: registry.cn-shenzhen.aliyuncs.com/geekmaster/xxl-job-admin:2.4.0
restart: always
ports:
- "8081:8080"
environment:
- PARAMS=--spring.config.location=/application.properties
volumes:
- ./logs/xxl-job:/data/applogs
- ./conf/xxl-job/application.properties:/application.properties
# xxl-job-admin:
# container_name: geekai-xxl-job-admin
# image: registry.cn-shenzhen.aliyuncs.com/geekmaster/xxl-job-admin:2.4.0
# restart: always
# ports:
# - "8081:8080"
# environment:
# - PARAMS=--spring.config.location=/application.properties
# volumes:
# - ./logs/xxl-job:/data/applogs
# - ./conf/xxl-job/application.properties:/application.properties
tika:
image: registry.cn-shenzhen.aliyuncs.com/geekmaster/tika:latest
@ -46,14 +46,14 @@ services:
ports:
- "9998:9998"
midjourney-proxy:
image: registry.cn-shenzhen.aliyuncs.com/geekmaster/midjourney-proxy:2.6.2
container_name: geekai-midjourney-proxy
restart: always
ports:
- "8082:8080"
volumes:
- ./conf/mj-proxy:/home/spring/config
# midjourney-proxy:
# image: registry.cn-shenzhen.aliyuncs.com/geekmaster/midjourney-proxy:2.6.2
# container_name: geekai-midjourney-proxy
# restart: always
# ports:
# - "8082:8080"
# volumes:
# - ./conf/mj-proxy:/home/spring/config
# 后端 API 程序

3
web/src/views/admin/Users.vue
View File

@ -11,7 +11,8 @@
<el-table :data="users.items" border class="table" :row-key="row => row.id"
@selection-change="handleSelectionChange" table-layout="auto">
<el-table-column type="selection" width="38"></el-table-column>
<el-table-column prop="mobile" label="账号">
<el-table-column prop="id" label="ID"/>
<el-table-column label="账号">
<template #default="scope">
<span>{{ scope.row.username }}</span>
<el-image v-if="scope.row.vip" :src="vipImg" style="height: 20px;position: relative; top:5px; left: 5px"/>