mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-05 09:17:47 +02:00
rewrite-state-graph
This commit is contained in:
Quanzheng Long
@@ -1,245 +0,0 @@
|
||||
package advancedgraph
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type StateNodeFunc[StateT any] func(ctx *Context, state StateT) (StateT, error)
|
||||
|
||||
type BasicStateGraph[StateT any] struct {
|
||||
nodes map[string]StateNodeFunc[StateT]
|
||||
edges map[string][]string
|
||||
interruptChannel string
|
||||
interruptSteps map[int]struct{}
|
||||
}
|
||||
|
||||
func NewBasicStateGraph[StateT any]() *BasicStateGraph[StateT] {
|
||||
return &BasicStateGraph[StateT]{
|
||||
nodes: make(map[string]StateNodeFunc[StateT]),
|
||||
edges: make(map[string][]string),
|
||||
interruptSteps: make(map[int]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) AddNode(name string, fn StateNodeFunc[StateT]) {
|
||||
if name == "" {
|
||||
panic("node name cannot be empty")
|
||||
}
|
||||
if _, exists := g.nodes[name]; exists {
|
||||
panic(fmt.Sprintf("node `%s` already exists", name))
|
||||
}
|
||||
g.nodes[name] = fn
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) AddEdge(from string, to string) {
|
||||
if _, ok := g.nodes[from]; !ok {
|
||||
panic(fmt.Sprintf("source node `%s` does not exist", from))
|
||||
}
|
||||
if _, ok := g.nodes[to]; !ok {
|
||||
panic(fmt.Sprintf("target node `%s` does not exist", to))
|
||||
}
|
||||
g.edges[from] = append(g.edges[from], to)
|
||||
}
|
||||
|
||||
// EnableInterruptOnSuperstep enables pause/resume before dispatching a superstep.
|
||||
// superstep=1 means "after first superstep has completed, before second starts".
|
||||
func (g *BasicStateGraph[StateT]) EnableInterruptOnSuperstep(superstep int, channel string) {
|
||||
if superstep <= 0 {
|
||||
panic("interrupt superstep must be >= 1")
|
||||
}
|
||||
if channel == "" {
|
||||
panic("interrupt channel cannot be empty")
|
||||
}
|
||||
if g.interruptChannel != "" && g.interruptChannel != channel {
|
||||
panic("all interrupts must use the same channel")
|
||||
}
|
||||
g.interruptChannel = channel
|
||||
g.interruptSteps[superstep] = struct{}{}
|
||||
}
|
||||
|
||||
type CompiledBasicStateGraph[StateT any] struct {
|
||||
inner *CompiledGraph[StateT]
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] {
|
||||
if len(g.nodes) == 0 {
|
||||
panic("graph has no nodes")
|
||||
}
|
||||
levels, err := g.computeSupersteps()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
adv := NewAdvancedStateGraph[StateT]()
|
||||
|
||||
const finalNodeName = "__stategraph_finish"
|
||||
finalNode := func(_ *Context, _ any, state StateT) (Command, error) {
|
||||
return Command{Update: state}, nil
|
||||
}
|
||||
adv.AddFinishNodeAs(finalNodeName, finalNode)
|
||||
|
||||
if g.interruptChannel != "" {
|
||||
adv.AddAsyncChannel(g.interruptChannel)
|
||||
}
|
||||
|
||||
for stepIdx, stepNodes := range levels {
|
||||
for _, nodeName := range stepNodes {
|
||||
userFn := g.nodes[nodeName]
|
||||
nextBarrier := fmt.Sprintf("__stategraph_barrier_%d", stepIdx+1)
|
||||
wrapper := func(ctx *Context, _ any, state StateT) (Command, error) {
|
||||
updated, err := userFn(ctx, state)
|
||||
if err != nil {
|
||||
return Command{}, err
|
||||
}
|
||||
return Command{
|
||||
Update: updated,
|
||||
Goto: []Send{{Node: nextBarrier}},
|
||||
}, nil
|
||||
}
|
||||
adv.AddNodeAs(fmt.Sprintf("__stategraph_node_%s", nodeName), wrapper)
|
||||
}
|
||||
}
|
||||
|
||||
type barrierCounter struct {
|
||||
mu sync.Mutex
|
||||
counts map[*RustEngine]int
|
||||
}
|
||||
counters := make(map[int]*barrierCounter)
|
||||
for barrierStep := 1; barrierStep <= len(levels); barrierStep++ {
|
||||
counters[barrierStep] = &barrierCounter{
|
||||
counts: make(map[*RustEngine]int),
|
||||
}
|
||||
}
|
||||
|
||||
lastBarrier := len(levels)
|
||||
for barrierStep := 0; barrierStep <= lastBarrier; barrierStep++ {
|
||||
barrierName := fmt.Sprintf("__stategraph_barrier_%d", barrierStep)
|
||||
nextStep := barrierStep
|
||||
barrier := func(ctx *Context, _ any, state StateT) (Command, error) {
|
||||
if nextStep > 0 {
|
||||
counter := counters[nextStep]
|
||||
counter.mu.Lock()
|
||||
counter.counts[ctx.engine]++
|
||||
current := counter.counts[ctx.engine]
|
||||
needed := len(levels[nextStep-1])
|
||||
if current < needed {
|
||||
counter.mu.Unlock()
|
||||
return Command{Update: state}, nil
|
||||
}
|
||||
delete(counter.counts, ctx.engine)
|
||||
counter.mu.Unlock()
|
||||
}
|
||||
if _, needsInterrupt := g.interruptSteps[nextStep]; needsInterrupt {
|
||||
cond := AnyOf(ChannelCondition{Channel: g.interruptChannel, N: 1})
|
||||
if _, err := ctx.WaitFor(cond); err != nil {
|
||||
return Command{}, err
|
||||
}
|
||||
}
|
||||
if nextStep >= len(levels) {
|
||||
return Command{
|
||||
Update: state,
|
||||
Goto: []Send{{Node: finalNodeName}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
sends := make([]Send, 0, len(levels[nextStep]))
|
||||
for _, nodeName := range levels[nextStep] {
|
||||
sends = append(sends, Send{
|
||||
Node: fmt.Sprintf("__stategraph_node_%s", nodeName),
|
||||
})
|
||||
}
|
||||
return Command{
|
||||
Update: state,
|
||||
Goto: sends,
|
||||
}, nil
|
||||
}
|
||||
if barrierStep == 0 {
|
||||
adv.AddEntryNodeAs(barrierName, barrier)
|
||||
} else {
|
||||
adv.AddNodeAs(barrierName, barrier)
|
||||
}
|
||||
}
|
||||
|
||||
return &CompiledBasicStateGraph[StateT]{
|
||||
inner: adv.Compile(),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *CompiledBasicStateGraph[StateT]) Start(initialState StateT) (*Handler[StateT], error) {
|
||||
return g.inner.Start(nil, initialState)
|
||||
}
|
||||
|
||||
func (g *CompiledBasicStateGraph[StateT]) Invoke(initialState StateT) (StateT, error) {
|
||||
handler, err := g.Start(initialState)
|
||||
if err != nil {
|
||||
var zero StateT
|
||||
return zero, err
|
||||
}
|
||||
return handler.WaitForResult()
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) computeSupersteps() ([][]string, error) {
|
||||
indegree := make(map[string]int, len(g.nodes))
|
||||
for name := range g.nodes {
|
||||
indegree[name] = 0
|
||||
}
|
||||
for from, tos := range g.edges {
|
||||
if _, ok := g.nodes[from]; !ok {
|
||||
return nil, fmt.Errorf("edge source `%s` does not exist", from)
|
||||
}
|
||||
for _, to := range tos {
|
||||
if _, ok := g.nodes[to]; !ok {
|
||||
return nil, fmt.Errorf("edge target `%s` does not exist", to)
|
||||
}
|
||||
indegree[to]++
|
||||
}
|
||||
}
|
||||
|
||||
queue := make([]string, 0, len(g.nodes))
|
||||
level := make(map[string]int, len(g.nodes))
|
||||
for name, deg := range indegree {
|
||||
if deg == 0 {
|
||||
queue = append(queue, name)
|
||||
}
|
||||
}
|
||||
if len(queue) == 0 {
|
||||
return nil, fmt.Errorf("graph has no entry nodes (cycle suspected)")
|
||||
}
|
||||
|
||||
processed := 0
|
||||
for len(queue) > 0 {
|
||||
curr := queue[0]
|
||||
queue = queue[1:]
|
||||
processed++
|
||||
currLevel := level[curr]
|
||||
for _, to := range g.edges[curr] {
|
||||
if level[to] < currLevel+1 {
|
||||
level[to] = currLevel + 1
|
||||
}
|
||||
indegree[to]--
|
||||
if indegree[to] == 0 {
|
||||
queue = append(queue, to)
|
||||
}
|
||||
}
|
||||
}
|
||||
if processed != len(g.nodes) {
|
||||
return nil, fmt.Errorf("graph contains a cycle")
|
||||
}
|
||||
|
||||
maxLevel := 0
|
||||
for _, lv := range level {
|
||||
if lv > maxLevel {
|
||||
maxLevel = lv
|
||||
}
|
||||
}
|
||||
levels := make([][]string, maxLevel+1)
|
||||
for nodeName := range g.nodes {
|
||||
lv := level[nodeName]
|
||||
levels[lv] = append(levels[lv], nodeName)
|
||||
}
|
||||
for i := range levels {
|
||||
slices.Sort(levels[i])
|
||||
}
|
||||
return levels, nil
|
||||
}
|
||||
@@ -0,0 +1,320 @@
|
||||
package stategraph
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type StateNodeFunc[StateT any] func(ctx *Context, state StateT) (StateT, error)
|
||||
|
||||
const (
|
||||
internalBarrierChannel = "__stategraph_barrier"
|
||||
internalInterruptChannel = "__stategraph_interrupt"
|
||||
)
|
||||
|
||||
type Context struct {
|
||||
inner *ag.Context
|
||||
}
|
||||
|
||||
func (c *Context) Interrupt(name string) (any, error) {
|
||||
if name == "" {
|
||||
return nil, fmt.Errorf("interrupt name cannot be empty")
|
||||
}
|
||||
event, err := c.inner.WaitFor(ag.AnyOf(ag.ChannelCondition{
|
||||
Channel: internalInterruptChannel,
|
||||
N: 1,
|
||||
}))
|
||||
if err != nil {
|
||||
if waitReq, ok := ag.AsErrWaitRequested(err); ok {
|
||||
return nil, errInterruptRequested{
|
||||
Name: name,
|
||||
Condition: waitReq.Condition,
|
||||
}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if len(event.Value) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var payload interruptPayload
|
||||
if err := json.Unmarshal(event.Value, &payload); err != nil {
|
||||
var value any
|
||||
if err := json.Unmarshal(event.Value, &value); err != nil {
|
||||
return nil, fmt.Errorf("decode interrupt `%s` value: %w", name, err)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
if payload.Name != "" && payload.Name != name {
|
||||
return nil, fmt.Errorf("interrupt name mismatch: expected `%s`, got `%s`", name, payload.Name)
|
||||
}
|
||||
if len(payload.Value) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var value any
|
||||
if err := json.Unmarshal(payload.Value, &value); err != nil {
|
||||
return nil, fmt.Errorf("decode interrupt `%s` payload: %w", name, err)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
type errInterruptRequested struct {
|
||||
Name string
|
||||
Condition ag.AnyOfCondition
|
||||
}
|
||||
|
||||
func (e errInterruptRequested) Error() string {
|
||||
if e.Name == "" {
|
||||
return "interrupt requested"
|
||||
}
|
||||
return fmt.Sprintf("interrupt requested: %s", e.Name)
|
||||
}
|
||||
|
||||
func asErrInterruptRequested(err error) (errInterruptRequested, bool) {
|
||||
var target errInterruptRequested
|
||||
if !errors.As(err, &target) {
|
||||
return target, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
|
||||
type BasicStateGraph[StateT any] struct {
|
||||
nodes map[string]StateNodeFunc[StateT]
|
||||
edges map[string][]string
|
||||
}
|
||||
|
||||
type interruptPayload struct {
|
||||
Name string `json:"name"`
|
||||
Value json.RawMessage `json:"value"`
|
||||
}
|
||||
|
||||
func NewBasicStateGraph[StateT any]() *BasicStateGraph[StateT] {
|
||||
return &BasicStateGraph[StateT]{
|
||||
nodes: make(map[string]StateNodeFunc[StateT]),
|
||||
edges: make(map[string][]string),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) AddNode(fn StateNodeFunc[StateT]) string {
|
||||
name := ag.NodeName(fn)
|
||||
if _, exists := g.nodes[name]; exists {
|
||||
panic(fmt.Sprintf("node `%s` already exists", name))
|
||||
}
|
||||
g.nodes[name] = fn
|
||||
return name
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) AddEdge(from StateNodeFunc[StateT], to StateNodeFunc[StateT]) {
|
||||
fromName := ag.NodeName(from)
|
||||
toName := ag.NodeName(to)
|
||||
if _, ok := g.nodes[fromName]; !ok {
|
||||
panic(fmt.Sprintf("source node `%s` does not exist", fromName))
|
||||
}
|
||||
if _, ok := g.nodes[toName]; !ok {
|
||||
panic(fmt.Sprintf("target node `%s` does not exist", toName))
|
||||
}
|
||||
g.edges[fromName] = append(g.edges[fromName], toName)
|
||||
}
|
||||
|
||||
type CompiledBasicStateGraph[StateT any] struct {
|
||||
inner *ag.CompiledGraph[StateT]
|
||||
}
|
||||
|
||||
type Handler[StateT any] struct {
|
||||
inner *ag.Handler[StateT]
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) WaitForResult() (StateT, error) {
|
||||
return h.inner.WaitForResult()
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) Interrupt(name string, value any) error {
|
||||
if name == "" {
|
||||
return fmt.Errorf("interrupt name cannot be empty")
|
||||
}
|
||||
return h.inner.PublishToChannel(internalInterruptChannel, map[string]any{
|
||||
"name": name,
|
||||
"value": value,
|
||||
})
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] {
|
||||
if len(g.nodes) == 0 {
|
||||
panic("graph has no nodes")
|
||||
}
|
||||
levels, err := g.computeSupersteps()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
adv := ag.NewAdvancedStateGraph[StateT]()
|
||||
adv.AddAsyncChannel(internalBarrierChannel)
|
||||
adv.AddAsyncChannel(internalInterruptChannel)
|
||||
|
||||
const finalNodeName = "__stategraph_finish"
|
||||
finalNode := func(_ *ag.Context, _ any, state StateT) (ag.Command, error) {
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
adv.AddFinishNodeAs(finalNodeName, finalNode)
|
||||
|
||||
for stepIdx, stepNodes := range levels {
|
||||
for _, nodeName := range stepNodes {
|
||||
userFn := g.nodes[nodeName]
|
||||
nextBarrier := fmt.Sprintf("__stategraph_barrier_%d", stepIdx+1)
|
||||
wrapper := func(ctx *ag.Context, _ any, state StateT) (ag.Command, error) {
|
||||
updated, err := userFn(&Context{inner: ctx}, state)
|
||||
if err != nil {
|
||||
if interruptReq, ok := asErrInterruptRequested(err); ok {
|
||||
cond := interruptReq.Condition
|
||||
if len(cond.Conditions) == 0 {
|
||||
cond = ag.AnyOf(ag.ChannelCondition{
|
||||
Channel: internalInterruptChannel,
|
||||
N: 1,
|
||||
})
|
||||
}
|
||||
return ag.Command{}, ag.ErrWaitRequested{Condition: cond}
|
||||
}
|
||||
return ag.Command{}, err
|
||||
}
|
||||
if err := ctx.PublishToChannel(internalBarrierChannel, map[string]any{
|
||||
"step": stepIdx,
|
||||
}); err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
return ag.Command{
|
||||
Update: updated,
|
||||
Goto: []ag.Send{{Node: nextBarrier}},
|
||||
}, nil
|
||||
}
|
||||
adv.AddNodeAs(fmt.Sprintf("__stategraph_node_%s", nodeName), wrapper)
|
||||
}
|
||||
}
|
||||
|
||||
lastBarrier := len(levels)
|
||||
for barrierStep := 0; barrierStep <= lastBarrier; barrierStep++ {
|
||||
barrierName := fmt.Sprintf("__stategraph_barrier_%d", barrierStep)
|
||||
nextStep := barrierStep
|
||||
barrier := func(ctx *ag.Context, _ any, state StateT) (ag.Command, error) {
|
||||
if nextStep > 0 {
|
||||
needed := len(levels[nextStep-1])
|
||||
_, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{
|
||||
Channel: internalBarrierChannel,
|
||||
N: needed,
|
||||
}))
|
||||
if err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
}
|
||||
if nextStep >= len(levels) {
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: []ag.Send{{Node: finalNodeName}},
|
||||
}, nil
|
||||
}
|
||||
sends := make([]ag.Send, 0, len(levels[nextStep]))
|
||||
for _, nodeName := range levels[nextStep] {
|
||||
sends = append(sends, ag.Send{
|
||||
Node: fmt.Sprintf("__stategraph_node_%s", nodeName),
|
||||
})
|
||||
}
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: sends,
|
||||
}, nil
|
||||
}
|
||||
if barrierStep == 0 {
|
||||
adv.AddEntryNodeAs(barrierName, barrier)
|
||||
} else {
|
||||
adv.AddNodeAs(barrierName, barrier)
|
||||
}
|
||||
}
|
||||
|
||||
return &CompiledBasicStateGraph[StateT]{
|
||||
inner: adv.Compile(),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *CompiledBasicStateGraph[StateT]) Start(initialState StateT) (*Handler[StateT], error) {
|
||||
raw, err := g.inner.Start(nil, initialState)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Handler[StateT]{inner: raw}, nil
|
||||
}
|
||||
|
||||
func (g *CompiledBasicStateGraph[StateT]) Invoke(initialState StateT) (StateT, error) {
|
||||
handler, err := g.Start(initialState)
|
||||
if err != nil {
|
||||
var zero StateT
|
||||
return zero, err
|
||||
}
|
||||
return handler.WaitForResult()
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) computeSupersteps() ([][]string, error) {
|
||||
indegree := make(map[string]int, len(g.nodes))
|
||||
for name := range g.nodes {
|
||||
indegree[name] = 0
|
||||
}
|
||||
for from, tos := range g.edges {
|
||||
if _, ok := g.nodes[from]; !ok {
|
||||
return nil, fmt.Errorf("edge source `%s` does not exist", from)
|
||||
}
|
||||
for _, to := range tos {
|
||||
if _, ok := g.nodes[to]; !ok {
|
||||
return nil, fmt.Errorf("edge target `%s` does not exist", to)
|
||||
}
|
||||
indegree[to]++
|
||||
}
|
||||
}
|
||||
|
||||
queue := make([]string, 0, len(g.nodes))
|
||||
level := make(map[string]int, len(g.nodes))
|
||||
for name, deg := range indegree {
|
||||
if deg == 0 {
|
||||
queue = append(queue, name)
|
||||
}
|
||||
}
|
||||
if len(queue) == 0 {
|
||||
return nil, fmt.Errorf("graph has no entry nodes (cycle suspected)")
|
||||
}
|
||||
|
||||
processed := 0
|
||||
for len(queue) > 0 {
|
||||
curr := queue[0]
|
||||
queue = queue[1:]
|
||||
processed++
|
||||
currLevel := level[curr]
|
||||
for _, to := range g.edges[curr] {
|
||||
if level[to] < currLevel+1 {
|
||||
level[to] = currLevel + 1
|
||||
}
|
||||
indegree[to]--
|
||||
if indegree[to] == 0 {
|
||||
queue = append(queue, to)
|
||||
}
|
||||
}
|
||||
}
|
||||
if processed != len(g.nodes) {
|
||||
return nil, fmt.Errorf("graph contains a cycle")
|
||||
}
|
||||
|
||||
maxLevel := 0
|
||||
for _, lv := range level {
|
||||
if lv > maxLevel {
|
||||
maxLevel = lv
|
||||
}
|
||||
}
|
||||
levels := make([][]string, maxLevel+1)
|
||||
for nodeName := range g.nodes {
|
||||
lv := level[nodeName]
|
||||
levels[lv] = append(levels[lv], nodeName)
|
||||
}
|
||||
for i := range levels {
|
||||
slices.Sort(levels[i])
|
||||
}
|
||||
return levels, nil
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
package stategraph_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
sg "github.com/langchain-ai/langgraph/langgraph-go/stategraph"
|
||||
)
|
||||
|
||||
type stateGraphState struct {
|
||||
Noop bool `json:"noop"`
|
||||
}
|
||||
|
||||
type orderRecorder struct {
|
||||
mu sync.Mutex
|
||||
orders map[string]int32
|
||||
seq int32
|
||||
}
|
||||
|
||||
func newOrderRecorder() *orderRecorder {
|
||||
return &orderRecorder{orders: make(map[string]int32)}
|
||||
}
|
||||
|
||||
func (f *orderRecorder) record(name string) {
|
||||
idx := atomic.AddInt32(&f.seq, 1)
|
||||
f.mu.Lock()
|
||||
f.orders[name] = idx
|
||||
f.mu.Unlock()
|
||||
}
|
||||
|
||||
type orderGraph struct {
|
||||
recorder *orderRecorder
|
||||
}
|
||||
|
||||
func newOrderGraph() *orderGraph {
|
||||
return &orderGraph{
|
||||
recorder: newOrderRecorder(),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *orderGraph) A(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("A")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) B1(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("B1")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) B2(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("B2")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) C1(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("C1")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) C2(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("C2")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) C3(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("C3")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) D(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("D")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func TestBasicStateGraphWithoutInterrupt(t *testing.T) {
|
||||
fixture := newOrderGraph()
|
||||
graph := sg.NewBasicStateGraph[stateGraphState]()
|
||||
graph.AddNode(fixture.A)
|
||||
graph.AddNode(fixture.B1)
|
||||
graph.AddNode(fixture.B2)
|
||||
graph.AddNode(fixture.C1)
|
||||
graph.AddNode(fixture.C2)
|
||||
graph.AddNode(fixture.C3)
|
||||
graph.AddNode(fixture.D)
|
||||
|
||||
graph.AddEdge(fixture.A, fixture.B1)
|
||||
graph.AddEdge(fixture.A, fixture.B2)
|
||||
graph.AddEdge(fixture.B1, fixture.C1)
|
||||
graph.AddEdge(fixture.B1, fixture.C2)
|
||||
graph.AddEdge(fixture.B2, fixture.C3)
|
||||
graph.AddEdge(fixture.C1, fixture.D)
|
||||
graph.AddEdge(fixture.C2, fixture.D)
|
||||
graph.AddEdge(fixture.C3, fixture.D)
|
||||
|
||||
_, err := graph.Compile().Invoke(stateGraphState{})
|
||||
if err != nil {
|
||||
t.Fatalf("invoke failed: %v", err)
|
||||
}
|
||||
|
||||
fixture.recorder.mu.Lock()
|
||||
orders := make(map[string]int32, len(fixture.recorder.orders))
|
||||
for k, v := range fixture.recorder.orders {
|
||||
orders[k] = v
|
||||
}
|
||||
fixture.recorder.mu.Unlock()
|
||||
|
||||
for _, name := range []string{"A", "B1", "B2", "C1", "C2", "C3", "D"} {
|
||||
if _, ok := orders[name]; !ok {
|
||||
t.Fatalf("node %s did not execute; orders=%v", name, orders)
|
||||
}
|
||||
}
|
||||
maxB := maxInt32(orders["B1"], orders["B2"])
|
||||
minC := minInt32(orders["C1"], minInt32(orders["C2"], orders["C3"]))
|
||||
maxC := maxInt32(orders["C1"], maxInt32(orders["C2"], orders["C3"]))
|
||||
if !(orders["A"] < orders["B1"] && orders["A"] < orders["B2"]) {
|
||||
t.Fatalf("A should run before B-step, orders=%v", orders)
|
||||
}
|
||||
if !(maxB < minC) {
|
||||
t.Fatalf("B-step should finish before C-step, orders=%v", orders)
|
||||
}
|
||||
if !(maxC < orders["D"]) {
|
||||
t.Fatalf("C-step should finish before D, orders=%v", orders)
|
||||
}
|
||||
}
|
||||
|
||||
type interruptState struct {
|
||||
A bool `json:"a"`
|
||||
B bool `json:"b"`
|
||||
}
|
||||
|
||||
type interruptFixture struct{}
|
||||
|
||||
func (f *interruptFixture) A(_ *sg.Context, state interruptState) (interruptState, error) {
|
||||
state.A = true
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (f *interruptFixture) B(ctx *sg.Context, state interruptState) (interruptState, error) {
|
||||
if !state.A {
|
||||
return state, fmt.Errorf("B should observe A=true")
|
||||
}
|
||||
value, err := ctx.Interrupt("resume_channel")
|
||||
if err != nil {
|
||||
return state, err
|
||||
}
|
||||
s, ok := value.(string)
|
||||
if !ok || s != "go" {
|
||||
return state, fmt.Errorf("unexpected interrupt payload: %#v", value)
|
||||
}
|
||||
state.B = true
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func TestBasicStateGraphWithInterrupt(t *testing.T) {
|
||||
fixture := &interruptFixture{}
|
||||
graph := sg.NewBasicStateGraph[interruptState]()
|
||||
graph.AddNode(fixture.A)
|
||||
graph.AddNode(fixture.B)
|
||||
graph.AddEdge(fixture.A, fixture.B)
|
||||
|
||||
handler, err := graph.Compile().Start(interruptState{})
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
doneCh := make(chan interruptState, 1)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
result, runErr := handler.WaitForResult()
|
||||
if runErr != nil {
|
||||
errCh <- runErr
|
||||
return
|
||||
}
|
||||
doneCh <- result
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-doneCh:
|
||||
t.Fatalf("run should pause for interrupt, but completed early")
|
||||
case err := <-errCh:
|
||||
t.Fatalf("run should pause for interrupt, but failed early: %v", err)
|
||||
case <-time.After(120 * time.Millisecond):
|
||||
// expected: paused
|
||||
}
|
||||
|
||||
if err := handler.Interrupt("resume_channel", "go"); err != nil {
|
||||
t.Fatalf("resume interrupt failed: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("run failed after interrupt: %v", err)
|
||||
case result := <-doneCh:
|
||||
if !(result.A && result.B) {
|
||||
t.Fatalf("unexpected final state: %#v", result)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatalf("timeout waiting for resumed run completion")
|
||||
}
|
||||
}
|
||||
|
||||
func minInt32(a int32, b int32) int32 {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func maxInt32(a int32, b int32) int32 {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -1,164 +0,0 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type stateGraphState struct {
|
||||
Noop bool `json:"noop"`
|
||||
}
|
||||
|
||||
func TestBasicStateGraphWithoutInterrupt(t *testing.T) {
|
||||
var (
|
||||
orderMu sync.Mutex
|
||||
orders = make(map[string]int32)
|
||||
seq int32
|
||||
)
|
||||
record := func(name string) {
|
||||
idx := atomic.AddInt32(&seq, 1)
|
||||
orderMu.Lock()
|
||||
orders[name] = idx
|
||||
orderMu.Unlock()
|
||||
}
|
||||
graph := ag.NewBasicStateGraph[stateGraphState]()
|
||||
graph.AddNode("A", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("A")
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("B1", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("B1")
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("B2", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("B2")
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("C1", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("C1")
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("C2", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("C2")
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("C3", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("C3")
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("D", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) {
|
||||
record("D")
|
||||
return state, nil
|
||||
})
|
||||
|
||||
graph.AddEdge("A", "B1")
|
||||
graph.AddEdge("A", "B2")
|
||||
graph.AddEdge("B1", "C1")
|
||||
graph.AddEdge("B1", "C2")
|
||||
graph.AddEdge("B2", "C3")
|
||||
graph.AddEdge("C1", "D")
|
||||
graph.AddEdge("C2", "D")
|
||||
graph.AddEdge("C3", "D")
|
||||
|
||||
_, err := graph.Compile().Invoke(stateGraphState{})
|
||||
if err != nil {
|
||||
t.Fatalf("invoke failed: %v", err)
|
||||
}
|
||||
for _, name := range []string{"A", "B1", "B2", "C1", "C2", "C3", "D"} {
|
||||
if _, ok := orders[name]; !ok {
|
||||
t.Fatalf("node %s did not execute; orders=%v", name, orders)
|
||||
}
|
||||
}
|
||||
maxB := maxInt32(orders["B1"], orders["B2"])
|
||||
minC := minInt32(orders["C1"], minInt32(orders["C2"], orders["C3"]))
|
||||
maxC := maxInt32(orders["C1"], maxInt32(orders["C2"], orders["C3"]))
|
||||
if !(orders["A"] < orders["B1"] && orders["A"] < orders["B2"]) {
|
||||
t.Fatalf("A should run before B-step, orders=%v", orders)
|
||||
}
|
||||
if !(maxB < minC) {
|
||||
t.Fatalf("B-step should finish before C-step, orders=%v", orders)
|
||||
}
|
||||
if !(maxC < orders["D"]) {
|
||||
t.Fatalf("C-step should finish before D, orders=%v", orders)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBasicStateGraphWithInterrupt(t *testing.T) {
|
||||
type interruptState struct {
|
||||
A bool `json:"a"`
|
||||
B bool `json:"b"`
|
||||
}
|
||||
graph := ag.NewBasicStateGraph[interruptState]()
|
||||
graph.AddNode("A", func(_ *ag.Context, state interruptState) (interruptState, error) {
|
||||
state.A = true
|
||||
return state, nil
|
||||
})
|
||||
graph.AddNode("B", func(_ *ag.Context, state interruptState) (interruptState, error) {
|
||||
if !state.A {
|
||||
return state, fmt.Errorf("B should observe A=true")
|
||||
}
|
||||
state.B = true
|
||||
return state, nil
|
||||
})
|
||||
graph.AddEdge("A", "B")
|
||||
graph.EnableInterruptOnSuperstep(1, "resume_channel")
|
||||
|
||||
handler, err := graph.Compile().Start(interruptState{})
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
doneCh := make(chan interruptState, 1)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
result, runErr := handler.WaitForResult()
|
||||
if runErr != nil {
|
||||
errCh <- runErr
|
||||
return
|
||||
}
|
||||
doneCh <- result
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-doneCh:
|
||||
t.Fatalf("run should pause for interrupt, but completed early")
|
||||
case err := <-errCh:
|
||||
t.Fatalf("run should pause for interrupt, but failed early: %v", err)
|
||||
case <-time.After(120 * time.Millisecond):
|
||||
// expected: paused
|
||||
}
|
||||
|
||||
if err := handler.PublishToChannel("resume_channel", "go"); err != nil {
|
||||
t.Fatalf("publish interrupt failed: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("run failed after interrupt: %v", err)
|
||||
case result := <-doneCh:
|
||||
if !(result.A && result.B) {
|
||||
t.Fatalf("unexpected final state: %#v", result)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatalf("timeout waiting for resumed run completion")
|
||||
}
|
||||
}
|
||||
|
||||
func minInt32(a int32, b int32) int32 {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func maxInt32(a int32, b int32) int32 {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
Reference in New Issue
Block a user