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

Merge branch 'yangjian102621:main' into main

Browse Source
This commit is contained in:
JingHong
2023-12-27 23:33:46 +08:00
committed by GitHub
parent f6b567d6fc d972e97c88
commit 131e051ddc
29 changed files with 413 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
api/core/app_server.go
View File

@@ -154,6 +154,7 @@ func authorizeMiddleware(s *AppServer, client *redis.Client) gin.HandlerFunc {
c.Request.URL.Path == "/api/chat/detail" ||
c.Request.URL.Path == "/api/role/list" ||
c.Request.URL.Path == "/api/mj/jobs" ||
c.Request.URL.Path == "/api/mj/client" ||
c.Request.URL.Path == "/api/invite/hits" ||
c.Request.URL.Path == "/api/sd/jobs" ||
c.Request.URL.Path == "/api/upload" ||

13
api/core/config.go
View File

@@ -14,13 +14,12 @@ var logger = logger2.GetLogger()
func NewDefaultConfig() *types.AppConfig {
return &types.AppConfig{
Listen: "0.0.0.0:5678",
ProxyURL: "",
Manager: types.Manager{Username: "admin", Password: "admin123"},
StaticDir: "./static",
StaticUrl: "http://localhost/5678/static",
Redis: types.RedisConfig{Host: "localhost", Port: 6379, Password: ""},
AesEncryptKey: utils.RandString(24),
Listen: "0.0.0.0:5678",
ProxyURL: "",
Manager: types.Manager{Username: "admin", Password: "admin123"},
StaticDir: "./static",
StaticUrl: "http://localhost/5678/static",
Redis: types.RedisConfig{Host: "localhost", Port: 6379, Password: ""},
Session: types.Session{
SecretKey: utils.RandString(64),
MaxAge: 86400,

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

@@ -5,22 +5,21 @@ import (
)
type AppConfig struct {
Path string `toml:"-"`
Listen string
Session Session
ProxyURL string
MysqlDns string // mysql 连接地址
Manager Manager // 后台管理员账户信息
StaticDir string // 静态资源目录
StaticUrl string // 静态资源 URL
Redis RedisConfig // redis 连接信息
ApiConfig ChatPlusApiConfig // ChatPlus API authorization configs
AesEncryptKey string
SmsConfig AliYunSmsConfig // AliYun send message service config
OSS OSSConfig // OSS config
MjConfigs []MidJourneyConfig // mj AI draw service pool
WeChatBot bool // 是否启用微信机器人
SdConfigs []StableDiffusionConfig // sd AI draw service pool
Path string `toml:"-"`
Listen string
Session Session
ProxyURL string
MysqlDns string // mysql 连接地址
Manager Manager // 后台管理员账户信息
StaticDir string // 静态资源目录
StaticUrl string // 静态资源 URL
Redis RedisConfig // redis 连接信息
ApiConfig ChatPlusApiConfig // ChatPlus API authorization configs
SmsConfig AliYunSmsConfig // AliYun send message service config
OSS OSSConfig // OSS config
MjConfigs []MidJourneyConfig // mj AI draw service pool
WeChatBot bool // 是否启用微信机器人
SdConfigs []StableDiffusionConfig // sd AI draw service pool
XXLConfig XXLConfig
AlipayConfig AlipayConfig
@@ -34,11 +33,15 @@ type ChatPlusApiConfig struct {
}
type MidJourneyConfig struct {
Enabled bool
UserToken string
BotToken string
GuildId string // Server ID
ChanelId string // Chanel ID
Enabled bool
UserToken string
BotToken string
GuildId string // Server ID
ChanelId string // Chanel ID
UseCDN bool
DiscordAPI string
DiscordCDN string
DiscordGateway string
}
type StableDiffusionConfig struct {

39
api/core/types/function.go
View File

@@ -32,46 +32,21 @@ const (
var InnerFunctions = []Function{
{
Name: FuncZaoBao,
Description: "每日早报，获取当天全球的热门新闻事件列表",
Description: "每日早报，获取当天新闻事件列表",
Parameters: Parameters{
Type: "object",
Properties: map[string]Property{
"text": {
Type: "string",
Description: "",
},
},
Required: []string{},
Type: "object",
Properties: map[string]Property{},
Required: []string{},
},
},
{
Name: FuncWeibo,
Description: "新浪微博热搜榜，微博当日热搜榜单",
Parameters: Parameters{
Type: "object",
Properties: map[string]Property{
"text": {
Type: "string",
Description: "",
},
},
Required: []string{},
},
},
{
Name: FuncHeadLine,
Description: "今日头条，给用户推荐当天的头条新闻，周榜热文",
Parameters: Parameters{
Type: "object",
Properties: map[string]Property{
"text": {
Type: "string",
Description: "",
},
},
Required: []string{},
Type: "object",
Properties: map[string]Property{},
Required: []string{},
},
},

2
api/core/types/locked_map.go
View File

@@ -6,7 +6,7 @@ import (
)
type MKey interface {
string | int
string | int | uint
}
type MValue interface {
*WsClient | *ChatSession | context.CancelFunc | []interface{}

3
api/go.mod
View File

@@ -6,7 +6,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/bwmarrin/discordgo v0.27.1
github.com/eatmoreapple/openwechat v1.2.1
github.com/gin-gonic/gin v1.9.1
github.com/go-redis/redis/v8 v8.11.5
@@ -26,6 +25,8 @@ require (
require github.com/xxl-job/xxl-job-executor-go v1.2.0
require github.com/bg5t/mydiscordgo v0.28.1
require (
github.com/andybalholm/brotli v1.0.4 // indirect
github.com/bytedance/sonic v1.9.1 // indirect

4
api/go.sum
View File

@@ -7,8 +7,8 @@ github.com/aliyun/aliyun-oss-go-sdk v2.2.9+incompatible/go.mod h1:T/Aws4fEfogEE9
github.com/andybalholm/brotli v1.0.4 h1:V7DdXeJtZscaqfNuAdSRuRFzuiKlHSC/Zh3zl9qY3JY=
github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
github.com/bwmarrin/discordgo v0.27.1 h1:ib9AIc/dom1E/fSIulrBwnez0CToJE113ZGt4HoliGY=
github.com/bwmarrin/discordgo v0.27.1/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
github.com/bg5t/mydiscordgo v0.28.1 h1:mVH0ZWstVdJffCi/EXJAYQDtXwIKAJYVXLmECu1hEK8=
github.com/bg5t/mydiscordgo v0.28.1/go.mod h1:n3aba73N18k1DzM0t0mGE8rwW3Z+vwTvI8pcsBgxN/8=
github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=

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

@@ -6,7 +6,11 @@ import (
"chatplus/handler"
"chatplus/store/model"
"chatplus/store/vo"
"chatplus/utils"
"chatplus/utils/resp"
"github.com/golang-jwt/jwt/v5"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
@@ -29,13 +33,65 @@ func (h *FunctionHandler) Save(c *gin.Context) {
return
}
logger.Info(data)
var f = model.Function{
Id: data.Id,
Name: data.Name,
Label: data.Label,
Description: data.Description,
Parameters: utils.JsonEncode(data.Parameters),
Required: utils.JsonEncode(data.Required),
Action: data.Action,
Token: data.Token,
Enabled: data.Enabled,
}
res := h.db.Save(&f)
if res.Error != nil {
resp.ERROR(c, "error with save data:"+res.Error.Error())
return
}
data.Id = f.Id
resp.SUCCESS(c, data)
}
func (h *FunctionHandler) 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
}
res := h.db.Model(&model.Function{}).Where("id = ?", data.Id).Update(data.Filed, data.Value)
if res.Error != nil {
resp.ERROR(c, "更新数据库失败！")
return
}
resp.SUCCESS(c)
}
func (h *FunctionHandler) List(c *gin.Context) {
var items []model.Function
res := h.db.Find(&items)
if res.Error != nil {
resp.ERROR(c, "No data found")
return
}
resp.SUCCESS(c)
functions := make([]vo.Function, 0)
for _, v := range items {
var f vo.Function
err := utils.CopyObject(v, &f)
if err != nil {
continue
}
functions = append(functions, f)
}
resp.SUCCESS(c, functions)
}
func (h *FunctionHandler) Remove(c *gin.Context) {
@@ -50,3 +106,20 @@ func (h *FunctionHandler) Remove(c *gin.Context) {
}
resp.SUCCESS(c)
}
// GenToken generate function api access token
func (h *FunctionHandler) GenToken(c *gin.Context) {
// 创建 token
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
"user_id": 0,
"expired": 0,
})
tokenString, err := token.SignedString([]byte(h.App.Config.Session.SecretKey))
if err != nil {
logger.Error("error with generate token", err)
resp.ERROR(c)
return
}
resp.SUCCESS(c, tokenString)
}

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

@@ -442,7 +442,7 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, platf
} else {
client = http.DefaultClient
}
logger.Infof("Sending %s request, KEY: %s, PROXY: %s, Model: %s", platform, *apiKey, proxyURL, req.Model)
logger.Infof("Sending %s request, ApiURL:%s, PROXY: %s, Model: %s", platform, apiURL, proxyURL, req.Model)
switch platform {
case types.Azure:
request.Header.Set("api-key", *apiKey)
@@ -452,7 +452,6 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, platf
if err != nil {
return nil, err
}
logger.Info(token)
request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
break
case types.Baidu:

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

@@ -50,7 +50,7 @@ type xunFeiResp struct {
}
var Model2URL = map[string]string{
"generalv1": "1.1",
"general": "v1.1",
"generalv2": "v2.1",
"generalv3": "v3.1",
}

38
api/handler/mj_handler.go
View File

@@ -13,7 +13,9 @@ import (
"encoding/base64"
"fmt"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"gorm.io/gorm"
"net/http"
"strings"
"time"
)
@@ -58,6 +60,27 @@ func (h *MidJourneyHandler) preCheck(c *gin.Context) bool {
}
// Client WebSocket 客户端，用于通知任务状态变更
func (h *MidJourneyHandler) Client(c *gin.Context) {
ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
if err != nil {
logger.Error(err)
c.Abort()
return
}
userId := h.GetInt(c, "user_id", 0)
if userId == 0 {
logger.Info("Invalid user ID")
c.Abort()
return
}
client := types.NewWsClient(ws)
h.pool.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
// Image 创建一个绘画任务
func (h *MidJourneyHandler) Image(c *gin.Context) {
var data struct {
@@ -147,6 +170,9 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
UserId: userId,
})
client := h.pool.Clients.Get(uint(job.UserId))
_ = client.Send([]byte("Task Updated"))
// update user's img calls
h.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
resp.SUCCESS(c)
@@ -205,6 +231,10 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
MessageId: data.MessageId,
MessageHash: data.MessageHash,
})
client := h.pool.Clients.Get(uint(job.UserId))
_ = client.Send([]byte("Task Updated"))
resp.SUCCESS(c)
}
@@ -226,6 +256,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
job := model.MidJourneyJob{
Type: types.TaskVariation.String(),
ChannelId: data.ChannelId,
ReferenceId: data.MessageId,
UserId: userId,
TaskId: data.TaskId,
@@ -250,6 +281,9 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
MessageHash: data.MessageHash,
})
client := h.pool.Clients.Get(uint(job.UserId))
_ = client.Send([]byte("Task Updated"))
// update user's img calls
h.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
resp.SUCCESS(c)
@@ -320,6 +354,7 @@ func (h *MidJourneyHandler) JobList(c *gin.Context) {
func (h *MidJourneyHandler) Remove(c *gin.Context) {
var data struct {
Id uint `json:"id"`
UserId uint `json:"user_id"`
ImgURL string `json:"img_url"`
}
if err := c.ShouldBindJSON(&data); err != nil {
@@ -340,5 +375,8 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
logger.Error("remove image failed: ", err)
}
client := h.pool.Clients.Get(data.UserId)
_ = client.Send([]byte("Task Updated"))
resp.SUCCESS(c)
}

4
api/main.go
View File

@@ -168,6 +168,7 @@ func main() {
fx.Invoke(func(pool *mj.ServicePool) {
if pool.HasAvailableService() {
pool.DownloadImages()
pool.CheckTaskNotify()
}
}),
@@ -234,6 +235,7 @@ func main() {
}),
fx.Invoke(func(s *core.AppServer, h *handler.MidJourneyHandler) {
group := s.Engine.Group("/api/mj/")
group.Any("client", h.Client)
group.POST("image", h.Image)
group.POST("upscale", h.Upscale)
group.POST("variation", h.Variation)
@@ -350,8 +352,10 @@ func main() {
fx.Invoke(func(s *core.AppServer, h *admin.FunctionHandler) {
group := s.Engine.Group("/api/admin/function/")
group.POST("save", h.Save)
group.POST("set", h.Set)
group.GET("list", h.List)
group.GET("remove", h.Remove)
group.GET("token", h.GenToken)
}),
fx.Provide(handler.NewTestHandler),

41
api/service/mj/bot.go
View File

@@ -4,7 +4,7 @@ import (
"chatplus/core/types"
logger2 "chatplus/logger"
"chatplus/utils"
"github.com/bwmarrin/discordgo"
discordgo "github.com/bg5t/mydiscordgo"
"github.com/gorilla/websocket"
"net/http"
"net/url"
@@ -17,33 +17,48 @@ import (
var logger = logger2.GetLogger()
type Bot struct {
config *types.MidJourneyConfig
config types.MidJourneyConfig
bot *discordgo.Session
name string
service *Service
}
func NewBot(name string, proxy string, config *types.MidJourneyConfig, service *Service) (*Bot, error) {
discord, err := discordgo.New("Bot " + config.BotToken)
func NewBot(name string, proxy string, config types.MidJourneyConfig, service *Service) (*Bot, error) {
bot, err := discordgo.New("Bot " + config.BotToken)
if err != nil {
logger.Error(err)
return nil, err
}
if proxy != "" {
proxy, _ := url.Parse(proxy)
discord.Client = &http.Client{
Transport: &http.Transport{
// use CDN reverse proxy
if config.UseCDN {
discordgo.SetEndpointDiscord(config.DiscordAPI)
discordgo.SetEndpointCDN(config.DiscordCDN)
discordgo.SetEndpointStatus(config.DiscordAPI + "/api/v2/")
bot.MjGateway = config.DiscordGateway + "/"
} else { // use proxy
discordgo.SetEndpointDiscord("https://discord.com")
discordgo.SetEndpointCDN("https://cdn.discordapp.com")
discordgo.SetEndpointStatus("https://discord.com/api/v2/")
bot.MjGateway = "wss://gateway.discord.gg"
if proxy != "" {
proxy, _ := url.Parse(proxy)
bot.Client = &http.Client{
Transport: &http.Transport{
Proxy: http.ProxyURL(proxy),
},
}
bot.Dialer = &websocket.Dialer{
Proxy: http.ProxyURL(proxy),
},
}
discord.Dialer = &websocket.Dialer{
Proxy: http.ProxyURL(proxy),
}
}
}
return &Bot{
config: config,
bot: discord,
bot: bot,
name: name,
service: service,
}, nil

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

@@ -12,25 +12,32 @@ import (
type Client struct {
client *req.Client
config types.MidJourneyConfig
Config types.MidJourneyConfig
apiURL string
}
func NewClient(config types.MidJourneyConfig, proxy string) *Client {
client := req.C().SetTimeout(10 * time.Second)
var apiURL string
// set proxy URL
if proxy != "" {
client.SetProxyURL(proxy)
if config.UseCDN {
apiURL = config.DiscordAPI + "/api/v9/interactions"
} else {
apiURL = "https://discord.com/api/v9/interactions"
if proxy != "" {
client.SetProxyURL(proxy)
}
}
logger.Info(config)
return &Client{client: client, config: config}
return &Client{client: client, Config: config, apiURL: apiURL}
}
func (c *Client) Imagine(prompt string) error {
interactionsReq := &InteractionsRequest{
Type: 2,
ApplicationID: ApplicationID,
GuildID: c.config.GuildId,
ChannelID: c.config.ChanelId,
GuildID: c.Config.GuildId,
ChannelID: c.Config.ChanelId,
SessionID: SessionID,
Data: map[string]any{
"version": "1166847114203123795",
@@ -68,11 +75,10 @@ func (c *Client) Imagine(prompt string) error {
},
}
url := "https://discord.com/api/v9/interactions"
r, err := c.client.R().SetHeader("Authorization", c.config.UserToken).
r, err := c.client.R().SetHeader("Authorization", c.Config.UserToken).
SetHeader("Content-Type", "application/json").
SetBody(interactionsReq).
Post(url)
Post(c.apiURL)
if err != nil || r.IsErrorState() {
return fmt.Errorf("error with http request: %w%v", err, r.Err)
@@ -87,8 +93,8 @@ func (c *Client) Upscale(index int, messageId string, hash string) error {
interactionsReq := &InteractionsRequest{
Type: 3,
ApplicationID: ApplicationID,
GuildID: c.config.GuildId,
ChannelID: c.config.ChanelId,
GuildID: c.Config.GuildId,
ChannelID: c.Config.ChanelId,
MessageFlags: &flags,
MessageID: &messageId,
SessionID: SessionID,
@@ -99,13 +105,12 @@ func (c *Client) Upscale(index int, messageId string, hash string) error {
Nonce: fmt.Sprintf("%d", time.Now().UnixNano()),
}
url := "https://discord.com/api/v9/interactions"
var res InteractionsResult
r, err := c.client.R().SetHeader("Authorization", c.config.UserToken).
r, err := c.client.R().SetHeader("Authorization", c.Config.UserToken).
SetHeader("Content-Type", "application/json").
SetBody(interactionsReq).
SetErrorResult(&res).
Post(url)
Post(c.apiURL)
if err != nil || r.IsErrorState() {
return fmt.Errorf("error with http request: %v%v%v", err, r.Err, res.Message)
}
@@ -119,8 +124,8 @@ func (c *Client) Variation(index int, messageId string, hash string) error {
interactionsReq := &InteractionsRequest{
Type: 3,
ApplicationID: ApplicationID,
GuildID: c.config.GuildId,
ChannelID: c.config.ChanelId,
GuildID: c.Config.GuildId,
ChannelID: c.Config.ChanelId,
MessageFlags: &flags,
MessageID: &messageId,
SessionID: SessionID,
@@ -131,13 +136,12 @@ func (c *Client) Variation(index int, messageId string, hash string) error {
Nonce: fmt.Sprintf("%d", time.Now().UnixNano()),
}
url := "https://discord.com/api/v9/interactions"
var res InteractionsResult
r, err := c.client.R().SetHeader("Authorization", c.config.UserToken).
r, err := c.client.R().SetHeader("Authorization", c.Config.UserToken).
SetHeader("Content-Type", "application/json").
SetBody(interactionsReq).
SetErrorResult(&res).
Post(url)
Post(c.apiURL)
if err != nil || r.IsErrorState() {
return fmt.Errorf("error with http request: %v%v%v", err, r.Err, res.Message)
}

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

@@ -6,9 +6,9 @@ import (
"chatplus/store"
"chatplus/store/model"
"fmt"
"github.com/go-redis/redis/v8"
"time"
"github.com/go-redis/redis/v8"
"gorm.io/gorm"
)
@@ -16,13 +16,16 @@ import (
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
}
func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, appConfig *types.AppConfig) *ServicePool {
services := make([]*Service, 0)
queue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
taskQueue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
notifyQueue := store.NewRedisQueue("MidJourney_Notify_Queue", redisCli)
// create mj client and service
for k, config := range appConfig.MjConfigs {
if config.Enabled == false {
@@ -33,9 +36,9 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
name := fmt.Sprintf("MjService-%d", k)
// create mj service
service := NewService(name, queue, 4, 600, db, client)
service := NewService(name, taskQueue, notifyQueue, 4, 600, db, client)
botName := fmt.Sprintf("MjBot-%d", k)
bot, err := NewBot(botName, appConfig.ProxyURL, &config, service)
bot, err := NewBot(botName, appConfig.ProxyURL, config, service)
if err != nil {
continue
}
@@ -54,13 +57,32 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
}
return &ServicePool{
taskQueue: queue,
taskQueue: taskQueue,
notifyQueue: notifyQueue,
services: services,
uploaderManager: manager,
db: db,
Clients: types.NewLMap[uint, *types.WsClient](),
}
}
func (p *ServicePool) CheckTaskNotify() {
go func() {
for {
var userId uint
err := p.notifyQueue.LPop(&userId)
if err != nil {
continue
}
client := p.Clients.Get(userId)
err = client.Send([]byte("Task Updated"))
if err != nil {
continue
}
}
}()
}
func (p *ServicePool) DownloadImages() {
go func() {
var items []model.MidJourneyJob
@@ -71,15 +93,21 @@ func (p *ServicePool) DownloadImages() {
}
// download images
for _, item := range items {
imgURL, err := p.uploaderManager.GetUploadHandler().PutImg(item.OrgURL, true)
for _, v := range items {
imgURL, err := p.uploaderManager.GetUploadHandler().PutImg(v.OrgURL, true)
if err != nil {
logger.Error("error with download image: ", err)
continue
}
item.ImgURL = imgURL
p.db.Updates(&item)
v.ImgURL = imgURL
p.db.Updates(&v)
client := p.Clients.Get(uint(v.UserId))
err = client.Send([]byte("Task Updated"))
if err != nil {
continue
}
}
time.Sleep(time.Second * 5)

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

@@ -15,6 +15,7 @@ type Service struct {
name string // service name
client *Client // MJ client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
maxHandleTaskNum int32 // max task number current service can handle
handledTaskNum int32 // already handled task number
@@ -22,11 +23,12 @@ type Service struct {
taskTimeout int64
}
func NewService(name string, queue *store.RedisQueue, maxTaskNum int32, timeout int64, db *gorm.DB, client *Client) *Service {
func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, maxTaskNum int32, timeout int64, db *gorm.DB, client *Client) *Service {
return &Service{
name: name,
db: db,
taskQueue: queue,
taskQueue: taskQueue,
notifyQueue: notifyQueue,
client: client,
taskTimeout: timeout,
maxHandleTaskNum: maxTaskNum,
@@ -53,9 +55,10 @@ func (s *Service) Run() {
}
// if it's reference message, check if it's this channel's message
if task.ChannelId != "" && task.ChannelId != s.client.config.ChanelId {
if task.ChannelId != "" && task.ChannelId != s.client.Config.ChanelId {
s.taskQueue.RPush(task)
s.db.Model(&model.MidJourneyJob{Id: uint(task.Id)}).UpdateColumn("progress", -1)
s.notifyQueue.RPush(task.UserId)
time.Sleep(time.Second)
continue
}
@@ -77,6 +80,7 @@ func (s *Service) Run() {
logger.Error("绘画任务执行失败：", err)
// update the task progress
s.db.Model(&model.MidJourneyJob{Id: uint(task.Id)}).UpdateColumn("progress", -1)
s.notifyQueue.RPush(task.UserId)
// restore img_call quota
s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1))
continue
@@ -134,6 +138,10 @@ func (s *Service) Notify(data CBReq) {
job.Prompt = data.Prompt
job.Hash = data.Image.Hash
job.OrgURL = data.Image.URL
if s.client.Config.UseCDN {
job.UseProxy = true
job.ImgURL = strings.ReplaceAll(data.Image.URL, "https://cdn.discordapp.com", s.client.Config.DiscordCDN)
}
res = s.db.Updates(&job)
if res.Error != nil {
@@ -146,4 +154,6 @@ func (s *Service) Notify(data CBReq) {
atomic.AddInt32(&s.handledTaskNum, -1)
}
s.notifyQueue.RPush(job.UserId)
}

2
api/store/model/function.go
View File

@@ -3,9 +3,11 @@ package model
type Function struct {
Id uint `gorm:"primarykey;column:id"`
Name string
Label string
Description string
Parameters string
Required string
Action string
Token string
Enabled bool
}

1
api/store/model/mj_job.go
View File

@@ -15,6 +15,7 @@ type MidJourneyJob struct {
Hash string // message hash
Progress int
Prompt string
UseProxy bool // 是否使用反代加载图片
CreatedAt time.Time
}

2
api/store/vo/function.go
View File

@@ -14,9 +14,11 @@ type Property struct {
type Function struct {
Id uint `json:"id"`
Name string `json:"name"`
Label string `json:"label"`
Description string `json:"description"`
Parameters Parameters `json:"parameters"`
Required []string `json:"required"`
Action string `json:"action"`
Token string `json:"token"`
Enabled bool `json:"enabled"`
}

1
api/store/vo/mj_job.go
View File

@@ -15,5 +15,6 @@ type MidJourneyJob struct {
Hash string `json:"hash"`
Progress int `json:"progress"`
Prompt string `json:"prompt"`
UseProxy bool `json:"use_proxy"`
CreatedAt time.Time `json:"created_at"`
}