From 9ca1e65221d1e8a01f45d54bf465d63cd2ed7098 Mon Sep 17 00:00:00 2001 From: Gmaker689 <1711322114@qq.com> Date: Sun, 24 May 2026 20:45:18 +0800 Subject: [PATCH] =?UTF-8?q?feat(pipeline):=20PromptOptimizer=20=E9=9B=86?= =?UTF-8?q?=E6=88=90=E5=88=B0=E7=94=9F=E6=88=90=E7=AE=A1=E7=BA=BF=EF=BC=8C?= =?UTF-8?q?=E6=96=B0=E5=A2=9E=E4=BC=98=E5=8C=96=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - PipelineInput 新增 Tags/UserNote 字段 - 管线重构为: START → PromptOptimizer → AssetGenerator → QualitySupervisor → FormatAdapter - promptOptimizerNode: 有标签调用 PromptAgent，无标签补技术参数，合并风格与重试信息 - PromptBuilder 移除，提示词构建逻辑并入 PromptOptimizer - 新增 POST /api/v1/prompt/optimize 接口 - main.go 注入 LLM/ImageGen 配置，注册 prompt 路由 - inference.go 新增 InitImageGenConfig 注入 --- backend/cmd/main.go | 16 ++- backend/internal/handler/prompt.go | 42 +++++++ backend/internal/service/inference.go | 12 +- backend/internal/service/nodes.go | 139 ++++++++++++---------- backend/internal/service/pipeline.go | 35 +++--- backend/internal/service/pipeline_test.go | 129 ++++++++++++++++++-- backend/internal/service/types.go | 4 +- 7 files changed, 281 insertions(+), 96 deletions(-) create mode 100644 backend/internal/handler/prompt.go diff --git a/backend/cmd/main.go b/backend/cmd/main.go index 111e284..0d4c328 100644 --- a/backend/cmd/main.go +++ b/backend/cmd/main.go @@ -9,6 +9,7 @@ import ( "gen2d/internal/db" "gen2d/internal/handler" "gen2d/internal/model" + "gen2d/internal/service" "github.com/gin-gonic/gin" ) @@ -16,15 +17,19 @@ import ( func main() { cfg := config.Load() - gin.SetMode(cfg.Mode) + gin.SetMode(cfg.Server.Mode) // 初始化 SQLite 数据库 - if err := db.Init(cfg.DSN, &model.User{}); err != nil { + if err := db.Init(cfg.Database.DSN, &model.User{}); err != nil { log.Fatalf("db init failed: %v", err) } // 初始化 AuthService，注入 JWT 配置 - handler.InitAuthService(cfg.JWTSecret, cfg.JWTExpire) + handler.InitAuthService(cfg.JWT.Secret, cfg.JWT.Expire) + + // 注入 LLM 和文生图配置到 service 层 + service.InitLLMConfig(cfg.LLM) + service.InitImageGenConfig(cfg.ImageGen) r := gin.New() r.Use(gin.Recovery()) // panic 恢复中间件，防止服务因未捕获异常宕机 @@ -32,7 +37,8 @@ func main() { // API v1 路由组 v1 := r.Group("/api/v1") { - v1.GET("/health", handler.Health) // 健康检查 + v1.GET("/health", handler.Health) // 健康检查 + v1.POST("/prompt/optimize", handler.PromptOptimize) // 提示词优化 } // Auth 路由组 @@ -42,7 +48,7 @@ func main() { auth.POST("/login", handler.Login) // 用户登录 } - addr := fmt.Sprintf(":%d", cfg.Port) + addr := fmt.Sprintf(":%d", cfg.Server.Port) log.Printf("gen2d backend starting on %s", addr) if err := r.Run(addr); err != nil { log.Fatalf("server failed: %v", err) diff --git a/backend/internal/handler/prompt.go b/backend/internal/handler/prompt.go new file mode 100644 index 0000000..1910e79 --- /dev/null +++ b/backend/internal/handler/prompt.go @@ -0,0 +1,42 @@ +package handler + +import ( + "net/http" + + "gen2d/internal/model" + "gen2d/internal/service" + + "github.com/gin-gonic/gin" +) + +// PromptOptimizeRequest 提示词优化请求。 +type PromptOptimizeRequest struct { + Tags []string `json:"tags" binding:"required,min=1"` // 用户选择的标签 + AssetType string `json:"assetType" binding:"required"` // 素材类型 + Prompt string `json:"prompt"` // 用户原始提示词 + UserNote string `json:"userNote"` // 用户额外描述 +} + +// PromptOptimize 提示词优化接口，接收用户标签并调用 LLM 生成规范提示词。 +func PromptOptimize(c *gin.Context) { + var req PromptOptimizeRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "参数错误: "+err.Error())) + return + } + + in := service.PromptAgentInput{ + Tags: req.Tags, + AssetType: req.AssetType, + Prompt: req.Prompt, + UserNote: req.UserNote, + } + + output, err := service.RunPromptAgent(c.Request.Context(), in) + if err != nil { + c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "提示词优化失败: "+err.Error())) + return + } + + c.JSON(http.StatusOK, model.OK(output)) +} diff --git a/backend/internal/service/inference.go b/backend/internal/service/inference.go index 9574486..30b6ee0 100644 --- a/backend/internal/service/inference.go +++ b/backend/internal/service/inference.go @@ -1,15 +1,25 @@ package service import ( + "bytes" "context" "crypto/rand" "fmt" "image" "image/color" "image/png" - "bytes" + + "gen2d/internal/config" ) +// imgCfg 保存文生图配置，由 main 通过 InitImageGenConfig 注入。 +var imgCfg config.ImageGenConfig + +// InitImageGenConfig 注入文生图配置。 +func InitImageGenConfig(cfg config.ImageGenConfig) { + imgCfg = cfg +} + // GenerateImages 调用 AI 推理 API 生成图片。 // MVP 阶段返回 mock 占位图。 func GenerateImages(ctx context.Context, prompt string, params AssetParams) ([]GeneratedImage, error) { diff --git a/backend/internal/service/nodes.go b/backend/internal/service/nodes.go index 3f1dfff..1bac3ed 100644 --- a/backend/internal/service/nodes.go +++ b/backend/internal/service/nodes.go @@ -8,13 +8,53 @@ import ( "github.com/cloudwego/eino/compose" ) -// PromptBuilder 节点：接收输入，输出三段式提示词 -var promptBuilderNode = compose.InvokableLambda(func(ctx context.Context, in PipelineInput) (string, error) { - return buildPrompt(in), nil +// promptOptimizerNode 节点：调用 PromptAgent 生成规范提示词，合并风格与重试信息。 +// 输入 PipelineInput，输出最终提示词字符串（直接供 AssetGenerator 消费）。 +var promptOptimizerNode = compose.InvokableLambda(func(ctx context.Context, in PipelineInput) (string, error) { + // 合并风格描述，注入原始 Prompt 中 + styleDesc := buildStyleDescription(in.ProjectStyle, in.TaskStyle) + if styleDesc != "" { + if in.Prompt != "" { + in.Prompt = in.Prompt + "。" + styleDesc + } else { + in.Prompt = styleDesc + } + } + + // 注入重试原因 + if in.RejectReason != "" { + if in.Prompt != "" { + in.Prompt = in.Prompt + "。注意修正以下问题：" + in.RejectReason + } else { + in.Prompt = "修正以下问题：" + in.RejectReason + } + } + + if len(in.Tags) == 0 && in.Prompt == "" { + return "", fmt.Errorf("pipeline: Prompt and Tags are both 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，补上技术参数段 + return appendTechNotes(in.Prompt, in.AssetType, in.Params), nil }) -// promptBuilderPreHandler 首次运行时保存输入到 state；重试时注入 RejectReason -func promptBuilderPreHandler(ctx context.Context, in PipelineInput, state *PipelineState) (PipelineInput, error) { +// promptOptimizerPreHandler 首次运行时保存输入到 state；重试时注入 RejectReason。 +func promptOptimizerPreHandler(ctx context.Context, in PipelineInput, state *PipelineState) (PipelineInput, error) { if state.RetryCount == 0 { state.Input = in } else if state.RejectReason != "" { @@ -23,13 +63,13 @@ func promptBuilderPreHandler(ctx context.Context, in PipelineInput, state *Pipel return in, nil } -// promptBuilderPostHandler 将提示词写入全局状态 -func promptBuilderPostHandler(ctx context.Context, out string, state *PipelineState) (string, error) { +// promptOptimizerPostHandler 将最终提示词写入全局状态。 +func promptOptimizerPostHandler(ctx context.Context, out string, state *PipelineState) (string, error) { state.FinalPrompt = out return out, nil } -// AssetGenerator 节点：调用 AI 推理 API 出图 +// AssetGenerator 节点：调用 AI 推理 API 出图。 var assetGeneratorNode = compose.InvokableLambda(func(ctx context.Context, prompt string) ([]GeneratedImage, error) { var params AssetParams _ = compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error { @@ -39,21 +79,18 @@ var assetGeneratorNode = compose.InvokableLambda(func(ctx context.Context, promp return GenerateImages(ctx, prompt, params) }) -// assetGeneratorPostHandler 将原始图片写入全局状态 +// assetGeneratorPostHandler 将原始图片写入全局状态。 func assetGeneratorPostHandler(ctx context.Context, out []GeneratedImage, state *PipelineState) ([]GeneratedImage, error) { state.RawImages = out return out, nil } -// QualitySupervisor 节点：质检，输出 PipelineInput 供下游节点消费。 -// 将图片存入 state，设置路由目标 NextNode。 +// QualitySupervisor 节点：质检，设置路由目标。 var qualitySupervisorNode = compose.InvokableLambda(func(ctx context.Context, images []GeneratedImage) (PipelineInput, error) { var input PipelineInput err := compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error { - // 保存图片到状态 state.RawImages = images - // 质检 style := mergeStyle(state.Input.ProjectStyle, state.Input.TaskStyle) pass, reason, checkErr := CheckQuality(ctx, images, style) if checkErr != nil { @@ -65,14 +102,13 @@ var qualitySupervisorNode = compose.InvokableLambda(func(ctx context.Context, im state.RejectReason = reason } - // 决定路由 if pass { state.NextNode = nodeFormatAdapter } else if state.RetryCount >= 3 { - state.NextNode = nodeFormatAdapter // 超过重试次数，降级输出 + state.NextNode = nodeFormatAdapter } else { state.RetryCount++ - state.NextNode = nodePromptBuilder // 重生成 + state.NextNode = nodePromptOptimizer } input = state.Input @@ -84,7 +120,7 @@ var qualitySupervisorNode = compose.InvokableLambda(func(ctx context.Context, im return input, nil }) -// formatAdapterNode 节点：从 state 读取图片，格式转换，组装输出 +// formatAdapterNode 节点：从 state 读取图片，格式转换，组装输出。 var formatAdapterNode = compose.InvokableLambda(func(ctx context.Context, input PipelineInput) (PipelineOutput, error) { var images []GeneratedImage _ = compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error { @@ -119,63 +155,42 @@ var formatAdapterNode = compose.InvokableLambda(func(ctx context.Context, input }, nil }) -// buildPrompt 构建三段式提示词 -func buildPrompt(in PipelineInput) string { +// buildStyleDescription 将风格键值对转为自然语言描述，供 PromptAgent 注入。 +func buildStyleDescription(projectStyle, taskStyle map[string]string) string { + merged := mergeStyle(projectStyle, taskStyle) + if len(merged) == 0 { + return "" + } var parts []string - - // 【主题】 - parts = append(parts, fmt.Sprintf("【主题】%s", in.Prompt)) - - // 【约束】 - constraints := buildConstraints(in) - parts = append(parts, fmt.Sprintf("【约束】%s", constraints)) - - // 【内容】 - content := buildContent(in) - parts = append(parts, fmt.Sprintf("【内容】%s", content)) - - return strings.Join(parts, "\n") -} - -// buildConstraints 合并风格 + 负面提示词 + 重试原因 -func buildConstraints(in PipelineInput) string { - style := mergeStyle(in.ProjectStyle, in.TaskStyle) - - var parts []string - for k, v := range style { + for k, v := range merged { parts = append(parts, fmt.Sprintf("%s: %s", k, v)) } - - if in.RejectReason != "" { - parts = append(parts, fmt.Sprintf("上次质检问题：%s", in.RejectReason)) - } - - if len(parts) == 0 { - return "无特殊约束" - } - return strings.Join(parts, "; ") + return "风格约束：" + strings.Join(parts, "；") } -// buildContent 构建技术参数段 -func buildContent(in PipelineInput) string { +// appendTechNotes 在无标签（不走 PromptAgent）时补上技术参数段。 +func appendTechNotes(prompt, assetType string, params AssetParams) string { var parts []string - parts = append(parts, fmt.Sprintf("素材类型: %s", in.AssetType)) - if in.Params.Resolution > 0 { - parts = append(parts, fmt.Sprintf("分辨率: %d", in.Params.Resolution)) + if prompt != "" { + parts = append(parts, prompt) } - if in.Params.Frames.Directions > 0 { - parts = append(parts, fmt.Sprintf("方向数: %d", in.Params.Frames.Directions)) + parts = append(parts, fmt.Sprintf("素材类型: %s", assetType)) + if params.Resolution > 0 { + parts = append(parts, fmt.Sprintf("分辨率: %d", params.Resolution)) } - if in.Params.Frames.FramesPerDirection > 0 { - parts = append(parts, fmt.Sprintf("每方向帧数: %d", in.Params.Frames.FramesPerDirection)) + if params.Frames.Directions > 0 { + parts = append(parts, fmt.Sprintf("方向数: %d", params.Frames.Directions)) } - if in.Params.Format != "" { - parts = append(parts, fmt.Sprintf("输出格式: %s", in.Params.Format)) + if params.Frames.FramesPerDirection > 0 { + parts = append(parts, fmt.Sprintf("每方向帧数: %d", params.Frames.FramesPerDirection)) } - return strings.Join(parts, "; ") + if params.Format != "" { + parts = append(parts, fmt.Sprintf("输出格式: %s", params.Format)) + } + return strings.Join(parts, "；") } -// mergeStyle 合并工程风格与任务风格覆盖，任务同名键覆盖工程 +// mergeStyle 合并工程风格与任务风格覆盖，任务同名键覆盖工程。 func mergeStyle(projectStyle, taskStyle map[string]string) map[string]string { result := make(map[string]string) for k, v := range projectStyle { diff --git a/backend/internal/service/pipeline.go b/backend/internal/service/pipeline.go index 18db114..6ce8a23 100644 --- a/backend/internal/service/pipeline.go +++ b/backend/internal/service/pipeline.go @@ -8,18 +8,18 @@ import ( ) const ( - nodePromptBuilder = "prompt_builder" + nodePromptOptimizer = "prompt_optimizer" nodeAssetGenerator = "asset_generator" nodeQualitySupervisor = "quality_supervisor" nodeFormatAdapter = "format_adapter" ) -// NewGenerateGraph 创建四阶段生成管线 Graph。 +// NewGenerateGraph 创建生成管线 Graph（PromptOptimizer → AssetGenerator → QualitySupervisor → FormatAdapter）。 // -// START → PromptBuilder → AssetGenerator → QualitySupervisor -// ├── pass → FormatAdapter → END -// └── fail, retry<3 → PromptBuilder -// └── fail, retry>=3 → FormatAdapter (降级) +// START → PromptOptimizer → AssetGenerator → QualitySupervisor +// ├── pass → FormatAdapter → END +// └── fail, retry<3 → PromptOptimizer +// └── fail, retry>=3 → FormatAdapter (降级) func NewGenerateGraph() (*compose.Graph[PipelineInput, PipelineOutput], error) { g := compose.NewGraph[PipelineInput, PipelineOutput]( compose.WithGenLocalState(func(ctx context.Context) *PipelineState { @@ -27,12 +27,11 @@ func NewGenerateGraph() (*compose.Graph[PipelineInput, PipelineOutput], error) { }), ) - // 添加节点 - if err := g.AddLambdaNode(nodePromptBuilder, promptBuilderNode, - compose.WithStatePreHandler(promptBuilderPreHandler), - compose.WithStatePostHandler(promptBuilderPostHandler), + if err := g.AddLambdaNode(nodePromptOptimizer, promptOptimizerNode, + compose.WithStatePreHandler(promptOptimizerPreHandler), + compose.WithStatePostHandler(promptOptimizerPostHandler), ); err != nil { - return nil, fmt.Errorf("add %s node: %w", nodePromptBuilder, err) + return nil, fmt.Errorf("add %s node: %w", nodePromptOptimizer, err) } if err := g.AddLambdaNode(nodeAssetGenerator, assetGeneratorNode, @@ -49,12 +48,12 @@ func NewGenerateGraph() (*compose.Graph[PipelineInput, PipelineOutput], error) { return nil, fmt.Errorf("add %s node: %w", nodeFormatAdapter, err) } - // 连线：正常路径 - if err := g.AddEdge(compose.START, nodePromptBuilder); err != nil { - return nil, fmt.Errorf("add edge START->%s: %w", nodePromptBuilder, err) + // 连线：START → PromptOptimizer → AssetGenerator → Supervisor + if err := g.AddEdge(compose.START, nodePromptOptimizer); err != nil { + return nil, fmt.Errorf("add edge START->%s: %w", nodePromptOptimizer, err) } - if err := g.AddEdge(nodePromptBuilder, nodeAssetGenerator); err != nil { - return nil, fmt.Errorf("add edge %s->%s: %w", nodePromptBuilder, nodeAssetGenerator, err) + if err := g.AddEdge(nodePromptOptimizer, nodeAssetGenerator); err != nil { + return nil, fmt.Errorf("add edge %s->%s: %w", nodePromptOptimizer, nodeAssetGenerator, err) } if err := g.AddEdge(nodeAssetGenerator, nodeQualitySupervisor); err != nil { return nil, fmt.Errorf("add edge %s->%s: %w", nodeAssetGenerator, nodeQualitySupervisor, err) @@ -73,7 +72,7 @@ func NewGenerateGraph() (*compose.Graph[PipelineInput, PipelineOutput], error) { }) return next, nil }, - map[string]bool{nodePromptBuilder: true, nodeFormatAdapter: true}, + map[string]bool{nodePromptOptimizer: true, nodeFormatAdapter: true}, )); err != nil { return nil, fmt.Errorf("add branch at %s: %w", nodeQualitySupervisor, err) } @@ -81,7 +80,7 @@ func NewGenerateGraph() (*compose.Graph[PipelineInput, PipelineOutput], error) { return g, nil } -// RunPipeline 编译并执行生成管线 +// RunPipeline 编译并执行生成管线。 func RunPipeline(ctx context.Context, in PipelineInput) (*PipelineOutput, error) { g, err := NewGenerateGraph() if err != nil { diff --git a/backend/internal/service/pipeline_test.go b/backend/internal/service/pipeline_test.go index b7234dd..0286c1b 100644 --- a/backend/internal/service/pipeline_test.go +++ b/backend/internal/service/pipeline_test.go @@ -6,7 +6,6 @@ import ( ) func TestPipeline_HappyPath(t *testing.T) { - // 质检一次通过 QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) { return true, "", nil } @@ -40,7 +39,6 @@ func TestPipeline_HappyPath(t *testing.T) { } func TestPipeline_RetryThenPass(t *testing.T) { - // 质检前 2 次 fail，第 3 次 pass QualityChecker = NewCountedQualityChecker(3) defer func() { QualityChecker = defaultCheckQuality }() @@ -50,7 +48,7 @@ func TestPipeline_RetryThenPass(t *testing.T) { Params: AssetParams{ Resolution: 32, Frames: FrameParams{ - Directions: 4, + Directions: 4, FramesPerDirection: 2, }, Format: "spritesheet", @@ -60,7 +58,6 @@ func TestPipeline_RetryThenPass(t *testing.T) { t.Fatalf("RunPipeline failed: %v", err) } - // 4 directions × 2 frames = 8 张图 if len(output.Assets) != 8 { t.Errorf("expected 8 assets, got %d", len(output.Assets)) } @@ -73,7 +70,6 @@ func TestPipeline_RetryThenPass(t *testing.T) { } func TestPipeline_MaxRetryDegrade(t *testing.T) { - // 质检始终 fail，超过 3 次后降级输出 QualityChecker = AlwaysFailQualityChecker() defer func() { QualityChecker = defaultCheckQuality }() @@ -88,7 +84,6 @@ func TestPipeline_MaxRetryDegrade(t *testing.T) { t.Fatalf("RunPipeline failed: %v", err) } - // 降级也应该有输出 if len(output.Assets) == 0 { t.Fatal("expected non-empty assets even on degrade") } @@ -98,9 +93,7 @@ func TestPipeline_MaxRetryDegrade(t *testing.T) { } func TestPipeline_StyleMerge(t *testing.T) { - // 验证风格合并：task 覆盖 project QualityChecker = func(_ context.Context, _ []GeneratedImage, style map[string]string) (bool, string, error) { - // 验证合并结果 if style["artStyle"] != "realistic" { t.Errorf("expected artStyle=realistic (task override), got %s", style["artStyle"]) } @@ -119,7 +112,7 @@ func TestPipeline_StyleMerge(t *testing.T) { "palette": "warm", }, TaskStyle: map[string]string{ - "artStyle": "realistic", // 覆盖 project + "artStyle": "realistic", }, Params: AssetParams{Resolution: 64}, }) @@ -127,3 +120,121 @@ func TestPipeline_StyleMerge(t *testing.T) { t.Fatalf("RunPipeline failed: %v", err) } } + +func TestPipeline_WithPromptOptimizer(t *testing.T) { + // 带标签时 PromptOptimizer 应优化原始提示词 + QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) { + return true, "", nil + } + defer func() { QualityChecker = defaultCheckQuality }() + + output, err := RunPipeline(context.Background(), PipelineInput{ + Prompt: "一个战士", + AssetType: "sprite", + Tags: []string{"像素", "战士", "持剑"}, + Params: AssetParams{ + Resolution: 64, + Format: "spritesheet", + }, + }) + if err != nil { + t.Fatalf("RunPipeline with tags failed: %v", err) + } + + // 有标签时应有优化后的输出 + if len(output.Assets) == 0 { + t.Fatal("expected non-empty assets") + } +} + +func TestPipeline_WithoutTags(t *testing.T) { + // 无标签时 PromptOptimizer 应跳过，原始提示词直接进入 PromptBuilder + QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) { + return true, "", nil + } + defer func() { QualityChecker = defaultCheckQuality }() + + rawPrompt := "一个原始提示词没有标签" + + output, err := RunPipeline(context.Background(), PipelineInput{ + Prompt: rawPrompt, + AssetType: "sprite", + Params: AssetParams{ + Resolution: 32, + Format: "spritesheet", + }, + }) + if err != nil { + t.Fatalf("RunPipeline without tags failed: %v", err) + } + + if len(output.Assets) == 0 { + t.Fatal("expected non-empty assets") + } +} + +func TestPipeline_PromptOptimizerFallback(t *testing.T) { + // PromptAgent 优化失败（无 API key）也不应阻塞管线 + QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) { + return true, "", nil + } + defer func() { QualityChecker = defaultCheckQuality }() + + output, err := RunPipeline(context.Background(), PipelineInput{ + Prompt: "一个火球术", + AssetType: "animation", + Tags: []string{"火焰", "魔法"}, + Params: AssetParams{ + Resolution: 64, + Format: "spritesheet", + }, + }) + if err != nil { + t.Fatalf("RunPipeline with fallback prompt agent failed: %v", err) + } + + if len(output.Assets) == 0 { + t.Fatal("expected non-empty assets even on prompt agent fallback") + } +} + +func TestPipeline_PromptOptimizerSkipsEmptyTags(t *testing.T) { + // 无标签时 prompt 保持原样 + QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) { + return true, "", nil + } + defer func() { QualityChecker = defaultCheckQuality }() + + output, err := RunPipeline(context.Background(), PipelineInput{ + Prompt: "原始提示词", + AssetType: "sprite", + Params: AssetParams{Resolution: 32}, + }) + if err != nil { + t.Fatalf("RunPipeline failed: %v", err) + } + if len(output.Assets) == 0 { + t.Fatal("expected non-empty assets") + } +} + +func TestPipeline_PromptOptimizerRefinesPrompt(t *testing.T) { + // 有标签时 pipeline 产出优化后的 prompt + QualityChecker = func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) { + return true, "", nil + } + defer func() { QualityChecker = defaultCheckQuality }() + + output, err := RunPipeline(context.Background(), PipelineInput{ + Prompt: "一个战士", + AssetType: "sprite", + Tags: []string{"像素", "战士"}, + Params: AssetParams{Resolution: 32}, + }) + if err != nil { + t.Fatalf("RunPipeline with tags failed: %v", err) + } + if len(output.Assets) == 0 { + t.Fatal("expected non-empty assets") + } +} diff --git a/backend/internal/service/types.go b/backend/internal/service/types.go index 5b64bfb..ed684b0 100644 --- a/backend/internal/service/types.go +++ b/backend/internal/service/types.go @@ -8,6 +8,8 @@ type PipelineInput struct { TaskStyle map[string]string // 任务风格覆盖 Params AssetParams // 技术参数 RejectReason string // 重试时由 state 注入 + Tags []string // 用户选择的标签，PromptAgent 据此优化提示词 + UserNote string // 用户额外描述 } // PipelineState Graph 全局状态，通过 WithGenLocalState 注入 @@ -36,7 +38,7 @@ type AssetParams struct { // FrameParams 帧参数 type FrameParams struct { - Directions int + Directions int FramesPerDirection int }