ai-agent/workflow/service/flow/flow_graph_util.go

package flow

import (
	"ai-agent/gateway"
	"ai-agent/workflow/consts/node"
	nodeDao "ai-agent/workflow/dao/node"
	flowDto "ai-agent/workflow/model/dto/flow"
	nodeDto "ai-agent/workflow/model/dto/node"
	"ai-agent/workflow/model/entity"
	"context"
	"fmt"
	"time"

	"github.com/cloudwego/eino/compose"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

// BuildNodeExecutionInput 构建节点执行入参，包含中断恢复逻辑
func BuildNodeExecutionInput(ctx context.Context, input any, flowNode entity.FlowNode) (*flowDto.FlowExecutionInput, *flowDto.NodeExecutionInput, error) {
	execInput := new(flowDto.FlowExecutionInput)

	wasInterrupted, _, _ := compose.GetInterruptState[any](ctx)
	if wasInterrupted {
		if err := compose.ProcessState(ctx, func(_ context.Context, s *flowDto.NodeExecutionState) error {
			execInput = s.SavedFlowInput
			return nil
		}); err != nil {
			return nil, nil, fmt.Errorf("节点：%v 进程状态读取失败: %v", flowNode.Name, err)
		}
		// 兼容旧 checkpoint（无 SavedFlowInput 时，降级使用 input 参数）
		if execInput == nil {
			var ok bool
			execInput, ok = input.(*flowDto.FlowExecutionInput)
			if !ok {
				return nil, nil, fmt.Errorf("节点：%v 进程状态为空，节点入参类型不匹配", flowNode.Name)
			}
			if g.IsEmpty(execInput) {
				return nil, nil, fmt.Errorf("节点：%v 进程状态为空，节点入参参数为空", flowNode.Name)
			}
		}
		// 续跑必定非全新执行：checkpoint 恢复的 SavedFlowInput 里 ForceNewRun 是上次（fresh）执行留下的 true，
		// 不清则 ModelLambda 误走"清段重生成"而非复用已成功段
		execInput.ForceNewRun = false
	} else {
		var ok bool
		execInput, ok = input.(*flowDto.FlowExecutionInput)
		if !ok {
			return nil, nil, fmt.Errorf("节点：%v 入参类型不匹配", flowNode.Name)
		}
		if g.IsEmpty(execInput) {
			return nil, nil, fmt.Errorf("节点：%v 入参参数为空", flowNode.Name)
		}
	}

	configMap := execInput.ConfigMap
	currentConfig := configMap[flowNode.Id]
	if currentConfig == nil {
		return nil, nil, fmt.Errorf("节点：%v 节点信息为空", flowNode.Name)
	}

	// 聚合输入来源
	//if len(flowNode.InputSource) > 0 {
	//	for _, inputSource := range currentConfig.InputSource {
	//		if sourceConfig, ok := configMap[inputSource.NodeId]; ok {
	//			currentConfig.OutputResult = append(currentConfig.OutputResult, sourceConfig.OutputResult...)
	//		}
	//	}
	//}

	// 构建节点执行入参
	realInput := &flowDto.NodeExecutionInput{
		Config: currentConfig,
		Global: execInput,
	}

	return execInput, realInput, nil
}

// HandleSuccessfulNodeExecution 处理节点执行成功的后续操作
func HandleSuccessfulNodeExecution(ctx context.Context, execInput *flowDto.FlowExecutionInput, realInput *flowDto.NodeExecutionInput, nodeExecutionId int64, flowNode entity.FlowNode, durationMs int64) error {
	// 上传输出到OSS
	ossResult, err := gateway.Upload(ctx, fmt.Sprintf("nodeInput:%v.txt", time.Now().UnixMilli()), gconv.Bytes(gconv.String(realInput)))
	if err != nil {
		return fmt.Errorf("节点：%v 上传OSS失败: %v", realInput.Config.Name, err)
	}

	// 更新执行记录为成功
	if err := UpdateNodeExecutionRecord(ctx, nodeExecutionId, durationMs, node.NodeExecutionStatusSuccess.Code(), ossResult, ""); err != nil {
		return fmt.Errorf("节点：%v 更新成功状态错误: %v", flowNode.Name, err)
	}

	// 记录成功到已执行列表
	RecordExecutionResult(execInput, flowNode.Id, node.NodeExecutionStatusSuccess.Code())
	return nil
}

// HandleFailedNodeExecution 处理节点执行失败的后续操作
func HandleFailedNodeExecution(ctx context.Context, execInput *flowDto.FlowExecutionInput, nodeExecutionId int64, flowNode entity.FlowNode, err error, durationMs int64) error {
	// 保存状态用于续跑
	if stateErr := compose.ProcessState(ctx, func(_ context.Context, s *flowDto.NodeExecutionState) error {
		s.CompletedNodes = append(s.CompletedNodes, flowNode.Name)
		s.SavedFlowInput = execInput
		s.ExecutionCount++
		return nil
	}); stateErr != nil {
		fmt.Printf("节点：%v 进程状态保存失败: %v", flowNode.Name, stateErr)
	}

	if !g.IsEmpty(nodeExecutionId) {
		// 更新执行记录为失败
		if updateErr := UpdateNodeExecutionRecord(ctx, nodeExecutionId, durationMs, node.NodeExecutionStatusFailed.Code(), "", err.Error()); updateErr != nil {
			fmt.Printf("节点：%v 更新失败状态错误: %v", flowNode.Name, updateErr)
		}
	}

	// 记录失败到已执行列表
	//RecordExecutionResult(execInput, flowNode.Id, node.NodeExecutionStatusFailed.Code())

	// 触发中断
	return compose.Interrupt(ctx, map[string]string{
		"node":  flowNode.Name,
		"error": err.Error(),
	})
}

// RecordExecutionResult 将节点执行结果写入 Global.ExecutedNodes
func RecordExecutionResult(execInput *flowDto.FlowExecutionInput, nodeId string, status node.NodeExecutionStatus) {
	execInput.ExecutedNodes = append(execInput.ExecutedNodes, flowDto.ExecutedNode{
		NodeId: nodeId,
		Status: status,
	})
}

// UpdateNodeExecutionRecord 更新节点执行记录
func UpdateNodeExecutionRecord(ctx context.Context, nodeExecutionId int64, durationMs int64, status node.NodeExecutionStatus, outputParamsPath string, errMsg string) error {
	if _, err := nodeDao.NodeExecutionDao.Update(ctx, &nodeDto.UpdateNodeExecutionReq{
		Id:               nodeExecutionId,
		DurationMs:       durationMs,
		Status:           status,
		OutputParamsPath: outputParamsPath,
		ErrorMessage:     errMsg,
	}); err != nil {
		return err
	}
	return nil
}

// CreateNodeExecutionRecord 创建节点执行记录，返回记录ID
func CreateNodeExecutionRecord(ctx context.Context, execInput *flowDto.FlowExecutionInput, flowNode entity.FlowNode, inputOssUrl string) (int64, error) {
	id, err := nodeDao.NodeExecutionDao.Insert(ctx, &nodeDto.CreateNodeExecutionReq{
		FlowExecutionId: execInput.ExecutionId,
		NodeId:          flowNode.Id,
		NodeName:        flowNode.Name,
		NodeGroupId:     execInput.NodeGroupId,
		InputParamsPath: inputOssUrl,
		Status:          node.NodeExecutionStatusRunning.Code(),
	})
	if err != nil {
		return 0, fmt.Errorf("节点：%v 创建节点执行记录失败: %v", flowNode.Name, err)
	}
	return id, nil
}