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

merge v4.1.2

Browse Source
This commit is contained in:
RockYang
2024-11-27 15:00:02 +08:00
parent c434f85045 740c3c1b00
commit 296eabe09a
111 changed files with 3072 additions and 2940 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
api/config.sample.toml
View File

@@ -3,8 +3,6 @@ ProxyURL = "" # 如 http://127.0.0.1:7777
MysqlDns = "root:12345678@tcp(172.22.11.200:3307)/chatgpt_plus?charset=utf8mb4&collation=utf8mb4_unicode_ci&parseTime=True&loc=Local"
StaticDir = "./static" # 静态资源的目录
StaticUrl = "/static" # 静态资源访问 URL
AesEncryptKey = ""
WeChatBot = false
TikaHost = "http://tika:9998"
[Session]
@@ -65,23 +63,6 @@ TikaHost = "http://tika:9998"
SubDir = ""
Domain = ""
[[MjProxyConfigs]]
Enabled = true
ApiURL = "http://midjourney-proxy:8082"
ApiKey = "sk-geekmaster"
[[MjPlusConfigs]]
Enabled = false
ApiURL = "https://api.chat-plus.net"
Mode = "fast" # MJ 绘画模式，可选值 relax/fast/turbo
ApiKey = "sk-xxx"
[[SdConfigs]]
Enabled = false
ApiURL = ""
ApiKey = ""
Txt2ImgJsonPath = "res/sd/text2img.json"
[XXLConfig] # xxl-job 配置，需要你部署 XXL-JOB 定时任务工具，用来定期清理未支付订单和清理过期 VIP，如果你没有启用支付服务，则该服务也无需启动
Enabled = false # 是否启用 XXL JOB 服务
ServerAddr = "http://172.22.11.47:8080/xxl-job-admin" # xxl-job-admin 管理地址

2
api/core/app_server.go
View File

@@ -225,7 +225,7 @@ func needLogin(c *gin.Context) bool {
c.Request.URL.Path == "/api/payment/doPay" ||
c.Request.URL.Path == "/api/payment/payWays" ||
c.Request.URL.Path == "/api/suno/client" ||
c.Request.URL.Path == "/api/suno/Detail" ||
c.Request.URL.Path == "/api/suno/detail" ||
c.Request.URL.Path == "/api/suno/play" ||
strings.HasPrefix(c.Request.URL.Path, "/api/test") ||
strings.HasPrefix(c.Request.URL.Path, "/api/user/clogin") ||

1
api/core/config.go
View File

@@ -38,7 +38,6 @@ func NewDefaultConfig() *types.AppConfig {
BasePath: "./static/upload",
},
},
WeChatBot: false,
AlipayConfig: types.AlipayConfig{Enabled: false, SandBox: false},
}
}

7
api/core/types/chat.go
View File

