From ec0aa273fda496187a3fde1adf23874d6331429a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 6 Jun 2024 17:53:14 -0700 Subject: [PATCH] Fix for update_state mistakenly resetting barrier channels - This introduces a new optional method BaseChannel.consume, which gets called when a channel triggers a node - This is an alternative place to clean up channel state, in addition to the existing pattern of clearing state in the update() call for the next step - Channels should implement consume() when they want to clean up state exactly if and only if the channel triggered a node - Channels should clean up state in update() when they instead need to guarantee their value is available for a single step, irrespective of whether it was actually read --- langgraph/channels/base.py | 10 ++- langgraph/channels/dynamic_barrier_value.py | 10 +-- langgraph/channels/named_barrier_value.py | 6 +- langgraph/pregel/__init__.py | 12 ++- tests/test_pregel.py | 81 +++++++++++++++++++++ 5 files changed, 109 insertions(+), 10 deletions(-) diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index bbabb5500..01d595464 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -65,7 +65,8 @@ class BaseChannel(Generic[Value, Update, C], ABC): def update(self, values: Sequence[Update]) -> None: """Update the channel's value with the given sequence of updates. The order of the updates in the sequence is arbitrary. - + This method is called by Pregel for all channels at the end of each step. + If there are no updates, it is called with an empty sequence. Raises InvalidUpdateError if the sequence of updates is invalid.""" @abstractmethod @@ -74,6 +75,13 @@ class BaseChannel(Generic[Value, Update, C], ABC): Raises EmptyChannelError if the channel is empty (never updated yet).""" + def consume(self) -> None: + """Mark the current value of the channel as consumed. By default, no-op. + This is called by Pregel before the start of the next step, for all + channels that triggered a node. + """ + pass + @contextmanager def ChannelsManager( diff --git a/langgraph/channels/dynamic_barrier_value.py b/langgraph/channels/dynamic_barrier_value.py index 40758bc0a..45cd83a7c 100644 --- a/langgraph/channels/dynamic_barrier_value.py +++ b/langgraph/channels/dynamic_barrier_value.py @@ -60,11 +60,6 @@ class DynamicBarrierValue( pass def update(self, values: Sequence[Union[Value, WaitForNames]]) -> None: - # switch to "priming" state after reading - if self.seen == self.names: - self.seen = set() - self.names = None - if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]: if len(wait_for_names) > 1: raise InvalidUpdateError( @@ -83,3 +78,8 @@ class DynamicBarrierValue( if self.seen != self.names: raise EmptyChannelError() return None + + def consume(self) -> None: + if self.seen == self.names: + self.seen = set() + self.names = None diff --git a/langgraph/channels/named_barrier_value.py b/langgraph/channels/named_barrier_value.py index 5c6072155..8181d0a02 100644 --- a/langgraph/channels/named_barrier_value.py +++ b/langgraph/channels/named_barrier_value.py @@ -42,8 +42,6 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): pass def update(self, values: Sequence[Value]) -> None: - if self.seen == self.names: - self.seen = set() for value in values: if value in self.names: self.seen.add(value) @@ -54,3 +52,7 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): if self.seen != self.names: raise EmptyChannelError() return None + + def consume(self) -> None: + if self.seen == self.names: + self.seen = set() diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 74b404f7b..4d53c6be9 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -1685,6 +1685,8 @@ def _prepare_next_tasks( tasks.append(PregelTaskDescription(packet.node, packet.arg)) if for_execution: checkpoint["pending_sends"].clear() + # Collect channels to consume + channels_to_consume = set() # Check if any processes should be run in next step # If so, prepare the values to be passed to them for name, proc in processes.items(): @@ -1698,6 +1700,7 @@ def _prepare_next_tasks( ) and checkpoint["channel_versions"][chan] > seen[chan] ]: + channels_to_consume.update(triggers) try: val = next(_proc_input(step, name, proc, managed, channels)) except StopIteration: @@ -1754,6 +1757,10 @@ def _prepare_next_tasks( ) else: tasks.append(PregelTaskDescription(name, val)) + # Consume all channels that were read + if for_execution: + for chan in channels_to_consume: + channels[chan].consume() return checkpoint, tasks @@ -1763,7 +1770,6 @@ def _proc_input( proc: PregelNode, managed: ManagedValueMapping, channels: Mapping[str, BaseChannel], - catch: bool = False, ) -> Iterator[Any]: # If all trigger channels subscribed by this process are not empty # then invoke the process with the values of all non-empty channels @@ -1771,7 +1777,9 @@ def _proc_input( try: val: dict = { k: read_channel( - channels, chan, catch=catch or chan not in proc.triggers + channels, + chan, + catch=chan not in proc.triggers, ) for k, chan in proc.channels.items() if isinstance(chan, str) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 24566598a..798d11241 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -6224,6 +6224,52 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) + with SqliteSaver.from_conn_string(":memory:") as saver: + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_before=["finish"] + ) + + thread1 = {"configurable": {"thread_id": "1"}} + + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared slow", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={ + "my_key": "value prepared slow", + "market": "DE", + }, + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + # update state + tool_two.update_state(thread1, {"my_key": "er"}) + assert tool_two.get_state(thread1) == StateSnapshot( + values={ + "my_key": "value prepared slower", + "market": "DE", + }, + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "update", + "step": 3, + "writes": {"tool_two_slow": {"my_key": "er"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + with SqliteSaver.from_conn_string(":memory:") as saver: tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_after=["prepare"] @@ -6440,6 +6486,41 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), + interrupt_before=["qa"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + app_w_interrupt.update_state(config, {"docs": ["doc5"]}) + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], + }, + next=("qa",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "source": "update", + "step": 4, + "writes": {"retriever_one": {"docs": ["doc5"]}}, + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, + ] + def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion,