20Graph / Chain / Workflow 与类型检查

// I: 编译产物的输入类型;O: 输出类型
type Graph[I, O any] struct {
    *graph        // 内嵌一个非泛型的内部 graph
}

// NewGraph create a directed graph that can compose components, lambda, chain, parallel etc.
// simultaneously provide flexible and multi-granular aspect governance capabilities.
// I: the input type of graph compiled product
// O: the output type of graph compiled product
//
// To share state between nodes, use WithGenLocalState option:
//
//  type testState struct {
//    UserInfo *UserInfo
//    KVs     map[string]any
//  }
//
//  genStateFunc := func(ctx context.Context) *testState {
//    return &testState{}
//  }
//
//  graph := compose.NewGraph[string, string](WithGenLocalState(genStateFunc))
//
//  // you can use WithStatePreHandler and WithStatePostHandler to do something with state
//  graph.AddNode("node1", someNode, compose.WithPreHandler(func(ctx context.Context, in string, state *testState) (string, error) {
//    // do something with state
//    return in, nil
//  }), compose.WithPostHandler(func(ctx context.Context, out string, state *testState) (string, error) {
//    // do something with state
//    return out, nil
//  }))
func NewGraph[I, O any](opts ...NewGraphOption) *Graph[I, O] {
  options := &newGraphOptions{}
  for _, opt := range opts {
    opt(options)
  }

  g := &Graph[I, O]{
    newGraphFromGeneric[I, O](
      ComponentOfGraph,
      options.withState,
      options.stateType,
      opts,
    ),
  }

  return g
}

// AddEdge adds an edge to the graph, edge means a data flow from startNode to endNode.
// the previous node's output type must be set to the next node's input type.
// NOTE: startNode and endNode must have been added to the graph before adding edge.
// e.g.
//
//  graph.AddNode("start_node_key", compose.NewPassthroughNode())
//  graph.AddNode("end_node_key", compose.NewPassthroughNode())
//
//  err := graph.AddEdge("start_node_key", "end_node_key")
func (g *Graph[I, O]) AddEdge(startNode, endNode string) (err error) {
  return g.graph.addEdgeWithMappings(startNode, endNode, false, false)
}

