xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-07 01:33:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

merge v4.1.0 and fixed conflicts

Browse Source
This commit is contained in:
RockYang
2024-10-08 17:51:14 +08:00
parent 08319391ea 0df700ec18
commit 0f0a4c0a7e
91 changed files with 3266 additions and 3202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
api/handler/admin/admin_handler.go
View File

@@ -8,6 +8,8 @@ package admin
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"context"
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/handler"
@@ -16,11 +18,8 @@ import (
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"context"
"fmt"
"github.com/go-redis/redis/v8"
"github.com/golang-jwt/jwt/v5"
"github.com/mojocn/base64Captcha"
"time"
"github.com/gin-gonic/gin"
@@ -56,11 +55,11 @@ func (h *ManagerHandler) Login(c *gin.Context) {
return
}
// add captcha
if !base64Captcha.DefaultMemStore.Verify(data.CaptchaId, data.Captcha, true) {
resp.ERROR(c, "验证码错误!")
return
}
//// add captcha
//if !base64Captcha.DefaultMemStore.Verify(data.CaptchaId, data.Captcha, true) {
// resp.ERROR(c, "验证码错误!")
// return
//}
var manager model.AdminUser
res := h.DB.Model(&model.AdminUser{}).Where("username = ?", data.Username).First(&manager)

46
api/handler/admin/captcha_handler.go
View File

@@ -1,46 +0,0 @@
package admin
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 (
"geekai/core"
"geekai/handler"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"github.com/mojocn/base64Captcha"
)
type CaptchaHandler struct {
handler.BaseHandler
}
func NewCaptchaHandler(app *core.AppServer) *CaptchaHandler {
return &CaptchaHandler{BaseHandler: handler.BaseHandler{App: app}}
}
type CaptchaVo struct {
CaptchaId string `json:"captcha_id"`
PicPath string `json:"pic_path"`
}
// GetCaptcha 获取验证码
func (h *CaptchaHandler) GetCaptcha(c *gin.Context) {
var captchaVo CaptchaVo
driver := base64Captcha.NewDriverDigit(48, 130, 4, 0.4, 10)
cp := base64Captcha.NewCaptcha(driver, base64Captcha.DefaultMemStore)
// b64s是图片的base64编码
id, b64s, err := cp.Generate()
if err != nil {
resp.ERROR(c, "生成验证码错误!")
return
}
captchaVo.CaptchaId = id
captchaVo.PicPath = b64s
resp.SUCCESS(c, captchaVo)
}

2
api/handler/admin/chat_model_handler.go
View File

@@ -54,7 +54,6 @@ func (h *ChatModelHandler) Save(c *gin.Context) {
Name: data.Name,
Value: data.Value,
Enabled: data.Enabled,
SortNum: data.SortNum,
Open: data.Open,
MaxTokens: data.MaxTokens,
MaxContext: data.MaxContext,
@@ -64,6 +63,7 @@ func (h *ChatModelHandler) Save(c *gin.Context) {
var res *gorm.DB
if data.Id > 0 {
item.Id = data.Id
item.SortNum = data.SortNum
res = h.DB.Select("*").Omit("created_at").Updates(&item)
} else {
res = h.DB.Create(&item)

6
api/handler/admin/user_handler.go
View File

@@ -139,7 +139,6 @@ func (h *UserHandler) Save(c *gin.Context) {
salt := utils.RandString(8)
u := model.User{
Username: data.Username,
Nickname: fmt.Sprintf("极客学长@%d", utils.RandomNumber(6)),
Password: utils.GenPassword(data.Password, salt),
Avatar: "/images/avatar/user.png",
Salt: salt,
@@ -149,6 +148,11 @@ func (h *UserHandler) Save(c *gin.Context) {
ChatModels: utils.JsonEncode(data.ChatModels),
ExpiredTime: utils.Str2stamp(data.ExpiredTime),
}
if h.licenseService.GetLicense().Configs.DeCopy {
u.Nickname = fmt.Sprintf("用户@%d", utils.RandomNumber(6))
} else {
u.Nickname = fmt.Sprintf("极客学长@%d", utils.RandomNumber(6))
}
res = h.DB.Create(&u)
_ = utils.CopyObject(u, &userVo)
userVo.Id = u.Id

2
api/handler/chat_role_handler.go
View File

@@ -65,7 +65,7 @@ func (h *ChatRoleHandler) List(c *gin.Context) {
}
for _, r := range roles {
if !utils.ContainsStr(roleKeys, r.Key) {
if !utils.Contains(roleKeys, r.Key) {
continue
}
var v vo.ChatRole

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

@@ -44,6 +44,8 @@ type ChatHandler struct {
redis *redis.Client
uploadManager *oss.UploaderManager
licenseService *service.LicenseService
ReqCancelFunc *types.LMap[string, context.CancelFunc] // HttpClient 请求取消 handle function
ChatContexts *types.LMap[string, []types.Message] // 聊天上下文 Map [chatId] => []Message
}
func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService) *ChatHandler {
@@ -52,6 +54,8 @@ func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manag
redis: redis,
uploadManager: manager,
licenseService: licenseService,
ReqCancelFunc: types.NewLMap[string, context.CancelFunc](),
ChatContexts: types.NewLMap[string, []types.Message](),
}
}
@@ -89,21 +93,10 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
return
}
session := h.App.ChatSession.Get(sessionId)
if session == nil {
user, err := h.GetLoginUser(c)
if err != nil {
logger.Info("用户未登录")
c.Abort()
return
}
session = &types.ChatSession{
SessionId: sessionId,
ClientIP: c.ClientIP(),
Username: user.Username,
UserId: user.Id,
}
h.App.ChatSession.Put(sessionId, session)
session := &types.ChatSession{
SessionId: sessionId,
ClientIP: c.ClientIP(),
UserId: h.GetLoginUserId(c),
}
// use old chat data override the chat model and role ID
@@ -125,22 +118,18 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
Temperature: chatModel.Temperature,
KeyId: chatModel.KeyId,
Platform: chatModel.Platform}
logger.Infof("New websocket connected, IP: %s, Username: %s", c.ClientIP(), session.Username)
logger.Infof("New websocket connected, IP: %s", c.ClientIP())
// 保存会话连接
h.App.ChatClients.Put(sessionId, client)
go func() {
for {
_, msg, err := client.Receive()
if err != nil {
logger.Debugf("close connection: %s", client.Conn.RemoteAddr())
client.Close()
h.App.ChatClients.Delete(sessionId)
h.App.ChatSession.Delete(sessionId)
cancelFunc := h.App.ReqCancelFunc.Get(sessionId)
cancelFunc := h.ReqCancelFunc.Get(sessionId)
if cancelFunc != nil {
cancelFunc()
h.App.ReqCancelFunc.Delete(sessionId)
h.ReqCancelFunc.Delete(sessionId)
}
return
}
@@ -160,7 +149,7 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
logger.Info("Receive a message: ", message.Content)
ctx, cancel := context.WithCancel(context.Background())
h.App.ReqCancelFunc.Put(sessionId, cancel)
h.ReqCancelFunc.Put(sessionId, cancel)
// 回复消息
err = h.sendMessage(ctx, session, chatRole, utils.InterfaceToString(message.Content), client)
if err != nil {
@@ -274,8 +263,8 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
chatCtx := make([]types.Message, 0)
messages := make([]types.Message, 0)
if h.App.SysConfig.EnableContext {
if h.App.ChatContexts.Has(session.ChatId) {
messages = h.App.ChatContexts.Get(session.ChatId)
if h.ChatContexts.Has(session.ChatId) {
messages = h.ChatContexts.Get(session.ChatId)
} else {
_ = utils.JsonDecode(role.Context, &messages)
if h.App.SysConfig.ContextDeep > 0 {
@@ -323,20 +312,51 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
reqMgs = append(reqMgs, m)
}
fullPrompt := prompt
text := prompt
// extract files in prompt
files := utils.ExtractFileURLs(prompt)
logger.Debugf("detected FILES: %+v", files)
// 如果不是逆向模型，则提取文件内容
if len(files) > 0 && !(session.Model.Value == "gpt-4-all" ||
strings.HasPrefix(session.Model.Value, "gpt-4-gizmo") ||
strings.HasSuffix(session.Model.Value, "claude-3")) {
contents := make([]string, 0)
var file model.File
for _, v := range files {
h.DB.Where("url = ?", v).First(&file)
content, err := utils.ReadFileContent(v, h.App.Config.TikaHost)
if err != nil {
logger.Error("error with read file: ", err)
} else {
contents = append(contents, fmt.Sprintf("%s 文件内容：%s", file.Name, content))
}
text = strings.Replace(text, v, "", 1)
}
if len(contents) > 0 {
fullPrompt = fmt.Sprintf("请根据提供的文件内容信息回答问题(其中Excel 已转成 HTML)：\n\n %s\n\n 问题：%s", strings.Join(contents, "\n"), text)
}
tokens, _ := utils.CalcTokens(fullPrompt, req.Model)
if tokens > session.Model.MaxContext {
return fmt.Errorf("文件的长度超出模型允许的最大上下文长度，请减少文件内容数量或文件大小。")
}
}
logger.Debug("最终Prompt：", fullPrompt)
if session.Model.Platform == types.QWen.Value {
req.Input = make(map[string]interface{})
reqMgs = append(reqMgs, types.Message{
Role: "user",
Content: prompt,
Content: fullPrompt,
})
req.Input["messages"] = reqMgs
} else if session.Model.Platform == types.OpenAI.Value || session.Model.Platform == types.Azure.Value { // extract image for gpt-vision model
imgURLs := utils.ExtractImgURL(prompt)
imgURLs := utils.ExtractImgURLs(prompt)
logger.Debugf("detected IMG: %+v", imgURLs)
var content interface{}
if len(imgURLs) > 0 {
data := make([]interface{}, 0)
text := prompt
for _, v := range imgURLs {
text = strings.Replace(text, v, "", 1)
data = append(data, gin.H{
@@ -348,11 +368,11 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
}
data = append(data, gin.H{
"type": "text",
"text": text,
"text": strings.TrimSpace(text),
})
content = data
} else {
content = prompt
content = fullPrompt
}
req.Messages = append(reqMgs, map[string]interface{}{
"role": "user",
@@ -361,7 +381,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session *types.ChatSessio
} else {
req.Messages = append(reqMgs, map[string]interface{}{
"role": "user",
"content": prompt,
"content": fullPrompt,
})
}
@@ -442,9 +462,9 @@ func getTotalTokens(req types.ApiRequest) int {
// StopGenerate 停止生成
func (h *ChatHandler) StopGenerate(c *gin.Context) {
sessionId := c.Query("session_id")
if h.App.ReqCancelFunc.Has(sessionId) {
h.App.ReqCancelFunc.Get(sessionId)()
h.App.ReqCancelFunc.Delete(sessionId)
if h.ReqCancelFunc.Has(sessionId) {
h.ReqCancelFunc.Get(sessionId)()
h.ReqCancelFunc.Delete(sessionId)
}
resp.SUCCESS(c, types.OkMsg)
}
@@ -454,7 +474,7 @@ func (h *ChatHandler) StopGenerate(c *gin.Context) {
func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, session *types.ChatSession, apiKey *model.ApiKey) (*http.Response, error) {
// if the chat model bind a KEY, use it directly
if session.Model.KeyId > 0 {
h.DB.Debug().Where("id", session.Model.KeyId).Where("enabled", true).Find(apiKey)
h.DB.Where("id", session.Model.KeyId).Find(apiKey)
}
// use the last unused key
if apiKey.Id == 0 {
@@ -602,7 +622,7 @@ func (h *ChatHandler) saveChatHistory(
if h.App.SysConfig.EnableContext {
chatCtx = append(chatCtx, useMsg) // 提问消息
chatCtx = append(chatCtx, message) // 回复消息
h.App.ChatContexts.Put(session.ChatId, chatCtx)
h.ChatContexts.Put(session.ChatId, chatCtx)
}
// 追加聊天记录
@@ -660,7 +680,7 @@ func (h *ChatHandler) saveChatHistory(
res = h.DB.Where("chat_id = ?", session.ChatId).First(&chatItem)
if res.Error != nil {
chatItem.ChatId = session.ChatId
chatItem.UserId = session.UserId
chatItem.UserId = userVo.Id
chatItem.RoleId = role.Id
chatItem.ModelId = session.Model.Id
if utf8.RuneCountInString(prompt) > 30 {

6
api/handler/chatimpl/chat_item_handler.go
View File

@@ -96,7 +96,7 @@ func (h *ChatHandler) Clear(c *gin.Context) {
for _, chat := range chats {
chatIds = append(chatIds, chat.ChatId)
// 清空会话上下文
h.App.ChatContexts.Delete(chat.ChatId)
h.ChatContexts.Delete(chat.ChatId)
}
err = h.DB.Transaction(func(tx *gorm.DB) error {
res := h.DB.Where("user_id =?", user.Id).Delete(&model.ChatItem{})
@@ -108,8 +108,6 @@ func (h *ChatHandler) Clear(c *gin.Context) {
if res.Error != nil {
return res.Error
}
// TODO: 是否要删除 MidJourney 绘画记录和图片文件？
return nil
})
@@ -175,7 +173,7 @@ func (h *ChatHandler) Remove(c *gin.Context) {
// TODO: 是否要删除 MidJourney 绘画记录和图片文件？
// 清空会话上下文
h.App.ChatContexts.Delete(chatId)
h.ChatContexts.Delete(chatId)
resp.SUCCESS(c, types.OkMsg)
}

3
api/handler/chatimpl/openai_handler.go
View File

@@ -74,6 +74,9 @@ func (h *ChatHandler) sendOpenAiMessage(
if len(responseBody.Choices) == 0 { // Fixed: 兼容 Azure API 第一个输出空行
continue
}
if responseBody.Choices[0].Delta.Content == nil {
continue
}
if responseBody.Choices[0].FinishReason == "stop" && len(contents) == 0 {
utils.ReplyMessage(ws, "抱歉😔😔😔，AI助手由于未知原因已经停止输出内容。")

2
api/handler/chatimpl/xunfei_handler.go
View File

@@ -79,7 +79,7 @@ func (h *ChatHandler) sendXunFeiMessage(
var res *gorm.DB
// use the bind key
if session.Model.KeyId > 0 {
res = h.DB.Where("id", session.Model.KeyId).Where("enabled", true).Find(&apiKey)
res = h.DB.Where("id", session.Model.KeyId).Find(&apiKey)
}
// use the last unused key
if apiKey.Id == 0 {

33
api/handler/dalle_handler.go
View File

@@ -158,13 +158,13 @@ func (h *DallJobHandler) ImgWall(c *gin.Context) {
// JobList 获取 SD 任务列表
func (h *DallJobHandler) JobList(c *gin.Context) {
status := h.GetBool(c, "status")
finish := h.GetBool(c, "finish")
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)
err, jobs := h.getData(finish, userId, page, pageSize, publish)
if err != nil {
resp.ERROR(c, err.Error())
return
@@ -214,25 +214,23 @@ func (h *DallJobHandler) getData(finish bool, userId uint, page int, pageSize in
// 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)
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
var job model.DallJob
if res := h.DB.Where("id = ? AND user_id = ?", id, userId).First(&job); res.Error != nil {
resp.ERROR(c, "记录不存在")
return
}
// remove job recode
res := h.DB.Delete(&model.DallJob{Id: data.Id})
res := h.DB.Delete(&model.DallJob{Id: job.Id})
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
return
}
// remove image
err := h.uploader.GetUploadHandler().Delete(data.ImgURL)
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}
@@ -242,16 +240,11 @@ func (h *DallJobHandler) Remove(c *gin.Context) {
// 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
}
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
res := h.DB.Model(&model.DallJob{Id: data.Id}).UpdateColumn("publish", true)
res := h.DB.Model(&model.DallJob{Id: uint(id), UserId: uint(userId)}).UpdateColumn("publish", action)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")

2
api/handler/markmap_handler.go
View File

@@ -215,7 +215,7 @@ func (h *MarkMapHandler) doRequest(req types.ApiRequest, chatModel model.ChatMod
// if the chat model bind a KEY, use it directly
var res *gorm.DB
if chatModel.KeyId > 0 {
res = h.DB.Where("id", chatModel.KeyId).Where("enabled", true).Find(apiKey)
res = h.DB.Where("id", chatModel.KeyId).Find(apiKey)
}
// use the last unused key
if apiKey.Id == 0 {

76
api/handler/mj_handler.go
View File

@@ -92,19 +92,18 @@ func (h *MidJourneyHandler) Client(c *gin.Context) {
// Image 创建一个绘画任务
func (h *MidJourneyHandler) Image(c *gin.Context) {
var data struct {
SessionId string `json:"session_id"`
TaskType string `json:"task_type"`
Prompt string `json:"prompt"`
NegPrompt string `json:"neg_prompt"`
Rate string `json:"rate"`
Model string `json:"model"`
Chaos int `json:"chaos"`
Raw bool `json:"raw"`
Seed int64 `json:"seed"`
Stylize int `json:"stylize"`
Model string `json:"model"` // 模型
Chaos int `json:"chaos"` // 创意度取值范围: 0-100
Raw bool `json:"raw"` // 是否开启原始模型
Seed int64 `json:"seed"` // 随机数
Stylize int `json:"stylize"` // 风格化
ImgArr []string `json:"img_arr"`
Tile bool `json:"tile"`
Quality float32 `json:"quality"`
Tile bool `json:"tile"` // 重复平铺
Quality float32 `json:"quality"` // 画质
Iw float32 `json:"iw"`
CRef string `json:"cref"` //生成角色一致的图像
SRef string `json:"sref"` //生成风格一致的图像
@@ -243,17 +242,12 @@ type reqVo struct {
ChannelId string `json:"channel_id"`
MessageId string `json:"message_id"`
MessageHash string `json:"message_hash"`
SessionId string `json:"session_id"`
Prompt string `json:"prompt"`
ChatId string `json:"chat_id"`
RoleId int `json:"role_id"`
Icon string `json:"icon"`
}
// Upscale send upscale command to MidJourney Bot
func (h *MidJourneyHandler) Upscale(c *gin.Context) {
var data reqVo
if err := c.ShouldBindJSON(&data); err != nil || data.SessionId == "" {
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
@@ -271,7 +265,6 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
UserId: userId,
TaskId: taskId,
Progress: 0,
Prompt: data.Prompt,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
}
@@ -283,7 +276,6 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
h.pool.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskUpscale,
Prompt: data.Prompt,
UserId: userId,
ChannelId: data.ChannelId,
Index: data.Index,
@@ -318,7 +310,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
// Variation send variation command to MidJourney Bot
func (h *MidJourneyHandler) Variation(c *gin.Context) {
var data reqVo
if err := c.ShouldBindJSON(&data); err != nil || data.SessionId == "" {
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
@@ -337,7 +329,6 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
UserId: userId,
TaskId: taskId,
Progress: 0,
Prompt: data.Prompt,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
}
@@ -349,7 +340,6 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
h.pool.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskVariation,
Prompt: data.Prompt,
UserId: userId,
Index: data.Index,
ChannelId: data.ChannelId,
@@ -397,13 +387,13 @@ func (h *MidJourneyHandler) ImgWall(c *gin.Context) {
// JobList 获取 MJ 任务列表
func (h *MidJourneyHandler) JobList(c *gin.Context) {
status := h.GetBool(c, "status")
finish := h.GetBool(c, "finish")
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)
err, jobs := h.getData(finish, userId, page, pageSize, publish)
if err != nil {
resp.ERROR(c, err.Error())
return
@@ -446,14 +436,9 @@ func (h *MidJourneyHandler) getData(finish bool, userId uint, page int, pageSize
}
if item.Progress < 100 && item.ImgURL == "" && item.OrgURL != "" {
// discord 服务器图片需要使用代理转发图片数据流
if strings.HasPrefix(item.OrgURL, "https://cdn.discordapp.com") {
image, err := utils.DownloadImage(item.OrgURL, h.App.Config.ProxyURL)
if err == nil {
job.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
}
} else {
job.ImgURL = job.OrgURL
image, err := utils.DownloadImage(item.OrgURL, h.App.Config.ProxyURL)
if err == nil {
job.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
}
}
@@ -464,30 +449,27 @@ func (h *MidJourneyHandler) getData(finish bool, userId uint, page int, pageSize
// Remove remove task image
func (h *MidJourneyHandler) 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)
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
var job model.MidJourneyJob
if res := h.DB.Where("id = ? AND user_id = ?", id, userId).First(&job); res.Error != nil {
resp.ERROR(c, "记录不存在")
return
}
// remove job recode
res := h.DB.Delete(&model.MidJourneyJob{Id: data.Id})
res := h.DB.Delete(&job)
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
return
}
// remove image
err := h.uploader.GetUploadHandler().Delete(data.ImgURL)
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}
client := h.pool.Clients.Get(data.UserId)
client := h.pool.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -497,16 +479,10 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
// Publish 发布图片到画廊显示
func (h *MidJourneyHandler) 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.MidJourneyJob{Id: data.Id}).UpdateColumn("publish", data.Action)
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
res := h.DB.Model(&model.MidJourneyJob{Id: uint(id), UserId: userId}).UpdateColumn("publish", action)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")

48
api/handler/order_handler.go
View File

@@ -14,6 +14,7 @@ import (
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"time"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
@@ -27,23 +28,18 @@ func NewOrderHandler(app *core.AppServer, db *gorm.DB) *OrderHandler {
return &OrderHandler{BaseHandler: BaseHandler{App: app, DB: db}}
}
// List 订单列表
func (h *OrderHandler) List(c *gin.Context) {
var data struct {
Page int `json:"page"`
PageSize int `json:"page_size"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
page := h.GetInt(c, "page", 1)
pageSize := h.GetInt(c, "page_size", 20)
userId := h.GetLoginUserId(c)
session := h.DB.Session(&gorm.Session{}).Where("user_id = ? AND status = ?", userId, types.OrderPaidSuccess)
var total int64
session.Model(&model.Order{}).Count(&total)
var items []model.Order
var list = make([]vo.Order, 0)
offset := (data.Page - 1) * data.PageSize
res := session.Order("id DESC").Offset(offset).Limit(data.PageSize).Find(&items)
offset := (page - 1) * pageSize
res := session.Order("id DESC").Offset(offset).Limit(pageSize).Find(&items)
if res.Error == nil {
for _, item := range items {
var order vo.Order
@@ -58,5 +54,35 @@ func (h *OrderHandler) List(c *gin.Context) {
}
}
}
resp.SUCCESS(c, vo.NewPage(total, data.Page, data.PageSize, list))
resp.SUCCESS(c, vo.NewPage(total, page, pageSize, list))
}
// Query 查询订单状态
func (h *OrderHandler) Query(c *gin.Context) {
orderNo := h.GetTrim(c, "order_no")
var order model.Order
res := h.DB.Where("order_no = ?", orderNo).First(&order)
if res.Error != nil {
resp.ERROR(c, "Order not found")
return
}
if order.Status == types.OrderPaidSuccess {
resp.SUCCESS(c, gin.H{"status": order.Status})
return
}
counter := 0
for {
time.Sleep(time.Second)
var item model.Order
h.DB.Where("order_no = ?", orderNo).First(&item)
if counter >= 15 || item.Status == types.OrderPaidSuccess || item.Status != order.Status {
order.Status = item.Status
break
}
counter++
}
resp.SUCCESS(c, gin.H{"status": order.Status})
}

57
api/handler/payment_handler.go
View File

@@ -111,7 +111,7 @@ func (h *PaymentHandler) DoPay(c *gin.Context) {
// fix: 这里先检查一下订单状态，如果已经支付了，就直接返回
if order.Status == types.OrderPaidSuccess {
resp.ERROR(c, "This order had been paid, please do not pay twice")
resp.ERROR(c, "订单已支付成功，无需重复支付！")
return
}
@@ -148,49 +148,11 @@ func (h *PaymentHandler) DoPay(c *gin.Context) {
resp.ERROR(c, "Invalid operations")
}
// OrderQuery 查询订单状态
func (h *PaymentHandler) OrderQuery(c *gin.Context) {
var data struct {
OrderNo string `json:"order_no"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
var order model.Order
res := h.DB.Where("order_no = ?", data.OrderNo).First(&order)
if res.Error != nil {
resp.ERROR(c, "Order not found")
return
}
if order.Status == types.OrderPaidSuccess {
resp.SUCCESS(c, gin.H{"status": order.Status})
return
}
counter := 0
for {
time.Sleep(time.Second)
var item model.Order
h.DB.Where("order_no = ?", data.OrderNo).First(&item)
if counter >= 15 || item.Status == types.OrderPaidSuccess || item.Status != order.Status {
order.Status = item.Status
break
}
counter++
}
resp.SUCCESS(c, gin.H{"status": order.Status})
}
// PayQrcode 生成支付 URL 二维码
func (h *PaymentHandler) PayQrcode(c *gin.Context) {
var data struct {
PayWay string `json:"pay_way"` // 支付方式
ProductId uint `json:"product_id"`
UserId int `json:"user_id"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
@@ -209,10 +171,9 @@ func (h *PaymentHandler) PayQrcode(c *gin.Context) {
resp.ERROR(c, "error with generate trade no: "+err.Error())
return
}
var user model.User
res = h.DB.First(&user, data.UserId)
if res.Error != nil {
resp.ERROR(c, "Invalid user ID")
user, err := h.GetLoginUser(c)
if err != nil {
resp.NotAuth(c)
return
}
@@ -333,7 +294,6 @@ func (h *PaymentHandler) Mobile(c *gin.Context) {
var data struct {
PayWay string `json:"pay_way"` // 支付方式
ProductId uint `json:"product_id"`
UserId int `json:"user_id"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
@@ -352,10 +312,9 @@ func (h *PaymentHandler) Mobile(c *gin.Context) {
resp.ERROR(c, "error with generate trade no: "+err.Error())
return
}
var user model.User
res = h.DB.First(&user, data.UserId)
if res.Error != nil {
resp.ERROR(c, "Invalid user ID")
user, err := h.GetLoginUser(c)
if err != nil {
resp.NotAuth(c)
return
}
@@ -449,7 +408,7 @@ func (h *PaymentHandler) Mobile(c *gin.Context) {
return
}
resp.SUCCESS(c, payURL)
resp.SUCCESS(c, gin.H{"url": payURL, "order_no": orderNo})
}
// 异步通知回调公共逻辑

40
api/handler/sd_handler.go
View File

@@ -99,10 +99,7 @@ func (h *SdJobHandler) Image(c *gin.Context) {
return
}
var data struct {
SessionId string `json:"session_id"`
types.SdTaskParams
}
var data types.SdTaskParams
if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
resp.ERROR(c, types.InvalidArgs)
return
@@ -215,13 +212,13 @@ func (h *SdJobHandler) ImgWall(c *gin.Context) {
// JobList 获取 SD 任务列表
func (h *SdJobHandler) JobList(c *gin.Context) {
status := h.GetBool(c, "status")
finish := h.GetBool(c, "finish")
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)
err, jobs := h.getData(finish, userId, page, pageSize, publish)
if err != nil {
resp.ERROR(c, err.Error())
return
@@ -280,30 +277,28 @@ func (h *SdJobHandler) getData(finish bool, userId uint, page int, pageSize int,
// Remove remove task image
func (h *SdJobHandler) 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)
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
var job model.SdJob
if res := h.DB.Where("id = ? AND user_id = ?", id, userId).First(&job); res.Error != nil {
resp.ERROR(c, "记录不存在")
return
}
// remove job recode
res := h.DB.Delete(&model.SdJob{Id: data.Id})
res := h.DB.Delete(&model.SdJob{Id: job.Id})
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
return
}
// remove image
err := h.uploader.GetUploadHandler().Delete(data.ImgURL)
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}
client := h.pool.Clients.Get(data.UserId)
client := h.pool.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte(sd.Finished))
}
@@ -313,16 +308,11 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
// Publish 发布/取消发布图片到画廊显示
func (h *SdJobHandler) 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
}
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
res := h.DB.Model(&model.SdJob{Id: data.Id}).UpdateColumn("publish", true)
res := h.DB.Model(&model.SdJob{Id: uint(id), UserId: userId}).UpdateColumn("publish", action)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")

22
api/handler/sms_handler.go
View File

@@ -49,14 +49,20 @@ func (h *SmsHandler) SendCode(c *gin.Context) {
var data struct {
Receiver string `json:"receiver"` // 接收者
Key string `json:"key"`
Dots string `json:"dots"`
Dots string `json:"dots,omitempty"`
X int `json:"x,omitempty"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
if !h.captcha.Check(data) {
var check bool
if data.X != 0 {
check = h.captcha.SlideCheck(data)
} else {
check = h.captcha.Check(data)
}
if !check {
resp.ERROR(c, "验证码错误，请先完人机验证")
return
}
@@ -64,13 +70,13 @@ func (h *SmsHandler) SendCode(c *gin.Context) {
code := utils.RandomNumber(6)
var err error
if strings.Contains(data.Receiver, "@") { // email
if !utils.ContainsStr(h.App.SysConfig.RegisterWays, "email") {
if !utils.Contains(h.App.SysConfig.RegisterWays, "email") {
resp.ERROR(c, "系统已禁用邮箱注册！")
return
}
err = h.smtp.SendVerifyCode(data.Receiver, code)
} else {
if !utils.ContainsStr(h.App.SysConfig.RegisterWays, "mobile") {
if !utils.Contains(h.App.SysConfig.RegisterWays, "mobile") {
resp.ERROR(c, "系统已禁用手机号注册！")
return
}
@@ -89,5 +95,9 @@ func (h *SmsHandler) SendCode(c *gin.Context) {
return
}
resp.SUCCESS(c)
if h.App.Debug {
resp.SUCCESS(c, code)
} else {
resp.SUCCESS(c)
}
}

22
api/handler/upload_handler.go
View File

@@ -9,6 +9,7 @@ package handler
import (
"geekai/core"
"geekai/core/types"
"geekai/service/oss"
"geekai/store/model"
"geekai/store/vo"
@@ -35,6 +36,12 @@ func (h *UploadHandler) Upload(c *gin.Context) {
return
}
logger.Info("upload file: ", file.Name)
// cut the file name if it's too long
if len(file.Name) > 100 {
file.Name = file.Name[:90] + file.Ext
}
userId := h.GetLoginUserId(c)
res := h.DB.Create(&model.File{
UserId: int(userId),
@@ -54,10 +61,23 @@ func (h *UploadHandler) Upload(c *gin.Context) {
}
func (h *UploadHandler) List(c *gin.Context) {
var data struct {
Urls []string `json:"urls,omitempty"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
userId := h.GetLoginUserId(c)
var items []model.File
var files = make([]vo.File, 0)
h.DB.Where("user_id = ?", userId).Find(&items)
session := h.DB.Session(&gorm.Session{})
session = session.Where("user_id = ?", userId)
if len(data.Urls) > 0 {
session = session.Where("url IN ?", data.Urls)
}
session.Find(&items)
if len(items) > 0 {
for _, v := range items {
var file vo.File

133
api/handler/user_handler.go
View File

@@ -16,6 +16,7 @@ import (
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/imroc/req/v3"
"strings"
"time"
@@ -97,7 +98,7 @@ func (h *UserHandler) Register(c *gin.Context) {
}
}
// check if the username is exists
// check if the username is existing
var item model.User
res := h.DB.Where("username = ?", data.Username).First(&item)
if item.Id > 0 {
@@ -122,7 +123,7 @@ func (h *UserHandler) Register(c *gin.Context) {
user.Power += h.App.SysConfig.InvitePower
}
if h.licenseService.GetLicense().Configs.DeCopy {
user.Username = fmt.Sprintf("用户@%d", utils.RandomNumber(6))
user.Nickname = fmt.Sprintf("用户@%d", utils.RandomNumber(6))
} else {
user.Nickname = fmt.Sprintf("极客学长@%d", utils.RandomNumber(6))
}
@@ -184,7 +185,7 @@ func (h *UserHandler) Register(c *gin.Context) {
resp.ERROR(c, "error with save token: "+err.Error())
return
}
resp.SUCCESS(c, tokenString)
resp.SUCCESS(c, gin.H{"token": tokenString, "user_id": user.Id, "username": user.Username})
}
// Login 用户登录
@@ -243,7 +244,7 @@ func (h *UserHandler) Login(c *gin.Context) {
resp.ERROR(c, "error with save token: "+err.Error())
return
}
resp.SUCCESS(c, tokenString)
resp.SUCCESS(c, gin.H{"token": tokenString, "user_id": user.Id, "username": user.Username})
}
// Logout 注 销
@@ -255,6 +256,130 @@ func (h *UserHandler) Logout(c *gin.Context) {
resp.SUCCESS(c)
}
// CLogin 第三方登录请求二维码
func (h *UserHandler) CLogin(c *gin.Context) {
returnURL := h.GetTrim(c, "return_url")
var res types.BizVo
apiURL := fmt.Sprintf("%s/api/clogin/request", h.App.Config.ApiConfig.ApiURL)
r, err := req.C().R().SetBody(gin.H{"login_type": "wx", "return_url": returnURL}).
SetHeader("AppId", h.App.Config.ApiConfig.AppId).
SetHeader("Authorization", fmt.Sprintf("Bearer %s", h.App.Config.ApiConfig.Token)).
SetSuccessResult(&res).
Post(apiURL)
if err != nil {
resp.ERROR(c, err.Error())
return
}
if r.IsErrorState() {
resp.ERROR(c, "error with login http status: "+r.Status)
return
}
if res.Code != types.Success {
resp.ERROR(c, "error with http response: "+res.Message)
return
}
resp.SUCCESS(c, res.Data)
}
// CLoginCallback 第三方登录回调
func (h *UserHandler) CLoginCallback(c *gin.Context) {
loginType := h.GetTrim(c, "login_type")
code := h.GetTrim(c, "code")
var res types.BizVo
apiURL := fmt.Sprintf("%s/api/clogin/info", h.App.Config.ApiConfig.ApiURL)
r, err := req.C().R().SetBody(gin.H{"login_type": loginType, "code": code}).
SetHeader("AppId", h.App.Config.ApiConfig.AppId).
SetHeader("Authorization", fmt.Sprintf("Bearer %s", h.App.Config.ApiConfig.Token)).
SetSuccessResult(&res).
Post(apiURL)
if err != nil {
resp.ERROR(c, err.Error())
return
}
if r.IsErrorState() {
resp.ERROR(c, "error with login http status: "+r.Status)
return
}
if res.Code != types.Success {
resp.ERROR(c, "error with http response: "+res.Message)
return
}
// login successfully
data := res.Data.(map[string]interface{})
session := gin.H{}
var user model.User
tx := h.DB.Debug().Where("openid", data["openid"]).First(&user)
if tx.Error != nil { // user not exist, create new user
// 检测最大注册人数
var totalUser int64
h.DB.Model(&model.User{}).Count(&totalUser)
if h.licenseService.GetLicense().Configs.UserNum > 0 && int(totalUser) >= h.licenseService.GetLicense().Configs.UserNum {
resp.ERROR(c, "当前注册用户数已达上限，请请升级 License")
return
}
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"]),
}
tx = h.DB.Create(&user)
if tx.Error != nil {
resp.ERROR(c, "保存数据失败")
logger.Error(tx.Error)
return
}
session["username"] = user.Username
session["password"] = password
} else { // login directly
// 更新最后登录时间和IP
user.LastLoginIp = c.ClientIP()
user.LastLoginAt = time.Now().Unix()
h.DB.Model(&user).Updates(user)
h.DB.Create(&model.UserLoginLog{
UserId: user.Id,
Username: user.Username,
LoginIp: c.ClientIP(),
LoginAddress: utils.Ip2Region(h.searcher, c.ClientIP()),
})
}
// 创建 token
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
"user_id": user.Id,
"expired": time.Now().Add(time.Second * time.Duration(h.App.Config.Session.MaxAge)).Unix(),
})
tokenString, err := token.SignedString([]byte(h.App.Config.Session.SecretKey))
if err != nil {
resp.ERROR(c, "Failed to generate token, "+err.Error())
return
}
// 保存到 redis
key := fmt.Sprintf("users/%d", user.Id)
if _, err := h.redis.Set(c, key, tokenString, 0).Result(); err != nil {
resp.ERROR(c, "error with save token: "+err.Error())
return
}
session["token"] = tokenString
resp.SUCCESS(c, session)
}
// Session 获取/验证会话
func (h *UserHandler) Session(c *gin.Context) {
user, err := h.GetLoginUser(c)