From 79bff13a47f05c7535b1be5c092bee1304e6b0fc Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 24 Apr 2024 14:06:07 -0700 Subject: [PATCH] Add `then=` arg for add_conditional_edge and set_conditional_entry_point - This makes it easy to implement graphs where you want to decide among N possible nodes, and then visit another node after whichever one you chose --- langgraph/channels/dynamic_barrier_value.py | 91 +++++ langgraph/graph/graph.py | 101 +++-- langgraph/graph/state.py | 39 +- langgraph/pregel/__init__.py | 8 + tests/__snapshots__/test_pregel.ambr | 392 ++++++++++++++++++++ tests/test_pregel.py | 255 +++++++++++++ tests/test_pregel_async.py | 226 +++++++++++ 7 files changed, 1065 insertions(+), 47 deletions(-) create mode 100644 langgraph/channels/dynamic_barrier_value.py diff --git a/langgraph/channels/dynamic_barrier_value.py b/langgraph/channels/dynamic_barrier_value.py new file mode 100644 index 000000000..bdc7b21ec --- /dev/null +++ b/langgraph/channels/dynamic_barrier_value.py @@ -0,0 +1,91 @@ +from contextlib import contextmanager +from typing import Generator, Generic, NamedTuple, Optional, Sequence, Type, Union + +from typing_extensions import Self + +from langgraph.channels.base import ( + BaseChannel, + EmptyChannelError, + InvalidUpdateError, + Value, +) + + +class WaitForNames(NamedTuple): + names: set[Value] + + +class DynamicBarrierValue( + 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, it can be read once, and it switches + back to the "priming" state. + """ + + names: Optional[set[Value]] + seen: set[Value] + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + self.names = None + self.seen = set() + + @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[Optional[set[Value]], set[Value]]: + return (self.names, self.seen) + + @contextmanager + def from_checkpoint( + self, checkpoint: Optional[tuple[Optional[set[Value]], set[Value]]] = None + ) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + names, seen = checkpoint + empty.names = names.copy() if names is not None else None + empty.seen = seen.copy() + + try: + yield empty + finally: + 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( + "Received multiple WaitForNames updates in the same step." + ) + self.names = wait_for_names[0].names + elif self.names is not None: + for value in values: + assert not isinstance(value, WaitForNames) + if value in self.names: + self.seen.add(value) + else: + raise InvalidUpdateError(f"Value {value} not in {self.names}") + + print(self.seen != self.names, self.seen, self.names, values) + + def get(self) -> Value: + if self.seen != self.names: + raise EmptyChannelError() + return None diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 4c06de11b..8150fa3f8 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -37,8 +37,9 @@ END = "__end__" class Branch(NamedTuple): - condition: Runnable[Any, Union[str, list[str]]] + path: Runnable[Any, Union[str, list[str]]] ends: Optional[dict[str, str]] + then: Optional[str] = None def run( self, @@ -64,7 +65,7 @@ class Branch(NamedTuple): reader: Optional[Callable[[], Any]], writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: - result = self.condition.invoke(reader(config) if reader else input, config) + result = self.path.invoke(reader(config) if reader else input, config) if not isinstance(result, list): result = [result] if self.ends: @@ -81,9 +82,7 @@ class Branch(NamedTuple): reader: Optional[Callable[[], Any]], writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: - result = await self.condition.ainvoke( - reader(config) if reader else input, config - ) + result = await self.path.ainvoke(reader(config) if reader else input, config) if not isinstance(result, list): result = [result] if self.ends: @@ -140,27 +139,27 @@ class Graph: def add_conditional_edges( self, - start_key: str, - condition: Union[ + source: str, + path: Union[ Callable[..., Union[str, list[str]]], Callable[..., Awaitable[Union[str, list[str]]]], Runnable[Any, Union[str, list[str]]], ], - conditional_edge_mapping: Optional[dict[str, str]] = None, + path_map: Optional[dict[str, str]] = None, + then: Optional[str] = None, ) -> None: """Add a conditional edge from the starting node to any number of destination nodes. Args: - start_key (str): The key of the starting node. - condition (Union[Callable, Runnable]): The condition that determines the destination of the edge. - conditional_edge_mapping (Optional[dict[str, str]]): A dictionary that maps the response of the condition to a name of - the destination node(s). If the condition returns a list, the response will be matched against the keys of the - dictionary. If the condition returns a string, the response will be matched against the values of the dictionary. - If the condition returns a string and the dictionary contains a key with the value of `END` ("__end__"`), - the graph will finish. - - Raises: - ValueError: If the starting node is not found in the graph or if the conditional edge mapping contains missing nodes. + source (str): The starting node. This conditional edge will run when + exiting this node. + path (Union[Callable, Runnable]): The callable that determines the next + node or nodes. If not specifying `path_map` it should return one or + more nodes. If it returns END, the graph will stop execution. + path_map (Optional[dict[str, str]]): Optional mapping of paths to node + names. If ommitted the paths returned by `path` should be node names. + then (Optional[str]): The name of a node to execute after the nodes + selected by `path`. Returns: None @@ -171,16 +170,15 @@ class Graph: "not be reflected in the compiled graph." ) # find a name for the condition - condition = coerce_to_runnable(condition) - name = condition.name or "condition" + path = coerce_to_runnable(path) + name = path.name or "condition" # validate the condition - if name in self.branches[start_key]: + if name in self.branches[source]: raise ValueError( - f"Branch with name `{condition.name}` already exists for node " - f"`{start_key}`" + f"Branch with name `{path.name}` already exists for node " f"`{source}`" ) # save it - self.branches[start_key][name] = Branch(condition, conditional_edge_mapping) + self.branches[source][name] = Branch(path, path_map, then) def set_entry_point(self, key: str) -> None: """Specifies the first node to be called in the graph. @@ -195,21 +193,27 @@ class Graph: def set_conditional_entry_point( self, - condition: Union[ + path: Union[ Callable[..., str], Callable[..., Awaitable[str]], Runnable[Any, str] ], - conditional_edge_mapping: Optional[Dict[str, str]] = None, + path_map: Optional[Dict[str, str]] = None, + then: Optional[str] = None, ) -> None: """Sets a conditional entry point in the graph. Args: - condition: A callable object that takes any number of arguments and returns a string or an awaitable string. - conditional_edge_mapping: A dictionary that maps condition names to edge names. + path (Union[Callable, Runnable]): The callable that determines the next + node or nodes. If not specifying `path_map` it should return one or + more nodes. If it returns END, the graph will stop execution. + path_map (Optional[dict[str, str]]): Optional mapping of paths to node + names. If ommitted the paths returned by `path` should be node names. + then (Optional[str]): The name of a node to execute after the nodes + selected by `path`. Returns: None """ - return self.add_conditional_edges(START, condition, conditional_edge_mapping) + return self.add_conditional_edges(START, path, path_map, then) def set_finish_point(self, key: str) -> None: """Marks a node as a finish point of the graph. @@ -226,9 +230,18 @@ class Graph: def validate(self, interrupt: Optional[Sequence[str]] = None) -> None: # assemble sources - all_sources = {src for src, _ in self._all_edges} | { - src for src in self.branches - } + all_sources = {src for src, _ in self._all_edges} + for start, branches in self.branches.items(): + for cond, branch in branches.items(): + all_sources.add(start) + if branch.then is not None: + if branch.ends is not None: + for end in branch.ends.values(): + all_sources.add(end) + else: + for node in self.nodes: + if node != start and node != branch.then: + all_sources.add(node) # validate sources for node in self.nodes: if node not in all_sources: @@ -241,6 +254,8 @@ class Graph: all_targets = {end for _, end in self._all_edges} for start, branches in self.branches.items(): for cond, branch in branches.items(): + if branch.then is not None: + all_targets.add(branch.then) if branch.ends is not None: for end in branch.ends.values(): if end not in self.nodes and end != END: @@ -251,7 +266,7 @@ class Graph: else: all_targets.add(END) for node in self.nodes: - if node != start: + if node != start and node != branch.then: all_targets.add(node) # validate targets for node in self.nodes: @@ -404,24 +419,34 @@ class CompiledGraph(Pregel): name for _, branches in self.graph.branches.items() for name in branches ) for start, branches in self.graph.branches.items(): + default_ends = { + **{k: k for k in self.graph.nodes if k != start}, + END: END, + } for name, branch in branches.items(): - ends = branch.ends or { - **{k: k for k in self.graph.nodes}, - END: END, - } + if branch.ends is not None: + ends = branch.ends + elif branch.then is not None: + ends = {k: k for k in default_ends if k not in (END, branch.then)} + else: + ends = default_ends if add_condition_nodes is True: cond = graph.add_node( - branch.condition, + branch.path, f"{start}_{name}" if branches_by_name[name] > 1 else name, ) graph.add_edge(start_nodes[start], cond) for label, end in ends.items(): graph.add_edge(cond, end_nodes[end], label, conditional=True) + if branch.then is not None: + graph.add_edge(start_nodes[end], end_nodes[branch.then]) else: for label, end in ends.items(): graph.add_edge( start_nodes[start], end_nodes[end], label, conditional=True ) + if branch.then is not None: + graph.add_edge(start_nodes[end], end_nodes[branch.then]) return graph diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index fb9380a1c..fc16e5a5a 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -8,6 +8,7 @@ from langchain_core.runnables.base import RunnableLike from langgraph.channels.base import BaseChannel, InvalidUpdateError from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitForNames from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue @@ -211,7 +212,7 @@ class CompiledStateGraph(CompiledGraph): writers=[ # publish to this channel and state keys ChannelWrite( - [ChannelWriteEntry(key)] + state_write_entries, + [ChannelWriteEntry(key, key)] + state_write_entries, tags=[TAG_HIDDEN], ), ], @@ -247,25 +248,45 @@ class CompiledStateGraph(CompiledGraph): def attach_branch(self, start: str, name: str, branch: Branch) -> None: def branch_writer(ends: list[str]) -> Optional[ChannelWrite]: if filtered_ends := [end for end in ends if end != END]: - return ChannelWrite( - [ - ChannelWriteEntry(f"branch:{start}:{name}:{end}", start) - for end in filtered_ends - ], - tags=[TAG_HIDDEN], - ) + writes = [ + ChannelWriteEntry(f"branch:{start}:{name}:{end}", start) + for end in filtered_ends + ] + if branch.then and branch.then != END: + writes.append( + ChannelWriteEntry( + f"branch:{start}:{name}:then", + WaitForNames(set(filtered_ends)), + ) + ) + return ChannelWrite(writes, tags=[TAG_HIDDEN]) # attach branch publisher self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.graph)) # attach branch subscribers - ends = branch.ends.values() if branch.ends else [node for node in self.nodes] + ends = ( + branch.ends.values() + if branch.ends + else [node for node in self.graph.nodes if node != branch.then] + ) for end in ends: if end != END: channel_name = f"branch:{start}:{name}:{end}" self.channels[channel_name] = EphemeralValue(Any) self.nodes[end].triggers.append(channel_name) + # attach then subscriber + if branch.then and branch.then != END: + channel_name = f"branch:{start}:{name}:then" + self.channels[channel_name] = DynamicBarrierValue(str) + self.nodes[branch.then].triggers.append(channel_name) + for end in ends: + if end != END: + self.nodes[end] |= ChannelWrite( + [ChannelWriteEntry(channel_name, end)], tags=[TAG_HIDDEN] + ) + def _get_state_reader(graph: StateGraph) -> ChannelRead: state_keys = list(graph.channels) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 9535017c0..0d18bfb2e 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -580,6 +580,10 @@ class Pregel( try: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") + if self.checkpointer and not config.get("configurable"): + raise ValueError( + f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in self.checkpointer.config_specs]}" + ) # assign defaults ( debug, @@ -789,6 +793,10 @@ class Pregel( try: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") + if self.checkpointer and not config.get("configurable"): + raise ValueError( + f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in self.checkpointer.config_specs]}" + ) # assign defaults ( debug, diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 7d9937b26..48cf24ecd 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1,4 +1,232 @@ # serializer version: 1 +# name: test_branch_then[end_of_run] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------+ + | prepare | + +---------+ + * + * + * + +-----------+ + | condition | + +-----------+ + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +--------+ + | finish | + +--------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_branch_then[end_of_run].1 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------+ + | prepare | + +---------+. + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +--------+ + | finish | + +--------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_branch_then[end_of_run].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + prepare([prepare]):::otherclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + finish([finish]):::otherclass; + condition([condition]):::otherclass; + __start__ --> prepare; + finish --> __end__; + prepare --> condition; + condition -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> finish; + condition -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> finish; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- +# name: test_branch_then[end_of_run].3 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + prepare([prepare]):::otherclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + finish([finish]):::otherclass; + __start__ --> prepare; + finish --> __end__; + prepare -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> finish; + prepare -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> finish; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- +# name: test_branch_then[end_of_step] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------+ + | prepare | + +---------+ + * + * + * + +-----------+ + | condition | + +-----------+ + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +--------+ + | finish | + +--------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_branch_then[end_of_step].1 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------+ + | prepare | + +---------+. + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +--------+ + | finish | + +--------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_branch_then[end_of_step].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + prepare([prepare]):::otherclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + finish([finish]):::otherclass; + condition([condition]):::otherclass; + __start__ --> prepare; + finish --> __end__; + prepare --> condition; + condition -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> finish; + condition -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> finish; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- +# name: test_branch_then[end_of_step].3 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + prepare([prepare]):::otherclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + finish([finish]):::otherclass; + __start__ --> prepare; + finish --> __end__; + prepare -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> finish; + prepare -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> finish; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- # name: test_conditional_entrypoint_graph '{"title": "LangGraphInput"}' # --- @@ -1933,3 +2161,167 @@ +---------+ ''' # --- +# name: test_start_branch_then[end_of_run] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-----------+ + | condition | + +-----------+ + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_start_branch_then[end_of_run].1 + ''' + +-----------+ + | __start__ | + +-----------+ + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_start_branch_then[end_of_run].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + condition([condition]):::otherclass; + __start__ --> condition; + condition -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> __end__; + condition -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> __end__; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- +# name: test_start_branch_then[end_of_run].3 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + __start__ -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> __end__; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- +# name: test_start_branch_then[end_of_step] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-----------+ + | condition | + +-----------+ + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_start_branch_then[end_of_step].1 + ''' + +-----------+ + | __start__ | + +-----------+ + .. .. + .. .. + .. .. + +---------------+ +---------------+ + | tool_two_slow | | tool_two_fast | + +---------------+ +---------------+ + ** ** + ** ** + ** ** + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_start_branch_then[end_of_step].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + condition([condition]):::otherclass; + __start__ --> condition; + condition -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> __end__; + condition -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> __end__; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- +# name: test_start_branch_then[end_of_step].3 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + tool_two_slow([tool_two_slow]):::otherclass; + tool_two_fast([tool_two_fast]):::otherclass; + __start__ -. tool_two_slow .-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -. tool_two_fast .-> tool_two_fast; + tool_two_fast --> __end__; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 4aeb543d4..f7d42269b 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -3004,6 +3004,261 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] +@pytest.mark.parametrize( + "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] +) +def test_start_branch_then( + snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt +) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + # this graph is invalid because there is no path to END + invalid_graph = StateGraph(State) + invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": "slow"}) + invalid_graph.add_node("tool_two_fast", lambda s: {"my_key": "fast"}) + invalid_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast" + ) + with pytest.raises(ValueError): + invalid_graph.compile() + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + tool_two_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END + ) + tool_two = tool_two_graph.compile() + assert tool_two.get_graph().draw_ascii() == snapshot + assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot + assert tool_two.get_graph().draw_mermaid() == snapshot + assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot + + assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { + "my_key": "value slow", + "market": "DE", + } + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value fast", + "market": "US", + } + + with SqliteSaver.from_conn_string(":memory:") as saver: + saver.at = checkpoint_at + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value slow", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value slow", "market": "DE"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + ) + + +@pytest.mark.parametrize( + "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] +) +def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + # this graph is invalid because there is no path to "finish" + invalid_graph = StateGraph(State) + invalid_graph.set_entry_point("prepare") + invalid_graph.set_finish_point("finish") + invalid_graph.add_conditional_edges( + source="prepare", + path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + ) + invalid_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) + invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + invalid_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + invalid_graph.add_node("finish", lambda s: {"my_key": " finished"}) + with pytest.raises(ValueError): + invalid_graph.compile() + + tool_two_graph = StateGraph(State) + tool_two_graph.set_entry_point("prepare") + tool_two_graph.set_finish_point("finish") + tool_two_graph.add_conditional_edges( + source="prepare", + path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + then="finish", + ) + tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) + tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) + tool_two = tool_two_graph.compile() + assert tool_two.get_graph().draw_ascii() == snapshot + assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot + assert tool_two.get_graph().draw_mermaid() == snapshot + assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot + + assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value prepared fast finished", + "market": "US", + } + + with SqliteSaver.from_conn_string(":memory:") as saver: + saver.at = checkpoint_at + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + 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", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + ) + + with SqliteSaver.from_conn_string(":memory:") as saver: + saver.at = checkpoint_at + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_after=["prepare"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + 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", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + ) + + @pytest.mark.parametrize( "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] ) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 7c3abddee..3140cbe1a 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2679,6 +2679,232 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] +@pytest.mark.parametrize( + "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] +) +async def test_start_branch_then( + snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt +) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + tool_two_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END + ) + tool_two = tool_two_graph.compile() + + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value fast", + "market": "US", + } + + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + saver.at = checkpoint_at + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value slow", "market": "DE"}, + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + ) + + +@pytest.mark.parametrize( + "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] +) +async def test_branch_then( + snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt +) -> None: + pass + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.set_entry_point("prepare") + tool_two_graph.set_finish_point("finish") + tool_two_graph.add_conditional_edges( + source="prepare", + path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + then="finish", + ) + tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) + tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) + tool_two = tool_two_graph.compile() + + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value prepared fast finished", + "market": "US", + } + + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + saver.at = checkpoint_at + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + ) + + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + saver.at = checkpoint_at + tool_two = tool_two_graph.compile( + checkpointer=saver, interrupt_after=["prepare"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + ) + + @pytest.mark.parametrize( "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] )