diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index 03dd566e7..32f21ca79 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -217,7 +217,10 @@ func compileNodeExecutor(fn any, expectedStateType reflect.Type) (nodeExecutor, if err != nil { return Command{}, fmt.Errorf("node `%s` state decode failed: %w", NodeName(fn), err) } - beforeState := stateArg.Interface() + beforeStateMap, ok := structValueToMap(stateArg) + if !ok { + return Command{}, fmt.Errorf("node `%s` failed to snapshot state fields", NodeName(fn)) + } args := []reflect.Value{ reflect.ValueOf(ctx), reflect.Zero(inputType), @@ -251,7 +254,7 @@ func compileNodeExecutor(fn any, expectedStateType reflect.Type) (nodeExecutor, ) } } - cmd.Update = reduceStructUpdateToChangedFields(beforeState, cmd.Update, stateType) + cmd.Update = reduceStructUpdateToChangedFields(beforeStateMap, cmd.Update, stateType) if out[1].IsNil() { return cmd, nil } @@ -306,38 +309,45 @@ func mustTypeOf[T any]() reflect.Type { return reflect.TypeOf((*T)(nil)).Elem() } -func reduceStructUpdateToChangedFields(before any, update any, stateType reflect.Type) any { +func reduceStructUpdateToChangedFields( + before map[string]any, + update any, + stateType reflect.Type, +) any { if update == nil { return nil } - beforeV := reflect.ValueOf(before) updateV := reflect.ValueOf(update) - if !beforeV.IsValid() || !updateV.IsValid() { - return update + if !updateV.IsValid() { + panic("internal invariant violated: update is non-nil but reflect value is invalid") } - if beforeV.Type() != stateType || updateV.Type() != stateType { - return update + if updateV.Type() != stateType { + panic(fmt.Sprintf( + "internal invariant violated: update type mismatch in reducer: got %s, expected %s", + updateV.Type().String(), + stateType.String(), + )) } if stateType.Kind() != reflect.Struct { - return update + panic(fmt.Sprintf( + "internal invariant violated: stateType must be struct in reducer, got %s", + stateType.Kind().String(), + )) + } + updateMap, ok := structValueToMap(updateV) + if !ok { + panic(fmt.Sprintf( + "internal invariant violated: failed to convert struct update to map for type %s", + stateType.String(), + )) } - changed := make(map[string]any) - for i := 0; i < stateType.NumField(); i++ { - field := stateType.Field(i) - if field.PkgPath != "" { + for key, updateValue := range updateMap { + prevValue, ok := before[key] + if ok && reflect.DeepEqual(prevValue, updateValue) { continue } - beforeField := beforeV.Field(i) - updateField := updateV.Field(i) - if reflect.DeepEqual(beforeField.Interface(), updateField.Interface()) { - continue - } - key := jsonFieldName(field) - if key == "-" { - continue - } - changed[key] = updateField.Interface() + changed[key] = updateValue } if len(changed) == 0 { return nil @@ -345,14 +355,49 @@ func reduceStructUpdateToChangedFields(before any, update any, stateType reflect return changed } -func jsonFieldName(field reflect.StructField) string { +func structValueToMap(value reflect.Value) (map[string]any, bool) { + if !value.IsValid() || value.Kind() != reflect.Struct { + return nil, false + } + out := make(map[string]any, value.NumField()) + typ := value.Type() + for i := 0; i < value.NumField(); i++ { + field := typ.Field(i) + if field.PkgPath != "" { + continue + } + name, omitEmpty, skip := parseJSONFieldTag(field) + if skip { + continue + } + fv := value.Field(i) + if omitEmpty && fv.IsZero() { + continue + } + out[name] = fv.Interface() + } + return out, true +} + +func parseJSONFieldTag(field reflect.StructField) (name string, omitEmpty bool, skip bool) { tag := field.Tag.Get("json") + if tag == "-" { + return "", false, true + } if tag == "" { - return field.Name + return field.Name, false, false } parts := strings.Split(tag, ",") - if len(parts) == 0 || parts[0] == "" { - return field.Name + fieldName := parts[0] + if fieldName == "" { + fieldName = field.Name } - return parts[0] + omit := false + for _, opt := range parts[1:] { + if opt == "omitempty" { + omit = true + break + } + } + return fieldName, omit, false } diff --git a/langgraph-go/tests/test_update_elision_test.go b/langgraph-go/tests/test_update_elision_test.go index b6b509de0..c12d30d23 100644 --- a/langgraph-go/tests/test_update_elision_test.go +++ b/langgraph-go/tests/test_update_elision_test.go @@ -1,6 +1,7 @@ package tests import ( + "reflect" "testing" "time" @@ -8,39 +9,103 @@ import ( ) type updateElisionState struct { - X int `json:"x"` + X int `json:"x"` + S updateStruct `json:"s"` + M map[string]int `json:"m"` + L []int `json:"l"` + PS *updateStruct `json:"ps"` + PM *map[string]int `json:"pm"` + PL *[]int `json:"pl"` +} + +type updateStruct struct { + V int `json:"v"` } type updateElisionWorkflow struct{} -func (w *updateElisionWorkflow) startNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) { +func makeState(v int) updateElisionState { + m := map[string]int{"n": v} + l := []int{v} + return updateElisionState{ + X: v, + S: updateStruct{V: v}, + M: map[string]int{"n": v}, + L: []int{v}, + PS: &updateStruct{V: v}, + PM: &m, + PL: &l, + } +} + +func (w *updateElisionWorkflow) startNoopNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) { return ag.Command{ Goto: []ag.Send{ {Node: w.fastNode}, - {Node: w.slowNode}, + {Node: w.slowNoopNode}, + }, + }, nil +} + +func (w *updateElisionWorkflow) startChangedNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) { + return ag.Command{ + Goto: []ag.Send{ + {Node: w.fastNode}, + {Node: w.slowChangedNode}, }, }, nil } func (w *updateElisionWorkflow) fastNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) { - state.X = 1 + _ = state + return ag.Command{Update: makeState(1)}, nil +} + +func (w *updateElisionWorkflow) slowNoopNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) { + time.Sleep(100 * time.Millisecond) + // Returns same state as initial snapshot; SDK should elide this update. return ag.Command{Update: state}, nil } -func (w *updateElisionWorkflow) slowNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) { +func (w *updateElisionWorkflow) slowChangedNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) { time.Sleep(100 * time.Millisecond) - // Returns same X as initial snapshot; SDK should elide this update. - return ag.Command{Update: state}, nil + // Real change should not be elided. + return ag.Command{Update: makeState(2)}, nil +} + +func assertStateEquals(t *testing.T, got updateElisionState, expected updateElisionState) { + t.Helper() + if got.X != expected.X { + t.Fatalf("unexpected X: got=%d want=%d", got.X, expected.X) + } + if got.S != expected.S { + t.Fatalf("unexpected S: got=%#v want=%#v", got.S, expected.S) + } + if !reflect.DeepEqual(got.M, expected.M) { + t.Fatalf("unexpected M: got=%#v want=%#v", got.M, expected.M) + } + if !reflect.DeepEqual(got.L, expected.L) { + t.Fatalf("unexpected L: got=%#v want=%#v", got.L, expected.L) + } + if got.PS == nil || expected.PS == nil || *got.PS != *expected.PS { + t.Fatalf("unexpected PS: got=%#v want=%#v", got.PS, expected.PS) + } + if got.PM == nil || expected.PM == nil || !reflect.DeepEqual(*got.PM, *expected.PM) { + t.Fatalf("unexpected PM: got=%#v want=%#v", got.PM, expected.PM) + } + if got.PL == nil || expected.PL == nil || !reflect.DeepEqual(*got.PL, *expected.PL) { + t.Fatalf("unexpected PL: got=%#v want=%#v", got.PL, expected.PL) + } } func TestNoopSlowUpdateDoesNotOverrideFastUpdate(t *testing.T) { workflow := &updateElisionWorkflow{} graph := ag.NewAdvancedStateGraph[updateElisionState]() - graph.AddEntryNode(workflow.startNode) + graph.AddEntryNode(workflow.startNoopNode) graph.AddNode(workflow.fastNode) - graph.AddFinishNode(workflow.slowNode) + graph.AddFinishNode(workflow.slowNoopNode) - handler, err := graph.Compile().Start(nil, updateElisionState{X: 0}) + handler, err := graph.Compile().Start(nil, makeState(0)) if err != nil { t.Fatalf("start failed: %v", err) } @@ -48,7 +113,23 @@ func TestNoopSlowUpdateDoesNotOverrideFastUpdate(t *testing.T) { if err != nil { t.Fatalf("result failed: %v", err) } - if result.X != 1 { - t.Fatalf("expected fast update to win, got x=%d", result.X) - } + assertStateEquals(t, result, makeState(1)) +} + +func TestChangedSlowUpdateOverridesFastUpdate(t *testing.T) { + workflow := &updateElisionWorkflow{} + graph := ag.NewAdvancedStateGraph[updateElisionState]() + graph.AddEntryNode(workflow.startChangedNode) + graph.AddNode(workflow.fastNode) + graph.AddFinishNode(workflow.slowChangedNode) + + handler, err := graph.Compile().Start(nil, makeState(0)) + if err != nil { + t.Fatalf("start failed: %v", err) + } + result, err := handler.WaitForResult() + if err != nil { + t.Fatalf("result failed: %v", err) + } + assertStateEquals(t, result, makeState(2)) } diff --git a/libs/langgraph/langgraph/advanced_graph/state.py b/libs/langgraph/langgraph/advanced_graph/state.py index cf43251e6..da00c047e 100644 --- a/libs/langgraph/langgraph/advanced_graph/state.py +++ b/libs/langgraph/langgraph/advanced_graph/state.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import copy import inspect from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass @@ -252,7 +253,7 @@ class _GraphEngineRun: def _execute_node_for_rust( self, node_name: str, node_input: Any, state: Any ) -> dict[str, Any]: - before_state_markers = _state_shallow_markers(state) + before_state_snapshot = copy.deepcopy(state) if isinstance(state, dict) else None if node_name not in self._nodes: raise ValueError(f"Unknown node `{node_name}`") @@ -270,11 +271,9 @@ class _GraphEngineRun: if update is None and isinstance(state, dict): # Preserve in-place state mutations for prototype nodes like wait_node. - if _has_shallow_state_change(before_state_markers, state): - update = state - elif isinstance(update, dict): - if not _has_shallow_update_change(before_state_markers, update): - update = None + update = state + if isinstance(update, dict): + update = _reduce_update_to_changed_fields(before_state_snapshot, update) return { "update": update, @@ -390,44 +389,16 @@ def _resolve_target_name(target: Any) -> str: raise ValueError(f"Unsupported node target type: {type(target)!r}") -def _state_shallow_markers(state: Any) -> dict[str, Any] | None: - if not isinstance(state, dict): - return None - return {k: _value_shallow_marker(v) for k, v in state.items()} - - -def _value_shallow_marker(value: Any) -> Any: - if value is None or isinstance(value, (bool, int, float, str, bytes)): - return ("primitive", value) - if isinstance(value, (list, tuple, set, dict)): - return ("container", type(value).__name__, id(value), len(value)) - return ("object", type(value).__name__, id(value)) - - -def _has_shallow_state_change(before: dict[str, Any] | None, state: Any) -> bool: - if before is None or not isinstance(state, dict): - return True - if len(before) != len(state): - return True - for key, old_marker in before.items(): - if key not in state: - return True - if old_marker != _value_shallow_marker(state[key]): - return True - return False - - -def _has_shallow_update_change( +def _reduce_update_to_changed_fields( before: dict[str, Any] | None, update: dict[str, Any] -) -> bool: +) -> dict[str, Any] | None: if before is None: - return True + return update + changed: dict[str, Any] = {} for key, new_value in update.items(): - if key not in before: - return True - if _value_shallow_marker(new_value) != before[key]: - return True - return False + if key not in before or before[key] != new_value: + changed[key] = new_value + return changed or None def _invoke_node(node: Callable[..., Any], ctx: Context, node_input: Any, state: Any) -> Any: diff --git a/libs/langgraph/tests/advanced-graph/test_update_elision.py b/libs/langgraph/tests/advanced-graph/test_update_elision.py index 63ebc3fa4..b56202db4 100644 --- a/libs/langgraph/tests/advanced-graph/test_update_elision.py +++ b/libs/langgraph/tests/advanced-graph/test_update_elision.py @@ -1,35 +1,125 @@ import asyncio +from dataclasses import dataclass +from pydantic import BaseModel from typing_extensions import TypedDict import pytest -from langgraph.advanced_graph import AdvancedStateGraph +from langgraph.advanced_graph import AdvancedStateGraph, CompiledGraphEngine from langgraph.types import Command, Send pytestmark = pytest.mark.anyio +@dataclass(frozen=True) +class DataClassPayload: + value: int + + +class PydanticPayload(BaseModel): + value: int + + +class InnerTypedDict(TypedDict): + flag: bool + n: int + + class UpdateElisionState(TypedDict): x: int + dc: DataClassPayload + model: PydanticPayload + td: InnerTypedDict + obj: dict[str, int] + items: list[int] + + +def _initial_state() -> UpdateElisionState: + return { + "x": 0, + "dc": DataClassPayload(0), + "model": PydanticPayload(value=0), + "td": {"flag": False, "n": 0}, + "obj": {"n": 0}, + "items": [0], + } + async def test_noop_slow_update_does_not_override_fast_update() -> None: - graph = AdvancedStateGraph(UpdateElisionState) + graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState) async def start_node(state: UpdateElisionState) -> Command: return Command(goto=[Send("fast_node", None), Send("slow_node", None)]) - async def fast_node(state: UpdateElisionState) -> dict[str, int]: - return {"x": 1} + async def fast_node(state: UpdateElisionState) -> UpdateElisionState: + state.x = 1 + state.dc.value = 1 + state.model.value = 1 + state.td["flag"] = True + state.td["n"] = 1 + state.obj["n"] = 1 + state.items.append(1) + return state - async def slow_node(state: UpdateElisionState) -> dict[str, int]: + async def slow_node(state: UpdateElisionState) -> UpdateElisionState: await asyncio.sleep(0.1) - # Returns the same value as initial snapshot. - return {"x": state["x"]} + # Returns the same values as the initial snapshot. + return state graph.add_entry_node(start_node) graph.add_node(fast_node) graph.add_finish_node(slow_node) - result = await graph.compile().ainvoke({"x": 0}) + compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile() + initial_state: UpdateElisionState = _initial_state() + result: UpdateElisionState = await compiled.ainvoke(initial_state) assert result["x"] == 1 + assert result["dc"] == DataClassPayload(1) + assert result["model"].value == 1 + assert result["td"] == {"flag": True, "n": 1} + assert result["obj"] == {"n": 1} + assert result["items"] == [1] + + +async def test_changed_slow_update_overrides_fast_update() -> None: + graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState) + + async def start_node(state: UpdateElisionState) -> Command: + return Command(goto=[Send("fast_node", None), Send("slow_node", None)]) + + async def fast_node(state: UpdateElisionState) -> UpdateElisionState: + state.x = 1 + state.dc.value = 1 + state.model.value = 1 + state.td["flag"] = True + state.td["n"] = 1 + state.obj["n"] = 1 + state.items.append(1) + return state + + async def slow_node(state: UpdateElisionState) -> UpdateElisionState: + await asyncio.sleep(0.1) + # Slow node makes real changes for all field types. + state.x = 2 + state.dc.value = 2 + state.model.value = 2 + state.td["flag"] = False + state.td["n"] = 2 + state.obj["n"] = 2 + state.items.append(2) + return state + + graph.add_entry_node(start_node) + graph.add_node(fast_node) + graph.add_finish_node(slow_node) + + compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile() + initial_state: UpdateElisionState = _initial_state() + result: UpdateElisionState = await compiled.ainvoke(initial_state) + assert result["x"] == 2 + assert result["dc"] == DataClassPayload(2) + assert result["model"].value == 2 + assert result["td"] == {"flag": False, "n": 2} + assert result["obj"] == {"n": 2} + assert result["items"] == [2]