Gmarker689
5d5fe60426
后端: - pipeline.go: ProgressReporter 回调类型 + WithProgressReporter 注入 context - nodes.go: 各节点 pre/post handler 调用 reportProgress() 上报阶段进度 - generate.go: TaskResponse 新增 stage 字段，runPipelineBg 注入进度回调 前端: - vite.config.ts: 添加 /generation 代理到后端静态文件服务 - generation.ts: 轮询读取 stage 字段，暴露 stage/retryCount/rejectReason - GeneratePage.tsx: ProgressBar 接收真实管线阶段数据
121 lines
4.1 KiB
Go
Executable File
121 lines
4.1 KiB
Go
Executable File
package service
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
|
||
"github.com/cloudwego/eino/compose"
|
||
)
|
||
|
||
const (
|
||
nodePromptOptimizer = "prompt_optimizer"
|
||
nodeAssetGenerator = "asset_generator"
|
||
nodeQualitySupervisor = "quality_supervisor"
|
||
nodeFormatAdapter = "format_adapter"
|
||
)
|
||
|
||
// ProgressReporter 管线进度回调：stage 为当前节点名，progress 为 0-100。
|
||
type ProgressReporter func(stage string, progress int)
|
||
|
||
type progressKeyType struct{}
|
||
|
||
var progressCtxKey progressKeyType
|
||
|
||
// WithProgressReporter 将进度回调注入 context。
|
||
func WithProgressReporter(ctx context.Context, r ProgressReporter) context.Context {
|
||
return context.WithValue(ctx, progressCtxKey, r)
|
||
}
|
||
|
||
// reportProgress 从 context 取出回调上报进度。
|
||
func reportProgress(ctx context.Context, stage string, progress int) {
|
||
if r, ok := ctx.Value(progressCtxKey).(ProgressReporter); ok {
|
||
r(stage, progress)
|
||
}
|
||
}
|
||
|
||
// NewGenerateGraph 创建生成管线 Graph（PromptOptimizer → AssetGenerator → QualitySupervisor → 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 {
|
||
return &PipelineState{}
|
||
}),
|
||
)
|
||
|
||
if err := g.AddLambdaNode(nodePromptOptimizer, promptOptimizerNode,
|
||
compose.WithStatePreHandler(promptOptimizerPreHandler),
|
||
compose.WithStatePostHandler(promptOptimizerPostHandler),
|
||
); err != nil {
|
||
return nil, fmt.Errorf("add %s node: %w", nodePromptOptimizer, err)
|
||
}
|
||
|
||
if err := g.AddLambdaNode(nodeAssetGenerator, assetGeneratorNode,
|
||
compose.WithStatePostHandler(assetGeneratorPostHandler),
|
||
); err != nil {
|
||
return nil, fmt.Errorf("add %s node: %w", nodeAssetGenerator, err)
|
||
}
|
||
|
||
if err := g.AddLambdaNode(nodeQualitySupervisor, qualitySupervisorNode); err != nil {
|
||
return nil, fmt.Errorf("add %s node: %w", nodeQualitySupervisor, err)
|
||
}
|
||
|
||
if err := g.AddLambdaNode(nodeFormatAdapter, formatAdapterNode); err != nil {
|
||
return nil, fmt.Errorf("add %s node: %w", nodeFormatAdapter, 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(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)
|
||
}
|
||
if err := g.AddEdge(nodeFormatAdapter, compose.END); err != nil {
|
||
return nil, fmt.Errorf("add edge %s->END: %w", nodeFormatAdapter, err)
|
||
}
|
||
|
||
// 连线：质检分支（从 state.NextNode 读取路由目标）
|
||
if err := g.AddBranch(nodeQualitySupervisor, compose.NewGraphBranch(
|
||
func(ctx context.Context, _ PipelineInput) (string, error) {
|
||
var next string
|
||
_ = compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error {
|
||
next = state.NextNode
|
||
return nil
|
||
})
|
||
return next, nil
|
||
},
|
||
map[string]bool{nodePromptOptimizer: true, nodeFormatAdapter: true},
|
||
)); err != nil {
|
||
return nil, fmt.Errorf("add branch at %s: %w", nodeQualitySupervisor, err)
|
||
}
|
||
|
||
return g, nil
|
||
}
|
||
|
||
// RunPipeline 编译并执行生成管线。
|
||
func RunPipeline(ctx context.Context, in PipelineInput) (*PipelineOutput, error) {
|
||
g, err := NewGenerateGraph()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create graph: %w", err)
|
||
}
|
||
|
||
r, err := g.Compile(ctx, compose.WithMaxRunSteps(20))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("compile graph: %w", err)
|
||
}
|
||
|
||
output, err := r.Invoke(ctx, in)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("invoke pipeline: %w", err)
|
||
}
|
||
|
||
return &output, nil
|
||
}
|