wonder/gen2d
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pipeline): PromptOptimizer 集成到生成管线，新增优化 API

Browse Source 
- 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 注入
This commit is contained in:
Gmaker689
2026-05-24 20:45:18 +08:00
parent 788cd021ef
commit 9ca1e65221
7 changed files with 281 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/inference.go
+11 -1
View File
@@ -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) {
backend/internal/service/nodes.go
+77 -62
View File
@@ -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 {
backend/internal/service/pipeline.go
+17 -18
View File
@@ -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 {
backend/internal/service/pipeline_test.go
+120 -9
View File
@@ -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")
}
}
backend/internal/service/types.go
+3 -1
View File
@@ -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
}