From 56e9fe1b10d123e2fc502dbbdd0e7d1a8ca424ee Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Fri, 13 Mar 2026 23:31:52 -0700 Subject: [PATCH] more --- langgraph-go/advancedgraph/graph.go | 19 +- langgraph-go/advancedgraph/rust_engine.go | 21 +- langgraph-go/advancedgraph/state_graph.go | 245 ++++++++++++++++++++++ langgraph-go/tests/state_graph_test.go | 164 +++++++++++++++ 4 files changed, 445 insertions(+), 4 deletions(-) create mode 100644 langgraph-go/advancedgraph/state_graph.go create mode 100644 langgraph-go/tests/state_graph_test.go diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index fcdcc26ad..57d645d85 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -31,6 +31,10 @@ func NewAdvancedStateGraph[StateT any]() *AdvancedStateGraph[StateT] { func (g *AdvancedStateGraph[StateT]) AddNode(fn any) string { name := NodeName(fn) + return g.AddNodeAs(name, fn) +} + +func (g *AdvancedStateGraph[StateT]) AddNodeAs(name string, fn any) string { if _, exists := g.nodes[name]; exists { panic(fmt.Sprintf("node `%s` already exists", name)) } @@ -47,13 +51,23 @@ func (g *AdvancedStateGraph[StateT]) AddAsyncChannel(name string) { } func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any) string { - name := g.AddNode(fn) + name := NodeName(fn) + return g.AddEntryNodeAs(name, fn) +} + +func (g *AdvancedStateGraph[StateT]) AddEntryNodeAs(name string, fn any) string { + name = g.AddNodeAs(name, fn) g.entryPoint = name return name } func (g *AdvancedStateGraph[StateT]) AddFinishNode(fn any) string { - name := g.AddNode(fn) + name := NodeName(fn) + return g.AddFinishNodeAs(name, fn) +} + +func (g *AdvancedStateGraph[StateT]) AddFinishNodeAs(name string, fn any) string { + name = g.AddNodeAs(name, fn) g.finishPoint = name return name } @@ -331,4 +345,3 @@ func unwrapResumeInput(input any) (any, *WaitEvent) { } return rawArg, &event } - diff --git a/langgraph-go/advancedgraph/rust_engine.go b/langgraph-go/advancedgraph/rust_engine.go index 163c0471d..75e1866b0 100644 --- a/langgraph-go/advancedgraph/rust_engine.go +++ b/langgraph-go/advancedgraph/rust_engine.go @@ -12,6 +12,7 @@ import "C" import ( "encoding/json" "fmt" + "reflect" "runtime/cgo" "unsafe" ) @@ -59,8 +60,12 @@ func goNodeCallback(userData C.ulong, node *C.char, argJSON *C.char, stateJSON * sends := make([]map[string]any, 0, len(cmd.Goto)) for _, send := range cmd.Goto { + targetNode, err := resolveSendTarget(send.Node) + if err != nil { + return cCallbackEnvelopeError(err.Error()) + } sends = append(sends, map[string]any{ - "node": NodeName(send.Node), + "node": targetNode, "arg": send.NodeInput, }) } @@ -215,6 +220,20 @@ func cCallbackEnvelopeSuspend(cond AnyOfCondition) *C.char { return C.CString(string(raw)) } +func resolveSendTarget(target any) (string, error) { + if name, ok := target.(string); ok { + if name == "" { + return "", fmt.Errorf("send target cannot be empty string") + } + return name, nil + } + rv := reflect.ValueOf(target) + if rv.IsValid() && rv.Kind() == reflect.Func { + return NodeName(target), nil + } + return "", fmt.Errorf("unsupported send target type %T", target) +} + func coerceJSONValue(v any) any { switch t := v.(type) { case map[string]any: diff --git a/langgraph-go/advancedgraph/state_graph.go b/langgraph-go/advancedgraph/state_graph.go new file mode 100644 index 000000000..d117c1d38 --- /dev/null +++ b/langgraph-go/advancedgraph/state_graph.go @@ -0,0 +1,245 @@ +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 +} diff --git a/langgraph-go/tests/state_graph_test.go b/langgraph-go/tests/state_graph_test.go new file mode 100644 index 000000000..002b6c99e --- /dev/null +++ b/langgraph-go/tests/state_graph_test.go @@ -0,0 +1,164 @@ +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 +}