xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-09-17 16:56:38 +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
commit 6cc05e9e27
29 changed files with 413 additions and 169 deletions
Show Stats 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"`
}

7
database/update-v3.2.3.sql
View File

@ -18,4 +18,9 @@ ALTER TABLE `chatgpt_functions` ADD UNIQUE(`name`);
ALTER TABLE `chatgpt_functions` ADD `enabled` TINYINT(1) NOT NULL DEFAULT '0' COMMENT '是否启用' AFTER `action`;
ALTER TABLE `chatgpt_functions` ADD `lebal` VARCHAR(30) NULL COMMENT '函数标签' AFTER `name`;
ALTER TABLE `chatgpt_functions` ADD `label` VARCHAR(30) NULL COMMENT '函数标签' AFTER `name`;
ALTER TABLE `chatgpt_mj_jobs` ADD `use_proxy` TINYINT(1) NOT NULL DEFAULT '0' COMMENT '是否使用反代' AFTER `progress`;
ALTER TABLE `chatgpt_mj_jobs` CHANGE `img_url` `img_url` VARCHAR(400) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL COMMENT '图片URL';
ALTER TABLE `chatgpt_functions` ADD `token` VARCHAR(255) NULL COMMENT 'API授权token' AFTER `action`;

BIN
web/public/images/mj/mj-v6.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 18 KiB

78
web/src/views/ImageMj.vue
View File

