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

Merge tag 'v4.1.7' of gitee.com:blackfox/geekai-plus

Browse Source
This commit is contained in:
RockYang
2025-03-24 11:27:11 +08:00
parent 5f0a62b63e 2e1bad387c
commit a228f603d2
58 changed files with 1146 additions and 494 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
api/handler/admin/redeem_handler.go
View File

@@ -8,6 +8,8 @@ package admin
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"encoding/csv"
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/handler"
@@ -35,12 +37,10 @@ func (h *RedeemHandler) List(c *gin.Context) {
session := h.DB.Session(&gorm.Session{})
if code != "" {
session.Where("code LIKE ?", "%"+code+"%")
session = session.Where("code LIKE ?", "%"+code+"%")
}
if status == 0 {
session.Where("redeem_at = ?", 0)
} else if status == 1 {
session.Where("redeem_at > ?", 0)
if status >= 0 {
session = session.Where("redeemed_at", status)
}
var total int64
@@ -80,6 +80,65 @@ func (h *RedeemHandler) List(c *gin.Context) {
resp.SUCCESS(c, vo.NewPage(total, page, pageSize, items))
}
// Export 导出 CVS 文件
func (h *RedeemHandler) Export(c *gin.Context) {
var data struct {
Status int `json:"status"`
Ids []int `json:"ids"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
}
session := h.DB.Session(&gorm.Session{})
if data.Status >= 0 {
session = session.Where("redeemed_at", data.Status)
}
if len(data.Ids) > 0 {
session = session.Where("id IN ?", data.Ids)
}
var items []model.Redeem
err := session.Order("id DESC").Find(&items).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
// 设置响应头，告诉浏览器这是一个附件，需要下载
c.Header("Content-Disposition", "attachment; filename=output.csv")
c.Header("Content-Type", "text/csv")
// 创建一个 CSV writer
writer := csv.NewWriter(c.Writer)
// 写入 CSV 文件的标题行
headers := []string{"名称", "兑换码", "算力", "创建时间"}
if err := writer.Write(headers); err != nil {
resp.ERROR(c, err.Error())
return
}
// 写入数据行
records := make([][]string, 0)
for _, item := range items {
records = append(records, []string{item.Name, item.Code, fmt.Sprintf("%d", item.Power), item.CreatedAt.Format("2006-01-02 15:04:05")})
}
for _, record := range records {
if err := writer.Write(record); err != nil {
resp.ERROR(c, err.Error())
return
}
}
// 确保所有数据都已写入响应
writer.Flush()
if err := writer.Error(); err != nil {
resp.ERROR(c, err.Error())
return
}
}
func (h *RedeemHandler) Create(c *gin.Context) {
var data struct {
Name string `json:"name"`

25
api/handler/chat_handler.go
View File

@@ -40,7 +40,7 @@ type ChatHandler struct {
uploadManager *oss.UploaderManager
licenseService *service.LicenseService
ReqCancelFunc *types.LMap[string, context.CancelFunc] // HttpClient 请求取消 handle function
ChatContexts *types.LMap[string, []types.Message] // 聊天上下文 Map [chatId] => []Message
ChatContexts *types.LMap[string, []interface{}] // 聊天上下文 Map [chatId] => []Message
userService *service.UserService
}
@@ -51,7 +51,7 @@ func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manag
uploadManager: manager,
licenseService: licenseService,
ReqCancelFunc: types.NewLMap[string, context.CancelFunc](),
ChatContexts: types.NewLMap[string, []types.Message](),
ChatContexts: types.NewLMap[string, []interface{}](),
userService: userService,
}
}
@@ -143,8 +143,8 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
}
// 加载聊天上下文
chatCtx := make([]types.Message, 0)
messages := make([]types.Message, 0)
chatCtx := make([]interface{}, 0)
messages := make([]interface{}, 0)
if h.App.SysConfig.EnableContext {
if h.ChatContexts.Has(session.ChatId) {
messages = h.ChatContexts.Get(session.ChatId)
@@ -174,7 +174,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
for i := len(messages) - 1; i >= 0; i-- {
v := messages[i]
tks, _ = utils.CalcTokens(v.Content, req.Model)
tks, _ = utils.CalcTokens(utils.JsonEncode(v), req.Model)
// 上下文 token 超出了模型的最大上下文长度
if tokens+tks >= session.Model.MaxContext {
break
@@ -192,8 +192,9 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
logger.Debugf("聊天上下文：%+v", chatCtx)
}
reqMgs := make([]interface{}, 0)
for _, m := range chatCtx {
reqMgs = append(reqMgs, m)
for i := len(chatCtx) - 1; i >= 0; i-- {
reqMgs = append(reqMgs, chatCtx[i])
}
fullPrompt := prompt
@@ -258,7 +259,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
logger.Debugf("%+v", req.Messages)
return h.sendOpenAiMessage(chatCtx, req, userVo, ctx, session, role, prompt, ws)
return h.sendOpenAiMessage(req, userVo, ctx, session, role, prompt, ws)
}
// Tokens 统计 token 数量
@@ -371,7 +372,7 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, sessi
} else {
client = http.DefaultClient
}
logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
logger.Infof("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey.Value))
// 更新API KEY 最后使用时间
h.DB.Model(&model.ApiKey{}).Where("id", apiKey.Id).UpdateColumn("last_used_at", time.Now().Unix())
@@ -399,17 +400,15 @@ func (h *ChatHandler) saveChatHistory(
req types.ApiRequest,
usage Usage,
message types.Message,
chatCtx []types.Message,
session *types.ChatSession,
role model.ChatRole,
userVo vo.User,
promptCreatedAt time.Time,
replyCreatedAt time.Time) {
useMsg := types.Message{Role: "user", Content: usage.Prompt}
// 更新上下文消息，如果是调用函数则不需要更新上下文
// 更新上下文消息
if h.App.SysConfig.EnableContext {
chatCtx = append(chatCtx, useMsg) // 提问消息
chatCtx := req.Messages // 提问消息
chatCtx = append(chatCtx, message) // 回复消息
h.ChatContexts.Put(session.ChatId, chatCtx)
}

20
api/handler/chat_model_handler.go
View File

@@ -30,29 +30,25 @@ func NewChatModelHandler(app *core.AppServer, db *gorm.DB) *ChatModelHandler {
func (h *ChatModelHandler) List(c *gin.Context) {
var items []model.ChatModel
var chatModels = make([]vo.ChatModel, 0)
var res *gorm.DB
session := h.DB.Session(&gorm.Session{}).Where("enabled", true)
t := c.Query("type")
if t != "" {
session = session.Where("type", t)
}
// 如果用户没有登录，则加载所有开放模型
if !h.IsLogin(c) {
res = session.Where("open", true).Order("sort_num ASC").Find(&items)
} else {
session = session.Where("open", true)
if h.IsLogin(c) {
user, _ := h.GetLoginUser(c)
var models []int
err := utils.JsonDecode(user.ChatModels, &models)
if err != nil {
resp.ERROR(c, "当前用户没有订阅任何模型")
return
}
// 查询用户有权限访问的模型以及所有开放的模型
res = h.DB.Where("enabled = ?", true).Where(
h.DB.Where("id IN ?", models).Or("open", true),
).Order("sort_num ASC").Find(&items)
if err == nil {
session = session.Or("id IN ?", models)
}
}
res := session.Order("sort_num ASC").Find(&items)
if res.Error == nil {
for _, item := range items {
var cm vo.ChatModel

51
api/handler/dalle_handler.go
View File

@@ -8,7 +8,6 @@ package handler
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
@@ -72,10 +71,21 @@ func (h *DallJobHandler) Image(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
task := types.DallTask{
ClientId: data.ClientId,
UserId: uint(userId),
Prompt: data.Prompt,
Quality: data.Quality,
Size: data.Size,
Style: data.Style,
Power: h.App.SysConfig.DallPower,
TranslateModelId: h.App.SysConfig.TranslateModelId,
}
job := model.DallJob{
UserId: uint(userId),
Prompt: data.Prompt,
Power: h.App.SysConfig.DallPower,
UserId: uint(userId),
Prompt: data.Prompt,
Power: task.Power,
TaskInfo: utils.JsonEncode(task),
}
res := h.DB.Create(&job)
if res.Error != nil {
@@ -83,16 +93,8 @@ func (h *DallJobHandler) Image(c *gin.Context) {
return
}
h.dallService.PushTask(types.DallTask{
ClientId: data.ClientId,
JobId: job.Id,
UserId: uint(userId),
Prompt: data.Prompt,
Quality: data.Quality,
Size: data.Size,
Style: data.Style,
Power: job.Power,
})
task.Id = job.Id
h.dallService.PushTask(task)
resp.SUCCESS(c)
}
@@ -179,25 +181,14 @@ func (h *DallJobHandler) Remove(c *gin.Context) {
}
// 删除任务
tx := h.DB.Begin()
tx.Delete(&job)
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
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())
return
}
err := h.DB.Delete(&job).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
err = h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}

9
api/handler/function_handler.go
View File

@@ -113,10 +113,13 @@ func (h *FunctionHandler) WeiBo(c *gin.Context) {
SetHeader("AppId", h.config.AppId).
SetHeader("Authorization", fmt.Sprintf("Bearer %s", h.config.Token)).
SetSuccessResult(&res).Get(url)
if err != nil || r.IsErrorState() {
resp.ERROR(c, fmt.Sprintf("%v%v", err, r.Err))
if err != nil {
resp.ERROR(c, fmt.Sprintf("%v", err))
return
}
if r.IsErrorState() {
resp.ERROR(c, fmt.Sprintf("error http code status: %v", r.Status))
}
if res.Code != types.Success {
resp.ERROR(c, res.Message)
@@ -209,7 +212,7 @@ func (h *FunctionHandler) Dall3(c *gin.Context) {
}
content, err := h.dallService.Image(types.DallTask{
JobId: job.Id,
Id: job.Id,
UserId: user.Id,
Prompt: job.Prompt,
N: 1,

2
api/handler/markmap_handler.go
View File

@@ -87,7 +87,7 @@ func (h *MarkMapHandler) Generate(c *gin.Context) {
请直接生成结果，不要任何解释性语句。
`})
messages = append(messages, types.Message{Role: "user", Content: fmt.Sprintf("请生成一份有关【%s】一份思维导图，要求结构清晰，有条理", data.Prompt)})
content, err := utils.SendOpenAIMessage(h.DB, messages, chatModel.Value, chatModel.KeyId)
content, err := utils.SendOpenAIMessage(h.DB, messages, data.ModelId)
if err != nil {
resp.ERROR(c, fmt.Sprintf("请求 OpenAI API 失败: %s", err))
return

122
api/handler/mj_handler.go
View File

@@ -152,10 +152,23 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
resp.ERROR(c, "error with generate task id: "+err.Error())
return
}
task := types.MjTask{
ClientId: data.ClientId,
TaskId: taskId,
Type: types.TaskType(data.TaskType),
Prompt: data.Prompt,
NegPrompt: data.NegPrompt,
Params: params,
UserId: userId,
ImgArr: data.ImgArr,
Mode: h.App.SysConfig.MjMode,
TranslateModelId: h.App.SysConfig.TranslateModelId,
}
job := model.MidJourneyJob{
Type: data.TaskType,
UserId: userId,
TaskId: taskId,
TaskInfo: utils.JsonEncode(task),
Progress: 0,
Prompt: fmt.Sprintf("%s %s", data.Prompt, params),
Power: h.App.SysConfig.MjPower,
@@ -175,18 +188,8 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
return
}
h.mjService.PushTask(types.MjTask{
Id: job.Id,
ClientId: data.ClientId,
TaskId: taskId,
Type: types.TaskType(data.TaskType),
Prompt: data.Prompt,
NegPrompt: data.NegPrompt,
Params: params,
UserId: userId,
ImgArr: data.ImgArr,
Mode: h.App.SysConfig.MjMode,
})
task.Id = job.Id
h.mjService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
@@ -225,22 +228,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
taskId, _ := h.snowflake.Next(true)
job := model.MidJourneyJob{
Type: types.TaskUpscale.String(),
ReferenceId: data.MessageId,
UserId: userId,
TaskId: taskId,
Progress: 0,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
}
if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
resp.ERROR(c, "添加任务失败："+res.Error.Error())
return
}
h.mjService.PushTask(types.MjTask{
Id: job.Id,
task := types.MjTask{
ClientId: data.ClientId,
Type: types.TaskUpscale,
UserId: userId,
@@ -249,7 +237,23 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
})
}
job := model.MidJourneyJob{
Type: types.TaskUpscale.String(),
UserId: userId,
TaskId: taskId,
TaskInfo: utils.JsonEncode(task),
Progress: 0,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
}
if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
resp.ERROR(c, "添加任务失败："+res.Error.Error())
return
}
task.Id = job.Id
h.mjService.PushTask(task)
// update user's power
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
@@ -280,23 +284,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
taskId, _ := h.snowflake.Next(true)
job := model.MidJourneyJob{
Type: types.TaskVariation.String(),
ChannelId: data.ChannelId,
ReferenceId: data.MessageId,
UserId: userId,
TaskId: taskId,
Progress: 0,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
}
if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
resp.ERROR(c, "添加任务失败："+res.Error.Error())
return
}
h.mjService.PushTask(types.MjTask{
Id: job.Id,
task := types.MjTask{
Type: types.TaskVariation,
ClientId: data.ClientId,
UserId: userId,
@@ -305,7 +293,24 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
})
}
job := model.MidJourneyJob{
Type: types.TaskVariation.String(),
ChannelId: data.ChannelId,
UserId: userId,
TaskId: taskId,
TaskInfo: utils.JsonEncode(task),
Progress: 0,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
}
if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
resp.ERROR(c, "添加任务失败："+res.Error.Error())
return
}
task.Id = job.Id
h.mjService.PushTask(task)
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
@@ -401,26 +406,15 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
return
}
// remove job recode
tx := h.DB.Begin()
tx.Delete(&job)
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
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())
return
}
// remove job
err := h.DB.Delete(&job).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
err = h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}

