more

langgraph-go/advancedgraph/graph.go

@@ -1,35 +1,38 @@
 package advancedgraph
 
 import (
+	"encoding/json"
 	"fmt"
 	"reflect"
 	"runtime"
 	"strings"
 )
 
-type NodeFunc[IN any] func(ctx *Context, input IN, state map[string]any) (Command, error)
+type NodeFunc[IN any, StateT any] func(ctx *Context, input IN, state StateT) (Command, error)
 
 type nodeExecutor func(ctx *Context, input any, state map[string]any) (Command, error)
 
-type AdvancedStateGraph struct {
+type AdvancedStateGraph[StateT any] struct {
 	nodes         map[string]nodeExecutor
 	asyncChannels []string
 	entryPoint    string
 	finishPoint   string
+	stateType     reflect.Type
 }
 
-func NewAdvancedStateGraph() *AdvancedStateGraph {
-	return &AdvancedStateGraph{
-		nodes: make(map[string]nodeExecutor),
+func NewAdvancedStateGraph[StateT any]() *AdvancedStateGraph[StateT] {
+	return &AdvancedStateGraph[StateT]{
+		nodes:     make(map[string]nodeExecutor),
+		stateType: mustTypeOf[StateT](),
 	}
 }
 
-func (g *AdvancedStateGraph) AddNode(fn any) string {
+func (g *AdvancedStateGraph[StateT]) AddNode(fn any) string {
 	name := NodeName(fn)
 	if _, exists := g.nodes[name]; exists {
 		panic(fmt.Sprintf("node `%s` already exists", name))
 	}
-	exec, err := compileNodeExecutor(fn)
+	exec, err := compileNodeExecutor(fn, g.stateType)
 	if err != nil {
 		panic(err)
 	}
@@ -37,44 +40,38 @@ func (g *AdvancedStateGraph) AddNode(fn any) string {
 	return name
 }
 
-func (g *AdvancedStateGraph) AddAsyncChannel(name string) {
+func (g *AdvancedStateGraph[StateT]) AddAsyncChannel(name string) {
 	g.asyncChannels = append(g.asyncChannels, name)
 }
 
-func (g *AdvancedStateGraph) SetEntryNode(fn any) {
-	g.entryPoint = NodeName(fn)
-}
-
-func (g *AdvancedStateGraph) SetFinishNode(fn any) {
-	g.finishPoint = NodeName(fn)
-}
-
-func (g *AdvancedStateGraph) AddEntryNode(fn any) string {
+func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any) string {
 	name := g.AddNode(fn)
 	g.entryPoint = name
 	return name
 }
 
-func (g *AdvancedStateGraph) AddFinishNode(fn any) string {
+func (g *AdvancedStateGraph[StateT]) AddFinishNode(fn any) string {
 	name := g.AddNode(fn)
 	g.finishPoint = name
 	return name
 }
 
-func (g *AdvancedStateGraph) Compile() *CompiledGraph {
-	return &CompiledGraph{
+func (g *AdvancedStateGraph[StateT]) Compile() *CompiledGraph[StateT] {
+	return &CompiledGraph[StateT]{
 		nodes:         g.nodes,
 		asyncChannels: g.asyncChannels,
 		entryPoint:    g.entryPoint,
 		finishPoint:   g.finishPoint,
+		stateType:     g.stateType,
 	}
 }
 
-type CompiledGraph struct {
+type CompiledGraph[StateT any] struct {
 	nodes         map[string]nodeExecutor
 	asyncChannels []string
 	entryPoint    string
 	finishPoint   string
+	stateType     reflect.Type
 }
 
 type Context struct {
@@ -89,26 +86,26 @@ func (c *Context) PublishToChannel(channel string, value any) error {
 	return c.engine.Publish(channel, value)
 }
 
-type Handler struct {
+type Handler[StateT any] struct {
 	engine *RustEngine
-	done   chan resultOrErr
+	done   chan resultOrErr[StateT]
 }
 
-type resultOrErr struct {
-	state map[string]any
+type resultOrErr[StateT any] struct {
+	state StateT
 	err   error
 }
 
-func (h *Handler) PublishToChannel(channel string, value any) error {
+func (h *Handler[StateT]) PublishToChannel(channel string, value any) error {
 	return h.engine.Publish(channel, value)
 }
 
-func (h *Handler) WaitForResult() (map[string]any, error) {
+func (h *Handler[StateT]) WaitForResult() (StateT, error) {
 	res := <-h.done
 	return res.state, res.err
 }
 
-func (g *CompiledGraph) Start(initialInput any, initialState map[string]any) (*Handler, error) {
+func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT) (*Handler[StateT], error) {
 	engine := NewRustEngine()
 	for _, ch := range g.asyncChannels {
 		if err := engine.AddAsyncChannel(ch); err != nil {
@@ -116,13 +113,13 @@ func (g *CompiledGraph) Start(initialInput any, initialState map[string]any) (*H
 		}
 	}
 
-	handler := &Handler{
+	handler := &Handler[StateT]{
 		engine: engine,
-		done:   make(chan resultOrErr, 1),
+		done:   make(chan resultOrErr[StateT], 1),
 	}
 	go func() {
 		defer engine.Close()
-		state, err := engine.RunGraph(
+		rawState, err := engine.RunGraph(
 			g.entryPoint,
 			g.finishPoint,
 			initialState,
@@ -138,7 +135,13 @@ func (g *CompiledGraph) Start(initialInput any, initialState map[string]any) (*H
 				return fn(&Context{engine: engine}, nodeInput, fallbackState)
 			},
 		)
-		handler.done <- resultOrErr{state: state, err: err}
+		if err != nil {
+			handler.done <- resultOrErr[StateT]{err: err}
+			close(handler.done)
+			return
+		}
+		state, err := mapToState[StateT](rawState)
+		handler.done <- resultOrErr[StateT]{state: state, err: err}
 		close(handler.done)
 	}()
 	return handler, nil
@@ -172,23 +175,19 @@ func NodeName(fn any) string {
 	return short
 }
 
-func compileNodeExecutor(fn any) (nodeExecutor, error) {
+func compileNodeExecutor(fn any, expectedStateType reflect.Type) (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))
+		return nil, fmt.Errorf("node `%s` must accept exactly 3 args: (*Context, input, state)", 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))
 	}
@@ -202,11 +201,24 @@ func compileNodeExecutor(fn any) (nodeExecutor, error) {
 	}
 
 	inputType := rt.In(1)
+	stateType := rt.In(2)
+	if stateType != expectedStateType {
+		return nil, fmt.Errorf(
+			"node `%s` state type mismatch: got %s, graph expects %s",
+			NodeName(fn),
+			stateType.String(),
+			expectedStateType.String(),
+		)
+	}
 	return func(ctx *Context, input any, state map[string]any) (Command, error) {
+		stateArg, err := convertStateArg(state, stateType)
+		if err != nil {
+			return Command{}, fmt.Errorf("node `%s` state decode failed: %w", NodeName(fn), err)
+		}
 		args := []reflect.Value{
 			reflect.ValueOf(ctx),
 			reflect.Zero(inputType),
-			reflect.ValueOf(state),
+			stateArg,
 		}
 		if input != nil {
 			inVal := reflect.ValueOf(input)
@@ -231,3 +243,50 @@ func compileNodeExecutor(fn any) (nodeExecutor, error) {
 		return cmd, out[1].Interface().(error)
 	}, nil
 }
+
+func convertStateArg(state map[string]any, stateType reflect.Type) (reflect.Value, error) {
+	if stateType == reflect.TypeOf(map[string]any{}) {
+		return reflect.ValueOf(state), nil
+	}
+	raw, err := json.Marshal(state)
+	if err != nil {
+		return reflect.Value{}, fmt.Errorf("marshal state: %w", err)
+	}
+	if stateType.Kind() == reflect.Ptr {
+		target := reflect.New(stateType.Elem())
+		if err := json.Unmarshal(raw, target.Interface()); err != nil {
+			return reflect.Value{}, fmt.Errorf("unmarshal state into %s: %w", stateType.String(), err)
+		}
+		return target, nil
+	}
+	target := reflect.New(stateType)
+	if err := json.Unmarshal(raw, target.Interface()); err != nil {
+		return reflect.Value{}, fmt.Errorf("unmarshal state into %s: %w", stateType.String(), err)
+	}
+	return target.Elem(), nil
+}
+
+func mapToState[StateT any](raw map[string]any) (StateT, error) {
+	var out StateT
+	if anyVal, ok := any(raw).(StateT); ok {
+		return anyVal, nil
+	}
+	payload, err := json.Marshal(raw)
+	if err != nil {
+		return out, fmt.Errorf("marshal state: %w", err)
+	}
+	if err := json.Unmarshal(payload, &out); err != nil {
+		return out, fmt.Errorf("unmarshal state: %w", err)
+	}
+	return out, nil
+}
+
+func mustTypeOf[T any]() reflect.Type {
+	var zero T
+	t := reflect.TypeOf(zero)
+	if t != nil {
+		return t
+	}
+	// Handles nil-able types where zero value has no dynamic type.
+	return reflect.TypeOf((*T)(nil)).Elem()
+}

langgraph-go/advancedgraph/rust_engine.go

@@ -138,7 +138,7 @@ func (e *RustEngine) WaitAnyOf(cond AnyOfCondition) (WaitEvent, error) {
 func (e *RustEngine) RunGraph(
 	entryPoint string,
 	finishPoint string,
-	initialState map[string]any,
+	initialState any,
 	initialInput any,
 	exec func(node string, nodeInput any, state map[string]any) (Command, error),
 ) (map[string]any, error) {

langgraph-go/advancedgraph/types.go

@@ -59,7 +59,7 @@ type Send struct {
 }
 
 type Command struct {
-	Update map[string]any
+	Update any
 	Goto   []Send
 }
 

langgraph-go/tests/sub_agents_test.go

@@ -34,28 +34,13 @@ type lunchWorkflow struct {
 	names   map[string]any
 }
 
-func outputSlice(state map[string]any) []string {
-	raw, ok := state["output"]
-	if !ok || raw == nil {
-		return []string{}
-	}
-	switch v := raw.(type) {
-	case []string:
-		return v
-	case []any:
-		out := make([]string, 0, len(v))
-		for _, item := range v {
-			if s, ok := item.(string); ok {
-				out = append(out, s)
-			}
-		}
-		return out
-	default:
-		return []string{}
-	}
+type lunchState struct {
+	Input  string   `json:"input"`
+	Output []string `json:"output"`
+	Done   string   `json:"done"`
 }
 
-func (w *lunchWorkflow) llmNode(ctx *ag.Context, _ any, state map[string]any) (ag.Command, error) {
+func (w *lunchWorkflow) llmNode(ctx *ag.Context, _ any, _ lunchState) (ag.Command, error) {
 	decisions := w.planner.invoke()
 	sends := make([]ag.Send, 0, 4)
 	for _, d := range decisions {
@@ -77,7 +62,7 @@ func (w *lunchWorkflow) llmNode(ctx *ag.Context, _ any, state map[string]any) (a
 	return ag.Command{Goto: sends}, nil
 }
 
-func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state map[string]any) (ag.Command, error) {
+func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state lunchState) (ag.Command, error) {
 	event, err := ctx.WaitFor(
 		ag.AnyOf(
 			ag.ChannelCondition{Channel: "tool_completion_channel", N: 1},
@@ -90,7 +75,7 @@ func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state map[string]any) (
 		return ag.Command{}, err
 	}
 
-	output := outputSlice(state)
+	output := append([]string(nil), state.Output...)
 	if event.Condition == "channel" {
 		payload := ag.DecodeString(event.Value)
 		switch event.Channel {
@@ -101,23 +86,23 @@ func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state map[string]any) (
 		case "user_input_channel":
 			output = append(output, "user_input: "+payload)
 		}
-		state["output"] = output
+		state.Output = output
 		return ag.Command{Goto: []ag.Send{{Node: w.names["llm"]}}, Update: state}, nil
 	}
 
 	output = append(output, "timer: no updates yet")
-	state["output"] = output
+	state.Output = output
 	return ag.Command{Goto: []ag.Send{{Node: w.names["wait"]}}, Update: state}, nil
 }
 
-func (w *lunchWorkflow) toolNode(ctx *ag.Context, input any, _ map[string]any) (ag.Command, error) {
+func (w *lunchWorkflow) toolNode(ctx *ag.Context, input any, _ lunchState) (ag.Command, error) {
 	toolInput, _ := input.(string)
 	time.Sleep(100 * time.Millisecond)
 	err := ctx.PublishToChannel("tool_completion_channel", "tool completed for: "+toolInput)
 	return ag.Command{}, err
 }
 
-func (w *lunchWorkflow) subAgentNode(ctx *ag.Context, input any, _ map[string]any) (ag.Command, error) {
+func (w *lunchWorkflow) subAgentNode(ctx *ag.Context, input any, _ lunchState) (ag.Command, error) {
 	subInput, _ := input.(string)
 	time.Sleep(5 * time.Second)
 	err := ctx.PublishToChannel(
@@ -127,12 +112,12 @@ func (w *lunchWorkflow) subAgentNode(ctx *ag.Context, input any, _ map[string]an
 	return ag.Command{}, err
 }
 
-func (w *lunchWorkflow) orderFoodNode(ctx *ag.Context, input any, state map[string]any) (ag.Command, error) {
+func (w *lunchWorkflow) orderFoodNode(ctx *ag.Context, input any, state lunchState) (ag.Command, error) {
 	complete, _ := input.(string)
-	output := outputSlice(state)
+	output := append([]string(nil), state.Output...)
 	output = append(output, "order_food: "+complete)
-	state["output"] = output
-	state["done"] = complete
+	state.Output = output
+	state.Done = complete
 	return ag.Command{Update: state}, nil
 }
 
@@ -154,31 +139,28 @@ func TestSubAgentsEquivalentFlow(t *testing.T) {
 		names:   make(map[string]any),
 	}
 
-	graph := ag.NewAdvancedStateGraph()
+	graph := ag.NewAdvancedStateGraph[lunchState]()
 	graph.AddAsyncChannel("tool_completion_channel")
 	graph.AddAsyncChannel("subagent_completion_channel")
 	graph.AddAsyncChannel("user_input_channel")
 
-	graph.AddNode(workflow.llmNode)
+	graph.AddEntryNode(workflow.llmNode)
 	graph.AddNode(workflow.waitNode)
 	graph.AddNode(workflow.toolNode)
 	graph.AddNode(workflow.subAgentNode)
-	graph.AddNode(workflow.orderFoodNode)
+	graph.AddFinishNode(workflow.orderFoodNode)
 	workflow.names["llm"] = workflow.llmNode
 	workflow.names["wait"] = workflow.waitNode
 	workflow.names["tool"] = workflow.toolNode
 	workflow.names["sub"] = workflow.subAgentNode
 	workflow.names["order"] = workflow.orderFoodNode
 
-	graph.SetEntryNode(workflow.llmNode)
-	graph.SetFinishNode(workflow.orderFoodNode)
-
 	handler, err := graph.Compile().Start(
 		nil,
-		map[string]any{
-			"input":  "help me get something for lunch",
-			"output": []string{},
-			"done":   nil,
+		lunchState{
+			Input:  "help me get something for lunch",
+			Output: []string{},
+			Done:   "",
 		},
 	)
 	if err != nil {
@@ -195,12 +177,12 @@ func TestSubAgentsEquivalentFlow(t *testing.T) {
 		t.Fatalf("result failed: %v", err)
 	}
 
-	output := outputSlice(result)
+	output := result.Output
 	if len(output) == 0 {
 		t.Fatalf("output is empty, full result=%#v", result)
 	}
-	if result["done"] != "order submitted" {
-		t.Fatalf("unexpected done: %v", result["done"])
+	if result.Done != "order submitted" {
+		t.Fatalf("unexpected done: %v", result.Done)
 	}
 	if !slices.Contains(output, "user_input: No spicy food please") {
 		t.Fatalf("missing user input output: %#v", output)

langgraph-go/tests/test_primitives_test.go

@@ -10,31 +10,14 @@ import (
 type primitiveWorkflow struct {
 }
 
-func logsSlice(state map[string]any) []string {
-	raw, ok := state["logs"]
-	if !ok || raw == nil {
-		return []string{}
-	}
-	switch v := raw.(type) {
-	case []string:
-		return v
-	case []any:
-		out := make([]string, 0, len(v))
-		for _, item := range v {
-			if s, ok := item.(string); ok {
-				out = append(out, s)
-			}
-		}
-		return out
-	default:
-		return []string{}
-	}
+type primitiveState struct {
+	Count int      `json:"count"`
+	Logs  []string `json:"logs"`
+	Done  string   `json:"done"`
 }
 
-func (w *primitiveWorkflow) startNode(ctx *ag.Context, input int, state map[string]any) (ag.Command, error) {
-	logs := logsSlice(state)
-	logs = append(logs, fmt.Sprintf("start:%d", input))
-	state["logs"] = logs
+func (w *primitiveWorkflow) startNode(ctx *ag.Context, input int, state primitiveState) (ag.Command, error) {
+	state.Logs = append(state.Logs, fmt.Sprintf("start:%d", input))
 	return ag.Command{
 		Update: state,
 		Goto: []ag.Send{
@@ -43,10 +26,8 @@ func (w *primitiveWorkflow) startNode(ctx *ag.Context, input int, state map[stri
 	}, nil
 }
 
-func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input string, state map[string]any) (ag.Command, error) {
-	logs := logsSlice(state)
-	logs = append(logs, "middle:"+input)
-	state["logs"] = logs
+func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input string, state primitiveState) (ag.Command, error) {
+	state.Logs = append(state.Logs, "middle:"+input)
 	return ag.Command{
 		Update: state,
 		Goto: []ag.Send{
@@ -55,25 +36,24 @@ func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input string, state map[
 	}, nil
 }
 
-func (w *primitiveWorkflow) finishNode(ctx *ag.Context, input string, state map[string]any) (ag.Command, error) {
-	logs := logsSlice(state)
-	logs = append(logs, "finish:"+input)
-	state["logs"] = logs
-	state["done"] = input
+func (w *primitiveWorkflow) finishNode(ctx *ag.Context, input string, state primitiveState) (ag.Command, error) {
+	state.Logs = append(state.Logs, "finish:"+input)
+	state.Done = input
 	return ag.Command{Update: state}, nil
 }
 
 func TestInputAndStatePrimitivesCompatible(t *testing.T) {
 	workflow := &primitiveWorkflow{}
-	graph := ag.NewAdvancedStateGraph()
+	graph := ag.NewAdvancedStateGraph[primitiveState]()
 
 	graph.AddEntryNode(workflow.startNode)
 	graph.AddNode(workflow.middleNode)
 	graph.AddFinishNode(workflow.finishNode)
 
-	handler, err := graph.Compile().Start(100, map[string]any{
-		"logs": []string{},
-		"done": nil,
+	handler, err := graph.Compile().Start(100, primitiveState{
+		Count: 1,
+		Logs:  []string{},
+		Done:  "",
 	})
 	if err != nil {
 		t.Fatalf("start failed: %v", err)
@@ -83,11 +63,13 @@ func TestInputAndStatePrimitivesCompatible(t *testing.T) {
 	if err != nil {
 		t.Fatalf("result failed: %v", err)
 	}
-	if result["done"] != "from_middle" {
-		t.Fatalf("unexpected done: %v", result["done"])
+	if result.Done != "from_middle" {
+		t.Fatalf("unexpected done: %v", result.Done)
 	}
-	logs := logsSlice(result)
-	if len(logs) != 3 || logs[0] != "start:100" || logs[1] != "middle:from_start" || logs[2] != "finish:from_middle" {
-		t.Fatalf("unexpected logs: %#v", logs)
+	if result.Count != 1 {
+		t.Fatalf("unexpected count: %v", result.Count)
+	}
+	if len(result.Logs) != 3 || result.Logs[0] != "start:100" || result.Logs[1] != "middle:from_start" || result.Logs[2] != "finish:from_middle" {
+		t.Fatalf("unexpected logs: %#v", result.Logs)
 	}
 }