@@ -61,7 +61,6 @@ type ChatSession struct {
type ChatModel struct {
Id uint `json:"id"`
Platform string `json:"platform"`
Name string `json:"name"`
Value string `json:"value"`
Power int `json:"power"`
@@ -91,7 +90,7 @@ const (
PowerConsume = PowerType(2) // 消费
PowerRefund = PowerType(3) // 任务（SD,MJ）执行失败，退款
PowerInvite = PowerType(4) // 邀请奖励
PowerReward = PowerType(5) // 众筹
PowerRedeem = PowerType(5) // 众筹
PowerGift = PowerType(6) // 系统赠送
)
@@ -103,8 +102,8 @@ func (t PowerType) String() string {
return "消费"
case PowerRefund:
return "退款"
case PowerReward:
return "众筹"
case PowerRedeem:
return "兑换"
}
return "其他"

61
api/core/types/config.go
View File

@@ -12,22 +12,18 @@ import (
)
type AppConfig struct {
Path string `toml:"-"`
Listen string
Session Session
AdminSession Session
ProxyURL string
MysqlDns string // mysql 连接地址
StaticDir string // 静态资源目录
StaticUrl string // 静态资源 URL
Redis RedisConfig // redis 连接信息
ApiConfig ApiConfig // ChatPlus API authorization configs
SMS SMSConfig // send mobile message config
OSS OSSConfig // OSS config
MjProxyConfigs []MjProxyConfig // MJ proxy config
MjPlusConfigs []MjPlusConfig // MJ plus config
WeChatBot bool // 是否启用微信机器人
SdConfigs []StableDiffusionConfig // sd AI draw service pool
Path string `toml:"-"`
Listen string
Session Session
AdminSession Session
ProxyURL string
MysqlDns string // mysql 连接地址
StaticDir string // 静态资源目录
StaticUrl string // 静态资源 URL
Redis RedisConfig // redis 连接信息
ApiConfig ApiConfig // ChatPlus API authorization configs
SMS SMSConfig // send mobile message config
OSS OSSConfig // OSS config
XXLConfig XXLConfig
AlipayConfig AlipayConfig // 支付宝支付渠道配置
@@ -53,27 +49,6 @@ type ApiConfig struct {
Token string
}
type MjProxyConfig struct {
Enabled bool
ApiURL string // api 地址
Mode string // 绘画模式，可选值：fast/turbo/relax
ApiKey string
}
type StableDiffusionConfig struct {
Enabled bool
Model string // 模型名称
ApiURL string
ApiKey string
}
type MjPlusConfig struct {
Enabled bool // 如果启用了 MidJourney Plus，将会自动禁用原生的MidJourney服务
ApiURL string // api 地址
Mode string // 绘画模式，可选值：fast/turbo/relax
ApiKey string
}
type AlipayConfig struct {
Enabled bool // 是否启用该支付通道
SandBox bool // 是否沙盒环境
@@ -168,10 +143,6 @@ type SystemConfig struct {
RegisterWays []string `json:"register_ways,omitempty"` // 注册方式：支持手机（mobile），邮箱注册（email），账号密码注册
EnabledRegister bool `json:"enabled_register,omitempty"` // 是否开放注册
RewardImg string `json:"reward_img,omitempty"` // 众筹收款二维码地址
EnabledReward bool `json:"enabled_reward,omitempty"` // 启用众筹功能
PowerPrice float64 `json:"power_price,omitempty"` // 算力单价
OrderPayTimeout int `json:"order_pay_timeout,omitempty"` //订单支付超时时间
VipInfoText string `json:"vip_info_text,omitempty"` // 会员页面充值说明
DefaultModels []int `json:"default_models,omitempty"` // 默认开通的 AI 模型
@@ -188,8 +159,10 @@ type SystemConfig struct {
ContextDeep int `json:"context_deep,omitempty"`
SdNegPrompt string `json:"sd_neg_prompt"` // SD 默认反向提示词
MjMode string `json:"mj_mode"` // midjourney 默认的API模式，relax, fast, turbo
IndexBgURL string `json:"index_bg_url"` // 前端首页背景图片
IndexNavs []int `json:"index_navs"` // 首页显示的导航菜单
Copyright string `json:"copyright"` // 版权信息
IndexBgURL string `json:"index_bg_url"` // 前端首页背景图片
IndexNavs []int `json:"index_navs"` // 首页显示的导航菜单
Copyright string `json:"copyright"` // 版权信息
MarkMapText string `json:"mark_map_text"` // 思维导入的默认文本
}

3
api/core/types/task.go
View File

@@ -27,7 +27,6 @@ type MjTask struct {
Id uint `json:"id"`
TaskId string `json:"task_id"`
ImgArr []string `json:"img_arr"`
ChannelId string `json:"channel_id"`
Type TaskType `json:"type"`
UserId int `json:"user_id"`
Prompt string `json:"prompt,omitempty"`
@@ -37,6 +36,8 @@ type MjTask struct {
MessageId string `json:"message_id,omitempty"`
MessageHash string `json:"message_hash,omitempty"`
RetryCount int `json:"retry_count"`
ChannelId string `json:"channel_id"` // 渠道ID，用来区分是哪个渠道创建的任务，一个任务的 create 和 action 操作必须要再同一个渠道
Mode string `json:"mode"` // 绘画模式，relax, fast, turbo
}
type SdTask struct {

3
api/go.mod
View File

@@ -8,7 +8,6 @@ require (
github.com/BurntSushi/toml v1.1.0
github.com/aliyun/alibaba-cloud-sdk-go v1.62.405
github.com/aliyun/aliyun-oss-go-sdk v2.2.9+incompatible
github.com/eatmoreapple/openwechat v1.2.1
github.com/gin-gonic/gin v1.9.1
github.com/go-redis/redis/v8 v8.11.5
github.com/golang-jwt/jwt/v5 v5.0.0
@@ -30,7 +29,6 @@ require (
github.com/go-pay/gopay v1.5.101
github.com/google/go-tika v0.3.1
github.com/microcosm-cc/bluemonday v1.0.26
github.com/mojocn/base64Captcha v1.3.6
github.com/shirou/gopsutil v3.21.11+incompatible
github.com/shopspring/decimal v1.3.1
github.com/syndtr/goleveldb v1.0.0
@@ -45,7 +43,6 @@ require (
github.com/go-pay/util v0.0.2 // indirect
github.com/go-pay/xlog v0.0.2 // indirect
github.com/go-pay/xtime v0.0.2 // indirect
github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db // indirect
github.com/gorilla/css v1.0.0 // indirect
github.com/tklauser/go-sysconf v0.3.13 // indirect

8
api/go.sum
View File

@@ -28,8 +28,6 @@ github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0
github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/eatmoreapple/openwechat v1.2.1 h1:ez4oqF/Y2NSEX/DbPV8lvj7JlfkYqvieeo4awx5lzfU=
github.com/eatmoreapple/openwechat v1.2.1/go.mod h1:61HOzTyvLobGdgWhL68jfGNwTJEv0mhQ1miCXQrvWU8=
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4=
github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
@@ -84,8 +82,6 @@ github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MG
github.com/goji/httpauth v0.0.0-20160601135302-2da839ab0f4d/go.mod h1:nnjvkQ9ptGaCkuDUx6wNykzzlUixGxvkme+H/lnzb+A=
github.com/golang-jwt/jwt/v5 v5.0.0 h1:1n1XNM9hk7O9mnQoNBGolZvzebBQ7p93ULHRc28XJUE=
github.com/golang-jwt/jwt/v5 v5.0.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 h1:DACJavvAHhabrF08vX0COfcOBJRhZ8lUbR+ZWIs0Y5g=
github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k=
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
@@ -157,8 +153,6 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ
github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
github.com/mojocn/base64Captcha v1.3.6 h1:gZEKu1nsKpttuIAQgWHO+4Mhhls8cAKyiV2Ew03H+Tw=
github.com/mojocn/base64Captcha v1.3.6/go.mod h1:i5CtHvm+oMbj1UzEPXaA8IH/xHFZ3DGY3Wh3dBpZ28E=
github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646 h1:zYyBkD/k9seD2A7fsi6Oo2LfFZAehjjQMERAvZLEDnQ=
github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646/go.mod h1:jpp1/29i3P1S/RLdc7JQKbRpFeM1dOBd8T9ki5s+AY8=
github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
@@ -267,7 +261,6 @@ golang.org/x/crypto v0.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI=
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 h1:vr/HnozRka3pE4EsMEg1lgkXJkTFJCVUX+S/ZT6wYzM=
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842/go.mod h1:XtvwrStGgqGPLc4cjQfWqZHG1YFdYs6swckp8vpsjnc=
golang.org/x/image v0.13.0/go.mod h1:6mmbMOeV28HuMTgA6OSRkdXKYw/t5W9Uwn2Yv1r3Yxk=
golang.org/x/image v0.15.0 h1:kOELfmgrmJlw4Cdb7g/QGuB3CvDrXbqEIww/pNtNBm8=
golang.org/x/image v0.15.0/go.mod h1:HUYqC05R2ZcZ3ejNQsIHQDQiwWM4JBqmm6MKANTp4LE=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
@@ -323,7 +316,6 @@ golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
golang.org/x/text v0.15.0 h1:h1V/4gjBv8v9cjcR6+AR5+/cIYK5N/WAgiv4xlsEtAk=
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=

26
api/handler/admin/api_key_handler.go
View File

@@ -8,6 +8,7 @@ package admin
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/handler"
@@ -53,17 +54,16 @@ func (h *ApiKeyHandler) Save(c *gin.Context) {
apiKey.Enabled = data.Enabled
apiKey.ProxyURL = data.ProxyURL
apiKey.Name = data.Name
res := h.DB.Save(&apiKey)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Save(&apiKey).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
var keyVo vo.ApiKey
err := utils.CopyObject(apiKey, &keyVo)
err = utils.CopyObject(apiKey, &keyVo)
if err != nil {
resp.ERROR(c, "数据拷贝失败！")
resp.ERROR(c, fmt.Sprintf("拷贝数据失败：%v", err))
return
}
keyVo.Id = apiKey.Id
@@ -119,10 +119,9 @@ func (h *ApiKeyHandler) Set(c *gin.Context) {
return
}
res := h.DB.Model(&model.ApiKey{}).Where("id = ?", data.Id).Update(data.Filed, data.Value)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.ApiKey{}).Where("id = ?", data.Id).Update(data.Filed, data.Value).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@@ -135,10 +134,9 @@ func (h *ApiKeyHandler) Remove(c *gin.Context) {
return
}
res := h.DB.Where("id", id).Delete(&model.ApiKey{})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Where("id", id).Delete(&model.ApiKey{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)

7
api/handler/admin/chat_handler.go
View File

@@ -259,10 +259,9 @@ func (h *ChatHandler) RemoveChat(c *gin.Context) {
// RemoveMessage 删除聊天记录
func (h *ChatHandler) RemoveMessage(c *gin.Context) {
id := h.GetInt(c, "id", 0)
tx := h.DB.Unscoped().Where("id = ?", id).Delete(&model.ChatMessage{})
if tx.Error != nil {
logger.Error("error with update database：", tx.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Unscoped().Where("id = ?", id).Delete(&model.ChatMessage{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)

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

@@ -147,10 +147,9 @@ func (h *ChatModelHandler) Set(c *gin.Context) {
return
}
res := h.DB.Model(&model.ChatModel{}).Where("id = ?", data.Id).Update(data.Filed, data.Value)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.ChatModel{}).Where("id = ?", data.Id).Update(data.Filed, data.Value).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@@ -168,10 +167,9 @@ func (h *ChatModelHandler) Sort(c *gin.Context) {
}
for index, id := range data.Ids {
res := h.DB.Model(&model.ChatModel{}).Where("id = ?", id).Update("sort_num", data.Sorts[index])
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.ChatModel{}).Where("id = ?", id).Update("sort_num", data.Sorts[index]).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}
@@ -186,10 +184,9 @@ func (h *ChatModelHandler) Remove(c *gin.Context) {
return
}
res := h.DB.Where("id = ?", id).Delete(&model.ChatModel{})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Where("id = ?", id).Delete(&model.ChatModel{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)

28
api/handler/admin/chat_role_handler.go
View File

@@ -8,6 +8,7 @@ package admin
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/handler"
@@ -45,11 +46,16 @@ func (h *ChatRoleHandler) Save(c *gin.Context) {
role.Id = data.Id
if data.CreatedAt > 0 {
role.CreatedAt = time.Unix(data.CreatedAt, 0)
} else {
err = h.DB.Where("marker", data.Key).First(&role).Error
if err == nil {
resp.ERROR(c, fmt.Sprintf("角色 %s 已存在", data.Key))
return
}
}
res := h.DB.Save(&role)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err = h.DB.Save(&role).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
// 填充 ID 数据
@@ -114,10 +120,9 @@ func (h *ChatRoleHandler) Sort(c *gin.Context) {
}
for index, id := range data.Ids {
res := h.DB.Model(&model.ChatRole{}).Where("id = ?", id).Update("sort_num", data.Sorts[index])
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.ChatRole{}).Where("id = ?", id).Update("sort_num", data.Sorts[index]).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}
@@ -137,10 +142,9 @@ func (h *ChatRoleHandler) Set(c *gin.Context) {
return
}
res := h.DB.Model(&model.ChatRole{}).Where("id = ?", data.Id).Update(data.Filed, data.Value)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.ChatRole{}).Where("id = ?", data.Id).Update(data.Filed, data.Value).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)

63
api/handler/admin/config_handler.go
View File

@@ -12,8 +12,6 @@ import (
"geekai/core/types"
"geekai/handler"
"geekai/service"
"geekai/service/mj"
"geekai/service/sd"
"geekai/store"
"geekai/store/model"
"geekai/utils"
@@ -28,16 +26,12 @@ type ConfigHandler struct {
handler.BaseHandler
levelDB *store.LevelDB
licenseService *service.LicenseService
mjServicePool *mj.ServicePool
sdServicePool *sd.ServicePool
}
func NewConfigHandler(app *core.AppServer, db *gorm.DB, levelDB *store.LevelDB, licenseService *service.LicenseService, mjPool *mj.ServicePool, sdPool *sd.ServicePool) *ConfigHandler {
func NewConfigHandler(app *core.AppServer, db *gorm.DB, levelDB *store.LevelDB, licenseService *service.LicenseService) *ConfigHandler {
return &ConfigHandler{
BaseHandler: handler.BaseHandler{App: app, DB: db},
levelDB: levelDB,
mjServicePool: mjPool,
sdServicePool: sdPool,
licenseService: licenseService,
}
}
@@ -146,58 +140,3 @@ func (h *ConfigHandler) GetLicense(c *gin.Context) {
license := h.licenseService.GetLicense()
resp.SUCCESS(c, license)
}
// GetAppConfig 获取内置配置
func (h *ConfigHandler) GetAppConfig(c *gin.Context) {
resp.SUCCESS(c, gin.H{
"mj_plus": h.App.Config.MjPlusConfigs,
"mj_proxy": h.App.Config.MjProxyConfigs,
"sd": h.App.Config.SdConfigs,
})
}
// SaveDrawingConfig 保存AI绘画配置
func (h *ConfigHandler) SaveDrawingConfig(c *gin.Context) {
var data struct {
Sd []types.StableDiffusionConfig `json:"sd"`
MjPlus []types.MjPlusConfig `json:"mj_plus"`
MjProxy []types.MjProxyConfig `json:"mj_proxy"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
changed := false
if configChanged(data.Sd, h.App.Config.SdConfigs) {
logger.Debugf("SD 配置变动了")
h.App.Config.SdConfigs = data.Sd
h.sdServicePool.InitServices(data.Sd)
changed = true
}
if configChanged(data.MjPlus, h.App.Config.MjPlusConfigs) || configChanged(data.MjProxy, h.App.Config.MjProxyConfigs) {
logger.Debugf("MidJourney 配置变动了")
h.App.Config.MjPlusConfigs = data.MjPlus
h.App.Config.MjProxyConfigs = data.MjProxy
h.mjServicePool.InitServices(data.MjPlus, data.MjProxy)
changed = true
}
if changed {
err := core.SaveConfig(h.App.Config)
if err != nil {
resp.ERROR(c, "更新配置文档失败！")
return
}
}
resp.SUCCESS(c)
}
func configChanged(c1 interface{}, c2 interface{}) bool {
encode1 := utils.JsonEncode(c1)
encode2 := utils.JsonEncode(c2)
return utils.Md5(encode1) != utils.Md5(encode2)
}

14
api/handler/admin/dashboard_handler.go
View File

@@ -60,13 +60,6 @@ func (h *DashboardHandler) Stats(c *gin.Context) {
stats.Tokens += item.Tokens
}
// 众筹收入
var rewards []model.Reward
res = h.DB.Where("created_at > ?", zeroTime).Find(&rewards)
for _, item := range rewards {
stats.Income += item.Amount
}
// 订单收入
var orders []model.Order
res = h.DB.Where("status = ?", types.OrderPaidSuccess).Where("created_at > ?", zeroTime).Find(&orders)
@@ -101,13 +94,6 @@ func (h *DashboardHandler) Stats(c *gin.Context) {
historyMessagesStatistic[item.CreatedAt.Format("2006-01-02")] += float64(item.Tokens)
}
// 浮点数相加？
// 统计最近7天的众筹
res = h.DB.Where("created_at > ?", startDate).Find(&rewards)
for _, item := range rewards {
incomeStatistic[item.CreatedAt.Format("2006-01-02")], _ = decimal.NewFromFloat(incomeStatistic[item.CreatedAt.Format("2006-01-02")]).Add(decimal.NewFromFloat(item.Amount)).Float64()
}
// 统计最近7天的订单
res = h.DB.Where("status = ?", types.OrderPaidSuccess).Where("created_at > ?", startDate).Find(&orders)
for _, item := range orders {

14
api/handler/admin/function_handler.go
View File

@@ -69,10 +69,9 @@ func (h *FunctionHandler) Set(c *gin.Context) {
return
}
res := h.DB.Model(&model.Function{}).Where("id = ?", data.Id).Update(data.Filed, data.Value)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.Function{}).Where("id = ?", data.Id).Update(data.Filed, data.Value).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@@ -102,10 +101,9 @@ func (h *FunctionHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
if id > 0 {
res := h.DB.Delete(&model.Function{Id: uint(id)})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Delete(&model.Function{Id: uint(id)}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}

30
api/handler/admin/menu_handler.go
View File

@@ -41,17 +41,16 @@ func (h *MenuHandler) Save(c *gin.Context) {
return
}
res := h.DB.Save(&model.Menu{
err := h.DB.Save(&model.Menu{
Id: data.Id,
Name: data.Name,
Icon: data.Icon,
URL: data.URL,
SortNum: data.SortNum,
Enabled: data.Enabled,
})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@@ -85,10 +84,9 @@ func (h *MenuHandler) Enable(c *gin.Context) {
return
}
res := h.DB.Model(&model.Menu{}).Where("id", data.Id).UpdateColumn("enabled", data.Enabled)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.Menu{}).Where("id", data.Id).UpdateColumn("enabled", data.Enabled).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@@ -106,10 +104,9 @@ func (h *MenuHandler) Sort(c *gin.Context) {
}
for index, id := range data.Ids {
res := h.DB.Model(&model.Menu{}).Where("id", id).Update("sort_num", data.Sorts[index])
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.Menu{}).Where("id", id).Update("sort_num", data.Sorts[index]).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}
@@ -121,10 +118,9 @@ func (h *MenuHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
if id > 0 {
res := h.DB.Where("id", id).Delete(&model.Menu{})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Where("id", id).Delete(&model.Menu{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}

17
api/handler/admin/order_handler.go
View File

@@ -92,12 +92,21 @@ func (h *OrderHandler) Remove(c *gin.Context) {
return
}
res = h.DB.Unscoped().Where("id = ?", id).Delete(&model.Order{})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Unscoped().Where("id = ?", id).Delete(&model.Order{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}
resp.SUCCESS(c)
}
func (h *OrderHandler) Clear(c *gin.Context) {
err := h.DB.Unscoped().Where("status <> ?", 2).Where("pay_time", 0).Delete(&model.Order{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
}

32
api/handler/admin/product_handler.go
View File

@@ -55,17 +55,16 @@ func (h *ProductHandler) Save(c *gin.Context) {
if item.Id > 0 {
item.CreatedAt = time.Unix(data.CreatedAt, 0)
}
res := h.DB.Save(&item)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Save(&item).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
var itemVo vo.Product
err := utils.CopyObject(item, &itemVo)
err = utils.CopyObject(item, &itemVo)
if err != nil {
resp.ERROR(c, "数据拷贝失败！")
resp.ERROR(c, "数据拷贝失败: "+err.Error())
return
}
itemVo.Id = item.Id
@@ -106,10 +105,9 @@ func (h *ProductHandler) Enable(c *gin.Context) {
return
}
res := h.DB.Model(&model.Product{}).Where("id", data.Id).UpdateColumn("enabled", data.Enabled)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.Product{}).Where("id", data.Id).UpdateColumn("enabled", data.Enabled).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
@@ -127,10 +125,9 @@ func (h *ProductHandler) Sort(c *gin.Context) {
}
for index, id := range data.Ids {
res := h.DB.Model(&model.Product{}).Where("id", id).Update("sort_num", data.Sorts[index])
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Model(&model.Product{}).Where("id", id).Update("sort_num", data.Sorts[index]).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}
@@ -142,10 +139,9 @@ func (h *ProductHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
if id > 0 {
res := h.DB.Where("id", id).Delete(&model.Product{})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err := h.DB.Where("id", id).Delete(&model.Product{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}

164
api/handler/admin/redeem_handler.go Normal file
View File

@@ -0,0 +1,164 @@
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/core/types"
"geekai/handler"
"geekai/store/model"
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
type RedeemHandler struct {
handler.BaseHandler
}
func NewRedeemHandler(app *core.AppServer, db *gorm.DB) *RedeemHandler {
return &RedeemHandler{BaseHandler: handler.BaseHandler{App: app, DB: db}}
}
func (h *RedeemHandler) List(c *gin.Context) {
page := h.GetInt(c, "page", 1)
pageSize := h.GetInt(c, "page_size", 20)
code := c.Query("code")
status := h.GetInt(c, "status", -1)
session := h.DB.Session(&gorm.Session{})
if code != "" {
session.Where("code LIKE ?", "%"+code+"%")
}
if status == 0 {
session.Where("redeem_at = ?", 0)
} else if status == 1 {
session.Where("redeem_at > ?", 0)
}
var total int64
session.Model(&model.Redeem{}).Count(&total)
var redeems []model.Redeem
offset := (page - 1) * pageSize
err := session.Order("id DESC").Offset(offset).Limit(pageSize).Find(&redeems).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
var items = make([]vo.Redeem, 0)
userIds := make([]uint, 0)
for _, v := range redeems {
userIds = append(userIds, v.UserId)
}
var users []model.User
h.DB.Where("id IN ?", userIds).Find(&users)
var userMap = make(map[uint]model.User)
for _, u := range users {
userMap[u.Id] = u
}
for _, v := range redeems {
var r vo.Redeem
err = utils.CopyObject(v, &r)
if err != nil {
continue
}
r.Id = v.Id
r.Username = userMap[v.UserId].Username
r.CreatedAt = v.CreatedAt.Unix()
items = append(items, r)
}
resp.SUCCESS(c, vo.NewPage(total, page, pageSize, items))
}
func (h *RedeemHandler) Create(c *gin.Context) {
var data struct {
Name string `json:"name"`
Power int `json:"power"`
Num int `json:"num"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
counter := 0
codes := make([]string, 0)
var errMsg = ""
if data.Num > 0 {
for i := 0; i < data.Num; i++ {
code, err := utils.GenRedeemCode(32)
if err != nil {
errMsg = err.Error()
continue
}
err = h.DB.Create(&model.Redeem{
Code: code,
Name: data.Name,
Power: data.Power,
Enabled: true,
}).Error
if err != nil {
errMsg = err.Error()
continue
}
codes = append(codes, code)
counter++
}
}
if counter == 0 {
resp.ERROR(c, errMsg)
return
}
resp.SUCCESS(c, gin.H{
"counter": counter,
})
}
func (h *RedeemHandler) Set(c *gin.Context) {
var data struct {
Id uint `json:"id"`
Filed string `json:"filed"`
Value interface{} `json:"value"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
err := h.DB.Model(&model.Redeem{}).Where("id = ?", data.Id).Update(data.Filed, data.Value).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c)
}
func (h *RedeemHandler) Remove(c *gin.Context) {
var data struct {
Id uint
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
if data.Id > 0 {
err := h.DB.Where("id", data.Id).Delete(&model.Redeem{}).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
}
resp.SUCCESS(c)
}

81
api/handler/admin/reward_handler.go
View File

@@ -1,81 +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/core/types"
"geekai/handler"
"geekai/store/model"
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
type RewardHandler struct {
handler.BaseHandler
}
func NewRewardHandler(app *core.AppServer, db *gorm.DB) *RewardHandler {
return &RewardHandler{BaseHandler: handler.BaseHandler{App: app, DB: db}}
}
func (h *RewardHandler) List(c *gin.Context) {
var items []model.Reward
res := h.DB.Order("id DESC").Find(&items)
var rewards = make([]vo.Reward, 0)
if res.Error == nil {
userIds := make([]uint, 0)
for _, v := range items {
userIds = append(userIds, v.UserId)
}
var users []model.User
h.DB.Where("id IN ?", userIds).Find(&users)
var userMap = make(map[uint]model.User)
for _, u := range users {
userMap[u.Id] = u
}
for _, v := range items {
var r vo.Reward
err := utils.CopyObject(v, &r)
if err != nil {
continue
}
r.Id = v.Id
r.Username = userMap[v.UserId].Username
r.CreatedAt = v.CreatedAt.Unix()
r.UpdatedAt = v.UpdatedAt.Unix()
rewards = append(rewards, r)
}
}
resp.SUCCESS(c, rewards)
}
func (h *RewardHandler) Remove(c *gin.Context) {
var data struct {
Id uint
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
if data.Id > 0 {
res := h.DB.Where("id = ?", data.Id).Delete(&model.Reward{})
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
return
}
}
resp.SUCCESS(c)
}

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

@@ -168,8 +168,7 @@ func (h *UserHandler) Save(c *gin.Context) {
}
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")
resp.ERROR(c, res.Error.Error())
return
}
@@ -226,7 +225,7 @@ func (h *UserHandler) Remove(c *gin.Context) {
// 删除算力日志
h.DB.Where("user_id = ?", id).Delete(&model.PowerLog{})
// 删除众筹日志
h.DB.Where("user_id = ?", id).Delete(&model.Reward{})
h.DB.Where("user_id = ?", id).Delete(&model.Redeem{})
// 删除绘图任务
h.DB.Where("user_id = ?", id).Delete(&model.MidJourneyJob{})
h.DB.Where("user_id = ?", id).Delete(&model.SdJob{})

7
api/handler/chat_role_handler.go
View File

@@ -81,10 +81,9 @@ func (h *ChatRoleHandler) UpdateRole(c *gin.Context) {
return
}
res := h.DB.Model(&model.User{}).Where("id = ?", user.Id).UpdateColumn("chat_roles_json", utils.JsonEncode(data.Keys))
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败！")
err = h.DB.Model(&model.User{}).Where("id = ?", user.Id).UpdateColumn("chat_roles_json", utils.JsonEncode(data.Keys)).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}

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

@@ -468,7 +468,7 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, sessi
} else {
client = http.DefaultClient
}
logger.Debugf("Sending %s request, Channel:%s, API KEY:%s, PROXY: %s, Model: %s", session.Model.Platform, apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
logger.Debugf("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())

57
api/handler/dalle_handler.go
View File

@@ -8,6 +8,7 @@ package handler
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service/dalle"
@@ -16,9 +17,9 @@ import (
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"net/http"
"github.com/gorilla/websocket"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
@@ -178,7 +179,7 @@ func (h *DallJobHandler) getData(finish bool, userId uint, page int, pageSize in
session := h.DB.Session(&gorm.Session{})
if finish {
session = session.Where("progress = ?", 100).Order("id DESC")
session = session.Where("progress >= ?", 100).Order("id DESC")
} else {
session = session.Where("progress < ?", 100).Order("id ASC")
}
@@ -215,20 +216,51 @@ func (h *DallJobHandler) getData(finish bool, userId uint, page int, pageSize in
// Remove remove task image
func (h *DallJobHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
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: job.Id})
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
// 删除任务
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 如果任务未完成，或者任务失败，则恢复用户算力
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
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
@@ -241,13 +273,12 @@ func (h *DallJobHandler) Remove(c *gin.Context) {
// Publish 发布/取消发布图片到画廊显示
func (h *DallJobHandler) Publish(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
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, "更新数据库失败")
err := h.DB.Model(&model.DallJob{Id: uint(id), UserId: userId}).UpdateColumn("publish", action).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}

9
api/handler/markmap_handler.go
View File

@@ -101,17 +101,13 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
return fmt.Errorf("error with query chat model: %v", res.Error)
}
if user.Status == false {
return errors.New("当前用户被禁用")
}
if user.Power < chatModel.Power {
return fmt.Errorf("您当前剩余算力（%d）已不足以支付当前模型算力（%d）！", user.Power, chatModel.Power)
}
messages := make([]interface{}, 0)
messages = append(messages, types.Message{Role: "system", Content: `
你是一位非常优秀的思维导图助手，你会把用户的所有提问都总结成思维导图，然后以 Markdown 格式输出。markdown 只需要输出一级标题，二级标题，三级标题，四级标题，最多输出四级，除此之外不要输出任何其他 markdown 标记。下面是一个合格的例子：
你是一位非常优秀的思维导图助手， 你能帮助用户整理思路，根据用户提供的主题或内容，快速生成结构清晰，有条理的思维导图，然后以 Markdown 格式输出。markdown 只需要输出一级标题，二级标题，三级标题，四级标题，最多输出四级，除此之外不要输出任何其他 markdown 标记。下面是一个合格的例子：
# Geek-AI 助手
## 完整的开源系统
@@ -130,7 +126,7 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
另外，除此之外不要任何解释性语句。
`})
messages = append(messages, types.Message{Role: "user", Content: prompt})
messages = append(messages, types.Message{Role: "user", Content: fmt.Sprintf("请生成一份有关【%s】一份思维导图，要求结构清晰，有条理", prompt)})
var req = types.ApiRequest{
Model: chatModel.Value,
Stream: true,
@@ -253,5 +249,6 @@ func (h *MarkMapHandler) doRequest(req types.ApiRequest, chatModel model.ChatMod
client = http.DefaultClient
}
request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey.Value))
logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
return client.Do(request)
}

89
api/handler/mj_handler.go
View File

@@ -30,15 +30,15 @@ import (
type MidJourneyHandler struct {
BaseHandler
pool *mj.ServicePool
service *mj.Service
snowflake *service.Snowflake
uploader *oss.UploaderManager
}
func NewMidJourneyHandler(app *core.AppServer, db *gorm.DB, snowflake *service.Snowflake, pool *mj.ServicePool, manager *oss.UploaderManager) *MidJourneyHandler {
func NewMidJourneyHandler(app *core.AppServer, db *gorm.DB, snowflake *service.Snowflake, service *mj.Service, manager *oss.UploaderManager) *MidJourneyHandler {
return &MidJourneyHandler{
snowflake: snowflake,
pool: pool,
service: service,
uploader: manager,
BaseHandler: BaseHandler{
App: app,
@@ -59,11 +59,6 @@ func (h *MidJourneyHandler) preCheck(c *gin.Context) bool {
return false
}
if !h.pool.HasAvailableService() {
resp.ERROR(c, "MidJourney 池子中没有没有可用的服务！")
return false
}
return true
}
@@ -85,7 +80,7 @@ func (h *MidJourneyHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.pool.Clients.Put(uint(userId), client)
h.service.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@@ -201,7 +196,7 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
return
}
h.pool.PushTask(types.MjTask{
h.service.PushTask(types.MjTask{
Id: job.Id,
TaskId: taskId,
Type: types.TaskType(data.TaskType),
@@ -210,9 +205,10 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
Params: params,
UserId: userId,
ImgArr: data.ImgArr,
Mode: h.App.SysConfig.MjMode,
})
client := h.pool.Clients.Get(uint(job.UserId))
client := h.service.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -273,7 +269,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
return
}
h.pool.PushTask(types.MjTask{
h.service.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskUpscale,
UserId: userId,
@@ -281,9 +277,10 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
Index: data.Index,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
})
client := h.pool.Clients.Get(uint(job.UserId))
client := h.service.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -337,7 +334,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
return
}
h.pool.PushTask(types.MjTask{
h.service.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskVariation,
UserId: userId,
@@ -345,9 +342,10 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
ChannelId: data.ChannelId,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
})
client := h.pool.Clients.Get(uint(job.UserId))
client := h.service.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -465,40 +463,42 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
return
}
// refund power
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
h.DB.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power + job.Power,
Mark: types.PowerAdd,
Model: "mid-journey",
Remark: fmt.Sprintf("绘画任务失败，退回算力。任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
// 如果任务未完成，或者任务失败，则恢复用户算力
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
if err != nil {
tx.Rollback()
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)
}
client := h.pool.Clients.Get(uint(job.UserId))
client := h.service.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -511,10 +511,9 @@ func (h *MidJourneyHandler) Publish(c *gin.Context) {
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, "更新数据库失败")
err := h.DB.Model(&model.MidJourneyJob{Id: uint(id), UserId: userId}).UpdateColumn("publish", action).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}

102
api/handler/redeem_handler.go Normal file
View File

@@ -0,0 +1,102 @@
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/store/model"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"sync"
"time"
)
type RedeemHandler struct {
BaseHandler
lock sync.Mutex
}
func NewRedeemHandler(app *core.AppServer, db *gorm.DB) *RedeemHandler {
return &RedeemHandler{BaseHandler: BaseHandler{App: app, DB: db}}
}
func (h *RedeemHandler) Verify(c *gin.Context) {
var data struct {
Code string `json:"code"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
userId := h.GetLoginUserId(c)
h.lock.Lock()
defer h.lock.Unlock()
var item model.Redeem
res := h.DB.Where("code", data.Code).First(&item)
if res.Error != nil {
resp.ERROR(c, "无效的兑换码！")
return
}
if !item.Enabled {
resp.ERROR(c, "当前兑换码已被禁用！")
return
}
if item.RedeemedAt > 0 {
resp.ERROR(c, "当前兑换码已使用，请勿重复使用！")
return
}
tx := h.DB.Begin()
err := tx.Model(&model.User{}).Where("id", userId).UpdateColumn("power", gorm.Expr("power + ?", item.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 更新核销状态
item.RedeemedAt = time.Now().Unix()
item.UserId = userId
err = tx.Updates(&item).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
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)
}

108
api/handler/reward_handler.go
View File

@@ -1,108 +0,0 @@
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/store/model"
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"math"
"strings"
"sync"
"time"
)
type RewardHandler struct {
BaseHandler
lock sync.Mutex
}
func NewRewardHandler(app *core.AppServer, db *gorm.DB) *RewardHandler {
return &RewardHandler{BaseHandler: BaseHandler{App: app, DB: db}}
}
// Verify 打赏码核销
func (h *RewardHandler) Verify(c *gin.Context) {
var data struct {
TxId string `json:"tx_id"`
}
if err := c.ShouldBindJSON(&data); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
user, err := h.GetLoginUser(c)
if err != nil {
resp.HACKER(c)
return
}
// 移除转账单号中间的空格，防止有人复制的时候多复制了空格
data.TxId = strings.ReplaceAll(data.TxId, " ", "")
h.lock.Lock()
defer h.lock.Unlock()
var item model.Reward
res := h.DB.Where("tx_id = ?", data.TxId).First(&item)
if res.Error != nil {
resp.ERROR(c, "无效的交易流水号！")
return
}
if item.Status {
resp.ERROR(c, "当前交易流水号已经被核销，请不要重复核销！")
return
}
tx := h.DB.Begin()
exchange := vo.RewardExchange{}
power := math.Ceil(item.Amount / h.App.SysConfig.PowerPrice)
exchange.Power = int(power)
res = tx.Model(&user).UpdateColumn("power", gorm.Expr("power + ?", exchange.Power))
if res.Error != nil {
tx.Rollback()
logger.Error("添加应用失败：", res.Error)
resp.ERROR(c, "更新数据库失败！")
return
}
// 更新核销状态
item.Status = true
item.UserId = user.Id
item.Exchange = utils.JsonEncode(exchange)
res = tx.Updates(&item)
if res.Error != nil {
tx.Rollback()
logger.Error("添加应用失败：", res.Error)
resp.ERROR(c, "更新数据库失败！")
return
}
// 记录算力充值日志
h.DB.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerReward,
Amount: exchange.Power,
Balance: user.Power + exchange.Power,
Mark: types.PowerAdd,
Model: "众筹支付",
Remark: fmt.Sprintf("充值算力，金额：%f，价格：%f", item.Amount, h.App.SysConfig.PowerPrice),
CreatedAt: time.Now(),
})
tx.Commit()
resp.SUCCESS(c)
}

73
api/handler/sd_handler.go
View File

@@ -32,15 +32,15 @@ import (
type SdJobHandler struct {
BaseHandler
redis *redis.Client
pool *sd.ServicePool
service *sd.Service
uploader *oss.UploaderManager
snowflake *service.Snowflake
leveldb *store.LevelDB
}
func NewSdJobHandler(app *core.AppServer, db *gorm.DB, pool *sd.ServicePool, 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, levelDB *store.LevelDB) *SdJobHandler {
return &SdJobHandler{
pool: pool,
service: service,
uploader: manager,
snowflake: snowflake,
leveldb: levelDB,
@@ -68,7 +68,7 @@ func (h *SdJobHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.pool.Clients.Put(uint(userId), client)
h.service.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@@ -79,11 +79,6 @@ func (h *SdJobHandler) preCheck(c *gin.Context) bool {
return false
}
if !h.pool.HasAvailableService() {
resp.ERROR(c, "Stable-Diffusion 池子中没有没有可用的服务！")
return false
}
if user.Power < h.App.SysConfig.SdPower {
resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
return false
@@ -164,14 +159,14 @@ func (h *SdJobHandler) Image(c *gin.Context) {
return
}
h.pool.PushTask(types.SdTask{
h.service.PushTask(types.SdTask{
Id: int(job.Id),
Type: types.TaskImage,
Params: params,
UserId: userId,
})
client := h.pool.Clients.Get(uint(job.UserId))
client := h.service.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -232,7 +227,7 @@ func (h *SdJobHandler) getData(finish bool, userId uint, page int, pageSize int,
session := h.DB.Session(&gorm.Session{})
if finish {
session = session.Where("progress = ?", 100).Order("id DESC")
session = session.Where("progress >= ?", 100).Order("id DESC")
} else {
session = session.Where("progress < ?", 100).Order("id ASC")
}
@@ -278,44 +273,68 @@ func (h *SdJobHandler) getData(finish bool, userId uint, page int, pageSize int,
// Remove remove task image
func (h *SdJobHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
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: job.Id})
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
// 删除任务
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 如果任务未完成，或者任务失败，则恢复用户算力
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
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}
client := h.pool.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte(sd.Finished))
}
resp.SUCCESS(c)
}
// Publish 发布/取消发布图片到画廊显示
func (h *SdJobHandler) Publish(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
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, "更新数据库失败")
err := h.DB.Model(&model.SdJob{Id: uint(id), UserId: int(userId)}).UpdateColumn("publish", action).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}

38
api/handler/suno_handler.go
View File

@@ -94,6 +94,7 @@ func (h *SunoHandler) Create(c *gin.Context) {
RefTaskId: data.RefTaskId,
ExtendSecs: data.ExtendSecs,
Power: h.App.SysConfig.SunoPower,
SongId: utils.RandString(32),
}
if data.Lyrics != "" {
job.Prompt = data.Lyrics
@@ -210,7 +211,42 @@ func (h *SunoHandler) Remove(c *gin.Context) {
return
}
// 删除任务
h.DB.Delete(&job)
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 如果任务未完成，或者任务失败，则恢复用户算力
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
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// 删除文件
_ = h.uploader.GetUploadHandler().Delete(job.CoverURL)
_ = h.uploader.GetUploadHandler().Delete(job.AudioURL)

37
api/handler/test_handler.go
View File

@@ -3,7 +3,9 @@ package handler
import (
"geekai/service"
"geekai/service/payment"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"net/http"
)
type TestHandler struct {
@@ -15,3 +17,38 @@ type TestHandler struct {
func NewTestHandler(db *gorm.DB, snowflake *service.Snowflake, js *payment.JPayService) *TestHandler {
return &TestHandler{db: db, snowflake: snowflake, js: js}
}
func (h *TestHandler) SseTest(c *gin.Context) {
//c.Header("Content-Type", "text/event-stream")
//c.Header("Cache-Control", "no-cache")
//c.Header("Connection", "keep-alive")
//
//
//// 模拟实时数据更新
//for i := 0; i < 10; i++ {
// // 发送 SSE 数据
// _, err := fmt.Fprintf(c.Writer, "data: %v\n\n", data)
// if err != nil {
// return
// }
// c.Writer.Flush() // 确保立即发送数据
// time.Sleep(1 * time.Second) // 每秒发送一次数据
//}
//c.Abort()
}
func (h *TestHandler) PostTest(c *gin.Context) {
var data struct {
Message string `json:"message"`
UserId uint `json:"user_id"`
}
if err := c.ShouldBindJSON(&data); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
// 将参数存储在上下文中
c.Set("data", data)
c.Next()
}

14
api/handler/user_handler.go
View File

@@ -481,10 +481,9 @@ func (h *UserHandler) UpdatePass(c *gin.Context) {
}
newPass := utils.GenPassword(data.Password, user.Salt)
res := h.DB.Model(&user).UpdateColumn("password", newPass)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")
err = h.DB.Model(&user).UpdateColumn("password", newPass).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}
@@ -562,10 +561,9 @@ func (h *UserHandler) BindUsername(c *gin.Context) {
return
}
res = h.DB.Model(&user).UpdateColumn("username", data.Username)
if res.Error != nil {
logger.Error(res.Error)
resp.ERROR(c, "更新数据库失败")
err = h.DB.Model(&user).UpdateColumn("username", data.Username).Error
if err != nil {
resp.ERROR(c, err.Error())
return
}

72
api/main.go
View File

@@ -24,7 +24,6 @@ import (
"geekai/service/sd"
"geekai/service/sms"
"geekai/service/suno"
"geekai/service/wx"
"geekai/store"
"io"
"log"
@@ -131,7 +130,7 @@ func main() {
fx.Provide(chatimpl.NewChatHandler),
fx.Provide(handler.NewUploadHandler),
fx.Provide(handler.NewSmsHandler),
fx.Provide(handler.NewRewardHandler),
fx.Provide(handler.NewRedeemHandler),
fx.Provide(handler.NewCaptchaHandler),
fx.Provide(handler.NewMidJourneyHandler),
fx.Provide(handler.NewChatModelHandler),
@@ -147,7 +146,7 @@ func main() {
fx.Provide(admin.NewApiKeyHandler),
fx.Provide(admin.NewUserHandler),
fx.Provide(admin.NewChatRoleHandler),
fx.Provide(admin.NewRewardHandler),
fx.Provide(admin.NewRedeemHandler),
fx.Provide(admin.NewDashboardHandler),
fx.Provide(admin.NewChatModelHandler),
fx.Provide(admin.NewProductHandler),
@@ -161,13 +160,12 @@ func main() {
return service.NewCaptchaService(config.ApiConfig)
}),
fx.Provide(oss.NewUploaderManager),
fx.Provide(mj.NewService),
fx.Provide(dalle.NewService),
fx.Invoke(func(service *dalle.Service) {
service.Run()
service.CheckTaskNotify()
service.DownloadImages()
service.CheckTaskStatus()
fx.Invoke(func(s *dalle.Service) {
s.Run()
s.CheckTaskNotify()
s.DownloadImages()
s.CheckTaskStatus()
}),
// 邮件服务
@@ -178,36 +176,22 @@ func main() {
licenseService.SyncLicense()
}),
// 微信机器人服务
fx.Provide(wx.NewWeChatBot),
fx.Invoke(func(config *types.AppConfig, bot *wx.Bot) {
if config.WeChatBot {
err := bot.Run()
if err != nil {
logger.Error("微信登录失败：", err)
}
}
}),
// MidJourney service pool
fx.Provide(mj.NewServicePool),
fx.Invoke(func(pool *mj.ServicePool, config *types.AppConfig) {
pool.InitServices(config.MjPlusConfigs, config.MjProxyConfigs)
if pool.HasAvailableService() {
pool.DownloadImages()
pool.CheckTaskNotify()
pool.SyncTaskProgress()
}
fx.Provide(mj.NewService),
fx.Provide(mj.NewClient),
fx.Invoke(func(s *mj.Service) {
s.Run()
s.SyncTaskProgress()
s.CheckTaskNotify()
s.DownloadImages()
}),
// Stable Diffusion 机器人
fx.Provide(sd.NewServicePool),
fx.Invoke(func(pool *sd.ServicePool, config *types.AppConfig) {
pool.InitServices(config.SdConfigs)
if pool.HasAvailableService() {
pool.CheckTaskNotify()
pool.CheckTaskStatus()
}
fx.Provide(sd.NewService),
fx.Invoke(func(s *sd.Service, config *types.AppConfig) {
s.Run()
s.CheckTaskStatus()
s.CheckTaskNotify()
}),
fx.Provide(suno.NewService),
@@ -280,8 +264,8 @@ func main() {
group.GET("slide/get", h.SlideGet)
group.POST("slide/check", h.SlideCheck)
}),
fx.Invoke(func(s *core.AppServer, h *handler.RewardHandler) {
group := s.Engine.Group("/api/reward/")
fx.Invoke(func(s *core.AppServer, h *handler.RedeemHandler) {
group := s.Engine.Group("/api/redeem/")
group.POST("verify", h.Verify)
}),
fx.Invoke(func(s *core.AppServer, h *handler.MidJourneyHandler) {
@@ -317,8 +301,6 @@ func main() {
group.GET("config/get", h.Get)
group.POST("active", h.Active)
group.GET("config/get/license", h.GetLicense)
group.GET("config/get/app", h.GetAppConfig)
group.POST("config/update/draw", h.SaveDrawingConfig)
}),
fx.Invoke(func(s *core.AppServer, h *admin.ManagerHandler) {
group := s.Engine.Group("/api/admin/")
@@ -354,9 +336,11 @@ func main() {
group.POST("set", h.Set)
group.GET("remove", h.Remove)
}),
fx.Invoke(func(s *core.AppServer, h *admin.RewardHandler) {
group := s.Engine.Group("/api/admin/reward/")
fx.Invoke(func(s *core.AppServer, h *admin.RedeemHandler) {
group := s.Engine.Group("/api/admin/redeem/")
group.GET("list", h.List)
group.POST("create", h.Create)
group.POST("set", h.Set)
group.POST("remove", h.Remove)
}),
fx.Invoke(func(s *core.AppServer, h *admin.DashboardHandler) {
@@ -398,6 +382,7 @@ func main() {
group := s.Engine.Group("/api/admin/order/")
group.POST("list", h.List)
group.GET("remove", h.Remove)
group.GET("clear", h.Clear)
}),
fx.Invoke(func(s *core.AppServer, h *handler.OrderHandler) {
group := s.Engine.Group("/api/order/")
@@ -497,6 +482,11 @@ func main() {
group.GET("play", h.Play)
group.POST("lyric", h.Lyric)
}),
fx.Provide(handler.NewTestHandler),
fx.Invoke(func(s *core.AppServer, h *handler.TestHandler) {
group := s.Engine.Group("/api/test")
group.Any("sse", h.PostTest, h.SseTest)
}),
fx.Invoke(func(s *core.AppServer, db *gorm.DB) {
go func() {
err := s.Run(db)

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

@@ -14,7 +14,6 @@ import (
logger2 "geekai/logger"
"geekai/service"
"geekai/service/oss"
"geekai/service/sd"
"geekai/store"
"geekai/store/model"
"geekai/utils"
@@ -70,10 +69,10 @@ func (s *Service) Run() {
if err != nil {
logger.Errorf("error with image task: %v", err)
s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
"progress": -1,
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Failed})
s.notifyQueue.RPush(service.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: service.TaskStatusFailed})
}
}
}()
@@ -148,7 +147,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
Where("enabled", true).
Order("last_used_at ASC").First(&apiKey)
if tx.Error != nil {
return "", fmt.Errorf("no available IMG api key: %v", tx.Error)
return "", fmt.Errorf("no available DALL-E api key: %v", tx.Error)
}
var res imgRes
@@ -191,7 +190,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
return "", fmt.Errorf("err with update database: %v", tx.Error)
}
s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Finished})
s.notifyQueue.RPush(service.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: service.TaskStatusFailed})
var content string
if sync {
imgURL, err := s.downloadImage(task.JobId, int(task.UserId), res.Data[0].Url)
@@ -208,7 +207,7 @@ func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running DALL-E task notify checking ...")
for {
var message sd.NotifyMessage
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
@@ -225,6 +224,30 @@ func (s *Service) CheckTaskNotify() {
}()
}
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running DALL-E task status checking ...")
for {
var jobs []model.DallJob
res := s.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
for _, job := range jobs {
// 超时的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
}
}
time.Sleep(time.Second * 10)
}
}()
}
func (s *Service) DownloadImages() {
go func() {
var items []model.DallJob
@@ -268,47 +291,6 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
if res.Error != nil {
return "", err
}
s.notifyQueue.RPush(sd.NotifyMessage{UserId: userId, JobId: int(jobId), Message: sd.Finished})
s.notifyQueue.RPush(service.NotifyMessage{UserId: userId, JobId: int(jobId), Message: service.TaskStatusFinished})
return imgURL, nil
}
// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running Stable-Diffusion task status checking ...")
for {
var jobs []model.DallJob
res := s.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
for _, job := range jobs {
// 5 分钟还没完成的任务直接删除
if time.Now().Sub(job.CreatedAt) > time.Minute*5 || job.Progress == -1 {
s.db.Delete(&job)
var user model.User
s.db.Where("id = ?", job.UserId).First(&user)
// 退回绘图次数
res = s.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
if res.Error == nil && res.RowsAffected > 0 {
s.db.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power + job.Power,
Mark: types.PowerAdd,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d", job.Id),
CreatedAt: time.Now(),
})
}
continue
}
}
time.Sleep(time.Second * 10)
}
}()
}

212
api/service/mj/client.go
View File

@@ -7,15 +7,28 @@ package mj
// * @Author yangjian102621@163.com
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import "geekai/core/types"
import (
"encoding/base64"
"errors"
"fmt"
"geekai/core/types"
logger2 "geekai/logger"
"geekai/service"
"geekai/store/model"
"geekai/utils"
"github.com/imroc/req/v3"
"gorm.io/gorm"
"io"
"time"
type Client interface {
Imagine(task types.MjTask) (ImageRes, error)
Blend(task types.MjTask) (ImageRes, error)
SwapFace(task types.MjTask) (ImageRes, error)
Upscale(task types.MjTask) (ImageRes, error)
Variation(task types.MjTask) (ImageRes, error)
QueryTask(taskId string) (QueryRes, error)
"github.com/gin-gonic/gin"
)
// Client MidJourney client
type Client struct {
client *req.Client
licenseService *service.LicenseService
db *gorm.DB
}
type ImageReq struct {
@@ -33,13 +46,8 @@ type ImageRes struct {
Description string `json:"description"`
Properties struct {
} `json:"properties"`
Result string `json:"result"`
}
type ErrRes struct {
Error struct {
Message string `json:"message"`
} `json:"error"`
Result string `json:"result"`
Channel string `json:"channel,omitempty"`
}
type QueryRes struct {
@@ -66,3 +74,177 @@ type QueryRes struct {
Status string `json:"status"`
SubmitTime int `json:"submitTime"`
}
var logger = logger2.GetLogger()
func NewClient(licenseService *service.LicenseService, db *gorm.DB) *Client {
return &Client{
client: req.C().SetTimeout(time.Minute).SetUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"),
licenseService: licenseService,
db: db,
}
}
func (c *Client) Imagine(task types.MjTask) (ImageRes, error) {
apiPath := fmt.Sprintf("mj-%s/mj/submit/imagine", task.Mode)
prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
if task.NegPrompt != "" {
prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
}
body := ImageReq{
BotType: "MID_JOURNEY",
Prompt: prompt,
Base64Array: make([]string, 0),
}
// 生成图片 Base64 编码
if len(task.ImgArr) > 0 {
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
}
}
return c.doRequest(body, apiPath, task.ChannelId)
}
// Blend 融图
func (c *Client) Blend(task types.MjTask) (ImageRes, error) {
apiPath := fmt.Sprintf("mj-%s/mj/submit/blend", task.Mode)
body := ImageReq{
BotType: "MID_JOURNEY",
Dimensions: "SQUARE",
Base64Array: make([]string, 0),
}
// 生成图片 Base64 编码
if len(task.ImgArr) > 0 {
for _, imgURL := range task.ImgArr {
imageData, err := utils.DownloadImage(imgURL, "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
}
}
}
return c.doRequest(body, apiPath, task.ChannelId)
}
// SwapFace 换脸
func (c *Client) SwapFace(task types.MjTask) (ImageRes, error) {
apiPath := fmt.Sprintf("mj-%s/mj/insight-face/swap", task.Mode)
// 生成图片 Base64 编码
if len(task.ImgArr) != 2 {
return ImageRes{}, errors.New("参数错误，必须上传2张图片")
}
var sourceBase64 string
var targetBase64 string
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
sourceBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
}
imageData, err = utils.DownloadImage(task.ImgArr[1], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
targetBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
}
body := gin.H{
"sourceBase64": sourceBase64,
"targetBase64": targetBase64,
"accountFilter": gin.H{
"instanceId": "",
},
"state": "",
}
return c.doRequest(body, apiPath, task.ChannelId)
}
// Upscale 放大指定的图片
func (c *Client) Upscale(task types.MjTask) (ImageRes, error) {
body := map[string]string{
"customId": fmt.Sprintf("MJ::JOB::upsample::%d::%s", task.Index, task.MessageHash),
"taskId": task.MessageId,
}
apiPath := fmt.Sprintf("mj-%s/mj/submit/action", task.Mode)
return c.doRequest(body, apiPath, task.ChannelId)
}
// Variation 以指定的图片的视角进行变换再创作，注意需要在对应的频道中关闭 Remix 变换，否则 Variation 指令将不会生效
func (c *Client) Variation(task types.MjTask) (ImageRes, error) {
body := map[string]string{
"customId": fmt.Sprintf("MJ::JOB::variation::%d::%s", task.Index, task.MessageHash),
"taskId": task.MessageId,
}
apiPath := fmt.Sprintf("mj-%s/mj/submit/action", task.Mode)
return c.doRequest(body, apiPath, task.ChannelId)
}
func (c *Client) doRequest(body interface{}, apiPath string, channel string) (ImageRes, error) {
var res ImageRes
session := c.db.Session(&gorm.Session{}).Where("type", "mj").Where("enabled", true)
if channel != "" {
session = session.Where("api_url", channel)
}
var apiKey model.ApiKey
err := session.Order("last_used_at ASC").First(&apiKey).Error
if err != nil {
return ImageRes{}, fmt.Errorf("no available MidJourney api key: %v", err)
}
if err = c.licenseService.IsValidApiURL(apiKey.ApiURL); err != nil {
return ImageRes{}, err
}
apiURL := fmt.Sprintf("%s/%s", apiKey.ApiURL, apiPath)
logger.Info("API URL: ", apiURL)
r, err := req.C().R().
SetHeader("Authorization", "Bearer "+apiKey.Value).
SetBody(body).
SetSuccessResult(&res).
Post(apiURL)
if err != nil {
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", err)
}
if r.IsErrorState() {
errMsg, _ := io.ReadAll(r.Body)
return ImageRes{}, fmt.Errorf("API 返回错误：%s", string(errMsg))
}
// update the api key last used time
if err = c.db.Model(&apiKey).Update("last_used_at", time.Now().Unix()).Error; err != nil {
logger.Error("update api key last used time error: ", err)
}
res.Channel = apiKey.ApiURL
return res, nil
}
func (c *Client) QueryTask(taskId string, channel string) (QueryRes, error) {
var apiKey model.ApiKey
err := c.db.Where("type", "mj").Where("enabled", true).Where("api_url", channel).First(&apiKey).Error
if err != nil {
return QueryRes{}, fmt.Errorf("no available MidJourney api key: %v", err)
}
apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", apiKey.ApiURL, taskId)
var res QueryRes
r, err := c.client.R().SetHeader("Authorization", "Bearer "+apiKey.Value).
SetSuccessResult(&res).
Get(apiURL)
if err != nil {
return QueryRes{}, err
}
if r.IsErrorState() {
return QueryRes{}, errors.New("error status:" + r.Status)
}
return res, nil
}

204
api/service/mj/plus_client.go
View File

@@ -1,204 +0,0 @@
package mj
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 (
"encoding/base64"
"errors"
"fmt"
"geekai/core/types"
"geekai/service"
"geekai/utils"
"github.com/imroc/req/v3"
"time"
"github.com/gin-gonic/gin"
)
// PlusClient MidJourney Plus ProxyClient
type PlusClient struct {
Config types.MjPlusConfig
apiURL string
client *req.Client
licenseService *service.LicenseService
}
func NewPlusClient(config types.MjPlusConfig, licenseService *service.LicenseService) *PlusClient {
return &PlusClient{
Config: config,
apiURL: config.ApiURL,
client: req.C().SetTimeout(time.Minute).SetUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"),
licenseService: licenseService,
}
}
func (c *PlusClient) preCheck() error {
return c.licenseService.IsValidApiURL(c.Config.ApiURL)
}
func (c *PlusClient) Imagine(task types.MjTask) (ImageRes, error) {
if err := c.preCheck(); err != nil {
return ImageRes{}, err
}
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/imagine", c.apiURL, c.Config.Mode)
prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
if task.NegPrompt != "" {
prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
}
body := ImageReq{
BotType: "MID_JOURNEY",
Prompt: prompt,
Base64Array: make([]string, 0),
}
// 生成图片 Base64 编码
if len(task.ImgArr) > 0 {
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
}
}
return c.doRequest(body, apiURL)
}
// Blend 融图
func (c *PlusClient) Blend(task types.MjTask) (ImageRes, error) {
if err := c.preCheck(); err != nil {
return ImageRes{}, err
}
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/blend", c.apiURL, c.Config.Mode)
logger.Info("API URL: ", apiURL)
body := ImageReq{
BotType: "MID_JOURNEY",
Dimensions: "SQUARE",
Base64Array: make([]string, 0),
}
// 生成图片 Base64 编码
if len(task.ImgArr) > 0 {
for _, imgURL := range task.ImgArr {
imageData, err := utils.DownloadImage(imgURL, "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
}
}
}
return c.doRequest(body, apiURL)
}
// SwapFace 换脸
func (c *PlusClient) SwapFace(task types.MjTask) (ImageRes, error) {
if err := c.preCheck(); err != nil {
return ImageRes{}, err
}
apiURL := fmt.Sprintf("%s/mj-%s/mj/insight-face/swap", c.apiURL, c.Config.Mode)
// 生成图片 Base64 编码
if len(task.ImgArr) != 2 {
return ImageRes{}, errors.New("参数错误，必须上传2张图片")
}
var sourceBase64 string
var targetBase64 string
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
sourceBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
}
imageData, err = utils.DownloadImage(task.ImgArr[1], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
targetBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
}
body := gin.H{
"sourceBase64": sourceBase64,
"targetBase64": targetBase64,
"accountFilter": gin.H{
"instanceId": "",
},
"state": "",
}
return c.doRequest(body, apiURL)
}
// Upscale 放大指定的图片
func (c *PlusClient) Upscale(task types.MjTask) (ImageRes, error) {
if err := c.preCheck(); err != nil {
return ImageRes{}, err
}
body := map[string]string{
"customId": fmt.Sprintf("MJ::JOB::upsample::%d::%s", task.Index, task.MessageHash),
"taskId": task.MessageId,
}
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/action", c.apiURL, c.Config.Mode)
return c.doRequest(body, apiURL)
}
// Variation 以指定的图片的视角进行变换再创作，注意需要在对应的频道中关闭 Remix 变换，否则 Variation 指令将不会生效
func (c *PlusClient) Variation(task types.MjTask) (ImageRes, error) {
if err := c.preCheck(); err != nil {
return ImageRes{}, err
}
body := map[string]string{
"customId": fmt.Sprintf("MJ::JOB::variation::%d::%s", task.Index, task.MessageHash),
"taskId": task.MessageId,
}
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/action", c.apiURL, c.Config.Mode)
return c.doRequest(body, apiURL)
}
func (c *PlusClient) doRequest(body interface{}, apiURL string) (ImageRes, error) {
var res ImageRes
var errRes ErrRes
logger.Info("API URL: ", apiURL)
r, err := req.C().R().
SetHeader("Authorization", "Bearer "+c.Config.ApiKey).
SetBody(body).
SetSuccessResult(&res).
SetErrorResult(&errRes).
Post(apiURL)
if err != nil {
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", err)
}
if r.IsErrorState() {
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
}
return res, nil
}
func (c *PlusClient) QueryTask(taskId string) (QueryRes, error) {
apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", c.apiURL, taskId)
var res QueryRes
r, err := c.client.R().SetHeader("Authorization", "Bearer "+c.Config.ApiKey).
SetSuccessResult(&res).
Get(apiURL)
if err != nil {
return QueryRes{}, err
}
if r.IsErrorState() {
return QueryRes{}, errors.New("error status:" + r.Status)
}
return res, nil
}
var _ Client = &PlusClient{}

207
api/service/mj/pool.go
View File

@@ -1,207 +0,0 @@
package mj
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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/types"
logger2 "geekai/logger"
"geekai/service"
"geekai/service/oss"
"geekai/service/sd"
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"strings"
"time"
"gorm.io/gorm"
)
// ServicePool Mj service pool
type ServicePool struct {
services []*Service
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
uploaderManager *oss.UploaderManager
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
licenseService *service.LicenseService
}
var logger = logger2.GetLogger()
func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService) *ServicePool {
services := make([]*Service, 0)
taskQueue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
notifyQueue := store.NewRedisQueue("MidJourney_Notify_Queue", redisCli)
return &ServicePool{
taskQueue: taskQueue,
notifyQueue: notifyQueue,
services: services,
uploaderManager: manager,
db: db,
Clients: types.NewLMap[uint, *types.WsClient](),
licenseService: licenseService,
}
}
func (p *ServicePool) InitServices(plusConfigs []types.MjPlusConfig, proxyConfigs []types.MjProxyConfig) {
// stop old service
for _, s := range p.services {
s.Stop()
}
p.services = make([]*Service, 0)
for _, config := range plusConfigs {
if config.Enabled == false {
continue
}
cli := NewPlusClient(config, p.licenseService)
name := utils.Md5(config.ApiURL)
plusService := NewService(name, p.taskQueue, p.notifyQueue, p.db, cli)
go func() {
plusService.Run()
}()
p.services = append(p.services, plusService)
}
// for mid-journey proxy
for _, config := range proxyConfigs {
if config.Enabled == false {
continue
}
cli := NewProxyClient(config)
name := utils.Md5(config.ApiURL)
proxyService := NewService(name, p.taskQueue, p.notifyQueue, p.db, cli)
go func() {
proxyService.Run()
}()
p.services = append(p.services, proxyService)
}
}
func (p *ServicePool) CheckTaskNotify() {
go func() {
for {
var message sd.NotifyMessage
err := p.notifyQueue.LPop(&message)
if err != nil {
continue
}
cli := p.Clients.Get(uint(message.UserId))
if cli == nil {
continue
}
err = cli.Send([]byte(message.Message))
if err != nil {
continue
}
}
}()
}
func (p *ServicePool) DownloadImages() {
go func() {
var items []model.MidJourneyJob
for {
res := p.db.Where("img_url = ? AND progress = ?", "", 100).Find(&items)
if res.Error != nil {
continue
}
// download images
for _, v := range items {
if v.OrgURL == "" {
continue
}
logger.Infof("try to download image: %s", v.OrgURL)
mjService := p.getService(v.ChannelId)
if mjService == nil {
logger.Errorf("Invalid task: %+v", v)
continue
}
task, _ := mjService.Client.QueryTask(v.TaskId)
if len(task.Buttons) > 0 {
v.Hash = GetImageHash(task.Buttons[0].CustomId)
}
// 如果是返回的是 discord 图片地址，则使用代理下载
proxy := false
if strings.HasPrefix(v.OrgURL, "https://cdn.discordapp.com") {
proxy = true
}
imgURL, err := p.uploaderManager.GetUploadHandler().PutUrlFile(v.OrgURL, proxy)
if err != nil {
logger.Errorf("error with download image %s, %v", v.OrgURL, err)
continue
} else {
logger.Infof("download image %s successfully.", v.OrgURL)
}
v.ImgURL = imgURL
p.db.Updates(&v)
cli := p.Clients.Get(uint(v.UserId))
if cli == nil {
continue
}
err = cli.Send([]byte(sd.Finished))
if err != nil {
continue
}
}
time.Sleep(time.Second * 5)
}
}()
}
// PushTask push a new mj task in to task queue
func (p *ServicePool) PushTask(task types.MjTask) {
logger.Debugf("add a new MidJourney task to the task list: %+v", task)
p.taskQueue.RPush(task)
}
// HasAvailableService check if it has available mj service in pool
func (p *ServicePool) HasAvailableService() bool {
return len(p.services) > 0
}
// SyncTaskProgress 异步拉取任务
func (p *ServicePool) SyncTaskProgress() {
go func() {
var jobs []model.MidJourneyJob
for {
res := p.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
continue
}
for _, job := range jobs {
if servicePlus := p.getService(job.ChannelId); servicePlus != nil {
_ = servicePlus.Notify(job)
}
}
time.Sleep(time.Second * 10)
}
}()
}
func (p *ServicePool) getService(name string) *Service {
for _, s := range p.services {
if s.Name == name {
return s
}
}
return nil
}

185
api/service/mj/proxy_client.go
View File

@@ -1,185 +0,0 @@
package mj
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 (
"encoding/base64"
"errors"
"fmt"
"geekai/core/types"
"geekai/utils"
"github.com/imroc/req/v3"
"io"
)
// ProxyClient MidJourney Proxy Client
type ProxyClient struct {
Config types.MjProxyConfig
apiURL string
}
func NewProxyClient(config types.MjProxyConfig) *ProxyClient {
return &ProxyClient{Config: config, apiURL: config.ApiURL}
}
func (c *ProxyClient) Imagine(task types.MjTask) (ImageRes, error) {
apiURL := fmt.Sprintf("%s/mj/submit/imagine", c.apiURL)
prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
if task.NegPrompt != "" {
prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
}
body := ImageReq{
Prompt: prompt,
Base64Array: make([]string, 0),
}
// 生成图片 Base64 编码
if len(task.ImgArr) > 0 {
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
}
}
logger.Info("API URL: ", apiURL)
var res ImageRes
var errRes ErrRes
r, err := req.C().R().
SetHeader("mj-api-secret", c.Config.ApiKey).
SetBody(body).
SetSuccessResult(&res).
SetErrorResult(&errRes).
Post(apiURL)
if err != nil {
return ImageRes{}, fmt.Errorf("请求 API %s 出错：%v", apiURL, err)
}
if r.IsErrorState() {
errStr, _ := io.ReadAll(r.Body)
return ImageRes{}, fmt.Errorf("API 返回错误：%s，%v", errRes.Error.Message, string(errStr))
}
return res, nil
}
// Blend 融图
func (c *ProxyClient) Blend(task types.MjTask) (ImageRes, error) {
apiURL := fmt.Sprintf("%s/mj/submit/blend", c.apiURL)
body := ImageReq{
Dimensions: "SQUARE",
Base64Array: make([]string, 0),
}
// 生成图片 Base64 编码
if len(task.ImgArr) > 0 {
for _, imgURL := range task.ImgArr {
imageData, err := utils.DownloadImage(imgURL, "")
if err != nil {
logger.Error("error with download image: ", err)
} else {
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
}
}
}
var res ImageRes
var errRes ErrRes
r, err := req.C().R().
SetHeader("mj-api-secret", c.Config.ApiKey).
SetBody(body).
SetSuccessResult(&res).
SetErrorResult(&errRes).
Post(apiURL)
if err != nil {
return ImageRes{}, fmt.Errorf("请求 API %s 出错：%v", apiURL, err)
}
if r.IsErrorState() {
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
}
return res, nil
}
// SwapFace 换脸
func (c *ProxyClient) SwapFace(_ types.MjTask) (ImageRes, error) {
return ImageRes{}, errors.New("MidJourney-Proxy暂未实现该功能，请使用 MidJourney-Plus")
}
// Upscale 放大指定的图片
func (c *ProxyClient) Upscale(task types.MjTask) (ImageRes, error) {
body := map[string]interface{}{
"action": "UPSCALE",
"index": task.Index,
"taskId": task.MessageId,
}
apiURL := fmt.Sprintf("%s/mj/submit/change", c.apiURL)
var res ImageRes
var errRes ErrRes
r, err := req.C().R().
SetHeader("mj-api-secret", c.Config.ApiKey).
SetBody(body).
SetSuccessResult(&res).
SetErrorResult(&errRes).
Post(apiURL)
if err != nil {
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", err)
}
if r.IsErrorState() {
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
}
return res, nil
}
// Variation 以指定的图片的视角进行变换再创作，注意需要在对应的频道中关闭 Remix 变换，否则 Variation 指令将不会生效
func (c *ProxyClient) Variation(task types.MjTask) (ImageRes, error) {
body := map[string]interface{}{
"action": "VARIATION",
"index": task.Index,
"taskId": task.MessageId,
}
apiURL := fmt.Sprintf("%s/mj/submit/change", c.apiURL)
var res ImageRes
var errRes ErrRes
r, err := req.C().R().
SetHeader("mj-api-secret", c.Config.ApiKey).
SetBody(body).
SetSuccessResult(&res).
SetErrorResult(&errRes).
Post(apiURL)
if err != nil {
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", err)
}
if r.IsErrorState() {
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
}
return res, nil
}
func (c *ProxyClient) QueryTask(taskId string) (QueryRes, error) {
apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", c.apiURL, taskId)
var res QueryRes
r, err := req.C().R().SetHeader("mj-api-secret", c.Config.ApiKey).
SetSuccessResult(&res).
Get(apiURL)
if err != nil {
return QueryRes{}, err
}
if r.IsErrorState() {
return QueryRes{}, errors.New("error status:" + r.Status)
}
return res, nil
}
var _ Client = &ProxyClient{}

380
api/service/mj/service.go
View File

@@ -11,10 +11,11 @@ import (
"fmt"
"geekai/core/types"
"geekai/service"
"geekai/service/sd"
"geekai/service/oss"
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"strings"
"time"
@@ -23,127 +24,112 @@ import (
// Service MJ 绘画服务
type Service struct {
Name string // service Name
Client Client // MJ Client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
running bool
retryCount map[uint]int
client *Client // MJ Client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
uploaderManager *oss.UploaderManager
}
func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, cli Client) *Service {
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager) *Service {
return &Service{
Name: name,
db: db,
taskQueue: taskQueue,
notifyQueue: notifyQueue,
Client: cli,
running: true,
retryCount: make(map[uint]int),
db: db,
taskQueue: store.NewRedisQueue("MidJourney_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("MidJourney_Notify_Queue", redisCli),
client: client,
Clients: types.NewLMap[uint, *types.WsClient](),
uploaderManager: manager,
}
}
const failedProgress = 101
func (s *Service) Run() {
logger.Infof("Starting MidJourney job consumer for %s", s.Name)
for s.running {
var task types.MjTask
err := s.taskQueue.LPop(&task)
if err != nil {
logger.Errorf("taking task with error: %v", err)
continue
}
// 如果配置了多个中转平台的 API KEY
// U,V 操作必须和 Image 操作属于同一个平台，否则找不到关联任务，需重新放回任务列表
if task.ChannelId != "" && task.ChannelId != s.Name {
if s.retryCount[task.Id] > 5 {
s.db.Model(model.MidJourneyJob{Id: task.Id}).Delete(&model.MidJourneyJob{})
logger.Info("Starting MidJourney job consumer for service")
go func() {
for {
var task types.MjTask
err := s.taskQueue.LPop(&task)
if err != nil {
logger.Errorf("taking task with error: %v", err)
continue
}
logger.Debugf("handle other service task, name: %s, channel_id: %s, drop it.", s.Name, task.ChannelId)
s.taskQueue.RPush(task)
s.retryCount[task.Id]++
time.Sleep(time.Second)
continue
}
// translate prompt
if utils.HasChinese(task.Prompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Prompt), "gpt-4o-mini")
if err == nil {
task.Prompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
// translate prompt
if utils.HasChinese(task.Prompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), "gpt-4o-mini")
if err == nil {
task.Prompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
}
}
}
// translate negative prompt
if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.NegPrompt), "gpt-4o-mini")
if err == nil {
task.NegPrompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
}
}
var job model.MidJourneyJob
tx := s.db.Where("id = ?", task.Id).First(&job)
if tx.Error != nil {
logger.Error("任务不存在，任务ID：", task.TaskId)
continue
}
logger.Infof("%s handle a new MidJourney task: %+v", s.Name, task)
var res ImageRes
switch task.Type {
case types.TaskImage:
res, err = s.Client.Imagine(task)
break
case types.TaskUpscale:
res, err = s.Client.Upscale(task)
break
case types.TaskVariation:
res, err = s.Client.Variation(task)
break
case types.TaskBlend:
res, err = s.Client.Blend(task)
break
case types.TaskSwapFace:
res, err = s.Client.SwapFace(task)
break
}
if err != nil || (res.Code != 1 && res.Code != 22) {
var errMsg string
if err != nil {
errMsg = err.Error()
} else {
errMsg = fmt.Sprintf("%v,%s", err, res.Description)
// translate negative prompt
if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.NegPrompt), "gpt-4o-mini")
if err == nil {
task.NegPrompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
}
}
logger.Error("绘画任务执行失败：", errMsg)
job.Progress = failedProgress
job.ErrMsg = errMsg
// update the task progress
// use fast mode as default
if task.Mode == "" {
task.Mode = "fast"
}
var job model.MidJourneyJob
tx := s.db.Where("id = ?", task.Id).First(&job)
if tx.Error != nil {
logger.Error("任务不存在，任务ID：", task.TaskId)
continue
}
logger.Infof("handle a new MidJourney task: %+v", task)
var res ImageRes
switch task.Type {
case types.TaskImage:
res, err = s.client.Imagine(task)
break
case types.TaskUpscale:
res, err = s.client.Upscale(task)
break
case types.TaskVariation:
res, err = s.client.Variation(task)
break
case types.TaskBlend:
res, err = s.client.Blend(task)
break
case types.TaskSwapFace:
res, err = s.client.SwapFace(task)
break
}
if err != nil || (res.Code != 1 && res.Code != 22) {
var errMsg string
if err != nil {
errMsg = err.Error()
} else {
errMsg = fmt.Sprintf("%v,%s", err, res.Description)
}
logger.Error("绘画任务执行失败：", errMsg)
job.Progress = service.FailTaskProgress
job.ErrMsg = errMsg
// update the task progress
s.db.Updates(&job)
// 任务失败，通知前端
s.notifyQueue.RPush(service.NotifyMessage{UserId: task.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
continue
}
logger.Infof("任务提交成功：%+v", res)
// 更新任务 ID/频道
job.TaskId = res.Result
job.MessageId = res.Result
job.ChannelId = res.Channel
s.db.Updates(&job)
// 任务失败，通知前端
s.notifyQueue.RPush(sd.NotifyMessage{UserId: task.UserId, JobId: int(job.Id), Message: sd.Failed})
continue
}
logger.Infof("任务提交成功：%+v", res)
// 更新任务 ID/频道
job.TaskId = res.Result
job.MessageId = res.Result
job.ChannelId = s.Name
s.db.Updates(&job)
}
}
func (s *Service) Stop() {
s.running = false
}()
}
type CBReq struct {
@@ -164,46 +150,6 @@ type CBReq struct {
} `json:"properties"`
}
func (s *Service) Notify(job model.MidJourneyJob) error {
task, err := s.Client.QueryTask(job.TaskId)
if err != nil {
return err
}
// 任务执行失败了
if task.FailReason != "" {
s.db.Model(&model.MidJourneyJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
"progress": failedProgress,
"err_msg": task.FailReason,
})
s.notifyQueue.RPush(sd.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: sd.Failed})
return fmt.Errorf("task failed: %v", task.FailReason)
}
if len(task.Buttons) > 0 {
job.Hash = GetImageHash(task.Buttons[0].CustomId)
}
oldProgress := job.Progress
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
job.Prompt = task.PromptEn
if task.ImageUrl != "" {
job.OrgURL = task.ImageUrl
}
tx := s.db.Updates(&job)
if tx.Error != nil {
return fmt.Errorf("error with update database: %v", tx.Error)
}
// 通知前端更新任务进度
if oldProgress != job.Progress {
message := sd.Running
if job.Progress == 100 {
message = sd.Finished
}
s.notifyQueue.RPush(sd.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: message})
}
return nil
}
func GetImageHash(action string) string {
split := strings.Split(action, "::")
if len(split) > 5 {
@@ -211,3 +157,143 @@ func GetImageHash(action string) string {
}
return split[len(split)-1]
}
func (s *Service) CheckTaskNotify() {
go func() {
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
cli := s.Clients.Get(uint(message.UserId))
if cli == nil {
continue
}
err = cli.Send([]byte(message.Message))
if err != nil {
continue
}
}
}()
}
func (s *Service) DownloadImages() {
go func() {
var items []model.MidJourneyJob
for {
res := s.db.Where("img_url = ? AND progress = ?", "", 100).Find(&items)
if res.Error != nil {
continue
}
// download images
for _, v := range items {
if v.OrgURL == "" {
continue
}
logger.Infof("try to download image: %s", v.OrgURL)
// 如果是返回的是 discord 图片地址，则使用代理下载
proxy := false
if strings.HasPrefix(v.OrgURL, "https://cdn.discordapp.com") {
proxy = true
}
imgURL, err := s.uploaderManager.GetUploadHandler().PutUrlFile(v.OrgURL, proxy)
if err != nil {
logger.Errorf("error with download image %s, %v", v.OrgURL, err)
continue
} else {
logger.Infof("download image %s successfully.", v.OrgURL)
}
v.ImgURL = imgURL
s.db.Updates(&v)
cli := s.Clients.Get(uint(v.UserId))
if cli == nil {
continue
}
err = cli.Send([]byte(service.TaskStatusFinished))
if err != nil {
continue
}
}
time.Sleep(time.Second * 5)
}
}()
}
// PushTask push a new mj task in to task queue
func (s *Service) PushTask(task types.MjTask) {
logger.Debugf("add a new MidJourney task to the task list: %+v", task)
s.taskQueue.RPush(task)
}
// SyncTaskProgress 异步拉取任务
func (s *Service) SyncTaskProgress() {
go func() {
var jobs []model.MidJourneyJob
for {
res := s.db.Where("progress < ?", 100).Where("channel_id <> ?", "").Find(&jobs)
if res.Error != nil {
continue
}
for _, job := range jobs {
// 10 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
continue
}
task, err := s.client.QueryTask(job.TaskId, job.ChannelId)
if err != nil {
logger.Errorf("error with query task: %v", err)
continue
}
// 任务执行失败了
if task.FailReason != "" {
s.db.Model(&model.MidJourneyJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": task.FailReason,
})
logger.Errorf("task failed: %v", task.FailReason)
s.notifyQueue.RPush(service.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
continue
}
if len(task.Buttons) > 0 {
job.Hash = GetImageHash(task.Buttons[0].CustomId)
}
oldProgress := job.Progress
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
job.Prompt = task.PromptEn
if task.ImageUrl != "" {
job.OrgURL = task.ImageUrl
}
err = s.db.Updates(&job).Error
if err != nil {
logger.Errorf("error with update database: %v", err)
continue
}
// 通知前端更新任务进度
if oldProgress != job.Progress {
message := service.TaskStatusRunning
if job.Progress == 100 {
message = service.TaskStatusFinished
}
s.notifyQueue.RPush(service.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: message})
}
}
time.Sleep(time.Second * 5)
}
}()
}

143
api/service/sd/pool.go
View File

@@ -1,143 +0,0 @@
package sd
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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/types"
"geekai/service/oss"
"geekai/store"
"geekai/store/model"
"time"
"github.com/go-redis/redis/v8"
"gorm.io/gorm"
)
type ServicePool struct {
services []*Service
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
uploader *oss.UploaderManager
levelDB *store.LevelDB
}
func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, levelDB *store.LevelDB) *ServicePool {
services := make([]*Service, 0)
taskQueue := store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli)
notifyQueue := store.NewRedisQueue("StableDiffusion_Queue", redisCli)
return &ServicePool{
taskQueue: taskQueue,
notifyQueue: notifyQueue,
services: services,
db: db,
Clients: types.NewLMap[uint, *types.WsClient](),
uploader: manager,
levelDB: levelDB,
}
}
func (p *ServicePool) InitServices(configs []types.StableDiffusionConfig) {
// stop old service
for _, s := range p.services {
s.Stop()
}
p.services = make([]*Service, 0)
for k, config := range configs {
if config.Enabled == false {
continue
}
// create sd service
name := fmt.Sprintf(" sd-service-%d", k)
service := NewService(name, config, p.taskQueue, p.notifyQueue, p.db, p.uploader, p.levelDB)
// run sd service
go func() {
service.Run()
}()
p.services = append(p.services, service)
}
}
// PushTask push a new mj task in to task queue
func (p *ServicePool) PushTask(task types.SdTask) {
logger.Debugf("add a new MidJourney task to the task list: %+v", task)
p.taskQueue.RPush(task)
}
func (p *ServicePool) CheckTaskNotify() {
go func() {
logger.Info("Running Stable-Diffusion task notify checking ...")
for {
var message NotifyMessage
err := p.notifyQueue.LPop(&message)
if err != nil {
continue
}
client := p.Clients.Get(uint(message.UserId))
if client == nil {
continue
}
err = client.Send([]byte(message.Message))
if err != nil {
continue
}
}
}()
}
// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (p *ServicePool) CheckTaskStatus() {
go func() {
logger.Info("Running Stable-Diffusion task status checking ...")
for {
var jobs []model.SdJob
res := p.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
for _, job := range jobs {
// 5 分钟还没完成的任务直接删除
if time.Now().Sub(job.CreatedAt) > time.Minute*5 || job.Progress == -1 {
p.db.Delete(&job)
var user model.User
p.db.Where("id = ?", job.UserId).First(&user)
// 退回绘图次数
res = p.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
if res.Error == nil && res.RowsAffected > 0 {
p.db.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power + job.Power,
Mark: types.PowerAdd,
Model: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
})
}
continue
}
}
time.Sleep(time.Second * 5)
}
}()
}
// HasAvailableService check if it has available mj service in pool
func (p *ServicePool) HasAvailableService() bool {
return len(p.services) > 0
}

184
api/service/sd/service.go
View File

@@ -10,95 +10,91 @@ package sd
import (
"fmt"
"geekai/core/types"
logger2 "geekai/logger"
"geekai/service"
"geekai/service/oss"
"geekai/store"
"geekai/store/model"
"geekai/utils"
"strings"
"github.com/go-redis/redis/v8"
"time"
"github.com/imroc/req/v3"
"gorm.io/gorm"
)
var logger = logger2.GetLogger()
// SD 绘画服务
type Service struct {
httpClient *req.Client
config types.StableDiffusionConfig
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
uploadManager *oss.UploaderManager
name string // service name
leveldb *store.LevelDB
running bool // 运行状态
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
}
func NewService(name string, config types.StableDiffusionConfig, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelDB) *Service {
config.ApiURL = strings.TrimRight(config.ApiURL, "/")
func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelDB, redisCli *redis.Client) *Service {
return &Service{
name: name,
config: config,
httpClient: req.C(),
taskQueue: taskQueue,
notifyQueue: notifyQueue,
taskQueue: store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("StableDiffusion_Queue", redisCli),
db: db,
leveldb: levelDB,
Clients: types.NewLMap[uint, *types.WsClient](),
uploadManager: manager,
running: true,
}
}
func (s *Service) Run() {
logger.Infof("Starting Stable-Diffusion job consumer for %s", s.name)
for s.running {
var task types.SdTask
err := s.taskQueue.LPop(&task)
if err != nil {
logger.Errorf("taking task with error: %v", err)
continue
}
logger.Infof("Starting Stable-Diffusion job consumer")
go func() {
for {
var task types.SdTask
err := s.taskQueue.LPop(&task)
if err != nil {
logger.Errorf("taking task with error: %v", err)
continue
}
// translate prompt
if utils.HasChinese(task.Params.Prompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Params.Prompt), "gpt-4o-mini")
if err == nil {
task.Params.Prompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
// translate prompt
if utils.HasChinese(task.Params.Prompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Params.Prompt), "gpt-4o-mini")
if err == nil {
task.Params.Prompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
}
}
// translate negative prompt
if task.Params.NegPrompt != "" && utils.HasChinese(task.Params.NegPrompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Params.NegPrompt), "gpt-4o-mini")
if err == nil {
task.Params.NegPrompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
}
}
logger.Infof("handle a new Stable-Diffusion task: %+v", task)
err = s.Txt2Img(task)
if err != nil {
logger.Error("绘画任务执行失败：", err.Error())
// update the task progress
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
// 通知前端，任务失败
s.notifyQueue.RPush(service.NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFailed})
continue
}
}
// translate negative prompt
if task.Params.NegPrompt != "" && utils.HasChinese(task.Params.NegPrompt) {
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Params.NegPrompt), "gpt-4o-mini")
if err == nil {
task.Params.NegPrompt = content
} else {
logger.Warnf("error with translate prompt: %v", err)
}
}
logger.Infof("%s handle a new Stable-Diffusion task: %+v", s.name, task)
err = s.Txt2Img(task)
if err != nil {
logger.Error("绘画任务执行失败：", err.Error())
// update the task progress
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumns(map[string]interface{}{
"progress": -1,
"err_msg": err.Error(),
})
// 通知前端，任务失败
s.notifyQueue.RPush(NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: Failed})
continue
}
}
}
func (s *Service) Stop() {
s.running = false
}()
}
// Txt2ImgReq 文生图请求实体
@@ -160,12 +156,19 @@ func (s *Service) Txt2Img(task types.SdTask) error {
}
var res Txt2ImgResp
var errChan = make(chan error)
apiURL := fmt.Sprintf("%s/sdapi/v1/txt2img", s.config.ApiURL)
var apiKey model.ApiKey
err := s.db.Where("type", "sd").Where("enabled", true).Order("last_used_at ASC").First(&apiKey).Error
if err != nil {
return fmt.Errorf("no available Stable-Diffusion api key: %v", err)
}
apiURL := fmt.Sprintf("%s/sdapi/v1/txt2img", apiKey.ApiURL)
logger.Debugf("send image request to %s", apiURL)
// send a request to sd api endpoint
go func() {
response, err := s.httpClient.R().
SetHeader("Authorization", s.config.ApiKey).
SetHeader("Authorization", apiKey.Value).
SetBody(body).
SetSuccessResult(&res).
Post(apiURL)
@@ -178,6 +181,10 @@ func (s *Service) Txt2Img(task types.SdTask) error {
return
}
// update the last used time
apiKey.LastUsedAt = time.Now().Unix()
s.db.Updates(&apiKey)
// 保存 Base64 图片
imgURL, err := s.uploadManager.GetUploadHandler().PutBase64(res.Images[0])
if err != nil {
@@ -206,17 +213,17 @@ func (s *Service) Txt2Img(task types.SdTask) error {
// task finished
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", 100)
s.notifyQueue.RPush(NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: Finished})
s.notifyQueue.RPush(service.NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFinished})
// 从 leveldb 中删除预览图片数据
_ = s.leveldb.Delete(task.Params.TaskId)
return nil
default:
err, resp := s.checkTaskProgress()
err, resp := s.checkTaskProgress(apiKey)
// 更新任务进度
if err == nil && resp.Progress > 0 {
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", int(resp.Progress*100))
// 发送更新状态信号
s.notifyQueue.RPush(NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: Running})
s.notifyQueue.RPush(service.NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusRunning})
// 保存预览图片数据
if resp.CurrentImage != "" {
_ = s.leveldb.Put(task.Params.TaskId, resp.CurrentImage)
@@ -229,11 +236,11 @@ func (s *Service) Txt2Img(task types.SdTask) error {
}
// 执行任务
func (s *Service) checkTaskProgress() (error, *TaskProgressResp) {
apiURL := fmt.Sprintf("%s/sdapi/v1/progress?skip_current_image=false", s.config.ApiURL)
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressResp) {
apiURL := fmt.Sprintf("%s/sdapi/v1/progress?skip_current_image=false", apiKey.ApiURL)
var res TaskProgressResp
response, err := s.httpClient.R().
SetHeader("Authorization", s.config.ApiKey).
SetHeader("Authorization", apiKey.Value).
SetSuccessResult(&res).
Get(apiURL)
if err != nil {
@@ -245,3 +252,54 @@ func (s *Service) checkTaskProgress() (error, *TaskProgressResp) {
return nil, &res
}
func (s *Service) PushTask(task types.SdTask) {
logger.Debugf("add a new MidJourney task to the task list: %+v", task)
s.taskQueue.RPush(task)
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Stable-Diffusion task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
client := s.Clients.Get(uint(message.UserId))
if client == nil {
continue
}
err = client.Send([]byte(message.Message))
if err != nil {
continue
}
}
}()
}
// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running Stable-Diffusion task status checking ...")
for {
var jobs []model.SdJob
res := s.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
for _, job := range jobs {
// 5 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
}
}
time.Sleep(time.Second * 5)
}
}()
}

