From c2776449fd72fdb9d40ed3f402fc8533b899c573 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 16:03:57 -0700 Subject: [PATCH] Deferred Node (#4269) --- libs/langgraph/langgraph/channels/base.py | 46 +- .../channels/dynamic_barrier_value.py | 113 +++- .../langgraph/channels/last_value.py | 70 +++ .../langgraph/channels/named_barrier_value.py | 86 +++ libs/langgraph/langgraph/graph/state.py | 36 +- libs/langgraph/langgraph/pregel/__init__.py | 56 +- libs/langgraph/langgraph/pregel/algo.py | 25 +- libs/langgraph/langgraph/pregel/draw.py | 17 +- libs/langgraph/langgraph/pregel/loop.py | 10 +- libs/langgraph/langgraph/types.py | 4 +- .../tests/__snapshots__/test_pregel.ambr | 47 +- libs/langgraph/tests/test_pregel.py | 557 ++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 188 ++++++ .../langgraph/scheduler/kafka/orchestrator.py | 2 + 14 files changed, 1207 insertions(+), 50 deletions(-) diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index f7b53238f..4d6335bc1 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -53,16 +53,7 @@ class BaseChannel(Generic[Value, Update, C], ABC): """Return a new identical channel, optionally initialized from a checkpoint. If the checkpoint contains complex data structures, they should be copied.""" - # state methods - - @abstractmethod - def update(self, values: Sequence[Update]) -> bool: - """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. - Returns True if the channel was updated, False otherwise.""" + # read methods @abstractmethod def get(self) -> Value: @@ -70,13 +61,6 @@ class BaseChannel(Generic[Value, Update, C], ABC): Raises EmptyChannelError if the channel is empty (never updated yet).""" - def consume(self) -> bool: - """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. If the channel was updated, return True. - """ - return False - def is_available(self) -> bool: """Return True if the channel is available (not empty), False otherwise. Subclasses should override this method to provide a more efficient @@ -88,6 +72,34 @@ class BaseChannel(Generic[Value, Update, C], ABC): except EmptyChannelError: return False + # write methods + + @abstractmethod + def update(self, values: Sequence[Update]) -> bool: + """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. + Returns True if the channel was updated, False otherwise.""" + + def consume(self) -> bool: + """Notify the channel that a subscribed task ran. By default, no-op. + A channel can use this method to modify its state, preventing the value + from being consumed again. + + Returns True if the channel was updated, False otherwise. + """ + return False + + def finish(self) -> bool: + """Notify the channel that the Pregel run is finishing. By default, no-op. + A channel can use this method to modify its state, preventing finish. + + Returns True if the channel was updated, False otherwise. + """ + return False + __all__ = [ "BaseChannel", diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py index c2c3c026f..c25551082 100644 --- a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py +++ b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py @@ -81,12 +81,9 @@ class DynamicBarrierValue( updated = False for value in values: assert not isinstance(value, WaitForNames) - if value in self.names: - if value not in self.seen: - self.seen.add(value) - updated = True - else: - raise InvalidUpdateError(f"Value {value} not in {self.names}") + if value in self.names and value not in self.seen: + self.seen.add(value) + updated = True return updated def get(self) -> Value: @@ -103,3 +100,107 @@ class DynamicBarrierValue( self.names = None return True return False + + +class DynamicBarrierValueAfterFinish( + Generic[Value], BaseChannel[Value, Union[Value, WaitForNames], Set[Value]] +): + """A channel that switches between two states + + - in the "priming" state it can't be read from. + - if it receives a WaitForNames update, it switches to the "waiting" state. + - in the "waiting" state it collects named values until all are received. + - once all named values are received, and the finished flag is set, it can be read once, and it switches + back to the "priming" state. + """ + + __slots__ = ("names", "seen", "finished") + + names: Optional[Set[Value]] + seen: set[Value] + finished: bool + + def __init__(self, typ: type[Value]) -> None: + super().__init__(typ) + self.names = None + self.seen = set() + self.finished = False + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, DynamicBarrierValueAfterFinish) + and value.names == self.names + ) + + @property + def ValueType(self) -> type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> type[Value]: + """The type of the update received by the channel.""" + return self.typ + + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ) + empty.key = self.key + empty.names = self.names + empty.seen = self.seen.copy() + empty.finished = self.finished + return empty + + def checkpoint(self) -> tuple[Optional[Set[Value]], set[Value], bool]: + return (self.names, self.seen, self.finished) + + def from_checkpoint( + self, checkpoint: tuple[Optional[Set[Value]], set[Value], bool] + ) -> Self: + empty = self.__class__(self.typ) + empty.key = self.key + if checkpoint is not MISSING: + names, seen, finished = checkpoint + empty.names = names if names is not None else None + empty.seen = seen + empty.finished = finished + return empty + + def update(self, values: Sequence[Union[Value, WaitForNames]]) -> bool: + if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]: + if len(wait_for_names) > 1: + raise InvalidUpdateError( + f"At key '{self.key}': Received multiple WaitForNames updates in the same step." + ) + self.names = wait_for_names[0].names + return True + elif self.names is not None: + updated = False + for value in values: + assert not isinstance(value, WaitForNames) + if value in self.names and value not in self.seen: + self.seen.add(value) + updated = True + return updated + + def get(self) -> Value: + if not self.finished and self.seen != self.names: + raise EmptyChannelError() + return None + + def is_available(self) -> bool: + return self.seen == self.names and self.finished + + def consume(self) -> bool: + if self.finished and self.seen == self.names: + self.seen = set() + self.names = None + return True + return False + + def finish(self) -> bool: + if not self.finished and self.seen == self.names: + self.finished = True + return True + else: + return False diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py index c067aeb00..59c8d3c1b 100644 --- a/libs/langgraph/langgraph/channels/last_value.py +++ b/libs/langgraph/langgraph/channels/last_value.py @@ -70,3 +70,73 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): def checkpoint(self) -> Value: return self.value + + +class LastValueAfterFinish( + Generic[Value], BaseChannel[Value, Value, tuple[Value, bool]] +): + """Stores the last value received, but only made available after finish(). + Once made available, clears the value.""" + + __slots__ = ("value", "finished") + + def __init__(self, typ: Any, key: str = "") -> None: + super().__init__(typ, key) + self.value = MISSING + self.finished = False + + def __eq__(self, value: object) -> bool: + return isinstance(value, LastValueAfterFinish) + + @property + def ValueType(self) -> type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> type[Value]: + """The type of the update received by the channel.""" + return self.typ + + def checkpoint(self) -> tuple[Value, bool]: + if self.value is MISSING: + return MISSING + return (self.value, self.finished) + + def from_checkpoint(self, checkpoint: tuple[Value, bool]) -> Self: + empty = self.__class__(self.typ) + empty.key = self.key + if checkpoint is not MISSING: + empty.value, empty.finished = checkpoint + return empty + + def update(self, values: Sequence[Value]) -> bool: + if len(values) == 0: + return False + + self.finished = False + self.value = values[-1] + return True + + def consume(self) -> bool: + if self.finished: + self.finished = False + self.value = MISSING + return True + + return False + + def finish(self) -> bool: + if not self.finished and self.value is not MISSING: + self.finished = True + return True + else: + return False + + def get(self) -> Value: + if self.value is MISSING or not self.finished: + raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING and self.finished diff --git a/libs/langgraph/langgraph/channels/named_barrier_value.py b/libs/langgraph/langgraph/channels/named_barrier_value.py index 628c0bcf2..e5e96a7fb 100644 --- a/libs/langgraph/langgraph/channels/named_barrier_value.py +++ b/libs/langgraph/langgraph/channels/named_barrier_value.py @@ -77,3 +77,89 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): self.seen = set() return True return False + + +class NamedBarrierValueAfterFinish( + Generic[Value], BaseChannel[Value, Value, set[Value]] +): + """A channel that waits until all named values are received before making the value ready to be made available. It is only made available after finish() is called.""" + + __slots__ = ("names", "seen", "finished") + + names: set[Value] + seen: set[Value] + + def __init__(self, typ: type[Value], names: set[Value]) -> None: + super().__init__(typ) + self.names = names + self.seen: set[str] = set() + self.finished = False + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, NamedBarrierValueAfterFinish) + and value.names == self.names + ) + + @property + def ValueType(self) -> type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> type[Value]: + """The type of the update received by the channel.""" + return self.typ + + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.names) + empty.key = self.key + empty.seen = self.seen.copy() + empty.finished = self.finished + return empty + + def checkpoint(self) -> tuple[set[Value], bool]: + return (self.seen, self.finished) + + def from_checkpoint(self, checkpoint: tuple[set[Value], bool]) -> Self: + empty = self.__class__(self.typ, self.names) + empty.key = self.key + if checkpoint is not MISSING: + empty.seen, empty.finished = checkpoint + return empty + + def update(self, values: Sequence[Value]) -> bool: + updated = False + for value in values: + if value in self.names: + if value not in self.seen: + self.seen.add(value) + updated = True + else: + raise InvalidUpdateError( + f"At key '{self.key}': Value {value} not in {self.names}" + ) + return updated + + def get(self) -> Value: + if not self.finished or self.seen != self.names: + raise EmptyChannelError() + return None + + def is_available(self) -> bool: + return self.finished and self.seen == self.names + + def consume(self) -> bool: + if self.finished and self.seen == self.names: + self.finished = False + self.seen = set() + return True + return False + + def finish(self) -> bool: + if not self.finished and self.seen == self.names: + self.finished = True + return True + else: + return False diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 59da4e355..c86fc6b14 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -28,10 +28,17 @@ from typing_extensions import Self from langgraph._api.deprecation import LangGraphDeprecationWarning from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitForNames +from langgraph.channels.dynamic_barrier_value import ( + DynamicBarrierValue, + DynamicBarrierValueAfterFinish, + WaitForNames, +) from langgraph.channels.ephemeral_value import EphemeralValue -from langgraph.channels.last_value import LastValue -from langgraph.channels.named_barrier_value import NamedBarrierValue +from langgraph.channels.last_value import LastValue, LastValueAfterFinish +from langgraph.channels.named_barrier_value import ( + NamedBarrierValue, + NamedBarrierValueAfterFinish, +) from langgraph.checkpoint.base import Checkpoint from langgraph.constants import ( EMPTY_SEQ, @@ -107,6 +114,7 @@ class StateNodeSpec(NamedTuple): input: type[Any] retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ + defer: bool = False class StateGraph(Graph): @@ -247,6 +255,7 @@ class StateGraph(Graph): self, node: RunnableLike, *, + defer: bool = False, metadata: Optional[dict[str, Any]] = None, input: Optional[type[Any]] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, @@ -263,6 +272,7 @@ class StateGraph(Graph): node: str, action: RunnableLike, *, + defer: bool = False, metadata: Optional[dict[str, Any]] = None, input: Optional[type[Any]] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, @@ -276,6 +286,7 @@ class StateGraph(Graph): node: Union[str, RunnableLike], action: Optional[RunnableLike] = None, *, + defer: bool = False, metadata: Optional[dict[str, Any]] = None, input: Optional[type[Any]] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, @@ -421,6 +432,7 @@ class StateGraph(Graph): input=input or self.schema, retry_policy=retry, ends=ends, + defer=defer, ) return self @@ -784,7 +796,11 @@ class CompiledStateGraph(CompiledGraph): self.schema_to_mapper[input_schema] = mapper branch_channel = CHANNEL_BRANCH_TO.format(key) - self.channels[branch_channel] = EphemeralValue(Any, guard=False) + self.channels[branch_channel] = ( + LastValueAfterFinish(Any) + if node.defer + else EphemeralValue(Any, guard=False) + ) self.nodes[key] = PregelNode( triggers=[branch_channel], # read state keys and managed values @@ -812,7 +828,12 @@ class CompiledStateGraph(CompiledGraph): elif end != END: channel_name = f"join:{'+'.join(starts)}:{end}" # register channel - self.channels[channel_name] = NamedBarrierValue(str, set(starts)) + if self.builder.nodes[end].defer: + self.channels[channel_name] = NamedBarrierValueAfterFinish( + str, set(starts) + ) + else: + self.channels[channel_name] = NamedBarrierValue(str, set(starts)) # subscribe to channel self.nodes[end].triggers.append(channel_name) # publish to channel @@ -889,7 +910,10 @@ class CompiledStateGraph(CompiledGraph): else [node for node in self.builder.nodes if node != branch.then] ) channel_name = f"branch:{start}:{name}::then" - self.channels[channel_name] = DynamicBarrierValue(str) + if self.builder.nodes[branch.then].defer: + self.channels[channel_name] = DynamicBarrierValueAfterFinish(str) + else: + self.channels[channel_name] = DynamicBarrierValue(str) self.nodes[branch.then].triggers.append(channel_name) for end in ends: if end != END: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 0663f8311..3abe2092f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -506,7 +506,7 @@ class Pregel(PregelProtocol): name: str = "LangGraph" - trigger_to_nodes: Mapping[str, Sequence[str]] | None = None + trigger_to_nodes: Mapping[str, Sequence[str]] def __init__( self, @@ -552,7 +552,7 @@ class Pregel(PregelProtocol): self.config_type = config_type self.input_model = input_model self.config = config - self.trigger_to_nodes = trigger_to_nodes + self.trigger_to_nodes = trigger_to_nodes or {} self.name = name if auto_validate: self.validate() @@ -949,6 +949,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, + self.trigger_to_nodes, ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: @@ -958,7 +959,9 @@ class Pregel(PregelProtocol): continue next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(saved.checkpoint, channels, tasks, None) + apply_writes( + saved.checkpoint, channels, tasks, None, self.trigger_to_nodes + ) tasks_with_writes = tasks_w_writes( next_tasks.values(), saved.pending_writes, @@ -1071,6 +1074,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, + self.trigger_to_nodes, ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: @@ -1080,7 +1084,9 @@ class Pregel(PregelProtocol): continue next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(saved.checkpoint, channels, tasks, None) + apply_writes( + saved.checkpoint, channels, tasks, None, self.trigger_to_nodes + ) tasks_with_writes = tasks_w_writes( next_tasks.values(), @@ -1407,6 +1413,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, + self.trigger_to_nodes, ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: @@ -1416,7 +1423,13 @@ class Pregel(PregelProtocol): continue next_tasks[tid].writes.append((k, v)) # clear all current tasks - apply_writes(checkpoint, channels, next_tasks.values(), None) + apply_writes( + checkpoint, + channels, + next_tasks.values(), + None, + self.trigger_to_nodes, + ) # save checkpoint next_config = checkpointer.put( checkpoint_config, @@ -1475,6 +1488,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, input_writes, [])], checkpointer.get_next_version, + self.trigger_to_nodes, ) # apply input write to channels @@ -1575,6 +1589,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, + self.trigger_to_nodes, ) # apply writes for tid, k, v in saved.pending_writes: @@ -1584,7 +1599,9 @@ class Pregel(PregelProtocol): continue next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(checkpoint, channels, tasks, None) + apply_writes( + checkpoint, channels, tasks, None, self.trigger_to_nodes + ) valid_updates: list[tuple[str, dict[str, Any] | None]] = [] if len(updates) == 1: values, as_node = updates[0] @@ -1672,7 +1689,11 @@ class Pregel(PregelProtocol): checkpointer.put_writes(checkpoint_config, channel_writes, task_id) # apply to checkpoint and save mv_writes, _ = apply_writes( - checkpoint, channels, run_tasks, checkpointer.get_next_version + checkpoint, + channels, + run_tasks, + checkpointer.get_next_version, + self.trigger_to_nodes, ) assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) @@ -1822,6 +1843,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, + self.trigger_to_nodes, ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: @@ -1831,7 +1853,13 @@ class Pregel(PregelProtocol): continue next_tasks[tid].writes.append((k, v)) # clear all current tasks - apply_writes(checkpoint, channels, next_tasks.values(), None) + apply_writes( + checkpoint, + channels, + next_tasks.values(), + None, + self.trigger_to_nodes, + ) # save checkpoint next_config = await checkpointer.aput( checkpoint_config, @@ -1890,6 +1918,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, input_writes, [])], checkpointer.get_next_version, + self.trigger_to_nodes, ) # apply input write to channels @@ -1990,6 +2019,7 @@ class Pregel(PregelProtocol): channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, + self.trigger_to_nodes, ) for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): @@ -1998,7 +2028,9 @@ class Pregel(PregelProtocol): continue next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(checkpoint, channels, tasks, None) + apply_writes( + checkpoint, channels, tasks, None, self.trigger_to_nodes + ) valid_updates: list[tuple[str, dict[str, Any] | None]] = [] if len(updates) == 1: values, as_node = updates[0] @@ -2084,7 +2116,11 @@ class Pregel(PregelProtocol): ) # apply to checkpoint and save mv_writes, _ = apply_writes( - checkpoint, channels, run_tasks, checkpointer.get_next_version + checkpoint, + channels, + run_tasks, + checkpointer.get_next_version, + self.trigger_to_nodes, ) assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index b6ce4433d..c51afd7d3 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -232,6 +232,7 @@ def apply_writes( channels: Mapping[str, BaseChannel], tasks: Iterable[WritesProtocol], get_next_version: Optional[GetNextVersion], + trigger_to_nodes: Mapping[str, Sequence[str]], ) -> tuple[dict[str, list[Any]], set[str]]: """Apply writes from a set of tasks (usually the tasks from a Pregel step) to the checkpoint and channels, and return managed values writes to be applied @@ -317,7 +318,9 @@ def apply_writes( max_version, channels[chan], ) - updated_channels.add(chan) + # unavailable channels can't trigger tasks, so don't add them + if channels[chan].is_available(): + updated_channels.add(chan) # Channels that weren't updated in this step are notified of a new step if bump_step: @@ -328,6 +331,26 @@ def apply_writes( max_version, channels[chan], ) + # unavailable channels can't trigger tasks, so don't add them + if channels[chan].is_available(): + updated_channels.add(chan) + + # If this is (tentatively) the last superstep, notify all channels of finish + if ( + bump_step + and not checkpoint["pending_sends"] + and updated_channels.isdisjoint(trigger_to_nodes) + ): + for chan in channels: + if channels[chan].finish() and get_next_version is not None: + checkpoint["channel_versions"][chan] = get_next_version( + max_version, + channels[chan], + ) + # unavailable channels can't trigger tasks, so don't add them + if channels[chan].is_available(): + updated_channels.add(chan) + # Return managed values writes to be applied externally return pending_writes_by_managed, updated_channels diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index 8df6dd458..6e6b7440c 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -31,7 +31,7 @@ def draw_graph( input_channels: Union[str, Sequence[str]], interrupt_after_nodes: Union[All, Sequence[str]], interrupt_before_nodes: Union[All, Sequence[str]], - trigger_to_nodes: Optional[Mapping[str, Sequence[str]]], + trigger_to_nodes: Mapping[str, Sequence[str]], checkpointer: Checkpointer, subgraphs: dict[str, Graph], limit: int = 250, @@ -79,6 +79,7 @@ def draw_graph( PregelTaskWrites((), INPUT, input_writes, []), ], get_next_version, + trigger_to_nodes, ) # prepare first tasks tasks = prepare_next_tasks( @@ -98,7 +99,7 @@ def draw_graph( ) start_tasks = tasks # run the pregel loop - for _ in range(limit): + for step in range(step, limit): if not tasks: break conditionals: dict[tuple[str, str, Any], Optional[str]] = {} @@ -144,7 +145,7 @@ def draw_graph( trigger_to_sources[trigger].add((src, cond, label)) # apply writes _, updated_channels = apply_writes( - checkpoint, channels, tasks.values(), get_next_version + checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes ) # prepare next tasks tasks = prepare_next_tasks( @@ -164,9 +165,19 @@ def draw_graph( ) # collect edges for task in tasks.values(): + added = False for trigger in task.triggers: for src, cond, label in sorted(trigger_to_sources[trigger]): edges.add((src, task.name, cond, label)) + # if the edge is from this step, skip adding the implicit edges + if (trigger, cond, label) in step_sources.get(src, set()): + added = True + else: + sources[src].discard((trigger, cond, label)) + # if no edges from this step, add implicit edges from all previous tasks + if not added: + for src in step_sources: + edges.add((src, task.name, True, None)) # assemble the graph graph = Graph() # add nodes diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f4a230b42..805779d47 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -211,13 +211,13 @@ class PregelLoop(LoopProtocol): specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], output_keys: Union[str, Sequence[str]], stream_keys: Union[str, Sequence[str]], + trigger_to_nodes: Mapping[str, Sequence[str]], interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, input_model: Optional[type[BaseModel]] = None, debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, - trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, checkpoint_during: bool = True, ) -> None: super().__init__( @@ -483,6 +483,7 @@ class PregelLoop(LoopProtocol): self.channels, self.tasks.values(), self.checkpointer_get_next_version, + self.trigger_to_nodes, ) # apply writes to managed values for key, values in mv_writes.items(): @@ -679,6 +680,7 @@ class PregelLoop(LoopProtocol): self.channels, [PregelTaskWrites((), INPUT, null_writes, [])], self.checkpointer_get_next_version, + self.trigger_to_nodes, ) for key, values in mv_writes.items(): self._update_mv(key, values) @@ -731,6 +733,7 @@ class PregelLoop(LoopProtocol): PregelTaskWrites((), INPUT, input_writes, []), ], self.checkpointer_get_next_version, + self.trigger_to_nodes, ) assert not mv_writes, "Can't write to SharedValues in graph input" # save input checkpoint @@ -868,6 +871,7 @@ class PregelLoop(LoopProtocol): self.channels, self.tasks.values(), self.checkpointer_get_next_version, + self.trigger_to_nodes, ) for key, values in mv_writes.items(): self._update_mv(key, values) @@ -966,6 +970,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + trigger_to_nodes: Mapping[str, Sequence[str]], manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, @@ -974,7 +979,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): input_model: Optional[type[BaseModel]] = None, debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, - trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, checkpoint_during: bool = True, ) -> None: super().__init__( @@ -1116,6 +1120,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + trigger_to_nodes: Mapping[str, Sequence[str]], interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, @@ -1124,7 +1129,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): input_model: Optional[type[BaseModel]] = None, debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, - trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, checkpoint_during: bool = True, ) -> None: super().__init__( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 1dad60a9e..3d34b35fd 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -313,8 +313,8 @@ class Command(Generic[N], ToolOutputMixin): graph: Optional[str] = None update: Optional[Any] = None - resume: Optional[Union[dict[str, Any], Any]] = None - goto: Union[Send, Sequence[Union[Send, str]], str] = () + resume: Optional[Union[Any, dict[str, Any]]] = None + goto: Union[Send, Sequence[Union[Send, N]], N] = () def __repr__(self) -> str: # get all non-None values diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 46955d25c..feef534c9 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -378,7 +378,7 @@ "target": "agent" }, { - "source": "human", + "source": "agent", "target": "__end__", "conditional": true } @@ -392,7 +392,50 @@ __start__ --> human; agent --> human; human --> agent; - human -.-> __end__; + agent -.-> __end__; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_defer_node[memory-False] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one -.-> qa; + retriever_one --> analyzer_one; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> retriever_one; + rewrite_query --> retriever_two; + qa --> __end__; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_defer_node[memory-True] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one -.-> qa; + retriever_one --> analyzer_one; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> retriever_one; + rewrite_query --> retriever_two; + qa --> __end__; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_then_defer_node[memory-True] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> qa; + analyzer_one --> retriever_one; + retriever_one -.-> qa; + retriever_two --> qa; + rewrite_query -.-> analyzer_one; + rewrite_query -.-> qa; + rewrite_query -.-> retriever_two; + qa --> __end__; ''' # --- diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a8fd2b5f3..a5baeeece 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2500,6 +2500,563 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] +@pytest.mark.parametrize("use_waiting_edge", (True, False)) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_in_one_fan_out_state_graph_defer_node( + snapshot: SnapshotAssertion, + request: pytest.FixtureRequest, + checkpointer_name: str, + use_waiting_edge: bool, +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + workflow = StateGraph(State) + + @workflow.add_node + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) # to ensure stream order + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow.add_node(analyzer_one) + workflow.add_node(retriever_one) + workflow.add_node(retriever_two) + workflow.add_node(qa, defer=True) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "retriever_one") + workflow.add_edge("retriever_one", "analyzer_one") + workflow.add_edge("rewrite_query", "retriever_two") + if use_waiting_edge: + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + else: + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + if checkpointer_name == "memory": + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert [*app.stream({"query": "what is weather in sf"}, stream_mode="debug")] == [ + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "input": {"query": "what is weather in sf", "docs": []}, + "triggers": ("branch:to:rewrite_query",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "error": None, + "result": [("query", "query: what is weather in sf")], + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_one",), + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_two",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "error": None, + "result": [("docs", ["doc1", "doc2"])], + "interrupts": [], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "error": None, + "result": [("docs", ["doc3", "doc4"])], + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "analyzer_one", + "input": { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:to:analyzer_one",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "analyzer_one", + "error": None, + "result": [("query", "analyzed: query: what is weather in sf")], + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "id": AnyStr(), + "name": "qa", + "input": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:to:qa", "join:retriever_one+retriever_two:qa") + if use_waiting_edge + else ("branch:to:qa",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "id": AnyStr(), + "name": "qa", + "error": None, + "result": [("answer", "doc1,doc2,doc3,doc4")], + "interrupts": [], + }, + }, + ] + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["analyzer_one"], + ) + 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"}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["qa"], + ) + config = {"configurable": {"thread_id": "2"}} + + 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"}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"__interrupt__": ()}, + ] + + app_w_interrupt.update_state(config, {"docs": ["doc5"]}) + expected_parent_config = ( + None + if "shallow" in checkpointer_name + else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ) + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], + }, + tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),), + next=("qa",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": {"analyzer_one": {"docs": ["doc5"]}}, + "thread_id": "2", + }, + parent_config=expected_parent_config, + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, + ] + + +@pytest.mark.parametrize("with_path_map", (True, False)) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_in_one_fan_out_state_graph_then_defer_node( + snapshot: SnapshotAssertion, + request: pytest.FixtureRequest, + checkpointer_name: str, + with_path_map: bool, +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + workflow = StateGraph(State) + + @workflow.add_node + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) # to ensure stream order + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow.add_node(analyzer_one) + workflow.add_node(retriever_one) + workflow.add_node(retriever_two) + workflow.add_node(qa, defer=True) + + workflow.set_entry_point("rewrite_query") + workflow.add_conditional_edges( + "rewrite_query", + lambda _: ["analyzer_one", "retriever_two"], + ["analyzer_one", "retriever_two"] if with_path_map else None, + then="qa", + ) + workflow.add_edge("analyzer_one", "retriever_one") + + app = workflow.compile() + + if checkpointer_name == "memory" and with_path_map: + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"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"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert [*app.stream({"query": "what is weather in sf"}, stream_mode="debug")] == [ + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "input": {"query": "what is weather in sf", "docs": []}, + "triggers": ("branch:to:rewrite_query",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "error": None, + "result": [("query", "query: what is weather in sf")], + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "analyzer_one", + "input": { + "query": "query: what is weather in sf", + "docs": [], + }, + "triggers": ("branch:to:analyzer_one",), + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_two",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "analyzer_one", + "error": None, + "result": [("query", "analyzed: query: what is weather in sf")], + "interrupts": [], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "error": None, + "result": [("docs", ["doc3", "doc4"])], + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "input": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc3", "doc4"], + }, + "triggers": ("branch:to:retriever_one",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "error": None, + "result": [("docs", ["doc1", "doc2"])], + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "id": AnyStr(), + "name": "qa", + "input": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:rewrite_query:condition::then", "branch:to:qa"), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "id": AnyStr(), + "name": "qa", + "error": None, + "result": [("answer", "doc1,doc2,doc3,doc4")], + "interrupts": [], + }, + }, + ] + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["analyzer_one"], + ) + 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"]}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["qa"], + ) + config = {"configurable": {"thread_id": "2"}} + + 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"]}}, + {"__interrupt__": ()}, + ] + + app_w_interrupt.update_state(config, {"docs": ["doc5"]}) + expected_parent_config = ( + None + if "shallow" in checkpointer_name + else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ) + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], + }, + tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),), + next=("qa",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": {"retriever_one": {"docs": ["doc5"]}}, + "thread_id": "2", + }, + parent_config=expected_parent_config, + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 661beec5b..38c3ed52d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4572,6 +4572,194 @@ async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) - ] +@pytest.mark.parametrize("use_waiting_edge", (True, False)) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_in_one_fan_out_state_graph_defer_node( + checkpointer_name: str, use_waiting_edge: bool +) -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa, defer=True) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + if use_waiting_edge: + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + else: + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"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"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"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"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +@pytest.mark.parametrize("with_path_map", (True, False)) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_in_one_fan_out_state_graph_then_defer_node( + checkpointer_name: str, with_path_map: bool +) -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa, defer=True) + + workflow.set_entry_point("rewrite_query") + workflow.add_conditional_edges( + "rewrite_query", + lambda _: ["analyzer_one", "retriever_two"], + ["analyzer_one", "retriever_two"] if with_path_map else None, + then="qa", + ) + workflow.add_edge("analyzer_one", "retriever_one") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"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"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"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"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, checkpointer_name: str diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 3e4499266..9f6d4110f 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -159,6 +159,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): stream_keys=graph.stream_channels, interrupt_after=graph.interrupt_after_nodes, interrupt_before=graph.interrupt_before_nodes, + trigger_to_nodes=graph.trigger_to_nodes, ) as loop: if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved @@ -345,6 +346,7 @@ class KafkaOrchestrator(AbstractContextManager): stream_keys=graph.stream_channels, interrupt_after=graph.interrupt_after_nodes, interrupt_before=graph.interrupt_before_nodes, + trigger_to_nodes=graph.trigger_to_nodes, ) as loop: if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved