diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 9fd2d3b75..687970a4d 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -58,6 +58,7 @@ from langgraph.channels.base import ( from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, + CheckpointMetadata, copy_checkpoint, empty_checkpoint, ) @@ -774,6 +775,52 @@ class Pregel( ) as executor, ManagedValuesManager( self.managed_values_dict, config, self ) as managed: + + def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: + nonlocal checkpoint, checkpoint_config, channels + + if self.checkpointer is None: + return + if debug: + print_step_checkpoint( + metadata["step"], channels, self.stream_channels_list + ) + + # create new checkpoint + checkpoint = create_checkpoint( + checkpoint, channels, metadata["step"] + ) + # save it, without blocking + bg.append( + executor.submit( + self.checkpointer.put, + checkpoint_config, + copy_checkpoint(checkpoint), + metadata, + ) + ) + # update checkpoint config + checkpoint_config = { + **checkpoint_config, + "configurable": { + **checkpoint_config["configurable"], + "thread_ts": checkpoint["id"], + }, + } + # yield debug checkpoint event + if "debug" in stream_modes: + yield from _with_mode( + "debug", + isinstance(stream_mode, list), + map_debug_checkpoint( + metadata["step"], + checkpoint_config, + channels, + self.stream_channels_asis, + metadata, + ), + ) + # map inputs to channel updates if input_writes := deque(map_input(input_keys, input)): # discard any unfinished tasks from previous checkpoint @@ -789,23 +836,13 @@ class Pregel( # apply input writes _apply_writes(checkpoint, channels, input_writes) # save input checkpoint - if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels, start) - bg.append( - executor.submit( - self.checkpointer.put, - checkpoint_config, - copy_checkpoint(checkpoint), - {"source": "input", "step": start, "writes": input}, - ) - ) - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, + yield from put_checkpoint( + { + "source": "input", + "step": start, + "writes": input, } + ) # increment start to 0 start += 1 else: @@ -890,9 +927,6 @@ class Pregel( # apply writes to channels _apply_writes(checkpoint, channels, pending_writes) - if debug: - print_step_checkpoint(step, channels, self.stream_channels_list) - # yield current value or updates if "updates" in stream_modes: yield from _with_mode( @@ -916,47 +950,21 @@ class Pregel( ) # save end of step checkpoint - if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels, step) - bg.append( - executor.submit( - self.checkpointer.put, - checkpoint_config, - copy_checkpoint(checkpoint), - { - "source": "loop", - "step": step, - "writes": single( - map_output_updates(output_keys, next_tasks) - ) - if self.stream_mode == "updates" - else single( - map_output_values( - output_keys, pending_writes, channels - ), - ), - }, + yield from put_checkpoint( + { + "source": "loop", + "step": step, + "writes": single( + map_output_updates(output_keys, next_tasks) ) - ) - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - } - # yield debug checkpoint - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - step, - checkpoint_config if self.checkpointer else None, - channels, - self.stream_channels_asis, + if self.stream_mode == "updates" + else single( + map_output_values( + output_keys, pending_writes, channels + ), ), - ) + } + ) # after execution, check if we should interrupt if _should_interrupt( @@ -1073,6 +1081,51 @@ class Pregel( ) as channels, AsyncManagedValuesManager( self.managed_values_dict, config, self ) as managed: + + def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: + nonlocal checkpoint, checkpoint_config, channels + + if self.checkpointer is None: + return + if debug: + print_step_checkpoint( + metadata["step"], channels, self.stream_channels_list + ) + + # create new checkpoint + checkpoint = create_checkpoint( + checkpoint, channels, metadata["step"] + ) + # save it, without blocking + bg.append( + asyncio.create_task( + self.checkpointer.aput( + checkpoint_config, copy_checkpoint(checkpoint), metadata + ) + ) + ) + # update checkpoint config + checkpoint_config = { + **checkpoint_config, + "configurable": { + **checkpoint_config["configurable"], + "thread_ts": checkpoint["id"], + }, + } + # yield debug checkpoint event + if "debug" in stream_modes: + yield from _with_mode( + "debug", + isinstance(stream_mode, list), + map_debug_checkpoint( + metadata["step"], + checkpoint_config, + channels, + self.stream_channels_asis, + metadata, + ), + ) + # map inputs to channel updates if input_writes := deque(map_input(input_keys, input)): # discard any unfinished tasks from previous checkpoint @@ -1088,24 +1141,10 @@ class Pregel( # apply input writes _apply_writes(checkpoint, channels, input_writes) # save input checkpoint - if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels, start) - bg.append( - asyncio.create_task( - self.checkpointer.aput( - checkpoint_config, - copy_checkpoint(checkpoint), - {"source": "input", "step": start, "writes": input}, - ) - ) - ) - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - } + for chunk in put_checkpoint( + {"source": "input", "step": start, "writes": input} + ): + yield chunk # increment start to 0 start += 1 else: @@ -1193,9 +1232,6 @@ class Pregel( # apply writes to channels _apply_writes(checkpoint, channels, pending_writes) - if debug: - print_step_checkpoint(step, channels, self.stream_channels_list) - # yield current value or updates if "updates" in stream_modes: for chunk in _with_mode( @@ -1222,49 +1258,20 @@ class Pregel( yield chunk # save end of step checkpoint - if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels, step) - bg.append( - asyncio.create_task( - self.checkpointer.aput( - checkpoint_config, - checkpoint, - { - "source": "loop", - "step": step, - "writes": single( - map_output_updates(output_keys, next_tasks) - ) - if self.stream_mode == "updates" - else single( - map_output_values( - output_keys, pending_writes, channels - ) - ), - }, - ) + for chunk in put_checkpoint( + { + "source": "loop", + "step": step, + "writes": single( + map_output_updates(output_keys, next_tasks) ) - ) - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - } - # yield debug checkpoint - if "debug" in stream_modes: - for chunk in _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - step, - checkpoint_config if self.checkpointer else None, - channels, - self.stream_channels_asis, + if self.stream_mode == "updates" + else single( + map_output_values(output_keys, pending_writes, channels) ), - ): - yield chunk + } + ): + yield chunk # after execution, check if we should interrupt if _should_interrupt( diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index 6dc86643f..80db8475f 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -9,6 +9,7 @@ from langchain_core.runnables.config import RunnableConfig from langchain_core.utils.input import get_bolded_text, get_colored_text from langgraph.channels.base import BaseChannel +from langgraph.checkpoint.base import CheckpointMetadata from langgraph.constants import TAG_HIDDEN from langgraph.pregel.io import read_channels from langgraph.pregel.types import PregelExecutableTask @@ -29,6 +30,7 @@ class TaskResultPayload(TypedDict): class CheckpointPayload(TypedDict): config: Optional[RunnableConfig] + metadata: CheckpointMetadata values: dict[str, Any] @@ -111,6 +113,7 @@ def map_debug_checkpoint( config: RunnableConfig, channels: Mapping[str, BaseChannel], stream_channels: Union[str, Sequence[str]], + metadata: CheckpointMetadata, ) -> Iterator[DebugOutputCheckpoint]: ts = datetime.now(timezone.utc).isoformat() yield { @@ -120,6 +123,7 @@ def map_debug_checkpoint( "payload": { "config": config, "values": read_channels(channels, stream_channels), + "metadata": metadata, }, } diff --git a/tests/test_pregel.py b/tests/test_pregel.py index dfa3bfd43..4da1ca2a5 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -584,12 +584,6 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "result": [("output", 13)], }, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": {"config": None, "values": {"output": 13, "inbox": [3]}}, - }, { "type": "task", "timestamp": AnyStr(), @@ -611,12 +605,6 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "result": [("output", 4)], }, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": {"config": None, "values": {"output": 4, "inbox": []}}, - }, ] @@ -5594,18 +5582,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ) ] == [ ("values", {"query": "what is weather in sf", "docs": []}), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": None, - "values": {"query": "what is weather in sf", "docs": []}, - }, - }, - ), ( "debug", { @@ -5639,18 +5615,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: }, ), ("values", {"query": "query: what is weather in sf", "docs": []}), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": None, - "values": {"query": "query: what is weather in sf", "docs": []}, - }, - }, - ), ( "debug", { @@ -5727,21 +5691,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "docs": ["doc1", "doc2", "doc3", "doc4"], }, ), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": None, - "values": { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - }, - }, - ), ( "debug", { @@ -5782,22 +5731,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "docs": ["doc1", "doc2", "doc3", "doc4"], }, ), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": None, - "values": { - "query": "query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - }, - }, - ), ] @@ -6009,6 +5942,30 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" ) ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "run_id": None, + "configurable": { + "thread_id": "10", + "thread_ts": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + }, + }, { "type": "checkpoint", "timestamp": AnyStr(), @@ -6025,7 +5982,15 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "thread_ts": AnyStr(), }, }, - "values": {"my_key": "value", "market": "DE"}, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 0, + "writes": None, + }, }, }, { @@ -6065,7 +6030,15 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "thread_ts": AnyStr(), }, }, - "values": {"my_key": "value prepared", "market": "DE"}, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, }, }, { @@ -6105,7 +6078,15 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "thread_ts": AnyStr(), }, }, - "values": {"my_key": "value prepared slow", "market": "DE"}, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, }, }, { @@ -6149,6 +6130,11 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slow finished", "market": "DE", }, + "metadata": { + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, }, }, ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4913e1f05..5eac3555d 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -544,12 +544,6 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "result": [("output", 13)], }, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": {"config": None, "values": {"output": 13, "inbox": [3]}}, - }, { "type": "task", "timestamp": AnyStr(), @@ -571,12 +565,6 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "result": [("output", 4)], }, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": {"config": None, "values": {"output": 4, "inbox": []}}, - }, ] @@ -4043,18 +4031,6 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ) ] == [ ("values", {"query": "what is weather in sf", "docs": []}), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": None, - "values": {"query": "what is weather in sf", "docs": []}, - }, - }, - ), ( "debug", { @@ -4088,18 +4064,6 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: }, ), ("values", {"query": "query: what is weather in sf", "docs": []}), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": None, - "values": {"query": "query: what is weather in sf", "docs": []}, - }, - }, - ), ( "debug", { @@ -4176,21 +4140,6 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "docs": ["doc1", "doc2", "doc3", "doc4"], }, ), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": None, - "values": { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - }, - }, - ), ( "debug", { @@ -4231,22 +4180,6 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "docs": ["doc1", "doc2", "doc3", "doc4"], }, ), - ( - "debug", - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": None, - "values": { - "query": "query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - }, - }, - ), ] @@ -4458,6 +4391,27 @@ async def test_branch_then() -> None: {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" ) ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "run_id": None, + "configurable": {"thread_id": "10", "thread_ts": AnyStr()}, + }, + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + }, + }, { "type": "checkpoint", "timestamp": AnyStr(), @@ -4474,7 +4428,11 @@ async def test_branch_then() -> None: "thread_ts": AnyStr(), }, }, - "values": {"my_key": "value", "market": "DE"}, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": {"source": "loop", "step": 0, "writes": None}, }, }, { @@ -4515,6 +4473,11 @@ async def test_branch_then() -> None: }, }, "values": {"my_key": "value prepared", "market": "DE"}, + "metadata": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, }, }, { @@ -4554,7 +4517,15 @@ async def test_branch_then() -> None: "thread_ts": AnyStr(), }, }, - "values": {"my_key": "value prepared slow", "market": "DE"}, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, }, }, { @@ -4598,6 +4569,11 @@ async def test_branch_then() -> None: "my_key": "value prepared slow finished", "market": "DE", }, + "metadata": { + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, }, }, ]