24
api/service/sd/types.go
View File

@@ -1,24 +0,0 @@
package sd
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 logger2 "geekai/logger"
var logger = logger2.GetLogger()
type NotifyMessage struct {
UserId int `json:"user_id"`
JobId int `json:"job_id"`
Message string `json:"message"`
}
const (
Running = "RUNNING"
Finished = "FINISH"
Failed = "FAIL"
)

15
api/service/suno/service.go
View File

@@ -13,8 +13,8 @@ import (
"fmt"
"geekai/core/types"
logger2 "geekai/logger"
"geekai/service"
"geekai/service/oss"
"geekai/service/sd"
"geekai/store"
"geekai/store/model"
"geekai/utils"
@@ -88,7 +88,7 @@ func (s *Service) Run() {
logger.Errorf("create task with error: %v", err)
s.db.Model(&model.SunoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
"err_msg": err.Error(),
"progress": 101,
"progress": service.FailTaskProgress,
})
continue
}
@@ -157,6 +157,9 @@ func (s *Service) Create(task types.SunoTask) (RespVo, error) {
if res.Code != "success" {
return RespVo{}, fmt.Errorf("API 返回失败：%s", res.Message)
}
// update the last_use_at for api key
apiKey.LastUsedAt = time.Now().Unix()
session.Updates(&apiKey)
res.Channel = apiKey.ApiURL
return res, nil
}
@@ -165,7 +168,7 @@ func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Suno task notify checking ...")
for {
var message sd.NotifyMessage
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
@@ -210,7 +213,7 @@ func (s *Service) DownloadImages() {
v.AudioURL = audioURL
v.Progress = 100
s.db.Updates(&v)
s.notifyQueue.RPush(sd.NotifyMessage{UserId: v.UserId, JobId: int(v.Id), Message: sd.Finished})
s.notifyQueue.RPush(service.NotifyMessage{UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished})
}
time.Sleep(time.Second * 10)
@@ -278,10 +281,10 @@ func (s *Service) SyncTaskProgress() {
tx.Commit()
} else if task.Data.FailReason != "" {
job.Progress = 101
job.Progress = service.FailTaskProgress
job.ErrMsg = task.Data.FailReason
s.db.Updates(&job)
s.notifyQueue.RPush(sd.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: sd.Failed})
s.notifyQueue.RPush(service.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
}
}

13
api/service/types.go
View File

@@ -1,4 +1,17 @@
package service
const FailTaskProgress = 101
const (
TaskStatusRunning = "RUNNING"
TaskStatusFinished = "FINISH"
TaskStatusFailed = "FAIL"
)
type NotifyMessage struct {
UserId int `json:"user_id"`
JobId int `json:"job_id"`
Message string `json:"message"`
}
const RewritePromptTemplate = "Please rewrite the following text into AI painting prompt words, and please try to add detailed description of the picture, painting style, scene, rendering effect, picture light and other creative elements. Just output the final prompt word directly. Do not output any explanation lines. The text to be rewritten is: [%s]"
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]"

101
api/service/wx/bot.go
View File

@@ -1,101 +0,0 @@
package wx
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 (
logger2 "geekai/logger"
"geekai/store/model"
"github.com/eatmoreapple/openwechat"
"github.com/skip2/go-qrcode"
"gorm.io/gorm"
"os"
"strconv"
)
// 微信收款机器人
var logger = logger2.GetLogger()
type Bot struct {
bot *openwechat.Bot
token string
db *gorm.DB
}
func NewWeChatBot(db *gorm.DB) *Bot {
bot := openwechat.DefaultBot(openwechat.Desktop)
return &Bot{
bot: bot,
db: db,
}
}
func (b *Bot) Run() error {
logger.Info("Starting WeChat Bot...")
// set message handler
b.bot.MessageHandler = func(msg *openwechat.Message) {
b.messageHandler(msg)
}
// scan code login callback
b.bot.UUIDCallback = b.qrCodeCallBack
debug, err := strconv.ParseBool(os.Getenv("APP_DEBUG"))
if debug {
reloadStorage := openwechat.NewJsonFileHotReloadStorage("storage.json")
err = b.bot.HotLogin(reloadStorage, true)
} else {
err = b.bot.Login()
}
if err != nil {
return err
}
logger.Info("微信登录成功！")
return nil
}
// message handler
func (b *Bot) messageHandler(msg *openwechat.Message) {
sender, err := msg.Sender()
if err != nil {
return
}
// 只处理微信支付的推送消息
if sender.NickName == "微信支付" ||
msg.MsgType == openwechat.MsgTypeApp ||
msg.AppMsgType == openwechat.AppMsgTypeUrl {
// 解析支付金额
message := parseTransactionMessage(msg.Content)
transaction := extractTransaction(message)
logger.Infof("解析到收款信息：%+v", transaction)
if transaction.TransId != "" {
var item model.Reward
res := b.db.Where("tx_id = ?", transaction.TransId).First(&item)
if item.Id > 0 {
logger.Error("当前交易 ID 己经存在！")
return
}
res = b.db.Create(&model.Reward{
TxId: transaction.TransId,
Amount: transaction.Amount,
Remark: transaction.Remark,
Status: false,
})
if res.Error != nil {
logger.Errorf("交易保存失败: %v", res.Error)
}
}
}
}
func (b *Bot) qrCodeCallBack(uuid string) {
logger.Info("请使用微信扫描下面二维码登录")
q, _ := qrcode.New("https://login.weixin.qq.com/l/"+uuid, qrcode.Medium)
logger.Info(q.ToString(true))
}

