upgrade to v4.0.4

api/service/captcha_service.go

@@ -9,11 +9,11 @@ import (
 )
 
 type CaptchaService struct {
-	config types.ChatPlusApiConfig
+	config types.ApiConfig
 	client *req.Client
 }
 
-func NewCaptchaService(config types.ChatPlusApiConfig) *CaptchaService {
+func NewCaptchaService(config types.ApiConfig) *CaptchaService {
 	return &CaptchaService{
 		config: config,
 		client: req.C().SetTimeout(10 * time.Second),

api/service/dalle/service.go

@@ -0,0 +1,300 @@
+package dalle
+
+import (
+	"chatplus/core/types"
+	logger2 "chatplus/logger"
+	"chatplus/service"
+	"chatplus/service/oss"
+	"chatplus/service/sd"
+	"chatplus/store"
+	"chatplus/store/model"
+	"chatplus/utils"
+	"errors"
+	"fmt"
+	"github.com/go-redis/redis/v8"
+	"time"
+
+	"github.com/imroc/req/v3"
+	"gorm.io/gorm"
+)
+
+var logger = logger2.GetLogger()
+
+// DALL-E 绘画服务
+
+type Service struct {
+	httpClient    *req.Client
+	db            *gorm.DB
+	uploadManager *oss.UploaderManager
+	taskQueue     *store.RedisQueue
+	notifyQueue   *store.RedisQueue
+	Clients       *types.LMap[uint, *types.WsClient] // UserId => Client
+}
+
+func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client) *Service {
+	return &Service{
+		httpClient:    req.C().SetTimeout(time.Minute * 3),
+		db:            db,
+		taskQueue:     store.NewRedisQueue("DallE_Task_Queue", redisCli),
+		notifyQueue:   store.NewRedisQueue("DallE_Notify_Queue", redisCli),
+		Clients:       types.NewLMap[uint, *types.WsClient](),
+		uploadManager: manager,
+	}
+}
+
+// PushTask push a new mj task in to task queue
+func (s *Service) PushTask(task types.DallTask) {
+	logger.Debugf("add a new MidJourney task to the task list: %+v", task)
+	s.taskQueue.RPush(task)
+}
+
+func (s *Service) Run() {
+	go func() {
+		for {
+			var task types.DallTask
+			err := s.taskQueue.LPop(&task)
+			if err != nil {
+				logger.Errorf("taking task with error: %v", err)
+				continue
+			}
+
+			_, err = s.Image(task, false)
+			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,
+					"err_msg":  err.Error(),
+				})
+				s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Failed})
+			}
+		}
+	}()
+}
+
+type imgReq struct {
+	Model   string `json:"model"`
+	Prompt  string `json:"prompt"`
+	N       int    `json:"n"`
+	Size    string `json:"size"`
+	Quality string `json:"quality"`
+	Style   string `json:"style"`
+}
+
+type imgRes struct {
+	Created int64 `json:"created"`
+	Data    []struct {
+		RevisedPrompt string `json:"revised_prompt"`
+		Url           string `json:"url"`
+	} `json:"data"`
+}
+
+type ErrRes struct {
+	Error struct {
+		Code    interface{} `json:"code"`
+		Message string      `json:"message"`
+		Param   interface{} `json:"param"`
+		Type    string      `json:"type"`
+	} `json:"error"`
+}
+
+func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
+	logger.Debugf("绘画参数：%+v", task)
+	prompt := task.Prompt
+	// translate prompt
+	if utils.HasChinese(task.Prompt) {
+		content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Prompt))
+		if err != nil {
+			return "", fmt.Errorf("error with translate prompt: %v", err)
+		}
+		prompt = content
+		logger.Debugf("重写后提示词：%s", prompt)
+	}
+
+	var user model.User
+	s.db.Where("id", task.UserId).First(&user)
+	if user.Power < task.Power {
+		return "", errors.New("insufficient of power")
+	}
+
+	// get image generation API KEY
+	var apiKey model.ApiKey
+	tx := s.db.Where("platform", types.OpenAI).
+		Where("type", "img").
+		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)
+	}
+
+	var res imgRes
+	var errRes ErrRes
+	if len(apiKey.ProxyURL) > 5 {
+		s.httpClient.SetProxyURL(apiKey.ProxyURL).R()
+	}
+	logger.Debugf("Sending %s request, ApiURL:%s, API KEY:%s, PROXY: %s", apiKey.Platform, apiKey.ApiURL, apiKey.Value, apiKey.ProxyURL)
+	r, err := s.httpClient.R().SetHeader("Content-Type", "application/json").
+		SetHeader("Authorization", "Bearer "+apiKey.Value).
+		SetBody(imgReq{
+			Model:   "dall-e-3",
+			Prompt:  prompt,
+			N:       1,
+			Size:    "1024x1024",
+			Style:   task.Style,
+			Quality: task.Quality,
+		}).
+		SetErrorResult(&errRes).
+		SetSuccessResult(&res).Post(apiKey.ApiURL)
+	if err != nil {
+		return "", fmt.Errorf("error with send request: %v", err)
+	}
+
+	if r.IsErrorState() {
+		return "", fmt.Errorf("error with send request: %v", errRes.Error)
+	}
+	// update the api key last use time
+	s.db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
+	// update task progress
+	s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
+		"progress": 100,
+		"org_url":  res.Data[0].Url,
+		"prompt":   prompt,
+	})
+
+	s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Finished})
+	var content string
+	if sync {
+		imgURL, err := s.downloadImage(task.JobId, int(task.UserId), res.Data[0].Url)
+		if err != nil {
+			return "", fmt.Errorf("error with download image: %v", err)
+		}
+		content = fmt.Sprintf("```\n%s\n```\n下面是我为你创作的图片：\n\n![](%s)\n", prompt, imgURL)
+	}
+
+	// 更新用户算力
+	tx = s.db.Model(&model.User{}).Where("id", user.Id).UpdateColumn("power", gorm.Expr("power - ?", task.Power))
+	// 记录算力变化日志
+	if tx.Error == nil && tx.RowsAffected > 0 {
+		var u model.User
+		s.db.Where("id", user.Id).First(&u)
+		s.db.Create(&model.PowerLog{
+			UserId:    user.Id,
+			Username:  user.Username,
+			Type:      types.PowerConsume,
+			Amount:    task.Power,
+			Balance:   u.Power,
+			Mark:      types.PowerSub,
+			Model:     "dall-e-3",
+			Remark:    fmt.Sprintf("绘画提示词：%s", utils.CutWords(task.Prompt, 10)),
+			CreatedAt: time.Now(),
+		})
+	}
+
+	return content, nil
+}
+
+func (s *Service) CheckTaskNotify() {
+	go func() {
+		logger.Info("Running DALL-E task notify checking ...")
+		for {
+			var message sd.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
+			}
+		}
+	}()
+}
+
+func (s *Service) DownloadImages() {
+	go func() {
+		var items []model.DallJob
+		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)
+				imgURL, err := s.downloadImage(v.Id, int(v.UserId), v.OrgURL)
+				if err != nil {
+					logger.Error("error with download image: %s, error: %v", imgURL, err)
+					continue
+				}
+
+			}
+
+			time.Sleep(time.Second * 5)
+		}
+	}()
+}
+
+func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string, error) {
+	// sava image
+	imgURL, err := s.uploadManager.GetUploadHandler().PutImg(orgURL, false)
+	if err != nil {
+		return "", err
+	}
+
+	// update img_url
+	res := s.db.Model(&model.DallJob{Id: jobId}).UpdateColumn("img_url", imgURL)
+	if res.Error != nil {
+		return "", err
+	}
+	s.notifyQueue.RPush(sd.NotifyMessage{UserId: userId, JobId: int(jobId), Message: sd.Failed})
+	return imgURL, nil
+}
+
+// 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 == -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：%s", job.TaskId),
+							CreatedAt: time.Now(),
+						})
+					}
+					continue
+				}
+			}
+			time.Sleep(time.Second * 10)
+		}
+	}()
+}

api/service/license_service.go

@@ -0,0 +1,108 @@
+package service
+
+import (
+	"chatplus/core"
+	"chatplus/core/types"
+	"chatplus/store"
+	"errors"
+	"fmt"
+	"github.com/imroc/req/v3"
+	"github.com/shirou/gopsutil/host"
+	"strings"
+	"time"
+)
+
+type LicenseService struct {
+	config types.ApiConfig
+	levelDB *store.LevelDB
+	license types.License
+	machineId string
+}
+
+func NewLicenseService(server *core.AppServer, levelDB *store.LevelDB) * LicenseService {
+	var license types.License
+	var machineId string
+	_ = levelDB.Get(types.LicenseKey, &license)
+	info, err := host.Info()
+	if err == nil {
+		machineId = info.HostID
+	}
+	return &LicenseService{
+		config: server.Config.ApiConfig,
+		levelDB: levelDB,
+		license: license,
+		machineId: machineId,
+	}
+}
+
+// ActiveLicense 激活 License
+func (s *LicenseService) ActiveLicense(license string, machineId string) error {
+	var res struct {
+		Code    types.BizCode `json:"code"`
+		Message string        `json:"message"`
+		Data    struct {
+			Name      string `json:"name"`
+			License   string `json:"license"`
+			Mid       string `json:"mid"`
+			ExpiredAt int64  `json:"expired_at"`
+			UserNum   int    `json:"user_num"`
+		}
+	}
+	apiURL := fmt.Sprintf("%s/%s", s.config.ApiURL, "api/license/active")
+	response, err := req.C().R().
+		SetBody(map[string]string{"license": license, "machine_id": machineId}).
+		SetSuccessResult(&res).Post(apiURL)
+	if err != nil {
+		return fmt.Errorf("发送激活请求失败: %v", err)
+	}
+
+	if response.IsErrorState() {
+		return fmt.Errorf( "发送激活请求失败：%v", response.Status)
+	}
+
+	if res.Code != types.Success {
+		return fmt.Errorf( "激活失败：%v", res.Message)
+	}
+
+	err = s.levelDB.Put(types.LicenseKey, types.License{
+		Key:       license,
+		MachineId: machineId,
+		UserNum:   res.Data.UserNum,
+		ExpiredAt: res.Data.ExpiredAt,
+		IsActive: true,
+	})
+	if err != nil {
+		return fmt.Errorf("保存许可证书失败：%v", err)
+	}
+
+	return nil
+}
+
+// GetLicense 获取许可信息
+func (s *LicenseService) GetLicense() types.License {
+	return s.license
+}
+
+// IsValidApiURL 判断是否合法的中转 URL
+func (s *LicenseService) IsValidApiURL(uri string) error {
+	// 获得许可授权的直接放行
+	if s.license.IsActive {
+		if s.license.MachineId != s.machineId {
+			return errors.New("系统使用了盗版的许可证书")
+		}
+
+		if time.Now().Unix() > s.license.ExpiredAt {
+			return errors.New("系统许可证书已经过期")
+		}
+		return nil
+	}
+
+	if !strings.HasPrefix(uri, "https://gpt.bemore.lol") &&
+		!strings.HasPrefix(uri, "https://api.openai.com") &&
+		!strings.HasPrefix(uri, "http://cdn.chat-plus.net") &&
+		!strings.HasPrefix(uri, "https://api.chat-plus.net") {
+		return fmt.Errorf("当前 API 地址 %s 不在白名单列表当中。",uri)
+	}
+
+	return nil
+}