@ -359,7 +359,7 @@
<template #default="scope">
<div class="job-item">
<el-image
:src="scope.item.type === 'upscale' ? scope.item['img_url'] + '?imageView2/1/w/480/h/600/q/75' : scope.item['img_url'] + '?imageView2/1/w/480/h/480/q/75'"
:src="scope.item['thumb_url']"
:class="scope.item.type === 'upscale' ? 'upscale' : ''" :zoom-rate="1.2"
:preview-src-list="[scope.item['img_url']]" fit="cover" :initial-index="scope.index"
loading="lazy" v-if="scope.item.progress > 0">
@ -477,7 +477,8 @@ const rates = [
{css: "size9-16", value: "9:16", text: "9:16", img: "/images/mj/rate_9_16.png"},
]
const models = [
{text: "最新模式MJ-5.2", value: " --v 5.2", img: "/images/mj/mj-v5.2.png"},
{text: "写实模式MJ-6.0", value: " --v 6", img: "/images/mj/mj-v6.png"},
{text: "优质模式MJ-5.2", value: " --v 5.2", img: "/images/mj/mj-v5.2.png"},
{text: "优质模式MJ-5.1", value: " --v 5.1", img: "/images/mj/mj-v5.1.jpg"},
{text: "虚幻模式MJ-5", value: " --v 5", img: "/images/mj/mj-v5.jpg"},
{text: "真实模式MJ-4", value: " --v 4", img: "/images/mj/mj-v4.jpg"},
@ -489,6 +490,10 @@ const models = [
]
const options = [
{
value: 0,
label: '默认'
},
{
value: 0.25,
label: '普通'
@ -515,7 +520,7 @@ const params = ref({
prompt: "",
neg_prompt: "",
tile: false,
quality: 0.5
quality: 0
})
const activeName = ref('图生图')
@ -527,6 +532,7 @@ const router = useRouter()
const socket = ref(null)
const imgCalls = ref(0)
const loading = ref(false)
const userId = ref(0)
const rewritePrompt = () => {
loading.value = true
@ -550,12 +556,40 @@ const translatePrompt = () => {
})
}
const connect = () => {
let host = process.env.VUE_APP_WS_HOST
if (host === '') {
if (location.protocol === 'https:') {
host = 'wss://' + location.host;
} else {
host = 'ws://' + location.host;
}
}
const _socket = new WebSocket(host + `/api/mj/client?user_id=${userId.value}`);
_socket.addEventListener('open', () => {
socket.value = _socket;
});
_socket.addEventListener('message', event => {
if (event.data instanceof Blob) {
fetchRunningJobs(userId.value)
fetchFinishJobs(userId.value)
}
});
_socket.addEventListener('close', () => {
connect()
});
}
onMounted(() => {
checkSession().then(user => {
imgCalls.value = user['img_calls']
userId.value = user.id
fetchRunningJobs(user.id)
fetchFinishJobs(user.id)
fetchRunningJobs(userId.value)
fetchFinishJobs(userId.value)
connect()
}).catch(() => {
router.push('/login')
@ -588,39 +622,29 @@ const fetchRunningJobs = (userId) => {
_jobs.push(jobs[i])
}
runningJobs.value = _jobs
setTimeout(() => fetchRunningJobs(userId), 1000)
}).catch(e => {
ElMessage.error("获取任务失败：" + e.message)
setTimeout(() => fetchRunningJobs(userId), 5000)
})
}
const fetchFinishJobs = (userId) => {
// 获取已完成的任务
httpGet(`/api/mj/jobs?status=1&user_id=${userId}`).then(res => {
if (finishedJobs.value.length === 0) {
finishedJobs.value = res.data
return
}
// check if the img url is changed
const list = res.data
let changed = false
for (let i = 0; i < list.length; i++) {
if (list[i]["img_url"] !== finishedJobs.value[i]["img_url"]) {
changed = true
break
const jobs = res.data
for (let i = 0; i < jobs.length; i++) {
if (jobs[i]['use_proxy']) {
jobs[i]['thumb_url'] = jobs[i]['img_url'] + '?x-oss-process=image/quality,q_60&format=webp'
} else {
if (jobs[i].type === 'upscale') {
jobs[i]['thumb_url'] = jobs[i]['img_url'] + '?imageView2/1/w/480/h/600/q/75'
} else {
jobs[i]['thumb_url'] = jobs[i]['img_url'] + '?imageView2/1/w/480/h/480/q/75'
}
}
}
if (changed) {
finishedJobs.value = list
}
setTimeout(() => fetchFinishJobs(userId), 1000)
finishedJobs.value = jobs
}).catch(e => {
ElMessage.error("获取任务失败：" + e.message)
setTimeout(() => fetchFinishJobs(userId), 5000)
})
}
@ -710,7 +734,7 @@ const removeImage = (item) => {
type: 'warning',
}
).then(() => {
httpPost("/api/mj/remove", {id: item.id, img_url: item.img_url}).then(() => {
httpPost("/api/mj/remove", {id: item.id, img_url: item.img_url, user_id: userId.value}).then(() => {
ElMessage.success("任务删除成功")
}).catch(e => {
ElMessage.error("任务删除失败：" + e.message)

3
web/src/views/ImageSd.vue
View File

@ -612,8 +612,9 @@ onMounted(() => {
// 获取已完成的任务
const fetchFinishJobs = (userId) => {
httpGet(`/api/sd/jobs?status=1&user_id=${userId}`).then(res => {
if (finishedJobs.value.length === 0) {
if (finishedJobs.value.length === 0 || res.data.length > finishedJobs.value.length) {
finishedJobs.value = res.data
setTimeout(() => fetchFinishJobs(userId), 1000)
return
}

2
web/src/views/admin/ApiKey.vue
View File

@ -32,7 +32,7 @@
<el-table-column label="操作" width="180">
<template #default="scope">
<el-button size="small" type="primary" @click="edit(scope.row)">编辑</el-button>
<el-popconfirm title="确定要删除当前记录吗?" @confirm="remove(scope.row)">
<el-popconfirm title="确定要删除当前记录吗?" @confirm="remove(scope.row)" :width="200">
<template #reference>
<el-button size="small" type="danger">删除</el-button>
</template>

2
web/src/views/admin/ChatModel.vue
View File

@ -35,7 +35,7 @@
<el-table-column label="操作" width="180">
<template #default="scope">
<el-button size="small" type="primary" @click="edit(scope.row)">编辑</el-button>
<el-popconfirm title="确定要删除当前记录吗?" @confirm="remove(scope.row)">
<el-popconfirm title="确定要删除当前记录吗?" @confirm="remove(scope.row)" :width="200">
<template #reference>
<el-button size="small" type="danger">删除</el-button>
</template>

102
web/src/views/admin/Functions.vue
View File

@ -1,7 +1,7 @@
<template>
<div class="container role-list" v-loading="loading">
<div class="handle-box">
<el-button type="primary" :icon="Plus" @click="addRole">新增</el-button>
<el-button type="primary" :icon="Plus" @click="addRow">新增</el-button>
</div>
<el-row>
<el-table :data="tableData" :border="parentBorder" style="width: 100%">
@ -10,18 +10,18 @@
<span class="sort" :data-id="scope.row.id">{{ scope.row.name }}</span>
</template>
</el-table-column>
<el-table-column label="功能描述" prop="key"/>
<el-table-column label="">
<el-table-column label="函数别名" prop="label"/>
<el-table-column label="功能描述" prop="description"/>
<el-table-column label="启用状态">
<template #default="scope">
<el-tag v-if="scope.row.enable" type="success">启用</el-tag>
<el-tag type="danger" v-else>禁用</el-tag>
<el-switch v-model="scope.row.enabled" @change="functionSet('enabled',scope.row)"/>
</template>
</el-table-column>
<el-table-column label="操作" width="150" align="right">
<template #default="scope">
<el-button size="small" type="primary" @click="rowEdit(scope.$index, scope.row)">编辑</el-button>
<el-popconfirm title="确定要删除当前函数吗?" @confirm="remove(scope.row)">
<el-popconfirm title="确定要删除当前函数吗?" @confirm="remove(scope.row)" :width="200">
<template #reference>
<el-button size="small" type="danger">删除</el-button>
</template>
@ -40,6 +40,7 @@
<el-form-item label="函数名称：" prop="name">
<el-input
v-model="item.name"
placeholder="函数名称最好为英文"
autocomplete="off"
/>
</el-form-item>
@ -114,6 +115,33 @@
</template>
</el-form-item>
<el-form-item label="API 地址：" prop="action">
<el-input
v-model="item.action"
autocomplete="off"
placeholder="该函数实现的API地址，可以是第三方服务API"
/>
</el-form-item>
<el-form-item label="API Token：" prop="token">
<el-input
v-model="item.token"
autocomplete="off"
placeholder="API授权Token"
>
<template #append>
<el-tooltip
class="box-item"
effect="dark"
content="只有本地服务才可以使用自动生成Token<br/>第三方服务请填写第三方服务API Token"
placement="top-end"
raw-content
>
<el-button @click="generateToken">生成Token</el-button>
</el-tooltip>
</template>
</el-input>
</el-form-item>
<el-form-item label="启用状态">
<el-switch v-model="item.enabled"/>
</el-form-item>
@ -131,21 +159,19 @@
<script setup>
import {Delete, Plus, RemoveFilled} from "@element-plus/icons-vue";
import {Delete, Plus} from "@element-plus/icons-vue";
import {onMounted, reactive, ref} from "vue";
import {httpGet, httpPost} from "@/utils/http";
import {ElMessage} from "element-plus";
import {copyObj, removeArrayItem} from "@/utils/libs";
import {Sortable} from "sortablejs"
import {arrayContains, copyObj} from "@/utils/libs";
const showDialog = ref(false)
const parentBorder = ref(true)
const childBorder = ref(true)
const tableData = ref([])
const item = ref({parameters: []})
const item = ref({})
const params = ref([])
const formRef = ref(null)
const editRow = ref({})
const loading = ref(true)
const title = ref("新增函数")
@ -163,7 +189,9 @@ onMounted(() => {
const fetch = () => {
httpGet('/api/admin/function/list').then((res) => {
tableData.value = res.data
if (res.data) {
tableData.value = res.data
}
loading.value = false
}).catch(() => {
ElMessage.error("获取数据失败");
@ -175,11 +203,25 @@ const curIndex = ref(0)
const rowEdit = function (index, row) {
curIndex.value = index
item.value = copyObj(row)
// initialize parameters
const props = item.value?.parameters?.properties
const required = item.value?.parameters?.required
const _params = []
for (let key in props) {
_params.push({
name: key,
type: props[key].type,
desc: props[key].description,
required: arrayContains(required, key)
})
}
params.value = _params
showDialog.value = true
}
const addRole = function () {
item.value = {parameters: []}
const addRow = function () {
item.value = {enabled:true}
params.value = []
showDialog.value = true
}
@ -199,12 +241,12 @@ const save = function () {
item.value.parameters = {type: "object", "properties": properties, "required": required}
httpPost('/api/admin/function/save', item.value).then((res) => {
ElMessage.success('操作成功')
// 更新当前数据行
// if (item.value.id) {
// tableData.value[curIndex.value] = item.value
// } else {
// tableData.value.push(res.data)
// }
console.log(res.data)
if (item.value.id > 0) {
tableData.value[curIndex.value] = item.value
} else {
tableData.value.push(res.data)
}
}).catch((e) => {
ElMessage.error('操作失败，' + e.message)
})
@ -213,7 +255,7 @@ const save = function () {
}
const remove = function (row) {
httpGet('/api/admin/role/remove?id=' + row.id).then(() => {
httpGet('/api/admin/function/remove?id=' + row.id).then(() => {
ElMessage.success("删除成功！")
fetch()
}).catch(() => {
@ -225,13 +267,29 @@ const addParam = function () {
if (!params.value) {
item.value = []
}
params.value.push({name: "", type: "", desc: "", required: false})
params.value.push({name: "", type: "string", desc: "", required: false})
}
const removeParam = function (index) {
params.value.splice(index, 1);
}
const functionSet = (filed,row) => {
httpPost('/api/admin/function/set', {id: row.id, filed: filed, value: row[filed]}).then(() => {
ElMessage.success("操作成功！")
}).catch(e => {
ElMessage.error("操作失败：" + e.message)
})
}
const generateToken = () => {
httpGet('/api/admin/function/token').then(res => {
item.value.token = res.data
}).catch(() => {
ElMessage.error("生成 Token 失败")
})
}
</script>
<style lang="stylus" scoped>

2
web/src/views/admin/Product.vue
View File

@ -38,7 +38,7 @@
<el-table-column label="操作" width="180">
<template #default="scope">
<el-button size="small" type="primary" @click="edit(scope.row)">编辑</el-button>
<el-popconfirm title="确定要删除当前记录吗?" @confirm="remove(scope.row)">
<el-popconfirm title="确定要删除当前记录吗?" @confirm="remove(scope.row)" :width="200">
<template #reference>
<el-button size="small" type="danger">删除</el-button>
</template>

2
web/src/views/admin/RoleList.vue
View File

@ -36,7 +36,7 @@
<el-table-column label="操作" width="150" align="right">
<template #default="scope">
<el-button size="small" type="primary" @click="rowEdit(scope.$index, scope.row)">编辑</el-button>
<el-popconfirm title="确定要删除当前角色吗?" @confirm="removeRole(scope.row)">
<el-popconfirm title="确定要删除当前角色吗?" @confirm="removeRole(scope.row)" :width="200">
<template #reference>
<el-button size="small" type="danger">删除</el-button>
</template>