112
api/service/wx/tranaction.go
View File

@@ -1,112 +0,0 @@
package wx
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 (
"encoding/xml"
"net/url"
"strconv"
"strings"
)
// Message 转账消息
type Message struct {
Des string
Url string
}
// Transaction 解析后的交易信息
type Transaction struct {
TransId string `json:"trans_id"` // 微信转账交易 ID
Amount float64 `json:"amount"` // 微信转账交易金额
Remark string `json:"remark"` // 转账备注
}
// 解析微信转账消息
func parseTransactionMessage(xmlData string) *Message {
decoder := xml.NewDecoder(strings.NewReader(xmlData))
message := Message{}
for {
token, err := decoder.Token()
if err != nil {
break
}
switch se := token.(type) {
case xml.StartElement:
var value string
if se.Name.Local == "des" && message.Des == "" {
if err := decoder.DecodeElement(&value, &se); err == nil {
message.Des = strings.TrimSpace(value)
}
break
}
if se.Name.Local == "weapp_path" || se.Name.Local == "url" {
if err := decoder.DecodeElement(&value, &se); err == nil {
if strings.Contains(value, "?trans_id=") || strings.Contains(value, "?id=") {
message.Url = value
}
}
break
}
}
}
// 兼容旧版消息记录
if message.Url == "" {
var msg struct {
XMLName xml.Name `xml:"msg"`
AppMsg struct {
Des string `xml:"des"`
Url string `xml:"url"`
} `xml:"appmsg"`
}
if err := xml.Unmarshal([]byte(xmlData), &msg); err == nil {
message.Url = msg.AppMsg.Url
}
}
return &message
}
// 导出交易信息
func extractTransaction(message *Message) Transaction {
var tx = Transaction{}
// 导出交易金额和备注
lines := strings.Split(message.Des, "\n")
for _, line := range lines {
line = strings.TrimSpace(line)
if len(line) == 0 {
continue
}
// 解析收款金额
prefix := "收款金额￥"
if strings.HasPrefix(line, prefix) {
if value, err := strconv.ParseFloat(line[len(prefix):], 64); err == nil {
tx.Amount = value
continue
}
}
// 解析收款备注
prefix = "付款方备注"
if strings.HasPrefix(line, prefix) {
tx.Remark = line[len(prefix):]
break
}
}
// 解析交易 ID
parse, err := url.Parse(message.Url)
if err == nil {
tx.TransId = parse.Query().Get("id")
if tx.TransId == "" {
tx.TransId = parse.Query().Get("trans_id")
}
}
return tx
}