api/service/mj/plus_client.go

@@ -73,6 +73,7 @@ func (c *PlusClient) Imagine(task types.MjTask) (ImageRes, error) {
 // Blend 融图
 func (c *PlusClient) Blend(task types.MjTask) (ImageRes, error) {
 	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",
@@ -164,6 +165,7 @@ func (c *PlusClient) Upscale(task types.MjTask) (ImageRes, error) {
 		"taskId":   task.MessageId,
 	}
 	apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/action", c.apiURL, c.Config.Mode)
+	logger.Info("API URL: ", apiURL)
 	var res ImageRes
 	var errRes ErrRes
 	r, err := c.client.R().
@@ -190,6 +192,7 @@ func (c *PlusClient) Variation(task types.MjTask) (ImageRes, error) {
 		"taskId":   task.MessageId,
 	}
 	apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/action", c.apiURL, c.Config.Mode)
+	logger.Info("API URL: ", apiURL)
 	var res ImageRes
 	var errRes ErrRes
 	r, err := req.C().R().

api/service/mj/pool.go

@@ -3,7 +3,9 @@ package mj
 import (
 	"chatplus/core/types"
 	logger2 "chatplus/logger"
+	"chatplus/service"
 	"chatplus/service/oss"
+	"chatplus/service/sd"
 	"chatplus/store"
 	"chatplus/store/model"
 	"fmt"
@@ -25,7 +27,7 @@ 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, appConfig *types.AppConfig, licenseService *service.LicenseService) *ServicePool {
 	services := make([]*Service, 0)
 	taskQueue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
 	notifyQueue := store.NewRedisQueue("MidJourney_Notify_Queue", redisCli)
@@ -34,13 +36,19 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 		if config.Enabled == false {
 			continue
 		}
+		err := licenseService.IsValidApiURL(config.ApiURL)
+		if err != nil {
+			logger.Error(err)
+			continue
+		}
+		
 		cli := NewPlusClient(config)
 		name := fmt.Sprintf("mj-plus-service-%d", k)
-		service := NewService(name, taskQueue, notifyQueue, db, cli)
+		plusService := NewService(name, taskQueue, notifyQueue, db, cli)
 		go func() {
-			service.Run()
+			plusService.Run()
 		}()
-		services = append(services, service)
+		services = append(services, plusService)
 	}
 
 	for k, config := range appConfig.MjProxyConfigs {
@@ -49,11 +57,11 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 		}
 		cli := NewProxyClient(config)
 		name := fmt.Sprintf("mj-proxy-service-%d", k)
-		service := NewService(name, taskQueue, notifyQueue, db, cli)
+		proxyService := NewService(name, taskQueue, notifyQueue, db, cli)
 		go func() {
-			service.Run()
+			proxyService.Run()
 		}()
-		services = append(services, service)
+		services = append(services, proxyService)
 	}
 
 	return &ServicePool{
@@ -69,16 +77,16 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 func (p *ServicePool) CheckTaskNotify() {
 	go func() {
 		for {
-			var userId uint
-			err := p.notifyQueue.LPop(&userId)
+			var message sd.NotifyMessage
+			err := p.notifyQueue.LPop(&message)
 			if err != nil {
 				continue
 			}
-			cli := p.Clients.Get(userId)
+			cli := p.Clients.Get(uint(message.UserId))
 			if cli == nil {
 				continue
 			}
-			err = cli.Send([]byte("Task Updated"))
+			err = cli.Send([]byte(message.Message))
 			if err != nil {
 				continue
 			}
@@ -127,7 +135,7 @@ func (p *ServicePool) DownloadImages() {
 				if cli == nil {
 					continue
 				}
-				err = cli.Send([]byte("Task Updated"))
+				err = cli.Send([]byte(sd.Finished))
 				if err != nil {
 					continue
 				}
@@ -162,7 +170,6 @@ func (p *ServicePool) SyncTaskProgress() {
 			for _, job := range items {
 				// 失败或者 30 分钟还没完成的任务删除并退回算力
 				if time.Now().Sub(job.CreatedAt) > time.Minute*30 || job.Progress == -1 {
-					// 删除任务
 					p.db.Delete(&job)
 					// 退回算力
 					tx := p.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
@@ -189,7 +196,7 @@ func (p *ServicePool) SyncTaskProgress() {
 				}
 			}
 
-			time.Sleep(time.Second)
+			time.Sleep(time.Second * 10)
 		}
 	}()
 }

api/service/mj/service.go

@@ -3,6 +3,7 @@ package mj
 import (
 	"chatplus/core/types"
 	"chatplus/service"
+	"chatplus/service/sd"
 	"chatplus/store"
 	"chatplus/store/model"
 	"chatplus/utils"
@@ -53,7 +54,7 @@ func (s *Service) Run() {
 
 		// translate prompt
 		if utils.HasChinese(task.Prompt) {
-			content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt))
+			content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Prompt))
 			if err == nil {
 				task.Prompt = content
 			} else {
@@ -62,7 +63,7 @@ func (s *Service) Run() {
 		}
 		// translate negative prompt
 		if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) {
-			content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.NegPrompt))
+			content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.NegPrompt))
 			if err == nil {
 				task.NegPrompt = content
 			} else {
@@ -105,7 +106,7 @@ func (s *Service) Run() {
 			// update the task progress
 			s.db.Updates(&job)
 			// 任务失败，通知前端
-			s.notifyQueue.RPush(task.UserId)
+			s.notifyQueue.RPush(sd.NotifyMessage{UserId: task.UserId, JobId: int(job.Id), Message: sd.Failed})
 			continue
 		}
 		logger.Infof("任务提交成功：%+v", res)
@@ -147,7 +148,7 @@ func (s *Service) Notify(job model.MidJourneyJob) error {
 			"progress": -1,
 			"err_msg":  task.FailReason,
 		})
-		s.notifyQueue.RPush(job.UserId)
+		s.notifyQueue.RPush(sd.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: sd.Failed})
 		return fmt.Errorf("task failed: %v", task.FailReason)
 	}
 
