mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-11-08 10:13:44 +08:00
dalle image page is ready
This commit is contained in:
RockYang
@@ -104,8 +104,10 @@ func (h *ChatHandler) sendOpenAiMessage(
|
||||
res := h.DB.Where("name = ?", tool.Function.Name).First(&function)
|
||||
if res.Error == nil {
|
||||
toolCall = true
|
||||
callMsg := fmt.Sprintf("正在调用工具 `%s` 作答 ...\n\n", function.Label)
|
||||
utils.ReplyChunkMessage(ws, types.WsMessage{Type: types.WsStart})
|
||||
utils.ReplyChunkMessage(ws, types.WsMessage{Type: types.WsMiddle, Content: fmt.Sprintf("正在调用工具 `%s` 作答 ...\n\n", function.Label)})
|
||||
utils.ReplyChunkMessage(ws, types.WsMessage{Type: types.WsMiddle, Content: callMsg})
|
||||
contents = append(contents, callMsg)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
260
api/handler/dalle_handler.go
Normal file
260
api/handler/dalle_handler.go
Normal file
@@ -0,0 +1,260 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"chatplus/core"
|
||||
"chatplus/core/types"
|
||||
"chatplus/service/dalle"
|
||||
"chatplus/service/oss"
|
||||
"chatplus/store/model"
|
||||
"chatplus/store/vo"
|
||||
"chatplus/utils"
|
||||
"chatplus/utils/resp"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type DallJobHandler struct {
|
||||
BaseHandler
|
||||
redis *redis.Client
|
||||
service *dalle.Service
|
||||
uploader *oss.UploaderManager
|
||||
}
|
||||
|
||||
func NewDallJobHandler(app *core.AppServer, db *gorm.DB, service *dalle.Service, manager *oss.UploaderManager) *DallJobHandler {
|
||||
return &DallJobHandler{
|
||||
service: service,
|
||||
uploader: manager,
|
||||
BaseHandler: BaseHandler{
|
||||
App: app,
|
||||
DB: db,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Client WebSocket 客户端，用于通知任务状态变更
|
||||
func (h *DallJobHandler) Client(c *gin.Context) {
|
||||
ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
logger.Error(err)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
userId := h.GetInt(c, "user_id", 0)
|
||||
if userId == 0 {
|
||||
logger.Info("Invalid user ID")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
client := types.NewWsClient(ws)
|
||||
h.service.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))
|
||||
return
|
||||
}
|
||||
|
||||
var message types.WsMessage
|
||||
err = utils.JsonDecode(string(msg), &message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// 心跳消息
|
||||
if message.Type == "heartbeat" {
|
||||
logger.Debug("收到 DallE 心跳消息：", message.Content)
|
||||
continue
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (h *DallJobHandler) preCheck(c *gin.Context) bool {
|
||||
user, err := h.GetLoginUser(c)
|
||||
if err != nil {
|
||||
resp.NotAuth(c)
|
||||
return false
|
||||
}
|
||||
|
||||
if user.Power < h.App.SysConfig.SdPower {
|
||||
resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
|
||||
}
|
||||
|
||||
// Image 创建一个绘画任务
|
||||
func (h *DallJobHandler) Image(c *gin.Context) {
|
||||
if !h.preCheck(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var data types.DallTask
|
||||
if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
|
||||
resp.ERROR(c, types.InvalidArgs)
|
||||
return
|
||||
}
|
||||
|
||||
idValue, _ := c.Get(types.LoginUserID)
|
||||
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
|
||||
job := model.DallJob{
|
||||
UserId: uint(userId),
|
||||
Prompt: data.Prompt,
|
||||
Power: h.App.SysConfig.DallPower,
|
||||
}
|
||||
res := h.DB.Create(&job)
|
||||
if res.Error != nil {
|
||||
resp.ERROR(c, "error with save job: "+res.Error.Error())
|
||||
return
|
||||
}
|
||||
|
||||
h.service.PushTask(types.DallTask{
|
||||
JobId: job.Id,
|
||||
UserId: uint(userId),
|
||||
Prompt: data.Prompt,
|
||||
Quality: data.Quality,
|
||||
Size: data.Size,
|
||||
Style: data.Style,
|
||||
Power: job.Power,
|
||||
})
|
||||
|
||||
client := h.service.Clients.Get(job.UserId)
|
||||
if client != nil {
|
||||
_ = client.Send([]byte("Task Updated"))
|
||||
}
|
||||
resp.SUCCESS(c)
|
||||
}
|
||||
|
||||
// ImgWall 照片墙
|
||||
func (h *DallJobHandler) ImgWall(c *gin.Context) {
|
||||
page := h.GetInt(c, "page", 0)
|
||||
pageSize := h.GetInt(c, "page_size", 0)
|
||||
err, jobs := h.getData(true, 0, page, pageSize, true)
|
||||
if err != nil {
|
||||
resp.ERROR(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
resp.SUCCESS(c, jobs)
|
||||
}
|
||||
|
||||
// JobList 获取 SD 任务列表
|
||||
func (h *DallJobHandler) JobList(c *gin.Context) {
|
||||
status := h.GetBool(c, "status")
|
||||
userId := h.GetLoginUserId(c)
|
||||
page := h.GetInt(c, "page", 0)
|
||||
pageSize := h.GetInt(c, "page_size", 0)
|
||||
publish := h.GetBool(c, "publish")
|
||||
|
||||
err, jobs := h.getData(status, userId, page, pageSize, publish)
|
||||
if err != nil {
|
||||
resp.ERROR(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
resp.SUCCESS(c, jobs)
|
||||
}
|
||||
|
||||
// JobList 获取任务列表
|
||||
func (h *DallJobHandler) getData(finish bool, userId uint, page int, pageSize int, publish bool) (error, []vo.DallJob) {
|
||||
|
||||
session := h.DB.Session(&gorm.Session{})
|
||||
if finish {
|
||||
session = session.Where("progress = ?", 100).Order("id DESC")
|
||||
} else {
|
||||
session = session.Where("progress < ?", 100).Order("id ASC")
|
||||
}
|
||||
if userId > 0 {
|
||||
session = session.Where("user_id = ?", userId)
|
||||
}
|
||||
if publish {
|
||||
session = session.Where("publish", publish)
|
||||
}
|
||||
if page > 0 && pageSize > 0 {
|
||||
offset := (page - 1) * pageSize
|
||||
session = session.Offset(offset).Limit(pageSize)
|
||||
}
|
||||
|
||||
var items []model.DallJob
|
||||
res := session.Find(&items)
|
||||
if res.Error != nil {
|
||||
return res.Error, nil
|
||||
}
|
||||
|
||||
var jobs = make([]vo.DallJob, 0)
|
||||
for _, item := range items {
|
||||
// delete failed or timeout tasks
|
||||
if (item.Progress < 100 && time.Now().Sub(item.CreatedAt) > time.Minute*5) || item.Progress == -1 {
|
||||
h.DB.Delete(&item)
|
||||
}
|
||||
var job vo.DallJob
|
||||
err := utils.CopyObject(item, &job)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
jobs = append(jobs, job)
|
||||
}
|
||||
|
||||
return nil, jobs
|
||||
}
|
||||
|
||||
// Remove remove task image
|
||||
func (h *DallJobHandler) Remove(c *gin.Context) {
|
||||
var data struct {
|
||||
Id uint `json:"id"`
|
||||
UserId uint `json:"user_id"`
|
||||
ImgURL string `json:"img_url"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&data); err != nil {
|
||||
resp.ERROR(c, types.InvalidArgs)
|
||||
return
|
||||
}
|
||||
|
||||
// remove job recode
|
||||
res := h.DB.Delete(&model.DallJob{Id: data.Id})
|
||||
if res.Error != nil {
|
||||
resp.ERROR(c, res.Error.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// remove image
|
||||
err := h.uploader.GetUploadHandler().Delete(data.ImgURL)
|
||||
if err != nil {
|
||||
logger.Error("remove image failed: ", err)
|
||||
}
|
||||
|
||||
resp.SUCCESS(c)
|
||||
}
|
||||
|
||||
// Publish 发布/取消发布图片到画廊显示
|
||||
func (h *DallJobHandler) Publish(c *gin.Context) {
|
||||
var data struct {
|
||||
Id uint `json:"id"`
|
||||
Action bool `json:"action"` // 发布动作，true => 发布，false => 取消分享
|
||||
}
|
||||
if err := c.ShouldBindJSON(&data); err != nil {
|
||||
resp.ERROR(c, types.InvalidArgs)
|
||||
return
|
||||
}
|
||||
|
||||
res := h.DB.Model(&model.DallJob{Id: data.Id}).UpdateColumn("publish", true)
|
||||
if res.Error != nil {
|
||||
resp.ERROR(c, "更新数据库失败")
|
||||
return
|
||||
}
|
||||
|
||||
resp.SUCCESS(c)
|
||||
}
|
||||
@@ -3,27 +3,35 @@ package handler
|
||||
import (
|
||||
"chatplus/core"
|
||||
"chatplus/core/types"
|
||||
"chatplus/service/dalle"
|
||||
"chatplus/service/oss"
|
||||
"chatplus/store/model"
|
||||
"chatplus/utils"
|
||||
"chatplus/utils/resp"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/imroc/req/v3"
|
||||
"gorm.io/gorm"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type FunctionHandler struct {
|
||||
BaseHandler
|
||||
config types.ChatPlusApiConfig
|
||||
uploadManager *oss.UploaderManager
|
||||
dallService *dalle.Service
|
||||
}
|
||||
|
||||
func NewFunctionHandler(server *core.AppServer, db *gorm.DB, config *types.AppConfig, manager *oss.UploaderManager) *FunctionHandler {
|
||||
func NewFunctionHandler(
|
||||
server *core.AppServer,
|
||||
db *gorm.DB,
|
||||
config *types.AppConfig,
|
||||
manager *oss.UploaderManager,
|
||||
dallService *dalle.Service) *FunctionHandler {
|
||||
return &FunctionHandler{
|
||||
BaseHandler: BaseHandler{
|
||||
App: server,
|
||||
@@ -31,6 +39,7 @@ func NewFunctionHandler(server *core.AppServer, db *gorm.DB, config *types.AppCo
|
||||
},
|
||||
config: config.ApiConfig,
|
||||
uploadManager: manager,
|
||||
dallService: dallService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,30 +160,6 @@ func (h *FunctionHandler) ZaoBao(c *gin.Context) {
|
||||
resp.SUCCESS(c, strings.Join(builder, "\n\n"))
|
||||
}
|
||||
|
||||
type imgReq struct {
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
N int `json:"n"`
|
||||
Size string `json:"size"`
|
||||
}
|
||||
|
||||
type imgRes struct {
|
||||
Created int64 `json:"created"`
|
||||
Data []struct {
|
||||
RevisedPrompt string `json:"revised_prompt"`
|
||||
Url string `json:"url"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type ErrRes struct {
|
||||
Error struct {
|
||||
Code interface{} `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Param interface{} `json:"param"`
|
||||
Type string `json:"type"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
// Dall3 DallE3 AI 绘图
|
||||
func (h *FunctionHandler) Dall3(c *gin.Context) {
|
||||
if err := h.checkAuth(c); err != nil {
|
||||
@@ -190,85 +175,40 @@ func (h *FunctionHandler) Dall3(c *gin.Context) {
|
||||
|
||||
logger.Debugf("绘画参数：%+v", params)
|
||||
var user model.User
|
||||
tx := h.DB.Where("id = ?", params["user_id"]).First(&user)
|
||||
if tx.Error != nil {
|
||||
res := h.DB.Where("id = ?", params["user_id"]).First(&user)
|
||||
if res.Error != nil {
|
||||
resp.ERROR(c, "当前用户不存在！")
|
||||
return
|
||||
}
|
||||
|
||||
if user.Power < h.App.SysConfig.DallPower {
|
||||
resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
|
||||
return
|
||||
}
|
||||
|
||||
// create dall task
|
||||
prompt := utils.InterfaceToString(params["prompt"])
|
||||
// get image generation API KEY
|
||||
var apiKey model.ApiKey
|
||||
tx = h.DB.Where("platform = ?", types.OpenAI).Where("type = ?", "img").Where("enabled = ?", true).Order("last_used_at ASC").First(&apiKey)
|
||||
if tx.Error != nil {
|
||||
resp.ERROR(c, "获取绘图 API KEY 失败: "+tx.Error.Error())
|
||||
job := model.DallJob{
|
||||
UserId: user.Id,
|
||||
Prompt: prompt,
|
||||
Power: h.App.SysConfig.DallPower,
|
||||
}
|
||||
res = h.DB.Create(&job)
|
||||
|
||||
if res.Error != nil {
|
||||
resp.ERROR(c, "创建 DALL-E 绘图任务失败："+res.Error.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// translate prompt
|
||||
const translatePromptTemplate = "Translate the following painting prompt words into English keyword phrases. Without any explanation, directly output the keyword phrases separated by commas. The content to be translated is: [%s]"
|
||||
pt, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(translatePromptTemplate, params["prompt"]))
|
||||
if err == nil {
|
||||
logger.Debugf("翻译绘画提示词，原文：%s，译文：%s", prompt, pt)
|
||||
prompt = pt
|
||||
}
|
||||
var res imgRes
|
||||
var errRes ErrRes
|
||||
var request *req.Request
|
||||
if len(apiKey.ProxyURL) > 5 {
|
||||
request = req.C().SetProxyURL(apiKey.ProxyURL).R()
|
||||
} else {
|
||||
request = req.C().R()
|
||||
}
|
||||
logger.Debugf("Sending %s request, ApiURL:%s, API KEY:%s, PROXY: %s", apiKey.Platform, apiKey.ApiURL, apiKey.Value, apiKey.ProxyURL)
|
||||
r, err := request.SetHeader("Content-Type", "application/json").
|
||||
SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||||
SetBody(imgReq{
|
||||
Model: "dall-e-3",
|
||||
Prompt: prompt,
|
||||
N: 1,
|
||||
Size: "1024x1024",
|
||||
}).
|
||||
SetErrorResult(&errRes).
|
||||
SetSuccessResult(&res).Post(apiKey.ApiURL)
|
||||
if r.IsErrorState() {
|
||||
resp.ERROR(c, "请求 OpenAI API 失败: "+errRes.Error.Message)
|
||||
return
|
||||
}
|
||||
// 更新 API KEY 的最后使用时间
|
||||
h.DB.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
|
||||
logger.Debugf("%+v", res)
|
||||
// 存储图片
|
||||
imgURL, err := h.uploadManager.GetUploadHandler().PutImg(res.Data[0].Url, false)
|
||||
content, err := h.dallService.Image(types.DallTask{
|
||||
JobId: job.Id,
|
||||
UserId: user.Id,
|
||||
Prompt: job.Prompt,
|
||||
N: 1,
|
||||
Quality: "standard",
|
||||
Size: "1024x1024",
|
||||
Style: "vivid",
|
||||
Power: job.Power,
|
||||
}, true)
|
||||
if err != nil {
|
||||
resp.ERROR(c, "下载图片失败: "+err.Error())
|
||||
resp.ERROR(c, "任务执行失败："+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
content := fmt.Sprintf("下面是根据您的描述创作的图片，它描绘了 【%s】 的场景。 \n\n\n", prompt, imgURL)
|
||||
// 更新用户算力
|
||||
tx = h.DB.Model(&model.User{}).Where("id", user.Id).UpdateColumn("power", gorm.Expr("power - ?", h.App.SysConfig.DallPower))
|
||||
// 记录算力变化日志
|
||||
if tx.Error == nil && tx.RowsAffected > 0 {
|
||||
var u model.User
|
||||
h.DB.Where("id", user.Id).First(&u)
|
||||
h.DB.Create(&model.PowerLog{
|
||||
UserId: user.Id,
|
||||
Username: user.Username,
|
||||
Type: types.PowerConsume,
|
||||
Amount: h.App.SysConfig.DallPower,
|
||||
Balance: u.Power,
|
||||
Mark: types.PowerSub,
|
||||
Model: "dall-e-3",
|
||||
Remark: fmt.Sprintf("绘画提示词：%s", utils.CutWords(prompt, 10)),
|
||||
CreatedAt: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
resp.SUCCESS(c, content)
|
||||
}
|
||||
|
||||
@@ -65,7 +65,7 @@ func (h *SdJobHandler) Client(c *gin.Context) {
|
||||
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
|
||||
}
|
||||
|
||||
func (h *SdJobHandler) checkLimits(c *gin.Context) bool {
|
||||
func (h *SdJobHandler) preCheck(c *gin.Context) bool {
|
||||
user, err := h.GetLoginUser(c)
|
||||
if err != nil {
|
||||
resp.NotAuth(c)
|
||||
@@ -88,7 +88,7 @@ func (h *SdJobHandler) checkLimits(c *gin.Context) bool {
|
||||
|
||||
// Image 创建一个绘画任务
|
||||
func (h *SdJobHandler) Image(c *gin.Context) {
|
||||
if !h.checkLimits(c) {
|
||||
if !h.preCheck(c) {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -298,7 +298,7 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
|
||||
|
||||
client := h.pool.Clients.Get(data.UserId)
|
||||
if client != nil {
|
||||
_ = client.Send([]byte("Task Updated"))
|
||||
_ = client.Send([]byte(sd.Finished))
|
||||
}
|
||||
|
||||
resp.SUCCESS(c)
|
||||
|
||||
Reference in New Issue
Block a user