From d7614999b0af94249f1f194fbb1fc60202a28e9c Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Fri, 13 Mar 2026 13:44:44 -0700 Subject: [PATCH] more --- langgraph-go/advancedgraph/graph.go | 90 +++++++++++++++++++--- langgraph-go/advancedgraph/types.go | 2 +- langgraph-go/tests/sub_agents_test.go | 4 +- langgraph-go/tests/test_primitives_test.go | 17 ++-- rust-core/src/lib_c.rs | 2 +- 5 files changed, 93 insertions(+), 22 deletions(-) diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index f98ead71c..f592f9578 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -7,10 +7,12 @@ import ( "strings" ) -type NodeFunc func(ctx *Context, input any, state map[string]any) (Command, error) +type NodeFunc[IN any] func(ctx *Context, input IN, state map[string]any) (Command, error) + +type nodeExecutor func(ctx *Context, input any, state map[string]any) (Command, error) type AdvancedStateGraph struct { - nodes map[string]NodeFunc + nodes map[string]nodeExecutor asyncChannels []string entryPoint string finishPoint string @@ -18,16 +20,20 @@ type AdvancedStateGraph struct { func NewAdvancedStateGraph() *AdvancedStateGraph { return &AdvancedStateGraph{ - nodes: make(map[string]NodeFunc), + nodes: make(map[string]nodeExecutor), } } -func (g *AdvancedStateGraph) AddNode(fn NodeFunc) string { +func (g *AdvancedStateGraph) AddNode(fn any) string { name := NodeName(fn) if _, exists := g.nodes[name]; exists { panic(fmt.Sprintf("node `%s` already exists", name)) } - g.nodes[name] = fn + exec, err := compileNodeExecutor(fn) + if err != nil { + panic(err) + } + g.nodes[name] = exec return name } @@ -35,11 +41,11 @@ func (g *AdvancedStateGraph) AddAsyncChannel(name string) { g.asyncChannels = append(g.asyncChannels, name) } -func (g *AdvancedStateGraph) SetEntryNode(fn NodeFunc) { +func (g *AdvancedStateGraph) SetEntryNode(fn any) { g.entryPoint = NodeName(fn) } -func (g *AdvancedStateGraph) SetFinishNode(fn NodeFunc) { +func (g *AdvancedStateGraph) SetFinishNode(fn any) { g.finishPoint = NodeName(fn) } @@ -53,7 +59,7 @@ func (g *AdvancedStateGraph) Compile() *CompiledGraph { } type CompiledGraph struct { - nodes map[string]NodeFunc + nodes map[string]nodeExecutor asyncChannels []string entryPoint string finishPoint string @@ -125,8 +131,12 @@ func (g *CompiledGraph) Start(initialState map[string]any) (*Handler, error) { return handler, nil } -func NodeName(fn NodeFunc) string { - pc := reflect.ValueOf(fn).Pointer() +func NodeName(fn any) string { + rv := reflect.ValueOf(fn) + if !rv.IsValid() || rv.Kind() != reflect.Func { + panic("cannot infer node name from non-function value") + } + pc := rv.Pointer() f := runtime.FuncForPC(pc) if f == nil { panic("cannot infer node name from nil function") @@ -148,3 +158,63 @@ func NodeName(fn NodeFunc) string { } return short } + +func compileNodeExecutor(fn any) (nodeExecutor, error) { + rv := reflect.ValueOf(fn) + if !rv.IsValid() || rv.Kind() != reflect.Func { + return nil, fmt.Errorf("node must be a function") + } + rt := rv.Type() + if rt.NumIn() != 3 { + return nil, fmt.Errorf("node `%s` must accept exactly 3 args: (*Context, input, map[string]any)", NodeName(fn)) + } + ctxType := reflect.TypeOf((*Context)(nil)) + if rt.In(0) != ctxType { + return nil, fmt.Errorf("node `%s` first arg must be *Context", NodeName(fn)) + } + stateType := reflect.TypeOf(map[string]any{}) + if rt.In(2) != stateType { + return nil, fmt.Errorf("node `%s` third arg must be map[string]any", NodeName(fn)) + } + if rt.NumOut() != 2 { + return nil, fmt.Errorf("node `%s` must return (Command, error)", NodeName(fn)) + } + cmdType := reflect.TypeOf(Command{}) + if rt.Out(0) != cmdType { + return nil, fmt.Errorf("node `%s` first return must be Command", NodeName(fn)) + } + errType := reflect.TypeOf((*error)(nil)).Elem() + if !rt.Out(1).Implements(errType) { + return nil, fmt.Errorf("node `%s` second return must be error", NodeName(fn)) + } + + inputType := rt.In(1) + return func(ctx *Context, input any, state map[string]any) (Command, error) { + args := []reflect.Value{ + reflect.ValueOf(ctx), + reflect.Zero(inputType), + reflect.ValueOf(state), + } + if input != nil { + inVal := reflect.ValueOf(input) + if inVal.Type().AssignableTo(inputType) { + args[1] = inVal + } else if inVal.Type().ConvertibleTo(inputType) { + args[1] = inVal.Convert(inputType) + } else { + return Command{}, fmt.Errorf( + "node `%s` input type mismatch: got %T, want %s", + NodeName(fn), + input, + inputType.String(), + ) + } + } + out := rv.Call(args) + cmd := out[0].Interface().(Command) + if out[1].IsNil() { + return cmd, nil + } + return cmd, out[1].Interface().(error) + }, nil +} diff --git a/langgraph-go/advancedgraph/types.go b/langgraph-go/advancedgraph/types.go index 4ef8fba3d..4af5a3cd5 100644 --- a/langgraph-go/advancedgraph/types.go +++ b/langgraph-go/advancedgraph/types.go @@ -54,7 +54,7 @@ type WaitEvent struct { } type Send struct { - Node NodeFunc + Node any NodeInput any } diff --git a/langgraph-go/tests/sub_agents_test.go b/langgraph-go/tests/sub_agents_test.go index 90635fc45..38f6467d5 100644 --- a/langgraph-go/tests/sub_agents_test.go +++ b/langgraph-go/tests/sub_agents_test.go @@ -31,7 +31,7 @@ func (m *mockLLM) invoke() []decision { type lunchWorkflow struct { planner *mockLLM - names map[string]ag.NodeFunc + names map[string]any } func outputSlice(state map[string]any) []string { @@ -151,7 +151,7 @@ func TestSubAgentsEquivalentFlow(t *testing.T) { } workflow := &lunchWorkflow{ planner: planner, - names: make(map[string]ag.NodeFunc), + names: make(map[string]any), } graph := ag.NewAdvancedStateGraph() diff --git a/langgraph-go/tests/test_primitives_test.go b/langgraph-go/tests/test_primitives_test.go index 3f8c0e3f3..e56bd6a1f 100644 --- a/langgraph-go/tests/test_primitives_test.go +++ b/langgraph-go/tests/test_primitives_test.go @@ -1,6 +1,7 @@ package tests import ( + "fmt" "testing" ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph" @@ -30,9 +31,9 @@ func logsSlice(state map[string]any) []string { } } -func (w *primitiveWorkflow) startNode(ctx *ag.Context, _ any, state map[string]any) (ag.Command, error) { +func (w *primitiveWorkflow) startNode(ctx *ag.Context, input int, state map[string]any) (ag.Command, error) { logs := logsSlice(state) - logs = append(logs, "start") + logs = append(logs, fmt.Sprintf("start:%d", input)) state["logs"] = logs return ag.Command{ Update: state, @@ -42,9 +43,9 @@ func (w *primitiveWorkflow) startNode(ctx *ag.Context, _ any, state map[string]a }, nil } -func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input any, state map[string]any) (ag.Command, error) { +func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input string, state map[string]any) (ag.Command, error) { logs := logsSlice(state) - logs = append(logs, "middle:"+input.(string)) + logs = append(logs, "middle:"+input) state["logs"] = logs return ag.Command{ Update: state, @@ -54,11 +55,11 @@ func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input any, state map[str }, nil } -func (w *primitiveWorkflow) finishNode(ctx *ag.Context, input any, state map[string]any) (ag.Command, error) { +func (w *primitiveWorkflow) finishNode(ctx *ag.Context, input string, state map[string]any) (ag.Command, error) { logs := logsSlice(state) - logs = append(logs, "finish:"+input.(string)) + logs = append(logs, "finish:"+input) state["logs"] = logs - state["done"] = input.(string) + state["done"] = input return ag.Command{Update: state}, nil } @@ -88,7 +89,7 @@ func TestInputAndStatePrimitivesCompatible(t *testing.T) { t.Fatalf("unexpected done: %v", result["done"]) } logs := logsSlice(result) - if len(logs) != 3 || logs[0] != "start" || logs[1] != "middle:from_start" || logs[2] != "finish:from_middle" { + if len(logs) != 3 || logs[0] != "start:0" || logs[1] != "middle:from_start" || logs[2] != "finish:from_middle" { t.Fatalf("unexpected logs: %#v", logs) } } diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs index b874fa545..8d79240ec 100644 --- a/rust-core/src/lib_c.rs +++ b/rust-core/src/lib_c.rs @@ -145,7 +145,7 @@ fn run_graph_scheduler_json( let tx_for_spawn = tx.clone(); let state_for_spawn = Arc::clone(&state); let state_for_merge = Arc::clone(&state); - let initial_arg = state.lock().expect("state mutex poisoned").clone(); + let initial_arg = Value::Null; run_scheduler_loop( entry_point, &finish_point,