merge v4.0.6

api/service/mj/pool.go

@@ -16,6 +16,7 @@ import (
 	"geekai/store"
 	"geekai/store/model"
 	"github.com/go-redis/redis/v8"
+	"strings"
 	"time"
 
 	"gorm.io/gorm"
@@ -33,38 +34,10 @@ type ServicePool struct {
 
 var logger = logger2.GetLogger()
 
-func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, appConfig *types.AppConfig) *ServicePool {
+func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager) *ServicePool {
 	services := make([]*Service, 0)
 	taskQueue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
 	notifyQueue := store.NewRedisQueue("MidJourney_Notify_Queue", redisCli)
-
-	for k, config := range appConfig.MjPlusConfigs {
-		if config.Enabled == false {
-			continue
-		}
-
-		cli := NewPlusClient(config)
-		name := fmt.Sprintf("mj-plus-service-%d", k)
-		plusService := NewService(name, taskQueue, notifyQueue, db, cli)
-		go func() {
-			plusService.Run()
-		}()
-		services = append(services, plusService)
-	}
-
-	for k, config := range appConfig.MjProxyConfigs {
-		if config.Enabled == false {
-			continue
-		}
-		cli := NewProxyClient(config)
-		name := fmt.Sprintf("mj-proxy-service-%d", k)
-		proxyService := NewService(name, taskQueue, notifyQueue, db, cli)
-		go func() {
-			proxyService.Run()
-		}()
-		services = append(services, proxyService)
-	}
-
 	return &ServicePool{
 		taskQueue:       taskQueue,
 		notifyQueue:     notifyQueue,
@@ -75,6 +48,42 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 	}
 }
 
+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 k, config := range plusConfigs {
+		if config.Enabled == false {
+			continue
+		}
+
+		cli := NewPlusClient(config)
+		name := fmt.Sprintf("mj-plus-service-%d", k)
+		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 k, config := range proxyConfigs {
+		if config.Enabled == false {
+			continue
+		}
+		cli := NewProxyClient(config)
+		name := fmt.Sprintf("mj-proxy-service-%d", k)
+		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 {
@@ -111,17 +120,23 @@ func (p *ServicePool) DownloadImages() {
 				}
 
 				logger.Infof("try to download image: %s", v.OrgURL)
-				var imgURL string
-				var err error
-				if servicePlus := p.getService(v.ChannelId); servicePlus != nil {
-					task, _ := servicePlus.Client.QueryTask(v.TaskId)
-					if len(task.Buttons) > 0 {
-						v.Hash = GetImageHash(task.Buttons[0].CustomId)
-					}
-					imgURL, err = p.uploaderManager.GetUploadHandler().PutImg(v.OrgURL, false)
-				} else {
-					imgURL, err = p.uploaderManager.GetUploadHandler().PutImg(v.OrgURL, true)
+				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().PutImg(v.OrgURL, proxy)
+
 				if err != nil {
 					logger.Errorf("error with download image %s, %v", v.OrgURL, err)
 					continue

api/service/mj/service.go

@@ -28,6 +28,7 @@ type Service struct {
 	taskQueue   *store.RedisQueue
 	notifyQueue *store.RedisQueue
 	db          *gorm.DB
+	running     bool
 }
 
 func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, cli Client) *Service {
@@ -37,12 +38,13 @@ func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.Red
 		taskQueue:   taskQueue,
 		notifyQueue: notifyQueue,
 		Client:      cli,
+		running:     true,
 	}
 }
 
 func (s *Service) Run() {
 	logger.Infof("Starting MidJourney job consumer for %s", s.Name)
-	for {
+	for s.running {
 		var task types.MjTask
 		err := s.taskQueue.LPop(&task)
 		if err != nil {
@@ -125,6 +127,10 @@ func (s *Service) Run() {
 	}
 }
 
+func (s *Service) Stop() {
+	s.running = false
+}
+
 type CBReq struct {
 	Id          string      `json:"id"`
 	Action      string      `json:"action"`