func (g *graph) addEdgeWithMappings(startNode, endNode string, noControl bool, noData bool, mappings ...*FieldMapping) (err error) {
  if g.buildError != nil {
    return g.buildError
  }
  if g.compiled {
    return ErrGraphCompiled
  }

  if noControl && noData {
    return fmt.Errorf("edge[%s]-[%s] cannot be both noDirectDependency and noDataFlow", startNode, endNode)
  }

  defer func() {
    if err != nil {
      g.buildError = err
    }
  }()
  if startNode == END {
    return errors.New("END cannot be a start node")
  }
  if endNode == START {
    return errors.New("START cannot be an end node")
  }

  if _, ok := g.nodes[startNode]; !ok && startNode != START {
    return fmt.Errorf("edge start node '%s' needs to be added to graph first", startNode)
  }
  if _, ok := g.nodes[endNode]; !ok && endNode != END {
    return fmt.Errorf("edge end node '%s' needs to be added to graph first", endNode)
  }

  if !noControl {
    for i := range g.controlEdges[startNode] {
      if g.controlEdges[startNode][i] == endNode {
        return fmt.Errorf("control edge[%s]-[%s] have been added yet", startNode, endNode)
      }
    }

    g.controlEdges[startNode] = append(g.controlEdges[startNode], endNode)
    if startNode == START {
      g.startNodes = append(g.startNodes, endNode)
    }
    if endNode == END {
      g.endNodes = append(g.endNodes, startNode)
    }
  }
  if !noData {
    for i := range g.dataEdges[startNode] {
// … 省略 15 行；完整声明 L232–294，点击上方「浏览完整文件」

type graph struct {
  nodes        map[string]*graphNode
  controlEdges map[string][]string
  dataEdges    map[string][]string
  branches     map[string][]*GraphBranch
  startNodes   []string
  endNodes     []string

  toValidateMap map[string][]struct {
    endNode  string
    mappings []*FieldMapping
  }

  stateType      reflect.Type
  stateGenerator func(ctx context.Context) any
  newOpts        []NewGraphOption

  expectedInputType, expectedOutputType reflect.Type

  *genericHelper

  fieldMappingRecords map[string][]*FieldMapping

  buildError error

  cmp component

  compiled bool

  handlerOnEdges   map[string]map[string][]handlerPair
  handlerPreNode   map[string][]handlerPair
  handlerPreBranch map[string][][]handlerPair
}

// updateToValidateMap after update node, check validate map
// check again if nodes in toValidateMap have been updated. because when there are multiple linked passthrough nodes, in the worst scenario, only one node can be updated at a time.
func (g *graph) updateToValidateMap() error {
  var startNodeOutputType, endNodeInputType reflect.Type
  for {
    hasChanged := false
    for startNode := range g.toValidateMap {
      startNodeOutputType = g.getNodeOutputType(startNode)

      for i := 0; i < len(g.toValidateMap[startNode]); i++ {
        endNode := g.toValidateMap[startNode][i]

        endNodeInputType = g.getNodeInputType(endNode.endNode)
        if startNodeOutputType == nil && endNodeInputType == nil {
          continue
        }

        // update toValidateMap
        g.toValidateMap[startNode] = append(g.toValidateMap[startNode][:i], g.toValidateMap[startNode][i+1:]...)
        i--

        hasChanged = true
        // assume that START and END type isn't empty
        if startNodeOutputType != nil && endNodeInputType == nil {
          g.nodes[endNode.endNode].cr.inputType = startNodeOutputType
          g.nodes[endNode.endNode].cr.outputType = g.nodes[endNode.endNode].cr.inputType
          g.nodes[endNode.endNode].cr.genericHelper = g.getNodeGenericHelper(startNode).forSuccessorPassthrough()
        } else if startNodeOutputType == nil /* redundant condition && endNodeInputType != nil */ {
          g.nodes[startNode].cr.inputType = endNodeInputType
          g.nodes[startNode].cr.outputType = g.nodes[startNode].cr.inputType
          g.nodes[startNode].cr.genericHelper = g.getNodeGenericHelper(endNode.endNode).forPredecessorPassthrough()
        } else if len(endNode.mappings) == 0 {
          // common node check
          result := checkAssignable(startNodeOutputType, endNodeInputType)
          if result == assignableTypeMustNot {
            return fmt.Errorf("graph edge[%s]-[%s]: start node's output type[%s] and end node's input type[%s] mismatch",
              startNode, endNode.endNode, startNodeOutputType.String(), endNodeInputType.String())
          } else if result == assignableTypeMay {
            // add runtime check edges
            if _, ok := g.handlerOnEdges[startNode]; !ok {
              g.handlerOnEdges[startNode] = make(map[string][]handlerPair)
            }
            g.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], g.getNodeGenericHelper(endNode.endNode).inputConverter)
          }
          continue
        }

        if len(endNode.mappings) > 0 {
// … 省略 31 行；完整声明 L559–637，点击上方「浏览完整文件」

// toValidateMap isn't empty means there are nodes that cannot infer type
for _, v := range g.toValidateMap {
    if len(v) > 0 {
        return nil, fmt.Errorf("some node's input or output types cannot be inferred: %v", g.toValidateMap)
    }
}

// NewChain create a chain with input/output type.
func NewChain[I, O any](opts ...NewGraphOption) *Chain[I, O] {
  ch := &Chain[I, O]{
    gg: NewGraph[I, O](opts...),
  }

  ch.gg.cmp = ComponentOfChain

  return ch
}

// Chain is a chain of components.
// Chain nodes can be parallel / branch / sequence components.
// Chain is designed to be used in a builder pattern (should Compile() before use).
// And the interface is `Chain style`, you can use it like: `chain.AppendXX(...).AppendXX(...)`
//
// Normal usage:
//  1. create a chain with input/output type: `chain := NewChain[inputType, outputType]()`
//  2. add components to chainable list:
//     2.1 add components: `chain.AppendChatTemplate(...).AppendChatModel(...).AppendToolsNode(...)`
//     2.2 add parallel or branch node if needed: `chain.AppendParallel()`, `chain.AppendBranch()`
//  3. compile: `r, err := c.Compile()`
//  4. run:
//     4.1 `one input & one output` use `r.Invoke(ctx, input)`
//     4.2 `one input & multi output chunk` use `r.Stream(ctx, input)`
//     4.3 `multi input chunk & one output` use `r.Collect(ctx, inputReader)`
//     4.4 `multi input chunk & multi output chunk` use `r.Transform(ctx, inputReader)`
//
// Using in graph or other chain:
// chain1 := NewChain[inputType, outputType]()
// graph := NewGraph[](runTypePregel)
// graph.AddGraph("key", chain1) // chain is an AnyGraph implementation
//
// // or in another chain:
// chain2 := NewChain[inputType, outputType]()
// chain2.AppendGraph(chain1)
type Chain[I, O any] struct {
  err error

  gg *Graph[I, O]

  nodeIdx int

  preNodeKeys []string

  hasEnd bool
}

// 自动生成 key,加节点
err := c.gg.addNode(nodeKey, node, options)
// ...第一个节点自动从 START 连过来
if len(c.preNodeKeys) == 0 {
    c.preNodeKeys = append(c.preNodeKeys, START)
}
// 把上一个节点连到当前节点
for _, preNodeKey := range c.preNodeKeys {
    c.gg.AddEdge(preNodeKey, nodeKey)
}
c.preNodeKeys = []string{nodeKey}   // 当前节点成为下一个的"上游"

// Workflow is wrapper of graph, replacing AddEdge with declaring dependencies
// and field mappings between nodes.
// Under the hood it uses NodeTriggerMode(AllPredecessor), so does not support cycles.
type Workflow[I, O any] struct {
    g                *graph
    workflowNodes    map[string]*WorkflowNode
    // ...
}

// AddInput creates both data and execution dependencies between nodes.
// It configures how data flows from the predecessor node (fromNodeKey) to the current node,
// and ensures the current node only executes after the predecessor completes.
//
// Parameters:
//   - fromNodeKey: the key of the predecessor node
//   - inputs: field mappings that specify how data should flow from the predecessor
//     to the current node. If no mappings are provided, the entire output of the
//     predecessor will be used as input.
//
// Example:
//
//  // Map between specific field
//  node.AddInput("userNode", MapFields("user.name", "displayName"))
//
//  // Use entire output
//  node.AddInput("dataNode")
//
// Returns the current node for method chaining.
func (n *WorkflowNode) AddInput(fromNodeKey string, inputs ...*FieldMapping) *WorkflowNode {
  return n.addDependencyRelation(fromNodeKey, inputs, &workflowAddInputOpts{})
}

// MapFields creates a FieldMapping that maps a single predecessor field to a single successor field.
// Field: either the field of a struct, or the key of a map.
func MapFields(from, to string) *FieldMapping {
  return &FieldMapping{
    from: from,
    to:   to,
  }
}

// FromField creates a FieldMapping that maps a single predecessor field to the entire successor input.
// This is an exclusive mapping - once set, no other field mappings can be added since the successor input
// has already been fully mapped.
// Field: either the field of a struct, or the key of a map.
func FromField(from string) *FieldMapping {
  return &FieldMapping{
    from: from,
  }
}

// ToField creates a FieldMapping that maps the entire predecessor output to a single successor field.
// Field: either the field of a struct, or the key of a map.
func ToField(to string, opts ...FieldMappingOption) *FieldMapping {
  fm := &FieldMapping{
    to: to,
  }
  for _, opt := range opts {
    opt(fm)
  }
  return fm
}

type FieldMapping struct {
  fromNodeKey string
  from        string
  to          string

  customExtractor func(input any) (any, error)
}

func (n *WorkflowNode) addDependencyRelation(fromNodeKey string, inputs []*FieldMapping, options *workflowAddInputOpts) *WorkflowNode {
  for _, input := range inputs {
    input.fromNodeKey = fromNodeKey
  }

  if options.noDirectDependency {
    n.addInputs = append(n.addInputs, func() error {
      var paths []FieldPath
      for _, input := range inputs {
        paths = append(paths, input.targetPath())
      }
      if err := n.checkAndAddMappedPath(paths); err != nil {
        return err
      }

      if err := n.g.addEdgeWithMappings(fromNodeKey, n.key, true, false, inputs...); err != nil {
        return err
      }
      n.dependencySetter(fromNodeKey, noDirectDependency)
      return nil
    })
  } else if options.dependencyWithoutInput {
    n.addInputs = append(n.addInputs, func() error {
      if len(inputs) > 0 {
        return fmt.Errorf("dependency without input should not have inputs. node: %s, fromNode: %s, inputs: %v", n.key, fromNodeKey, inputs)
      }
      if err := n.g.addEdgeWithMappings(fromNodeKey, n.key, false, true); err != nil {
        return err
      }
      n.dependencySetter(fromNodeKey, normalDependency)
      return nil
    })
  } else {
    n.addInputs = append(n.addInputs, func() error {
      var paths []FieldPath
      for _, input := range inputs {
        paths = append(paths, input.targetPath())
      }
      if err := n.checkAndAddMappedPath(paths); err != nil {
        return err
      }

      if err := n.g.addEdgeWithMappings(fromNodeKey, n.key, false, false, inputs...); err != nil {
        return err
      }
      n.dependencySetter(fromNodeKey, normalDependency)
      return nil
    })
// … 省略 4 行；完整声明 L316–367，点击上方「浏览完整文件」