71 lines
2.0 KiB
Go
71 lines
2.0 KiB
Go
package node
|
||
|
||
import (
|
||
"context"
|
||
"strconv"
|
||
|
||
"ai-agent/gateway"
|
||
"ai-agent/workflow/consts/model"
|
||
"ai-agent/workflow/consts/node"
|
||
nodeDto "ai-agent/workflow/model/dto/node"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
var NodeLibraryService = &nodeLibraryService{}
|
||
|
||
type nodeLibraryService struct{}
|
||
|
||
func (s *nodeLibraryService) GetNodeLibrary(ctx context.Context, req *nodeDto.WorkflowNodeTreeReq) (*nodeDto.WorkflowNodeTreeRes, error) {
|
||
tree := node.GetFilterNodeTree([]node.NodeGroup{node.NodeGroupBase})
|
||
if opts, err := videoModelOptions(ctx); err != nil {
|
||
g.Log().Warningf(ctx, "加载视频模型选项失败，节点库降级返回: %v", err)
|
||
} else {
|
||
applyVideoModelOptions(tree, opts)
|
||
}
|
||
return &nodeDto.WorkflowNodeTreeRes{
|
||
Groups: tree,
|
||
}, nil
|
||
}
|
||
|
||
// videoModelOptions 按视频模型类型（600）查模型网关，转成 modelId 下拉选项：key=模型ID，value=模型名称。
|
||
func videoModelOptions(ctx context.Context) ([]node.SelectOption, error) {
|
||
res, err := gateway.ListModelManage(ctx, &gateway.ListModelManageReq{ModelType: model.TypeVideo})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
opts := make([]node.SelectOption, 0, len(res.List))
|
||
for _, item := range res.List {
|
||
if item == nil || item.Id <= 0 || item.ModelName == "" {
|
||
continue
|
||
}
|
||
opts = append(opts, node.SelectOption{
|
||
Key: strconv.FormatInt(item.Id, 10),
|
||
Value: item.ModelName,
|
||
})
|
||
}
|
||
return opts, nil
|
||
}
|
||
|
||
// applyVideoModelOptions 把 script_transcribe 节点的 modelId 预设下拉替换为视频模型选项。
|
||
// 深拷贝 PresetOption 后再改，避免改写全局 NodeTypeMetaList。
|
||
func applyVideoModelOptions(tree []node.NodeGroupTree, opts []node.SelectOption) {
|
||
for gi := range tree {
|
||
for ni := range tree[gi].Nodes {
|
||
n := &tree[gi].Nodes[ni]
|
||
if n.Key != node.NodeTypeScriptTranscribe {
|
||
continue
|
||
}
|
||
presets := append([]node.NodePresetField(nil), n.PresetOption...)
|
||
for pi := range presets {
|
||
if presets[pi].Field == "modelId" {
|
||
presets[pi].Options = opts
|
||
break
|
||
}
|
||
}
|
||
n.PresetOption = presets
|
||
return
|
||
}
|
||
}
|
||
}
|