WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-29 11:19:54 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add then= arg for add_conditional_edge and set_conditional_entry_point

Browse Source 
- 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
This commit is contained in:
Nuno Campos
2024-04-24 14:06:07 -07:00
parent 9ee5699f9a
commit 79bff13a47
7 changed files with 1065 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/channels/dynamic_barrier_value.py
+91
View File
@@ -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
langgraph/graph/graph.py
+63 -38
View File
@@ -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
langgraph/graph/state.py
+30 -9
View File
@@ -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)
langgraph/pregel/__init__.py
+8
View File
@@ -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,
tests/__snapshots__/test_pregel.ambr
+392
View File
@@ -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;
'''
# ---
tests/test_pregel.py
+255
View File
@@ -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]
)
tests/test_pregel_async.py
+226
View File
@@ -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]
)