!68 !67 fix: 修复管线context取消、添加FIFO任务队列、移除重复提示词优化

Merge pull request !68 from 何朝晖/develop
This commit is contained in:
2026-05-25 11:13:54 +00:00
committed by Gitee
10 changed files with 210 additions and 43 deletions
+5
View File
@@ -44,6 +44,11 @@ func main() {
// 初始化项目服务
service.InitProjectService(db.GetDB())
// 初始化生成任务队列(FIFO 串行执行)
generateQueue := service.NewTaskQueue()
handler.SetGenerateQueue(generateQueue)
defer generateQueue.Stop()
r := gin.New()
r.Use(mildware.Logger())
r.Use(mildware.Recovery())
+3
View File
@@ -71,6 +71,7 @@ type QiniuConfig struct {
Bucket string `mapstructure:"bucket"`
CDNHost string `mapstructure:"cdn_host"`
UseHTTPS bool `mapstructure:"use_https"`
URLExpire int64 `mapstructure:"url_expire"` // 私有下载签名 URL 有效期(秒),默认 3600
}
// Load 从 YAML 配置文件和环境变量加载配置。
@@ -142,6 +143,7 @@ func setDefaults(v *viper.Viper) {
v.SetDefault("qiniu.bucket", "")
v.SetDefault("qiniu.cdn_host", "")
v.SetDefault("qiniu.use_https", true)
v.SetDefault("qiniu.url_expire", int64(3600))
}
func bindEnvVars(v *viper.Viper) {
@@ -179,4 +181,5 @@ func bindEnvVars(v *viper.Viper) {
v.BindEnv("qiniu.bucket", "GEN2D_QINIU_BUCKET")
v.BindEnv("qiniu.cdn_host", "GEN2D_QINIU_CDN_HOST")
v.BindEnv("qiniu.use_https", "GEN2D_QINIU_USE_HTTPS")
v.BindEnv("qiniu.url_expire", "GEN2D_QINIU_URL_EXPIRE")
}
+8
View File
@@ -31,3 +31,11 @@ image_gen:
width: 1024
height: 1024
quality: "low"
qiniu:
access_key: ""
secret_key: ""
bucket: ""
cdn_host: ""
use_https: true
url_expire: 3600 # 私有下载签名 URL 有效期(秒)
+30 -10
View File
@@ -16,6 +16,14 @@ import (
"github.com/gin-gonic/gin"
)
// generateQueue 全局生成任务队列,由 main 通过 SetGenerateQueue 注入。
var generateQueue *service.TaskQueue
// SetGenerateQueue 设置生成任务队列。
func SetGenerateQueue(q *service.TaskQueue) {
generateQueue = q
}
// GenerateRequest 素材生成请求。
type GenerateRequest struct {
ProjectID string `json:"projectId"`
@@ -43,7 +51,7 @@ type AssetsResponse struct {
}
// Generate 素材生成接口(异步)。
// 立即返回 taskId,后台执行管线,前端通过 GET /tasks/:taskId 轮询进度。
// 立即返回 taskId,任务进入 FIFO 队列串行执行,前端通过 GET /tasks/:taskId 轮询进度。
func Generate(c *gin.Context) {
var req GenerateRequest
if err := c.ShouldBindJSON(&req); err != nil {
@@ -57,18 +65,29 @@ func Generate(c *gin.Context) {
}
taskID := fmt.Sprintf("task-%d", time.Now().UnixMilli())
// 保存任务到数据库
// 保存任务到数据库,初始状态为 pending
if err := saveTaskToDB(c.Request.Context(), projectID, taskID, req); err != nil {
logger.FromCtx(c.Request.Context()).Error("failed to save task", "error", err)
c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "创建任务失败"))
return
}
// 返回 taskId
c.JSON(http.StatusOK, model.OK(GenerateResponse{TaskID: taskID}))
// 加入 FIFO 队列
queuePos := 1
if generateQueue != nil {
queuePos = generateQueue.Enqueue(&service.TaskJob{
Ctx: context.Background(),
ProjectID: projectID,
TaskID: taskID,
Execute: func(ctx context.Context) error {
runPipelineBg(ctx, projectID, taskID, req)
return nil
},
})
}
// 后台执行管线
go runPipelineBg(c.Request.Context(), projectID, taskID, req)
c.JSON(http.StatusOK, model.OK(GenerateResponse{TaskID: taskID}))
_ = queuePos
}
// runPipelineBg 后台执行生成管线,更新任务状态。
@@ -82,6 +101,7 @@ func runPipelineBg(ctx context.Context, projectID, taskID string, req GenerateRe
updateTaskInDB(ctx, taskID, "running", stage, "", progress)
})
// 队列调度后才标记为 running,初始写入时是 pending
updateTaskInDB(ctx, taskID, "running", "prompt_builder", "", 5)
in := service.PipelineInput{
@@ -207,7 +227,7 @@ func GetAssets(c *gin.Context) {
for _, a := range assets {
response = append(response, model.AssetResponse{
Key: a.Key,
URL: a.URL,
URL: storageSvc.GetSignedURL(a.Key),
Format: a.Format,
})
}
@@ -240,9 +260,9 @@ func saveTaskToDB(ctx context.Context, projectID, taskID string, req GenerateReq
ProjectID: uint(projectIDUint),
Prompt: req.Prompt,
AssetType: req.AssetType,
Status: "running",
Progress: 5,
Stage: "prompt_builder",
Status: "pending",
Progress: 0,
Stage: "",
RetryCount: 0,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
+6 -21
View File
@@ -14,10 +14,10 @@ import (
"github.com/cloudwego/eino/compose"
)
// promptOptimizerNode 节点:调用 PromptAgent 生成规范提示词,合并风格与重试信息。
// 输入 PipelineInput,输出最终提示词字符串(直接供 AssetGenerator 消费)。
// promptOptimizerNode 节点:合并风格描述、重试信息与技术参数,输出最终提示词。
// 提示词优化已由前端在提交前完成,管线内不再重复调用 PromptAgent。
var promptOptimizerNode = compose.InvokableLambda(func(ctx context.Context, in PipelineInput) (string, error) {
// 合并风格描述,注入原始 Prompt 中
// 合并风格描述
styleDesc := buildStyleDescription(in.ProjectStyle, in.TaskStyle)
if styleDesc != "" {
if in.Prompt != "" {
@@ -36,26 +36,11 @@ var promptOptimizerNode = compose.InvokableLambda(func(ctx context.Context, in P
}
}
if len(in.Tags) == 0 && in.Prompt == "" {
return "", fmt.Errorf("pipeline: Prompt and Tags are both empty")
if in.Prompt == "" {
return "", fmt.Errorf("pipeline: prompt is empty")
}
// 有标签时调用 PromptAgent 优化提示词
if len(in.Tags) > 0 {
agentIn := PromptAgentInput{
Tags: in.Tags,
AssetType: in.AssetType,
Prompt: in.Prompt,
UserNote: in.UserNote,
}
output, err := RunPromptAgent(ctx, agentIn)
if err != nil {
return "", fmt.Errorf("prompt agent: %w", err)
}
return output.Prompt, nil
}
// 无标签时直接使用原始 Prompt,补上技术参数段
// 追加技术参数段,不再调用 PromptAgent
return appendTechNotes(in.Prompt, in.AssetType, in.Params), nil
})
+104
View File
@@ -0,0 +1,104 @@
package service
import (
"context"
"sync"
"gen2d/internal/logger"
)
// TaskJob 队列中的任务。
type TaskJob struct {
Ctx context.Context
ProjectID string
TaskID string
Execute func(ctx context.Context) error
}
// TaskQueue 串行 FIFO 任务队列。
type TaskQueue struct {
mu sync.Mutex
jobs []*TaskJob
ready chan struct{}
stop chan struct{}
stopped bool
}
// NewTaskQueue 创建任务队列并启动调度协程。
func NewTaskQueue() *TaskQueue {
q := &TaskQueue{
jobs: make([]*TaskJob, 0),
ready: make(chan struct{}, 1),
stop: make(chan struct{}),
}
go q.run()
return q
}
// Enqueue 将任务加入队尾,返回队列中的位置(1-based)。
func (q *TaskQueue) Enqueue(job *TaskJob) int {
q.mu.Lock()
defer q.mu.Unlock()
q.jobs = append(q.jobs, job)
pos := len(q.jobs)
select {
case q.ready <- struct{}{}:
default:
}
return pos
}
// QueueLen 返回当前队列长度。
func (q *TaskQueue) QueueLen() int {
q.mu.Lock()
defer q.mu.Unlock()
return len(q.jobs)
}
// Stop 优雅关闭队列。
func (q *TaskQueue) Stop() {
q.mu.Lock()
defer q.mu.Unlock()
if !q.stopped {
q.stopped = true
close(q.stop)
}
}
func (q *TaskQueue) run() {
for {
select {
case <-q.stop:
return
case <-q.ready:
q.processNext()
}
}
}
func (q *TaskQueue) processNext() {
q.mu.Lock()
if len(q.jobs) == 0 {
q.mu.Unlock()
return
}
job := q.jobs[0]
q.jobs = q.jobs[1:]
// 如果队列还有任务,重新发信号
if len(q.jobs) > 0 {
select {
case q.ready <- struct{}{}:
default:
}
}
q.mu.Unlock()
l := logger.With("task_id", job.TaskID, "project_id", job.ProjectID)
l.Info("task queue executing job", "queue_remaining", len(q.jobs))
if err := job.Execute(job.Ctx); err != nil {
l.Error("task job failed", "error", err)
} else {
l.Info("task job completed")
}
}
+44 -5
View File
@@ -3,6 +3,9 @@ package service
import (
"bytes"
"context"
"crypto/hmac"
"crypto/sha1"
"encoding/base64"
"fmt"
"net/url"
"path/filepath"
@@ -65,15 +68,43 @@ func (s *StorageService) Upload(ctx context.Context, key string, data []byte) (s
return "", fmt.Errorf("qiniu upload: %w", err)
}
cdnURL := s.buildURL(key)
cdnURL := s.buildBaseURL(key)
l.Info("upload succeeded", "key", key, "url", cdnURL)
return cdnURL, nil
}
// GetDownloadURL 生成素材的下载 URL。
// 公开 bucket 直接返回 CDN URL;私有 bucket 生成带签名的临时 URL。
// 私有 bucket 生成带签名和过期时间的临时 URL。
func (s *StorageService) GetDownloadURL(_ context.Context, key string) (string, error) {
return s.buildURL(key), nil
return s.signDownloadURL(key), nil
}
// GetSignedURL 获取私有资源签名 URL(公开方法,供 handler 层使用)。
func (s *StorageService) GetSignedURL(key string) string {
return s.signDownloadURL(key)
}
// signDownloadURL 生成七牛云私有资源下载签名 URL。
// 格式:http://<domain>/<key>?e=<deadline>&token=<downloadToken>
func (s *StorageService) signDownloadURL(key string) string {
baseURL := s.buildBaseURL(key)
expire := s.cfg.URLExpire
if expire <= 0 {
expire = 3600
}
deadline := time.Now().Unix() + expire
signStr := fmt.Sprintf("%s?e=%d", baseURL, deadline)
mac := hmac.New(sha1.New, []byte(s.cfg.SecretKey))
mac.Write([]byte(signStr))
signature := mac.Sum(nil)
encodedSign := base64.URLEncoding.EncodeToString(signature)
token := fmt.Sprintf("%s:%s", s.cfg.AccessKey, encodedSign)
return fmt.Sprintf("%s?e=%d&token=%s", baseURL, deadline, token)
}
// Delete 从七牛云删除指定对象。
@@ -91,9 +122,17 @@ func (s *StorageService) Delete(ctx context.Context, key string) error {
return nil
}
// buildURL 根据配置构建完整的 CDN URL。
func (s *StorageService) buildURL(key string) string {
// buildBaseURL 根据配置构建完整的基础 CDN URL(不含签名参数)。
func (s *StorageService) buildBaseURL(key string) string {
host := strings.TrimRight(s.cfg.CDNHost, "/")
scheme := "http"
if s.cfg.UseHTTPS {
scheme = "https"
}
// 如果 host 已经包含 scheme,直接使用;否则添加
if !strings.HasPrefix(host, "http://") && !strings.HasPrefix(host, "https://") {
host = scheme + "://" + host
}
segments := strings.Split(key, "/")
for i, seg := range segments {
segments[i] = url.PathEscape(seg)
+2 -2
View File
@@ -516,14 +516,14 @@ POST /api/v1/prompt/optimize
### GET /api/v1/assets/download
公开接口(无需认证)。返回 302 重定向到七牛云 CDN URL。前端可通过此接口加载素材图片,作为 CDN 直接访问的替代方案。
公开接口(无需认证)。返回 302 重定向到七牛云带签名的临时下载 URL(私有 bucket 签名有效时长默认为 3600 秒)。前端可通过此接口加载素材图片。
| 参数 | 类型 | 说明 |
|------|------|------|
| `key` | string | 对象存储 Key,查询参数 |
响应:
- HTTP 302,`Location` 头指向 CDN URL
- HTTP 302,`Location` 头指向带签名的临时下载 URL
- 400:缺少 key 参数
- 500:生成下载链接失败
+7 -4
View File
@@ -150,16 +150,18 @@ users/user_a1B2c3/projects/proj_abc/tasks/task_xyz/output/spritesheet.png
### 访问方式
素材上传后返回 CDN URL,前端直接通过该 URL 访问图片。
素材上传后 asset 表存储对象 Key 和基础 CDN URL。由于 bucket 为私有读,前端通过 API 获取的 `url` 字段为带签名参数的临时下载链接(有效期可配置,默认 3600 秒)。
下载接口 `GET /api/v1/assets/download?key=...` 返回 302 重定向到 CDN URL。
格式:`https://cdn.example.com/<key>?e=<deadline>&token=<downloadToken>`
下载接口 `GET /api/v1/assets/download?key=...` 返回 302 重定向到带签名的临时 URL。
### asset 表的 url 字段
存储七牛云 CDN 完整 URL,如:
存储七牛云基础 CDN URL(不含签名),如:
`https://cdn.example.com/users/user_a1B2c3/projects/proj_abc/tasks/task_xyz/output/spritesheet.png`
前端直接使用该 URL 加载图片,无需拼接路径。
注意:该字段仅在数据库中使用,API 返回给前端的 URL 会根据对象 Key 实时生成带签名的临时下载链接。
---
@@ -190,6 +192,7 @@ users/user_a1B2c3/projects/proj_abc/tasks/task_xyz/output/spritesheet.png
| `GEN2D_QINIU_BUCKET` | — | 七牛云 Bucket 名称 |
| `GEN2D_QINIU_CDN_HOST` | — | CDN 域名(如 `https://cdn.example.com`) |
| `GEN2D_QINIU_USE_HTTPS` | `true` | 是否使用 HTTPS |
| `GEN2D_QINIU_URL_EXPIRE` | `3600` | 私有下载签名 URL 有效期(秒) |
在 `config.go` 中扩展字段即可,无需额外依赖。
+1 -1
View File
@@ -120,7 +120,7 @@ export const useGenerationStore = create<GenerationState>((set) => ({
function stageLabel(stage?: string): string {
switch (stage) {
case 'prompt_builder': return '优化提示词...'
case 'prompt_builder': return '构建提示词...'
case 'asset_generator': return '生成素材中...'
case 'quality_supervisor': return '质检中...'
case 'format_adapter': return '格式转换中...'