16
api/store/model/redeem.go Normal file
View File

@@ -0,0 +1,16 @@
package model
import "time"
// 兑换码
type Redeem struct {
Id uint `gorm:"primarykey;column:id"`
UserId uint // 用户 ID
Name string // 名称
Power int // 算力
Code string // 兑换码
Enabled bool // 启用状态
RedeemedAt int64 // 兑换时间
CreatedAt time.Time
}

13
api/store/model/reward.go
View File

@@ -1,13 +0,0 @@
package model
// 用户打赏
type Reward struct {
BaseModel
UserId uint // 用户 ID
TxId string // 交易ID
Amount float64 // 打赏金额
Remark string // 打赏备注
Status bool // 核销状态
Exchange string // 众筹兑换详情，JSON
}

13
api/store/vo/redeem.go Normal file
View File

@@ -0,0 +1,13 @@
package vo
type Redeem struct {
Id uint `json:"id"`
UserId uint `json:"user_id"` // 用户 ID
Name string `json:"name"`
Username string `json:"username"`
Power int `json:"power"` // 算力
Code string `json:"code"` // 兑换码
Enabled bool `json:"enabled"`
RedeemedAt int64 `json:"redeemed_at"` // 兑换时间
CreatedAt int64 `json:"created_at"`
}

16
api/store/vo/reward.go
View File