5
api/handler/openai_handler.go
View File

@@ -51,7 +51,6 @@ type OpenAIResVo struct {
// OPenAI 消息发送实现
func (h *ChatHandler) sendOpenAiMessage(
chatCtx []types.Message,
req types.ApiRequest,
userVo vo.User,
ctx context.Context,
@@ -201,7 +200,7 @@ func (h *ChatHandler) sendOpenAiMessage(
TotalTokens: 0,
}
message.Content = usage.Content
h.saveChatHistory(req, usage, message, chatCtx, session, role, userVo, promptCreatedAt, replyCreatedAt)
h.saveChatHistory(req, usage, message, session, role, userVo, promptCreatedAt, replyCreatedAt)
}
} else { // 非流式输出
var respVo OpenAIResVo
@@ -220,7 +219,7 @@ func (h *ChatHandler) sendOpenAiMessage(
utils.SendChunkMsg(ws, content)
respVo.Usage.Prompt = prompt
respVo.Usage.Content = content
h.saveChatHistory(req, respVo.Usage, respVo.Choices[0].Message, chatCtx, session, role, userVo, promptCreatedAt, time.Now())
h.saveChatHistory(req, respVo.Usage, respVo.Choices[0].Message, session, role, userVo, promptCreatedAt, time.Now())
}
return nil

123
api/handler/prompt_handler.go Normal file
View File

@@ -0,0 +1,123 @@
package handler
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
// * Use of this source code is governed by a Apache-2.0 license
// * that can be found in the LICENSE file.
// * @Author yangjian102621@163.com
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service"
"geekai/service/oss"
"geekai/service/suno"
"geekai/utils"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"strings"
)
// 提示词生成 handler
// 使用 AI 生成绘画指令，歌词，视频生成指令等
type PromptHandler struct {
BaseHandler
sunoService *suno.Service
uploader *oss.UploaderManager
userService *service.UserService
}
func NewPromptHandler(app *core.AppServer, db *gorm.DB, userService *service.UserService) *PromptHandler {
return &PromptHandler{
BaseHandler: BaseHandler{
App: app,
DB: db,
},
userService: userService,
}
}
// Lyric 生成歌词
func (h *PromptHandler) Lyric(c *gin.Context) {
var data struct {
Prompt string `json:"prompt"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
content, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(service.LyricPromptTemplate, data.Prompt), h.App.SysConfig.TranslateModelId)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, content)
}
// Image 生成 AI 绘画提示词
func (h *PromptHandler) Image(c *gin.Context) {
var data struct {
Prompt string `json:"prompt"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
content, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(service.ImagePromptOptimizeTemplate, data.Prompt), h.App.SysConfig.TranslateModelId)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, strings.Trim(content, `"`))
}
// Video 生成视频提示词
func (h *PromptHandler) Video(c *gin.Context) {
var data struct {
Prompt string `json:"prompt"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
content, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(service.VideoPromptTemplate, data.Prompt), h.App.SysConfig.TranslateModelId)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, strings.Trim(content, `"`))
}
// MetaPrompt 生成元提示词
func (h *PromptHandler) MetaPrompt(c *gin.Context) {
var data struct {
Prompt string `json:"prompt"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
messages := make([]interface{}, 0)
messages = append(messages, types.Message{
Role: "system",
Content: service.MetaPromptTemplate,
})
messages = append(messages, types.Message{
Role: "user",
Content: "Task, Goal, or the Role to actor is:\n" + data.Prompt,
})
content, err := utils.SendOpenAIMessage(h.DB, messages, 0)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, strings.Trim(content, `"`))
}

74
api/handler/sd_handler.go
View File

@@ -109,29 +109,37 @@ func (h *SdJobHandler) Image(c *gin.Context) {
resp.ERROR(c, "error with generate task id: "+err.Error())
return
}
params := types.SdTaskParams{
TaskId: taskId,
Prompt: data.Prompt,
NegPrompt: data.NegPrompt,
Steps: data.Steps,
Sampler: data.Sampler,
FaceFix: data.FaceFix,
CfgScale: data.CfgScale,
Seed: data.Seed,
Height: data.Height,
Width: data.Width,
HdFix: data.HdFix,
HdRedrawRate: data.HdRedrawRate,
HdScale: data.HdScale,
HdScaleAlg: data.HdScaleAlg,
HdSteps: data.HdSteps,
task := types.SdTask{
ClientId: data.ClientId,
Type: types.TaskImage,
Params: types.SdTaskParams{
TaskId: taskId,
Prompt: data.Prompt,
NegPrompt: data.NegPrompt,
Steps: data.Steps,
Sampler: data.Sampler,
FaceFix: data.FaceFix,
CfgScale: data.CfgScale,
Seed: data.Seed,
Height: data.Height,
Width: data.Width,
HdFix: data.HdFix,
HdRedrawRate: data.HdRedrawRate,
HdScale: data.HdScale,
HdScaleAlg: data.HdScaleAlg,
HdSteps: data.HdSteps,
},
UserId: userId,
TranslateModelId: h.App.SysConfig.TranslateModelId,
}
job := model.SdJob{
UserId: userId,
Type: types.TaskImage.String(),
TaskId: params.TaskId,
Params: utils.JsonEncode(params),
TaskId: taskId,
Params: utils.JsonEncode(task.Params),
TaskInfo: utils.JsonEncode(task),
Prompt: data.Prompt,
Progress: 0,
Power: h.App.SysConfig.SdPower,
@@ -143,13 +151,8 @@ func (h *SdJobHandler) Image(c *gin.Context) {
return
}
h.sdService.PushTask(types.SdTask{
Id: int(job.Id),
ClientId: data.ClientId,
Type: types.TaskImage,
Params: params,
UserId: userId,
})
task.Id = int(job.Id)
h.sdService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
@@ -249,25 +252,14 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
}
// 删除任务
tx := h.DB.Begin()
tx.Delete(&job)
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := h.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d， Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
err := h.DB.Delete(&job).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
err = h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}

90
api/handler/suno_handler.go
View File

@@ -89,13 +89,29 @@ func (h *SunoHandler) Create(c *gin.Context) {
data.Prompt = fmt.Sprintf("%s\n%s", song.Prompt, refSong.Prompt)
}
}
task := types.SunoTask{
ClientId: data.ClientId,
UserId: int(h.GetLoginUserId(c)),
Type: data.Type,
Title: data.Title,
RefTaskId: data.RefTaskId,
RefSongId: data.RefSongId,
ExtendSecs: data.ExtendSecs,
Prompt: data.Prompt,
Tags: data.Tags,
Model: data.Model,
Instrumental: data.Instrumental,
SongId: data.SongId,
AudioURL: data.AudioURL,
}
// 插入数据库
job := model.SunoJob{
UserId: int(h.GetLoginUserId(c)),
UserId: task.UserId,
Prompt: data.Prompt,
Instrumental: data.Instrumental,
ModelName: data.Model,
TaskInfo: utils.JsonEncode(task),
Tags: data.Tags,
Title: data.Title,
Type: data.Type,
@@ -115,26 +131,13 @@ func (h *SunoHandler) Create(c *gin.Context) {
}
// 创建任务
h.sunoService.PushTask(types.SunoTask{
ClientId: data.ClientId,
Id: job.Id,
UserId: job.UserId,
Type: job.Type,
Title: job.Title,
RefTaskId: data.RefTaskId,
RefSongId: data.RefSongId,
ExtendSecs: data.ExtendSecs,
Prompt: job.Prompt,
Tags: data.Tags,
Model: data.Model,
Instrumental: data.Instrumental,
SongId: data.SongId,
AudioURL: data.AudioURL,
})
task.Id = job.Id
h.sunoService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 文生歌曲，%s", job.ModelName),
CreatedAt: time.Now(),
})
@@ -219,25 +222,11 @@ func (h *SunoHandler) Remove(c *gin.Context) {
}
// 删除任务
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 恢复用户算力
err = h.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
})
err = h.DB.Delete(&job).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
tx.Commit()
// 删除文件
_ = h.uploader.GetUploadHandler().Delete(job.CoverURL)
@@ -334,40 +323,3 @@ func (h *SunoHandler) Play(c *gin.Context) {
}
h.DB.Model(&model.SunoJob{}).Where("song_id", songId).UpdateColumn("play_times", gorm.Expr("play_times + ?", 1))
}
const genLyricTemplate = `
你是一位才华横溢的作曲家，拥有丰富的情感和细腻的笔触，你对文字有着独特的感悟力，能将各种情感和意境巧妙地融入歌词中。
请以【%s】为主题创作一首歌曲，歌曲时间不要太短，3分钟左右，不要输出任何解释性的内容。
输出格式如下：
歌曲名称
第一节：
{{歌词内容}}
副歌：
{{歌词内容}}
第二节：
{{歌词内容}}
副歌：
{{歌词内容}}
尾声：
{{歌词内容}}
`
// Lyric 生成歌词
func (h *SunoHandler) Lyric(c *gin.Context) {
var data struct {
Prompt string `json:"prompt"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
content, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(genLyricTemplate, data.Prompt), "gpt-4o-mini", 0)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, content)
}

50
api/handler/user_handler.go
View File

@@ -130,15 +130,28 @@ func (h *UserHandler) Register(c *gin.Context) {
}
}
salt := utils.RandString(8)
user := model.User{
Username: data.Username,
Password: utils.GenPassword(data.Password, salt),
Avatar: "/images/avatar/user.png",
Salt: salt,
Status: true,
ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色
Power: h.App.SysConfig.InitPower,
}
// check if the username is existing
var item model.User
session := h.DB.Session(&gorm.Session{})
if data.Mobile != "" {
session = session.Where("mobile = ?", data.Mobile)
data.Username = data.Mobile
user.Username = data.Mobile
user.Mobile = data.Mobile
} else if data.Email != "" {
session = session.Where("email = ?", data.Email)
data.Username = data.Email
user.Username = data.Email
user.Email = data.Email
} else if data.Username != "" {
session = session.Where("username = ?", data.Username)
}
@@ -148,20 +161,6 @@ func (h *UserHandler) Register(c *gin.Context) {
return
}
salt := utils.RandString(8)
user := model.User{
Username: data.Username,
Mobile: data.Mobile,
Email: data.Email,
Password: utils.GenPassword(data.Password, salt),
Avatar: "/images/avatar/user.png",
Salt: salt,
Status: true,
ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色
ChatModels: utils.JsonEncode(h.App.SysConfig.DefaultModels), // 默认开通的模型
Power: h.App.SysConfig.InitPower,
}
// 被邀请人也获得赠送算力
if data.InviteCode != "" {
user.Power += h.App.SysConfig.InvitePower
@@ -417,16 +416,15 @@ func (h *UserHandler) CLoginCallback(c *gin.Context) {
salt := utils.RandString(8)
password := fmt.Sprintf("%d", utils.RandomNumber(8))
user = model.User{
Username: fmt.Sprintf("%s@%d", loginType, utils.RandomNumber(10)),
Password: utils.GenPassword(password, salt),
Avatar: fmt.Sprintf("%s", data["avatar"]),
Salt: salt,
Status: true,
ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色
ChatModels: utils.JsonEncode(h.App.SysConfig.DefaultModels), // 默认开通的模型
Power: h.App.SysConfig.InitPower,
OpenId: fmt.Sprintf("%s", data["openid"]),
Nickname: fmt.Sprintf("%s", data["nickname"]),
Username: fmt.Sprintf("%s@%d", loginType, utils.RandomNumber(10)),
Password: utils.GenPassword(password, salt),
Avatar: fmt.Sprintf("%s", data["avatar"]),
Salt: salt,
Status: true,
ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色
Power: h.App.SysConfig.InitPower,
OpenId: fmt.Sprintf("%s", data["openid"]),
Nickname: fmt.Sprintf("%s", data["nickname"]),
}
tx = h.DB.Create(&user)

44
api/handler/video_handler.go
View File

@@ -80,13 +80,21 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) {
StartImgURL: data.FirstFrameImg,
EndImgURL: data.EndFrameImg,
}
task := types.VideoTask{
ClientId: data.ClientId,
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Params: params,
TranslateModelId: h.App.SysConfig.TranslateModelId,
}
// 插入数据库
job := model.VideoJob{
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Power: h.App.SysConfig.LumaPower,
Params: utils.JsonEncode(params),
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Power: h.App.SysConfig.LumaPower,
TaskInfo: utils.JsonEncode(task),
}
tx := h.DB.Create(&job)
if tx.Error != nil {
@@ -95,14 +103,8 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) {
}
// 创建任务
h.videoService.PushTask(types.VideoTask{
ClientId: data.ClientId,
Id: job.Id,
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Params: params,
})
task.Id = job.Id
h.videoService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
@@ -181,25 +183,11 @@ func (h *VideoHandler) Remove(c *gin.Context) {
}
// 删除任务
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 恢复算力
err = h.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "luma",
Remark: fmt.Sprintf("Luma 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
})
err = h.DB.Delete(&job).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
tx.Commit()
// 删除文件
_ = h.uploader.GetUploadHandler().Delete(job.CoverURL)