diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 01d53feab..5436d426d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -484,7 +484,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): patch_checkpoint_map(saved.config, saved.metadata), saved.metadata, saved.checkpoint["ts"], - saved.parent_config, + patch_checkpoint_map(saved.parent_config, saved.metadata), tasks_w_writes( next_tasks.values(), saved.pending_writes, @@ -565,7 +565,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): patch_checkpoint_map(saved.config, saved.metadata), saved.metadata, saved.checkpoint["ts"], - saved.parent_config, + patch_checkpoint_map(saved.parent_config, saved.metadata), tasks_w_writes( next_tasks.values(), saved.pending_writes, diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 1fc922761..d772e7cba 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -179,8 +179,7 @@ def map_debug_checkpoint( "step": step, "payload": { "config": patch_checkpoint_map(config, metadata), - "parent_config": parent_config, - # "parent_config": patch_checkpoint_map(parent_config, metadata), + "parent_config": patch_checkpoint_map(parent_config, metadata), "values": read_channels(channels, stream_channels), "metadata": metadata, "next": [t.name for t in tasks], diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index febc85f24..a73083290 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9092,6 +9092,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, ), @@ -9250,6 +9253,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), @@ -9279,6 +9285,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, tasks=( @@ -9707,6 +9716,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, ) @@ -9770,6 +9786,15 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), } }, ), @@ -9798,6 +9823,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, ), @@ -10056,6 +10084,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, tasks=(), @@ -10085,6 +10116,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, tasks=( @@ -10170,6 +10204,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, tasks=(), @@ -10208,6 +10249,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, tasks=( @@ -10253,6 +10301,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, tasks=( @@ -10444,6 +10499,12 @@ def test_send_to_nested_graphs( "thread_id": "1", "checkpoint_ns": AnyStr("generate_joke:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("generate_joke:"): AnyStr(), + } + ), } }, tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), @@ -10476,6 +10537,12 @@ def test_send_to_nested_graphs( "thread_id": "1", "checkpoint_ns": AnyStr("generate_joke:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("generate_joke:"): AnyStr(), + } + ), } }, tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), @@ -10931,6 +10998,12 @@ def test_weather_subgraph( "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), } }, tasks=( @@ -11020,6 +11093,12 @@ def test_weather_subgraph( "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), } }, tasks=(), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 928cfc5eb..44bee3f72 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7738,6 +7738,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, ), @@ -7903,6 +7906,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), } }, tasks=( @@ -7934,6 +7940,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), } }, tasks=( @@ -8328,6 +8337,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, ).tasks[0] @@ -8374,6 +8386,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, ) @@ -8439,6 +8458,15 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), } }, ), @@ -8467,6 +8495,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, ), @@ -8732,6 +8763,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, tasks=(), @@ -8761,6 +8795,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, tasks=( @@ -8850,6 +8887,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, tasks=(), @@ -8888,6 +8932,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, tasks=( @@ -8933,6 +8984,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, tasks=(