geekai/api/service/mj/pool.go

package mj

import (
	"chatplus/core/types"
	"chatplus/service/oss"
	"chatplus/store"
	"chatplus/store/model"
	"fmt"
	"time"

	"github.com/go-redis/redis/v8"
	"gorm.io/gorm"
)

// ServicePool Mj service pool
type ServicePool struct {
	services        []*Service
	taskQueue       *store.RedisQueue
	db              *gorm.DB
	uploaderManager *oss.UploaderManager
}

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)
	// create mj client and service
	for k, config := range appConfig.MjConfigs {
		if config.Enabled == false {
			continue
		}
		// create mj client
		client := NewClient(config, appConfig.ProxyURL)

		name := fmt.Sprintf("MjService-%d", k)
		// create mj service
		service := NewService(name, queue, 4, 600, db, client)
		botName := fmt.Sprintf("MjBot-%d", k)
		bot, err := NewBot(botName, appConfig.ProxyURL, &config, service)
		if err != nil {
			continue
		}

		err = bot.Run()
		if err != nil {
			continue
		}

		// run mj service
		go func() {
			service.Run()
		}()

		services = append(services, service)
	}

	return &ServicePool{
		taskQueue:       queue,
		services:        services,
		uploaderManager: manager,
		db:              db,
	}
}

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 _, item := range items {
				imgURL, err := p.uploaderManager.GetUploadHandler().PutImg(item.OrgURL, true)
				if err != nil {
					logger.Error("error with download image: ", err)
					continue
				}

				item.ImgURL = imgURL
				p.db.Updates(&item)
			}

			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
}