@@ -166,7 +167,11 @@ func (s *Service) Notify(job model.MidJourneyJob) error {
 	}
 	// 通知前端更新任务进度
 	if oldProgress != job.Progress {
-		s.notifyQueue.RPush(job.UserId)
+		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
 }

api/service/sd/pool.go

@@ -60,16 +60,16 @@ func (p *ServicePool) CheckTaskNotify() {
 	go func() {
 		logger.Info("Running Stable-Diffusion task notify checking ...")
 		for {
-			var userId uint
-			err := p.notifyQueue.LPop(&userId)
+			var message NotifyMessage
+			err := p.notifyQueue.LPop(&message)
 			if err != nil {
 				continue
 			}
-			client := p.Clients.Get(userId)
+			client := p.Clients.Get(uint(message.UserId))
 			if client == nil {
 				continue
 			}
-			err = client.Send([]byte("Task Updated"))
+			err = client.Send([]byte(message.Message))
 			if err != nil {
 				continue
 			}
@@ -113,7 +113,7 @@ func (p *ServicePool) CheckTaskStatus() {
 					continue
 				}
 			}
-
+			time.Sleep(time.Second * 10)
 		}
 	}()
 }

api/service/sd/service.go

@@ -8,10 +8,11 @@ import (
 	"chatplus/store/model"
 	"chatplus/utils"
 	"fmt"
-	"github.com/imroc/req/v3"
-	"gorm.io/gorm"
 	"strings"
 	"time"
+
+	"github.com/imroc/req/v3"
+	"gorm.io/gorm"
 )
 
 // SD 绘画服务
@@ -80,7 +81,7 @@ func (s *Service) Run() {
 				"err_msg":  err.Error(),
 			})
 			// 通知前端，任务失败
-			s.notifyQueue.RPush(task.UserId)
+			s.notifyQueue.RPush(NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: Failed})
 			continue
 		}
 	}
@@ -145,8 +146,13 @@ func (s *Service) Txt2Img(task types.SdTask) error {
 	var errChan = make(chan error)
 	apiURL := fmt.Sprintf("%s/sdapi/v1/txt2img", s.config.ApiURL)
 	logger.Debugf("send image request to %s", apiURL)
+	// send a request to sd api endpoint
 	go func() {
-		response, err := s.httpClient.R().SetBody(body).SetSuccessResult(&res).Post(apiURL)
+		response, err := s.httpClient.R().
+			SetHeader("Authorization", s.config.ApiKey).
+			SetBody(body).
+			SetSuccessResult(&res).
+			Post(apiURL)
 		if err != nil {
 			errChan <- err
 			return
@@ -174,14 +180,17 @@ func (s *Service) Txt2Img(task types.SdTask) error {
 		errChan <- nil
 	}()
 
+	// waiting for task finish
 	for {
 		select {
-		case err := <-errChan: // 任务完成
+		case err := <-errChan:
 			if err != nil {
 				return err
 			}
+
+			// task finished
 			s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", 100)
-			s.notifyQueue.RPush(task.UserId)
+			s.notifyQueue.RPush(NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: Finished})
 			// 从 leveldb 中删除预览图片数据
 			_ = s.leveldb.Delete(task.Params.TaskId)
 			return nil
@@ -191,7 +200,7 @@ func (s *Service) Txt2Img(task types.SdTask) error {
 			if err == nil && resp.Progress > 0 {
 				s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", int(resp.Progress*100))
 				// 发送更新状态信号
-				s.notifyQueue.RPush(task.UserId)
+				s.notifyQueue.RPush(NotifyMessage{UserId: task.UserId, JobId: task.Id, Message: Running})
 				// 保存预览图片数据
 				if resp.CurrentImage != "" {
 					_ = s.leveldb.Put(task.Params.TaskId, resp.CurrentImage)
@@ -207,7 +216,10 @@ 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)
 	var res TaskProgressResp
-	response, err := s.httpClient.R().SetSuccessResult(&res).Get(apiURL)
+	response, err := s.httpClient.R().
+		SetHeader("Authorization", s.config.ApiKey).
+		SetSuccessResult(&res).
+		Get(apiURL)
 	if err != nil {
 		return err, nil
 	}

api/service/sd/types.go

@@ -4,44 +4,14 @@ import logger2 "chatplus/logger"
 
 var logger = logger2.GetLogger()
 
-type TaskInfo struct {
-	UserId      uint          `json:"user_id"`
-	SessionId   string        `json:"session_id"`
-	JobId       int           `json:"job_id"`
-	TaskId      string        `json:"task_id"`
-	Data        []interface{} `json:"data"`
-	EventData   interface{}   `json:"event_data"`
-	FnIndex     int           `json:"fn_index"`
-	SessionHash string        `json:"session_hash"`
+type NotifyMessage struct {
+	UserId  int    `json:"user_id"`
+	JobId   int    `json:"job_id"`
+	Message string `json:"message"`
 }
 
-type CBReq struct {
-	UserId    uint
-	SessionId string
-	JobId     int
-	TaskId    string
-	ImageName string
-	ImageData string
-	Progress  int
-	Seed      int64
-	Success   bool
-	Message   string
-}
-
-var ParamKeys = map[string]int{
-	"task_id":         0,
-	"prompt":          1,
-	"negative_prompt": 2,
-	"steps":           4,
-	"sampler":         5,
-	"face_fix":        7, // 面部修复
-	"cfg_scale":       8,
-	"seed":            27,
-	"height":          10,
-	"width":           9,
-	"hd_fix":          11,
-	"hd_redraw_rate":  12, //高清修复重绘幅度
-	"hd_scale":        13, // 高清修复放大倍数
-	"hd_scale_alg":    14, // 高清修复放大算法
-	"hd_sample_num":   15, // 高清修复采样次数
-}
+const (
+	Running  = "RUNNING"
+	Finished = "FINISH"
+	Failed   = "FAIL"
+)

api/service/smtp_sms_service.go

@@ -3,9 +3,11 @@ package service
 import (
 	"bytes"
 	"chatplus/core/types"
+	"crypto/tls"
 	"fmt"
 	"mime"
 	"net/smtp"
+	"net/textproto"
 )
 
 type SmtpService struct {
@@ -19,12 +21,18 @@ func NewSmtpService(appConfig *types.AppConfig) *SmtpService {
 }
 
 func (s *SmtpService) SendVerifyCode(to string, code int) error {
-	subject := "ChatPlus注册验证码"
-	body := fmt.Sprintf("您正在注册 ChatPlus AI 助手账户，注册验证码为 %d，请不要告诉他人。如非本人操作，请忽略此邮件。", code)
+	subject := "Geek-AI 注册验证码"
+	body := fmt.Sprintf("您正在注册 Geek-AI 助手账户，注册验证码为 %d，请不要告诉他人。如非本人操作，请忽略此邮件。", code)
 
-	// 设置SMTP客户端配置
 	auth := smtp.PlainAuth("", s.config.From, s.config.Password, s.config.Host)
+	if s.config.UseTls {
+		return s.sendTLS(auth, to, subject, body)
+	} else {
+		return s.send(auth, to, subject, body)
+	}
+}
 
+func (s *SmtpService) send(auth smtp.Auth, to string, subject string, body string) error {
 	// 对主题进行MIME编码
 	encodedSubject := mime.QEncoding.Encode("UTF-8", subject)
 	// 组装邮件
@@ -34,11 +42,83 @@ func (s *SmtpService) SendVerifyCode(to string, code int) error {
 	message.WriteString(fmt.Sprintf("Subject: %s\r\n", encodedSubject))
 	message.WriteString("\r\n" + body)
 
-	// 发送邮件
 	// 发送邮件
 	err := smtp.SendMail(s.config.Host+":"+fmt.Sprint(s.config.Port), auth, s.config.From, []string{to}, message.Bytes())
 	if err != nil {
 		return fmt.Errorf("error sending email: %v", err)
 	}
+
+	return err
+
+}
+
+func (s *SmtpService) sendTLS(auth smtp.Auth, to string, subject string, body string) error {
+	// TLS配置
+	tlsConfig := &tls.Config{
+		ServerName: s.config.Host,
+	}
+
+	// 建立TLS连接
+	conn, err := tls.Dial("tcp", fmt.Sprintf("%s:%d", s.config.Host, s.config.Port), tlsConfig)
+	if err != nil {
+		return fmt.Errorf("error connecting to SMTP server: %v", err)
+	}
+	defer conn.Close()
+
+	client, err := smtp.NewClient(conn, s.config.Host)
+	if err != nil {
+		return fmt.Errorf("error creating SMTP client: %v", err)
+	}
+	defer client.Quit()
+
+	// 身份验证
+	if err = client.Auth(auth); err != nil {
+		return fmt.Errorf("error authenticating: %v", err)
+	}
+
+	// 设置寄件人
+	if err = client.Mail(s.config.From); err != nil {
+		return fmt.Errorf("error setting sender: %v", err)
+	}
+
+	// 设置收件人
+	if err = client.Rcpt(to); err != nil {
+		return fmt.Errorf("error setting recipient: %v", err)
+	}
+
+	// 发送邮件内容
+	wc, err := client.Data()
+	if err != nil {
+		return fmt.Errorf("error getting data writer: %v", err)
+	}
+	defer wc.Close()
+
+	header := make(textproto.MIMEHeader)
+	header.Set("From", s.config.From)
+	header.Set("To", to)
+	header.Set("Subject", subject)
+
+	// 将邮件头写入
+	for key, values := range header {
+		for _, value := range values {
+			_, err = fmt.Fprintf(wc, "%s: %s\r\n", key, value)
+			if err != nil {
+				return fmt.Errorf("error sending email header: %v", err)
+			}
+		}
+	}
+	_, _ = fmt.Fprintln(wc)
+	// 将邮件内容写入
+	_, err = fmt.Fprintf(wc, body)
+	if err != nil {
+		return fmt.Errorf("error sending email: %v", err)
+	}
+
+	// 发送完毕
+	err = wc.Close()
+	if err != nil {
+		return fmt.Errorf("error closing data writer: %v", err)
+	}
+
 	return nil
 }

api/service/types.go

@@ -1,4 +1,4 @@
 package service
 
-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 elements. Please output directly in English without any explanation, within 150 words. The text to be rewritten is: [%s]"
+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]"