69 lines
2.5 KiB
Go
69 lines
2.5 KiB
Go
package flow
|
||
|
||
import (
|
||
"ai-agent/workflow/consts/node"
|
||
flowDto "ai-agent/workflow/model/dto/flow"
|
||
"ai-agent/workflow/model/entity"
|
||
"context"
|
||
"encoding/json"
|
||
"time"
|
||
|
||
flowDao "ai-agent/workflow/dao/flow"
|
||
|
||
"github.com/cloudwego/eino/compose"
|
||
"github.com/cloudwego/eino/schema"
|
||
)
|
||
|
||
// 注册 checkpoint 序列化类型
|
||
func init() {
|
||
// ========== 1. Eino 断点核心状态（根类型） ==========
|
||
schema.RegisterName[*flowDto.NodeExecutionState]("flow.NodeExecutionState")
|
||
schema.RegisterName[*flowDto.NodeExecutionInput]("flow.NodeExecutionInput")
|
||
schema.RegisterName[*flowDto.FlowExecutionInput]("flow.FlowExecutionInput")
|
||
|
||
// ========== 2. 原有第三方类型 ==========
|
||
schema.RegisterName[json.Number]("json.Number")
|
||
schema.RegisterName[time.Time]("time.Time")
|
||
schema.RegisterName[time.Duration]("time.Duration")
|
||
|
||
// ========== 3. flowDto 内部嵌套类型 ==========
|
||
schema.RegisterName[flowDto.ExecutedNode]("flow.ExecutedNode")
|
||
|
||
// ========== 4. entity 核心链路类型（递归自 *entity.FlowNode） ==========
|
||
schema.RegisterName[*entity.FlowNode]("entity.FlowNode")
|
||
schema.RegisterName[node.NodeType]("node.NodeType")
|
||
schema.RegisterName[*entity.SubFlowConfig]("entity.SubFlowConfig")
|
||
//schema.RegisterName[node.NodeFormField]("node.NodeFormField")
|
||
schema.RegisterName[entity.ModelItem]("node.ModelItem")
|
||
|
||
// ========== 5. FlowInfo 相关（流程拓扑结构） ==========
|
||
schema.RegisterName[entity.FlowInfo]("entity.FlowInfo")
|
||
schema.RegisterName[entity.FlowEdge]("entity.FlowEdge")
|
||
}
|
||
|
||
// DbCheckPointStore 数据库存储实现
|
||
type DbCheckPointStore struct{}
|
||
|
||
func NewDbCheckPointStore() compose.CheckPointStore {
|
||
return &DbCheckPointStore{}
|
||
}
|
||
|
||
func (d *DbCheckPointStore) Get(ctx context.Context, id string) ([]byte, bool, error) {
|
||
// 去掉取消信号：断连/取消场景下 graph ctx 已取消，仅保留值，避免恢复时被阻断
|
||
record, err := flowDao.FlowCheckpointDao.Get(context.WithoutCancel(ctx), id)
|
||
if err != nil {
|
||
return nil, false, err
|
||
}
|
||
if record == nil || record.Data == "" {
|
||
return nil, false, nil
|
||
}
|
||
return []byte(record.Data), true, nil
|
||
}
|
||
|
||
func (d *DbCheckPointStore) Set(ctx context.Context, id string, val []byte) error {
|
||
// 关键：Eino 在节点失败（Interrupt）时用 graph ctx 写 checkpoint。WS 断连/用户终止会使
|
||
// 该 ctx 已取消，直接透传会因 "context canceled" 落库失败，断点丢失、续跑失效。
|
||
// 去掉取消信号只保留 ctx 值，保证中断时断点必达。
|
||
return flowDao.FlowCheckpointDao.SaveOrUpdate(context.WithoutCancel(ctx), id, string(val))
|
||
}
|