more

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
+}

langgraph-go/advancedgraph/types.go

@@ -54,7 +54,7 @@ type WaitEvent struct {
 }
 
 type Send struct {
-	Node      NodeFunc
+	Node      any
 	NodeInput any
 }
 

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()

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)
 	}
 }

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,