From 75c8f99d3cc01d1e1d72d021c57d8ce7e6a32c0e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Oct 2024 14:45:09 -0700 Subject: [PATCH] Add "node finished" callback --- libs/langgraph/langgraph/constants.py | 2 ++ libs/langgraph/langgraph/pregel/__init__.py | 3 +++ libs/langgraph/langgraph/pregel/runner.py | 8 +++++++- libs/langgraph/tests/test_pregel.py | 16 ++++++++++++++-- libs/langgraph/tests/test_pregel_async.py | 16 ++++++++++++++-- 5 files changed, 40 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 52eb72c38..36e375751 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -65,6 +65,8 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id") # holds the current checkpoint_id, if any CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns") # holds the current checkpoint_ns, "" for root graph +CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished") +# callback to be called when a node is finished # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c6018cbc7..e6a5638d2 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -56,6 +56,7 @@ from langgraph.constants import ( CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, + CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, @@ -1288,6 +1289,7 @@ class Pregel(PregelProtocol): runner = PregelRunner( submit=loop.submit, put_writes=loop.put_writes, + node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming if subgraphs: @@ -1509,6 +1511,7 @@ class Pregel(PregelProtocol): submit=loop.submit, put_writes=loop.put_writes, use_astream=do_stream is not None, + node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming if subgraphs: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index b7ad68884..a2ca1c6cb 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -14,7 +14,7 @@ from typing import ( cast, ) -from langgraph.constants import ERROR, INTERRUPT, NO_WRITES +from langgraph.constants import ERROR, INTERRUPT, NO_WRITES, TAG_HIDDEN from langgraph.errors import GraphDelegate, GraphInterrupt from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry @@ -32,10 +32,12 @@ class PregelRunner: submit: Submit, put_writes: Callable[[str, Sequence[tuple[str, Any]]], None], use_astream: bool = False, + node_finished: Optional[Callable[[str], None]] = None, ) -> None: self.submit = submit self.put_writes = put_writes self.use_astream = use_astream + self.node_finished = node_finished def tick( self, @@ -209,6 +211,10 @@ class PregelRunner: # save error to checkpointer self.put_writes(task.id, [(ERROR, exception)]) else: + if self.node_finished and ( + task.config is None or TAG_HIDDEN not in task.config.get("tags", []) + ): + self.node_finished(task.name) if not task.writes: # add no writes marker task.writes.append((NO_WRITES, None)) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ded80e20d..71b7f4a4f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -54,7 +54,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import ERROR, PULL, PUSH +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt from langgraph.graph import END, Graph from langgraph.graph.graph import START @@ -8917,15 +8917,27 @@ def test_doubly_nested_graph_interrupts( } # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} + nodes: list[str] = [] + config = { + "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} + } assert [*app.stream({"my_key": "my value"}, config)] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] + assert nodes == ["parent_1", "grandchild_1"] assert [*app.stream(None, config)] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] + assert nodes == [ + "parent_1", + "grandchild_1", + "grandchild_2", + "child_1", + "child", + "parent_2", + ] # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 82368be78..03b119d0d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -50,7 +50,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import ERROR, PULL, PUSH +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt from langgraph.graph import END, Graph, StateGraph from langgraph.graph.graph import START @@ -7605,15 +7605,27 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: } # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} + nodes: list[str] = [] + config = { + "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} + } assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] + assert nodes == ["parent_1", "grandchild_1"] assert [c async for c in app.astream(None, config)] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] + assert nodes == [ + "parent_1", + "grandchild_1", + "grandchild_2", + "child_1", + "child", + "parent_2", + ] # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}}