@@ -1,16 +0,0 @@
package vo
type Reward struct {
BaseVo
UserId uint `json:"user_id"` // 用户 ID
Username string `json:"username"`
TxId string `json:"tx_id"` // 交易ID
Amount float64 `json:"amount"` // 打赏金额
Remark string `json:"remark"` // 打赏备注
Status bool `json:"status"` // 核销状态
Exchange RewardExchange `json:"exchange"`
}
type RewardExchange struct {
Power int `json:"power"`
}

55
api/test/test.go
View File

@@ -1,16 +1,55 @@
package main
import (
"crypto/rand"
"encoding/hex"
"fmt"
"geekai/utils"
"sync"
)
func main() {
file := "http://nk.img.r9it.com/chatgpt-plus/1719389335351828.xlsx"
content, err := utils.ReadFileContent(file, "http://172.22.11.69:9998")
if err != nil {
panic(err)
}
const (
codeLength = 32 // 兑换码长度
)
fmt.Println(content)
var (
codeMap = make(map[string]bool)
mapMutex = &sync.Mutex{}
)
// GenerateUniqueCode 生成唯一兑换码
func GenerateUniqueCode() (string, error) {
for {
code, err := generateCode()
if err != nil {
return "", err
}
mapMutex.Lock()
if !codeMap[code] {
codeMap[code] = true
mapMutex.Unlock()
return code, nil
}
mapMutex.Unlock()
}
}
// generateCode 生成兑换码
func generateCode() (string, error) {
bytes := make([]byte, codeLength/2) // 因为 hex 编码会使长度翻倍
if _, err := rand.Read(bytes); err != nil {
return "", err
}
return hex.EncodeToString(bytes), nil
}
func main() {
for i := 0; i < 10; i++ {
code, err := GenerateUniqueCode()
if err != nil {
fmt.Println("Error generating code:", err)
return
}
fmt.Println("Generated code:", code)
}
}

6
api/utils/openai.go
View File

@@ -64,6 +64,7 @@ func OpenAIRequest(db *gorm.DB, prompt string, modelName string) (string, error)
client.SetProxyURL(apiKey.ApiURL)
}
apiURL := fmt.Sprintf("%s/v1/chat/completions", apiKey.ApiURL)
logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, modelName)
r, err := client.R().SetHeader("Content-Type", "application/json").
SetHeader("Authorization", "Bearer "+apiKey.Value).
SetBody(types.ApiRequest{
@@ -76,6 +77,11 @@ func OpenAIRequest(db *gorm.DB, prompt string, modelName string) (string, error)
if err != nil {
return "", fmt.Errorf("请求 OpenAI API失败：%v", err)
}
if r.IsErrorState() {
return "", fmt.Errorf("请求 OpenAI API失败：%v", r.Status)
}
body, _ := io.ReadAll(r.Body)
err = json.Unmarshal(body, &response)
if err != nil {

18
api/utils/strings.go
View File

@@ -8,14 +8,16 @@ package utils
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"math/rand"
"strings"
"time"
"unicode"
"golang.org/x/crypto/sha3"
rand2 "math/rand"
)
// RandString generate rand string with specified length
@@ -23,7 +25,7 @@ func RandString(length int) string {
str := "0123456789abcdefghijklmnopqrstuvwxyz"
data := []byte(str)
var result []byte
r := rand.New(rand.NewSource(time.Now().UnixNano()))
r := rand2.New(rand2.NewSource(time.Now().UnixNano()))
for i := 0; i < length; i++ {
result = append(result, data[r.Intn(len(data))])
}
@@ -34,8 +36,8 @@ func RandomNumber(bit int) int {
minNum := intPow(10, bit-1)
maxNum := intPow(10, bit) - 1
rand.NewSource(time.Now().UnixNano())
return rand.Intn(maxNum-minNum+1) + minNum
rand2.NewSource(time.Now().UnixNano())
return rand2.Intn(maxNum-minNum+1) + minNum
}
func intPow(x, y int) int {
@@ -124,3 +126,11 @@ func HasChinese(text string) bool {
}
return false
}
func GenRedeemCode(codeLength int) (string, error) {
bytes := make([]byte, codeLength/2)
if _, err := rand.Read(bytes); err != nil {
return "", err
}
return hex.EncodeToString(bytes), nil
}