81 lines
1.9 KiB
Go
81 lines
1.9 KiB
Go
package einograph
|
||
|
||
import (
|
||
"context"
|
||
|
||
"github.com/cloudwego/eino/compose"
|
||
)
|
||
|
||
// 节点名称常量
|
||
const (
|
||
NodeInputAdapt = "input_adapt" // 输入适配
|
||
NodeBuildPrompt = "build_prompt" // 构建提示词
|
||
NodeGenText = "gen_text" // 大模型生成原始文案
|
||
NodeGenImage = "gen_image" // 生成图片 + HTML
|
||
)
|
||
|
||
// ==================== 构建流式图（已加入优化节点） ====================
|
||
func BuildCreationGraph(ctx context.Context) (compose.Runnable[any, any], error) {
|
||
g := compose.NewGraph[any, any]()
|
||
|
||
// 1. 输入适配
|
||
err := g.AddLambdaNode(
|
||
NodeInputAdapt,
|
||
compose.InvokableLambda(InputAdaptLambda),
|
||
compose.WithOutputKey("input"),
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 2. 构建提示词
|
||
err = g.AddLambdaNode(
|
||
NodeBuildPrompt,
|
||
compose.InvokableLambda(BuildPromptLambda),
|
||
compose.WithInputKey("input"),
|
||
compose.WithOutputKey("messages"),
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 3. LLM 生成原始文案
|
||
chatModel, err := NewChatModel(ctx)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
err = g.AddChatModelNode(
|
||
NodeGenText,
|
||
chatModel,
|
||
compose.WithInputKey("messages"),
|
||
compose.WithOutputKey("text_result"),
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 5. 生成图片 + HTML（使用优化后的文案）
|
||
err = g.AddLambdaNode(
|
||
NodeGenImage,
|
||
compose.InvokableLambda(GenerateImageLambda),
|
||
compose.WithInputKey("text_result"),
|
||
compose.WithOutputKey("output"),
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// ✅ 正确连线：优化节点必须放在生成文案之后、生成图片之前
|
||
_ = g.AddEdge(compose.START, NodeInputAdapt)
|
||
_ = g.AddEdge(NodeInputAdapt, NodeBuildPrompt)
|
||
_ = g.AddEdge(NodeBuildPrompt, NodeGenText)
|
||
_ = g.AddEdge(NodeGenText, NodeGenImage)
|
||
_ = g.AddEdge(NodeGenImage, compose.END)
|
||
|
||
// 编译
|
||
return g.Compile(ctx,
|
||
compose.WithGraphName("xiaohongshu_content_creation"),
|
||
compose.WithNodeTriggerMode(compose.AnyPredecessor),
|
||
)
|
||
}
|