Rewrite Graph/StateGraph.compile()

- Each iteration of the graph is one single iteration in Pregel (ie. no more "{node}:edges" nodes)
- .update_state() now acts exactly as one of the nodes in the graph (which can be chosen), which makes human-in-the-loop scenarios where you want to override the actions of a specific node a lot easier to build
- Graph/StateGraph.compile() now delegate adding nodes/edges/etc to dedicated methods that operate on the Pregel object, which is 90% of the way towards dynamic graphs where nodes and edges can be added during execution
This commit is contained in:
Nuno Campos
2024-03-31 17:17:18 -07:00
parent 4fd90f8ad6
commit b668bf8eab
13 changed files with 1082 additions and 842 deletions
+100 -104
View File
@@ -19,6 +19,8 @@ from langchain_core.runnables.graph import (
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import ChannelInvoke
from langgraph.pregel.write import ChannelWrite
logger = logging.getLogger(__name__)
@@ -31,11 +33,13 @@ class Branch(NamedTuple):
ends: Optional[dict[str, str]]
def run(self, writer: Callable[[str], Optional[Runnable]]) -> None:
return RunnableLambda(
self._route,
self._aroute,
name=self.condition.name,
).bind(writer=writer)
return ChannelWrite.register_writer(
RunnableLambda(
self._route,
self._aroute,
name=self.condition.name,
).bind(writer=writer)
)
def _route(
self, input: Any, *, writer: Callable[[str], Optional[Runnable]]
@@ -62,11 +66,9 @@ class Graph:
def __init__(self) -> None:
self.nodes: dict[str, Runnable] = {}
self.edges = set[tuple[str, str]]()
self.branches: defaultdict[str, list[Branch]] = defaultdict(list)
self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict)
self.support_multiple_edges = False
self.compiled = False
self.entry_point: Optional[str] = None
self.entry_point_branch: Optional[Branch] = None
@property
def _all_edges(self) -> set[tuple[str, str]]:
@@ -93,7 +95,9 @@ class Graph:
)
if start_key == END:
raise ValueError("END cannot be a start node")
if start_key not in self.nodes:
if end_key == START:
raise ValueError("START cannot be an end node")
if start_key not in self.nodes and start_key != START:
raise ValueError(f"Need to add_node `{start_key}` first")
if end_key not in self.nodes and end_key != END:
raise ValueError(f"Need to add_node `{end_key}` first")
@@ -118,7 +122,7 @@ class Graph:
"Adding an edge to a graph that has already been compiled. This will "
"not be reflected in the compiled graph."
)
if start_key not in self.nodes:
if start_key not in self.nodes and start_key != START:
raise ValueError(f"Need to add_node `{start_key}` first")
if conditional_edge_mapping and set(
conditional_edge_mapping.values()
@@ -131,18 +135,20 @@ class Graph:
)
if not isinstance(condition, Runnable):
condition = RunnableLambda(condition)
if condition.name is None:
condition.name = f"{start_key}_condition"
if condition.name in self.branches[start_key]:
raise ValueError(
f"Branch with name `{condition.name}` already exists for node "
f"`{start_key}`"
)
self.branches[start_key].append(Branch(condition, conditional_edge_mapping))
self.branches[start_key][condition.name] = Branch(
condition, conditional_edge_mapping
)
def set_entry_point(self, key: str) -> None:
if self.compiled:
logger.warning(
"Setting the entry point of a graph that has already been compiled. "
"This will not be reflected in the compiled graph."
)
if key not in self.nodes:
raise ValueError(f"Need to add_node `{key}` first")
self.entry_point = key
return self.add_edge(START, key)
def set_conditional_entry_point(
self,
@@ -151,23 +157,7 @@ class Graph:
],
conditional_edge_mapping: Optional[Dict[str, str]] = None,
) -> None:
if self.compiled:
logger.warning(
"Setting the entry point of a graph that has already been compiled. "
"This will not be reflected in the compiled graph."
)
if conditional_edge_mapping and set(
conditional_edge_mapping.values()
).difference([END]).difference(self.nodes):
raise ValueError(
f"Missing nodes which are in conditional edge mapping. Mapping "
f"contains possible destinations: "
f"{list(conditional_edge_mapping.values())}. Possible nodes are "
f"{list(self.nodes.keys())}."
)
if not isinstance(condition, Runnable):
condition = RunnableLambda(condition)
self.entry_point_branch = Branch(condition, conditional_edge_mapping)
return self.add_conditional_edges(START, condition, conditional_edge_mapping)
def set_finish_point(self, key: str) -> None:
return self.add_edge(key, END)
@@ -180,19 +170,14 @@ class Graph:
if node not in all_starts:
raise ValueError(f"Node `{node}` is a dead-end")
branches = [
branch for branch_list in self.branches.values() for branch in branch_list
all_branches = [
branch
for branches in self.branches.values()
for branch in branches.values()
]
if self.entry_point_branch is not None:
branches.append(self.entry_point_branch)
all_hard_ends = {end for _, end in self._all_edges}
if self.entry_point is not None:
all_hard_ends.add(self.entry_point)
if all(branch.ends is not None for branch in branches):
all_ends = all_hard_ends | {
end for branch in branches for end in branch.ends.values()
if all(branch.ends is not None for branch in all_branches):
all_ends = {end for _, end in self._all_edges} | {
end for branch in all_branches for end in branch.ends.values()
}
for node in self.nodes:
@@ -213,65 +198,89 @@ class Graph:
interrupt_after: Optional[Sequence[str]] = None,
debug: bool = False,
) -> "CompiledGraph":
# assign default values
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
# validate the graph
self.validate(interrupt=interrupt_before + interrupt_after)
outgoing_edges = defaultdict(list)
for start, end in self.edges:
outgoing_edges[start].append(f"{end}:inbox" if end != END else END)
nodes = {
key: (Channel.subscribe_to(f"{key}:inbox") | node | Channel.write_to(key))
for key, node in self.nodes.items()
}
node_outboxes = {
# we clear outbox channels after each step
key: EphemeralValue(Any)
for key in self.nodes
}
def branch_writer(dest: str) -> Optional[Runnable]:
return Channel.write_to(f"{dest}:inbox" if dest != END else END)
for key in self.nodes:
outgoing = outgoing_edges[key]
edges_key = f"{key}:edges"
if outgoing or key in self.branches:
nodes[edges_key] = Channel.subscribe_to(key, tags=["langsmith:hidden"])
if outgoing:
nodes[edges_key] |= Channel.write_to(*[dest for dest in outgoing])
if key in self.branches:
for branch in self.branches[key]:
nodes[edges_key] |= branch.run(branch_writer)
if self.entry_point_branch:
nodes[f"{START}:edges"] = Channel.subscribe_to(
START, tags=["langsmith:hidden"]
) | self.entry_point_branch.run(branch_writer)
elif self.entry_point is None:
raise ValueError("No entry point set")
return CompiledGraph(
# create empty compiled graph
compiled = CompiledGraph(
graph=self,
nodes=nodes,
channels={**node_outboxes},
input_channels=f"{self.entry_point}:inbox" if self.entry_point else START,
nodes={},
channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
input_channels=START,
output_channels=END,
stream_channels=list(self.nodes),
stream_mode="values",
stream_channels=[],
checkpointer=checkpointer,
interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],
interrupt_before_nodes=interrupt_before,
interrupt_after_nodes=interrupt_after,
auto_validate=False,
debug=debug,
)
# attach nodes, edges, and branches
for key, node in self.nodes.items():
compiled.attach_node(key, node)
for start, end in self.edges:
compiled.attach_edge(start, end)
for start, branches in self.branches.items():
for name, branch in branches.items():
compiled.attach_branch(start, name, branch)
# validate the compiled graph
return compiled.validate()
class CompiledGraph(Pregel):
graph: Graph
def attach_node(self, key: str, node: Runnable) -> None:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = (
ChannelInvoke(channels=[], triggers=[]) | node | Channel.write_to(key)
)
self.stream_channels.append(key)
def attach_edge(self, start: str, end: str) -> None:
if end == END:
# publish to end channel
self.nodes[start].writers.append(Channel.write_to(END))
else:
# subscribe to start channel
self.nodes[end].triggers.append(start)
self.nodes[end].channels.append(start)
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(end: str) -> Optional[ChannelWrite]:
return Channel.write_to(
f"branch:{start}:{name}:{end}" if end != END else END
)
# add hidden start node
if start == START and start not in self.nodes:
self.nodes[start] = Channel.subscribe_to(START, tags=["langsmith:hidden"])
# attach branch writer
self.nodes[start] |= branch.run(branch_writer)
# attach branch readers
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
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)
self.nodes[end].channels.append(channel_name)
def get_graph(
self, config: Optional[RunnableConfig] = None, *, xray: bool = False
) -> RunnableGraph:
"""Returns a drawable representation of the computation graph."""
graph = RunnableGraph()
start_nodes: dict[str, RunnableGraphNode] = {
START: graph.add_node(self.get_input_schema(config), START)
@@ -304,10 +313,8 @@ class CompiledGraph(Pregel):
for start, end in self.graph._all_edges:
graph.add_edge(start_nodes[start], end_nodes[end])
for start, branches in self.graph.branches.items():
for i, branch in enumerate(branches):
name = f"{start}_{branch.condition.name or 'condition'}"
if i > 0:
name += f"_{i}"
for name, branch in branches.items():
name = f"{start}_{name}"
cond = graph.add_node(branch.condition, name)
graph.add_edge(start_nodes[start], cond)
ends = branch.ends or {
@@ -316,16 +323,5 @@ class CompiledGraph(Pregel):
}
for label, end in ends.items():
graph.add_edge(cond, end_nodes[end], label)
if entry_point_branch := self.graph.entry_point_branch:
cond = graph.add_node(
entry_point_branch.condition,
entry_point_branch.condition.name or f"{START}_condition",
)
graph.add_edge(start_nodes[START], cond)
ends = entry_point_branch.ends or {k: k for k in self.graph.nodes}
for label, end in ends.items():
graph.add_edge(cond, end_nodes[end], label)
elif self.graph.entry_point:
graph.add_edge(start_nodes[START], end_nodes[self.graph.entry_point])
return graph
+5 -6
View File
@@ -11,13 +11,9 @@ Messages = Union[list[AnyMessage], AnyMessage]
def add_messages(left: Messages, right: Messages) -> Messages:
# coerce to list
if not isinstance(left, list):
left = [message_chunk_to_message(left)]
else:
left = [message_chunk_to_message(m) for m in left]
left = [left]
if not isinstance(right, list):
right = [message_chunk_to_message(right)]
else:
right = [message_chunk_to_message(m) for m in right]
right = [right]
# assign missing ids
for m in left:
if m.id is None:
@@ -25,6 +21,9 @@ def add_messages(left: Messages, right: Messages) -> Messages:
for m in right:
if m.id is None:
m.id = str(uuid.uuid4())
# coerce to message
left = [message_chunk_to_message(m) for m in left]
right = [message_chunk_to_message(m) for m in right]
# merge
left_idx_by_id = {m.id: i for i, m in enumerate(left)}
merged = left.copy()
+126 -119
View File
@@ -1,22 +1,20 @@
import logging
from collections import defaultdict
from functools import partial
from inspect import signature
from typing import Any, Optional, Sequence, Type, Union
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables import Runnable, RunnableLambda
from langchain_core.runnables.base import RunnableLike
from langgraph.channels.any_value import AnyValue
from langgraph.channels.base import BaseChannel, InvalidUpdateError
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.named_barrier_value import NamedBarrierValue
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.graph.graph import END, START, CompiledGraph, Graph
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
from langgraph.pregel import Channel
from langgraph.pregel.read import ChannelInvoke
from langgraph.pregel.read import ChannelInvoke, ChannelRead
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
logger = logging.getLogger(__name__)
@@ -73,134 +71,143 @@ class StateGraph(Graph):
interrupt_after: Optional[Sequence[str]] = None,
debug: bool = False,
) -> CompiledGraph:
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
self.validate(interrupt=interrupt_before + interrupt_after)
# prepare state channels
state_keys = list(self.channels)
state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
state_channels = (
{chan: chan for chan in state_keys}
if isinstance(state_keys_read, list)
else {None: state_keys_read}
)
update_channels = (
[ChannelWriteEntry("__root__", None, True)]
if not isinstance(state_keys_read, list)
else [
ChannelWriteEntry(
key, RunnableLambda(partial(_dict_getter, state_keys, key)), False
)
for key in state_keys_read
]
)
coerce_state = (
partial(_coerce_state, self.schema)
if isinstance(state_keys_read, list)
else None
)
waiting_edges = {
(f"{starts}:{end}", starts, end) for starts, end in self.w_edges
}
waiting_edge_channels = {
key: NamedBarrierValue(str, set(starts)) for key, starts, _ in waiting_edges
}
# assign default values
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
outgoing_edges = defaultdict(list)
for start, end in self.edges:
if end != END:
outgoing_edges[start].append(f"{end}:inbox")
for key, starts, end in waiting_edges:
for start in starts:
outgoing_edges[start].append(key)
# validate the graph
self.validate(interrupt=interrupt_before + interrupt_after)
nodes = {
key: (
ChannelInvoke(
triggers=[
f"{key}:inbox",
*[chan for chan, _, end in waiting_edges if end == key],
],
channels=state_channels,
mapper=coerce_state,
)
| node
| ChannelWrite(
channels=[ChannelWriteEntry(key, None, False)] + update_channels
)
)
for key, node in self.nodes.items()
}
node_inboxes = {
# we take any value written to channel because all writers
# write the entire state as of that step, which is equal for all
f"{key}:inbox": AnyValue(self.schema)
for key in list(self.nodes) + [START]
}
node_outboxes = {
# we clear outbox channels after each step
key: EphemeralValue(Any)
for key in list(self.nodes) + [START]
}
def branch_writer(src: str, dest: str) -> Optional[ChannelWrite]:
if dest != END:
return ChannelWrite(
channels=[ChannelWriteEntry(f"{dest}:inbox", src, False)]
)
for key in self.nodes:
outgoing = outgoing_edges[key]
edges_key = f"{key}:edges"
if outgoing or key in self.branches:
nodes[edges_key] = ChannelInvoke(
triggers=[key], tags=["langsmith:hidden"], channels=state_channels
)
if outgoing:
nodes[edges_key] |= ChannelWrite(
channels=[ChannelWriteEntry(dest, key, True) for dest in outgoing]
)
if key in self.branches:
for branch in self.branches[key]:
nodes[edges_key] |= branch.run(partial(branch_writer, key))
nodes[START] = Channel.subscribe_to(
f"{START}:inbox", tags=["langsmith:hidden"]
) | ChannelWrite(
channels=[ChannelWriteEntry(START, None, False)] + update_channels
)
nodes[f"{START}:edges"] = ChannelInvoke(
triggers=[START], tags=["langsmith:hidden"], channels=state_channels
)
if self.entry_point:
nodes[f"{START}:edges"] |= Channel.write_to(f"{self.entry_point}:inbox")
elif self.entry_point_branch:
nodes[f"{START}:edges"] |= self.entry_point_branch.run(
partial(branch_writer, START)
)
else:
raise ValueError("No entry point set")
return CompiledGraph(
compiled = CompiledStateGraph(
graph=self,
nodes=nodes,
channels={
**self.channels,
**node_inboxes,
**node_outboxes,
**waiting_edge_channels,
},
input_channels=f"{START}:inbox",
nodes={},
channels={**self.channels, START: EphemeralValue(self.schema)},
input_channels=START,
stream_mode="updates",
output_channels=state_keys_read,
stream_channels=state_keys_read,
checkpointer=checkpointer,
interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],
interrupt_before_nodes=interrupt_before,
interrupt_after_nodes=interrupt_after,
auto_validate=False,
debug=debug,
)
compiled.attach_node(START, None)
for key, node in self.nodes.items():
compiled.attach_node(key, node)
for start, end in self.edges:
compiled.attach_edge(start, end)
for starts, end in self.w_edges:
compiled.attach_edge(starts, end)
for start, branches in self.branches.items():
for name, branch in branches.items():
compiled.attach_branch(start, name, branch)
return compiled.validate()
class CompiledStateGraph(CompiledGraph):
graph: StateGraph
def attach_node(self, key: str, node: Optional[Runnable]) -> None:
# prepare state
state_keys = list(self.graph.channels)
state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
state_write_entries = [
ChannelWriteEntry(
key,
RunnableLambda(partial(_dict_getter, state_keys, key)),
False,
)
if key != "__root__"
else ChannelWriteEntry(key, None, True)
for key in state_keys
]
# add node and output channel
if key == START:
self.nodes[key] = Channel.subscribe_to(
START, tags=["langsmith:hidden"]
).pipe(ChannelWrite(channels=state_write_entries))
else:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = ChannelInvoke(
triggers=[],
# read state keys
channels=(
{chan: chan for chan in state_keys}
if isinstance(state_keys_read, list)
else [state_keys_read]
),
# coerce state dict to schema class (eg. pydantic model)
mapper=(
partial(_coerce_state, self.graph.schema)
if isinstance(state_keys_read, list)
else None
),
# publish to this channel and state keys
writers=[
ChannelWrite(
[ChannelWriteEntry(key, None, False)] + state_write_entries
),
# read back state with updates applied
ChannelRead(state_keys_read, fresh=True),
],
).pipe(node)
def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None:
if isinstance(starts, str):
if starts == START:
channel_name = f"start:{end}"
# register channel
self.channels[channel_name] = EphemeralValue(Any)
# subscribe to channel
self.nodes[end].triggers.append(channel_name)
# publish to channel
self.nodes[START] |= ChannelWrite(
[ChannelWriteEntry(channel_name, START, False)]
)
elif end != END:
# subscribe to start channel
self.nodes[end].triggers.append(starts)
else:
channel_name = f"join:{starts}:{end}"
# register channel
self.channels[channel_name] = NamedBarrierValue(str, set(starts))
# subscribe to channel
self.nodes[end].triggers.append(channel_name)
# publish to channel
for start in starts:
self.nodes[start] |= ChannelWrite(
[ChannelWriteEntry(channel_name, start, False)]
)
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(end: str) -> Optional[ChannelWrite]:
if end != END:
return ChannelWrite(
[ChannelWriteEntry(f"branch:{start}:{name}:{end}", start, False)]
)
# attach branch publisher
self.nodes[start] |= branch.run(branch_writer)
# attach branch subscribers
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
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)
def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
return schema(**input)
+271 -139
View File
@@ -14,6 +14,7 @@ from typing import (
Mapping,
NamedTuple,
Optional,
Self,
Sequence,
Type,
Union,
@@ -29,6 +30,7 @@ from langchain_core.globals import get_debug
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain_core.runnables import (
Runnable,
RunnableSequence,
RunnableSerializable,
)
from langchain_core.runnables.base import Input, Output, coerce_to_runnable
@@ -52,6 +54,7 @@ from langgraph.channels.base import (
InvalidUpdateError,
create_checkpoint,
)
from langgraph.channels.last_value import LastValue
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
@@ -59,12 +62,15 @@ from langgraph.checkpoint.base import (
copy_checkpoint,
empty_checkpoint,
)
from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, INTERRUPT
from langgraph.constants import (
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
INTERRUPT,
)
from langgraph.pregel.debug import print_checkpoint, print_step_start
from langgraph.pregel.io import map_input, map_output_updates, map_output_values
from langgraph.pregel.log import logger
from langgraph.pregel.read import ChannelInvoke
from langgraph.pregel.reserved import AllReservedChannels, ReservedChannels
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
@@ -94,7 +100,6 @@ class Channel:
channels: str,
*,
key: Optional[str] = None,
when: Optional[Callable[[Any], bool]] = None,
tags: Optional[list[str]] = None,
) -> ChannelInvoke:
...
@@ -106,7 +111,6 @@ class Channel:
channels: Sequence[str],
*,
key: None = None,
when: Optional[Callable[[Any], bool]] = None,
tags: Optional[list[str]] = None,
) -> ChannelInvoke:
...
@@ -117,7 +121,6 @@ class Channel:
channels: Union[str, Sequence[str]],
*,
key: Optional[str] = None,
when: Optional[Callable[[Any], bool]] = None,
tags: Optional[list[str]] = None,
) -> ChannelInvoke:
"""Runs process.invoke() each time channels are updated,
@@ -130,11 +133,12 @@ class Channel:
channels=cast(
Union[Mapping[None, str], Mapping[str, str]],
{key: channels}
if isinstance(channels, str) and key is not None
else [channels]
if isinstance(channels, str)
else {chan: chan for chan in channels},
),
triggers=[channels] if isinstance(channels, str) else channels,
when=when,
tags=tags,
)
@@ -146,13 +150,11 @@ class Channel:
) -> ChannelWrite:
"""Writes to channels the result of the lambda, or None to skip writing."""
return ChannelWrite(
channels=(
[ChannelWriteEntry(c, None, False) for c in channels]
+ [
ChannelWriteEntry(k, _coerce_write_value(v), True)
for k, v in kwargs.items()
]
)
[ChannelWriteEntry(c, None, False) for c in channels]
+ [
ChannelWriteEntry(k, _coerce_write_value(v), True)
for k, v in kwargs.items()
]
)
@@ -177,6 +179,10 @@ class Pregel(
channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
default_channel_cls: Type[BaseChannel] = Field(default=LastValue)
auto_validate: bool = True
stream_mode: StreamMode = "values"
output_channels: Union[str, Sequence[str]] = "output"
@@ -203,7 +209,9 @@ class Pregel(
arbitrary_types_allowed = True
@root_validator(skip_on_failure=True)
def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]:
if not values["auto_validate"]:
return values
validate_graph(
values["nodes"],
values["channels"],
@@ -212,12 +220,29 @@ class Pregel(
values["stream_channels"],
values["interrupt_after_nodes"],
values["interrupt_before_nodes"],
values["default_channel_cls"],
)
if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]:
if not values["checkpointer"]:
raise ValueError("Interrupts require a checkpointer")
return values
def validate(self) -> Self:
validate_graph(
self.nodes,
self.channels,
self.input_channels,
self.output_channels,
self.stream_channels,
self.interrupt_after_nodes,
self.interrupt_before_nodes,
self.default_channel_cls,
)
if self.interrupt_after_nodes or self.interrupt_before_nodes:
if not self.checkpointer:
raise ValueError("Interrupts require a checkpointer")
return self
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
return [
@@ -274,8 +299,7 @@ class Pregel(
return (
[self.stream_channels]
if isinstance(self.stream_channels, str)
else self.stream_channels
or [k for k in self.channels if k not in AllReservedChannels]
else self.stream_channels or [k for k in self.channels]
)
def get_state(self, config: RunnableConfig) -> StateSnapshot:
@@ -290,10 +314,13 @@ class Pregel(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
k: _read_channel(channels, k, return_exception=True)
for k in channels
if k in self.snapshot_channels_list
}
values = {
k: v for k, v in values.items() if not isinstance(v, EmptyChannelError)
}
return StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
@@ -314,10 +341,13 @@ class Pregel(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
k: _read_channel(channels, k, return_exception=True)
for k in channels
if k in self.snapshot_channels_list
}
values = {
k: v for k, v in values.items() if not isinstance(v, EmptyChannelError)
}
return StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
@@ -336,10 +366,15 @@ class Pregel(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
k: _read_channel(channels, k, return_exception=True)
for k in channels
if k in self.snapshot_channels_list
}
values = {
k: v
for k, v in values.items()
if not isinstance(v, EmptyChannelError)
}
yield StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
@@ -361,10 +396,15 @@ class Pregel(
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
k: _read_channel(channels, k, return_exception=True)
for k in channels
if k in self.snapshot_channels_list
}
values = {
k: v
for k, v in values.items()
if not isinstance(v, EmptyChannelError)
}
yield StateSnapshot(
values[self.stream_channels]
if isinstance(self.stream_channels, str)
@@ -375,49 +415,117 @@ class Pregel(
)
def update_state(
self, config: RunnableConfig, values: dict[str, Any] | Any
self,
config: RunnableConfig,
values: dict[str, Any] | Any,
as_node: Optional[str] = None,
) -> RunnableConfig:
"""Update the state of the graph with the given values, as if they came from
node `as_node`. If `as_node` is not provided, it will be set to the last node
that updated the state, if not ambiguous.
"""
if not self.checkpointer:
raise ValueError("No checkpointer set")
values = (
{self.stream_channels: values}
if isinstance(self.stream_channels, str)
else values
)
# get last checkpoint
checkpoint = self.checkpointer.get(config)
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
# find last node that updated the state, if not provided
if as_node is None:
last_seen_by_node = sorted(
(v, n)
for n, seen in checkpoint["versions_seen"].items()
for v in seen.values()
)
# if two nodes updated the state at the same time, it's ambiguous
if last_seen_by_node:
if last_seen_by_node[-1][0] != last_seen_by_node[-2][0]:
as_node = last_seen_by_node[-1][1]
if as_node is None:
raise InvalidUpdateError("Ambiguous update, specify as_node")
# update channels
with ChannelsManager(self.channels, checkpoint) as channels:
for k, v in values.items():
channels[k].update([v])
checkpoint["channel_versions"][k] += 1
for k in self.snapshot_channels_list:
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
# create task to run all writers of the chosen node
task = PregelExecutableTask(
RunnableSequence(*self.nodes[as_node].writers)
if len(self.nodes[as_node].writers) > 1
else self.nodes[as_node].writers[0],
values,
as_node,
deque(),
)
# execute task
task.proc.invoke(
task.input,
patch_config(
config,
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, task.writes
),
},
),
)
# apply to checkpoint and save
_apply_writes(checkpoint, channels, task.writes)
return self.checkpointer.put(
config, create_checkpoint(checkpoint, channels)
)
async def aupdate_state(
self, config: RunnableConfig, values: dict[str, Any] | Any
self,
config: RunnableConfig,
values: dict[str, Any] | Any,
as_node: Optional[str] = None,
) -> RunnableConfig:
if not self.checkpointer:
raise ValueError("No checkpointer set")
values = (
{self.stream_channels: values}
if isinstance(self.stream_channels, str)
else values
)
# get last checkpoint
checkpoint = await self.checkpointer.aget(config)
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
# find last node that updated the state, if not provided
if as_node is None:
last_seen_by_node = sorted(
(v, n)
for n, seen in checkpoint["versions_seen"].items()
for v in seen.values()
)
# if two nodes updated the state at the same time, it's ambiguous
if last_seen_by_node:
if last_seen_by_node[-1][0] != last_seen_by_node[-2][0]:
as_node = last_seen_by_node[-1][1]
if as_node is None:
raise InvalidUpdateError("Ambiguous update, specify as_node")
# update channels, acting as the chosen node
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
for k, v in values.items():
channels[k].update([v])
checkpoint["channel_versions"][k] += 1
for k in self.stream_channels or self.channels:
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
# create task to run all writers of the chosen node
task = PregelExecutableTask(
RunnableSequence(*self.nodes[as_node].writers)
if len(self.nodes[as_node].writers) > 1
else self.nodes[as_node].writers[0],
values,
as_node,
deque(),
)
# execute task
await task.proc.ainvoke(
task.input,
patch_config(
config,
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, task.writes
),
},
),
)
# apply to checkpoint and save
_apply_writes(checkpoint, channels, task.writes)
return await self.checkpointer.aput(
config, create_checkpoint(checkpoint, channels)
)
@@ -503,13 +611,7 @@ class Pregel(
checkpoint, processes, channels, for_execution=True
)
# apply input writes
_apply_writes(
checkpoint,
channels,
input_writes,
config,
0,
)
_apply_writes(checkpoint, channels, input_writes)
else:
# if received no input, take that as signal to proceed
# past previous interrupt, if any
@@ -518,8 +620,6 @@ class Pregel(
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
read = partial(_read_channel, channels)
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1
@@ -543,6 +643,15 @@ class Pregel(
"limit by setting the `recursion_limit` config key."
)
# before execution, check if we should interrupt
if _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
next_tasks,
):
break
if debug:
print_step_start(step, next_tasks)
@@ -558,7 +667,9 @@ class Pregel(
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: read,
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, writes
),
},
),
)
@@ -587,9 +698,7 @@ class Pregel(
pending_writes.extend(writes)
# apply writes to channels
_apply_writes(
checkpoint, channels, pending_writes, config, step + 1
)
_apply_writes(checkpoint, channels, pending_writes)
if debug:
print_checkpoint(step, channels)
@@ -604,31 +713,21 @@ class Pregel(
if step_output := map_output_updates(output_keys, next_tasks):
yield step_output
# with previous step's checkpoint
if _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
pending_writes,
):
break
# save end of step checkpoint
if self.checkpointer is not None and (
self.checkpointer.at == CheckpointAt.END_OF_STEP
or interrupt_before_nodes
):
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint_config = self.checkpointer.put(
checkpoint_config, checkpoint
)
# with this step's checkpoint,
# after execution, check if we should interrupt
if _should_interrupt(
checkpoint,
interrupt_after_nodes,
self.snapshot_channels_list,
pending_writes,
next_tasks,
):
break
@@ -636,7 +735,6 @@ class Pregel(
if (
self.checkpointer is not None
and self.checkpointer.at == CheckpointAt.END_OF_RUN
and not interrupt_before_nodes
):
checkpoint = create_checkpoint(checkpoint, channels)
self.checkpointer.put(checkpoint_config, checkpoint)
@@ -697,13 +795,7 @@ class Pregel(
checkpoint, processes, channels, for_execution=True
)
# apply input writes
_apply_writes(
checkpoint,
channels,
input_writes,
config,
0,
)
_apply_writes(checkpoint, channels, input_writes)
else:
# if received no input, take that as signal to proceed
# past previous interrupt, if any
@@ -712,8 +804,6 @@ class Pregel(
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
read = partial(_read_channel, channels)
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1,
@@ -726,7 +816,10 @@ class Pregel(
# if no more tasks, we're done
if not next_tasks:
break
if step == 0:
raise ValueError("No tasks to run in graph.")
else:
break
elif step == config["recursion_limit"]:
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached"
@@ -734,6 +827,15 @@ class Pregel(
"by setting the `recursion_limit` config key."
)
# before execution, check if we should interrupt
if _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
next_tasks,
):
break
if debug:
print_step_start(step, next_tasks)
@@ -749,7 +851,9 @@ class Pregel(
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: read,
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, writes
),
},
),
)
@@ -785,9 +889,7 @@ class Pregel(
pending_writes.extend(writes)
# apply writes to channels
_apply_writes(
checkpoint, channels, pending_writes, config, step + 1
)
_apply_writes(checkpoint, channels, pending_writes)
if debug:
print_checkpoint(step, channels)
@@ -802,31 +904,21 @@ class Pregel(
if step_output := map_output_updates(output_keys, next_tasks):
yield step_output
# with previous step's checkpoint
if _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
pending_writes,
):
break
# save end of step checkpoint
if self.checkpointer is not None and (
self.checkpointer.at == CheckpointAt.END_OF_STEP
or interrupt_before_nodes
):
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint_config = await self.checkpointer.aput(
checkpoint_config, checkpoint
)
# with this step's checkpoint
# after execution, check if we should interrupt
if _should_interrupt(
checkpoint,
interrupt_after_nodes,
self.snapshot_channels_list,
pending_writes,
next_tasks,
):
break
@@ -834,7 +926,6 @@ class Pregel(
if (
self.checkpointer is not None
and self.checkpointer.at == CheckpointAt.END_OF_RUN
and not interrupt_before_nodes
):
checkpoint = create_checkpoint(checkpoint, channels)
await self.checkpointer.aput(checkpoint_config, checkpoint)
@@ -1041,7 +1132,7 @@ def _should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Sequence[str],
snapshot_channels: Sequence[str],
pending_writes: Sequence[tuple[str, Any]],
tasks: list[PregelExecutableTask],
) -> bool:
return (
# interrupt if any of snapshopt_channels has been updated since last interrupt
@@ -1051,10 +1142,32 @@ def _should_interrupt(
for chan in snapshot_channels
)
# and any channel written to is in interrupt_nodes list
and any(chan for chan, _ in pending_writes if chan in interrupt_nodes)
and any(node for _, _, node, _ in tasks if node in interrupt_nodes)
)
def _local_read(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
writes: Sequence[tuple[str, Any]],
select: Union[list[str], str],
fresh: bool = False,
) -> Union[dict[str, Any], Any]:
if fresh:
checkpoint = create_checkpoint(checkpoint, channels)
with ChannelsManager(channels, checkpoint) as channels:
_apply_writes(copy_checkpoint(checkpoint), channels, writes)
if isinstance(select, str):
return _read_channel(channels, select)
else:
return {k: _read_channel(channels, k) for k in select}
else:
if isinstance(select, str):
return _read_channel(channels, select)
else:
return {k: _read_channel(channels, k) for k in select}
def _read_channel(
channels: Mapping[str, BaseChannel],
chan: str,
@@ -1077,20 +1190,17 @@ def _apply_writes(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
pending_writes: Sequence[tuple[str, Any]],
config: RunnableConfig,
for_step: int,
) -> None:
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
# Group writes by channel
for chan, val in pending_writes:
if chan in AllReservedChannels:
raise ValueError(f"Can't write to reserved channel {chan}")
pending_writes_by_channel[chan].append(val)
# Update reserved channels
pending_writes_by_channel[ReservedChannels.is_last_step] = [
for_step + 1 == config["recursion_limit"]
]
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = 0
updated_channels: set[str] = set()
# Apply writes to channels
@@ -1102,7 +1212,7 @@ def _apply_writes(
raise InvalidUpdateError(
f"Invalid update for channel {chan}: {e}"
) from e
checkpoint["channel_versions"][chan] += 1
checkpoint["channel_versions"][chan] = max_version + 1
updated_channels.add(chan)
else:
logger.warning(f"Skipping write for channel {chan} which has no readers")
@@ -1112,13 +1222,26 @@ def _apply_writes(
channels[chan].update([])
class PregelTask(NamedTuple):
proc: Runnable
input: Any
name: str
class PregelExecutableTask(NamedTuple):
proc: Runnable
input: Any
name: str
writes: deque[tuple[str, Any]]
@overload
def _prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, ChannelInvoke],
channels: Mapping[str, BaseChannel],
for_execution: Literal[False],
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]:
) -> tuple[Checkpoint, list[PregelTask]]:
...
@@ -1128,7 +1251,7 @@ def _prepare_next_tasks(
processes: Mapping[str, ChannelInvoke],
channels: Mapping[str, BaseChannel],
for_execution: Literal[True],
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]]]:
) -> tuple[Checkpoint, list[PregelExecutableTask]]:
...
@@ -1138,18 +1261,9 @@ def _prepare_next_tasks(
channels: Mapping[str, BaseChannel],
*,
for_execution: bool,
) -> tuple[
Checkpoint,
Union[
list[tuple[Runnable, Any, str]],
list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
],
]:
) -> tuple[Checkpoint, Union[list[PregelTask], list[PregelExecutableTask]]]:
checkpoint = copy_checkpoint(checkpoint)
tasks: Union[
list[tuple[Runnable, Any, str]],
list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
] = []
tasks: Union[list[PregelTask], list[PregelExecutableTask]] = []
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
for name, proc in processes.items():
@@ -1164,23 +1278,34 @@ def _prepare_next_tasks(
):
# If all trigger channels subscribed by this process are not empty
# then invoke the process with the values of all non-empty channels
try:
val: Any = {
k: _read_channel(channels, chan, catch=chan not in proc.triggers)
for k, chan in proc.channels.items()
}
except EmptyChannelError:
continue
if isinstance(proc.channels, dict):
try:
val: Any = {
k: _read_channel(
channels, chan, catch=chan not in proc.triggers
)
for k, chan in proc.channels.items()
}
except EmptyChannelError:
continue
elif isinstance(proc.channels, list):
for chan in proc.channels:
try:
val = _read_channel(channels, chan, catch=False)
break
except EmptyChannelError:
pass
else:
continue
else:
raise RuntimeError(
"Invalid channels type, expected list or dict, got {proc.channels}"
)
# If the process has a mapper, apply it to the value
if proc.mapper is not None:
val = proc.mapper(val)
# Processes that subscribe to a single keyless channel get
# the value directly, instead of a dict
if list(proc.channels.keys()) == [None]:
val = val[None]
# update seen versions
if for_execution:
seen.update(
@@ -1190,12 +1315,19 @@ def _prepare_next_tasks(
}
)
# skip if condition is not met
if proc.when is None or proc.when(val):
if for_execution:
tasks.append((proc, val, name, deque()))
else:
tasks.append((proc, val, name))
if for_execution:
tasks.append(
PregelExecutableTask(
RunnableSequence(proc, *proc.writers, name=name)
if proc.writers
else proc,
val,
name,
deque(),
)
)
else:
tasks.append(PregelTask(proc, val, name))
return checkpoint, tasks
+1 -1
View File
@@ -25,7 +25,7 @@ def print_checkpoint(step: int, channels: Mapping[str, BaseChannel]) -> None:
print(
f"{get_colored_text('[langgraph/checkpoint]', color='blue')} "
+ get_bolded_text(f"Finishing step {step}. Channel values:\n")
+ pformat({name: val for name, val in _read_channels(channels)}, depth=2)
+ pformat({name: val for name, val in _read_channels(channels)}, depth=3)
)
+45 -25
View File
@@ -15,11 +15,16 @@ from langchain_core.runnables.config import merge_configs
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_READ
from langgraph.pregel.write import ChannelWrite
READ_TYPE = Callable[[str, bool], Union[Any, dict[str, Any]]]
class ChannelRead(RunnableLambda):
channel: Union[str, list[str]]
fresh: bool = False
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
return [
@@ -32,51 +37,44 @@ class ChannelRead(RunnableLambda):
),
]
def __init__(self, channel: Union[str, list[str]]) -> None:
def __init__(self, channel: Union[str, list[str]], fresh: bool = False) -> None:
super().__init__(func=self._read, afunc=self._aread)
self.fresh = fresh
self.channel = channel
self.name = f"ChannelRead<{channel}>"
def _read(self, _: Any, config: RunnableConfig) -> Any:
try:
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
read: READ_TYPE = config["configurable"][CONFIG_KEY_READ]
except KeyError:
raise RuntimeError(
f"Runnable {self} is not configured with a read function"
"Make sure to call in the context of a Pregel process"
)
return (
read(self.channel)
if isinstance(self.channel, str)
else {chan: read(chan) for chan in self.channel}
)
return read(self.channel, self.fresh)
async def _aread(self, _: Any, config: RunnableConfig) -> Any:
try:
read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
read: READ_TYPE = config["configurable"][CONFIG_KEY_READ]
except KeyError:
raise RuntimeError(
f"Runnable {self} is not configured with a read function"
"Make sure to call in the context of a Pregel process"
)
return (
read(self.channel)
if isinstance(self.channel, str)
else {chan: read(chan) for chan in self.channel}
)
return read(self.channel, self.fresh)
default_bound: RunnablePassthrough = RunnablePassthrough()
class ChannelInvoke(RunnableBindingBase):
channels: Union[Mapping[None, str], Mapping[str, str]]
channels: Union[list[str], Mapping[str, str]]
triggers: list[str] = Field(default_factory=list)
mapper: Optional[Callable[[Any], Any]] = None
when: Optional[Callable[[Any], bool]] = None
writers: list[Runnable] = Field(default_factory=list)
bound: Runnable[Any, Any] = Field(default=default_bound)
@@ -84,12 +82,12 @@ class ChannelInvoke(RunnableBindingBase):
def __init__(
self,
channels: Mapping[None, str] | Mapping[str, str],
*,
channels: Union[list[str], Mapping[str, str]],
triggers: Sequence[str],
mapper: Optional[Callable[[Any], Any]] = None,
when: Optional[Callable[[Any], bool]] = None,
writers: Optional[list[Runnable]] = None,
tags: Optional[list[str]] = None,
*,
bound: Optional[Runnable[Any, Any]] = None,
kwargs: Optional[Mapping[str, Any]] = None,
config: Optional[RunnableConfig] = None,
@@ -99,19 +97,22 @@ class ChannelInvoke(RunnableBindingBase):
channels=channels,
triggers=triggers,
mapper=mapper,
when=when,
writers=writers or [],
bound=bound or default_bound,
kwargs=kwargs or {},
config=merge_configs(config, {"tags": tags or []}),
**other_kwargs,
)
def __repr_args__(self) -> Any:
return [(k, v) for k, v in super().__repr_args__() if k != "bound"]
def join(self, channels: Sequence[str]) -> ChannelInvoke:
assert isinstance(channels, list) or isinstance(
channels, tuple
), "channels must be a list or tuple"
assert all(
k is not None for k in self.channels.keys()
assert isinstance(
self.channels, dict
), "all channels must be named when using .join()"
return ChannelInvoke(
channels={
@@ -120,7 +121,7 @@ class ChannelInvoke(RunnableBindingBase):
},
triggers=self.triggers,
mapper=self.mapper,
when=self.when,
writers=self.writers,
bound=self.bound,
kwargs=self.kwargs,
config=self.config,
@@ -134,12 +135,22 @@ class ChannelInvoke(RunnableBindingBase):
Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
],
) -> ChannelInvoke:
if self.bound is default_bound:
if ChannelWrite.is_writer(other):
return ChannelInvoke(
channels=self.channels,
triggers=self.triggers,
mapper=self.mapper,
when=self.when,
writers=[*self.writers, other],
bound=self.bound,
kwargs=self.kwargs,
config=self.config,
)
elif self.bound is default_bound:
return ChannelInvoke(
channels=self.channels,
triggers=self.triggers,
mapper=self.mapper,
writers=self.writers,
bound=coerce_to_runnable(other),
kwargs=self.kwargs,
config=self.config,
@@ -149,13 +160,22 @@ class ChannelInvoke(RunnableBindingBase):
channels=self.channels,
triggers=self.triggers,
mapper=self.mapper,
when=self.when,
writers=self.writers,
# delegate to __or__ in self.bound
bound=self.bound | other,
kwargs=self.kwargs,
config=self.config,
)
def pipe(
self,
*others: Runnable[Any, Other] | Callable[[Any], Other],
name: Optional[str] = None,
) -> RunnableSerializable[Any, Other]:
for other in others:
self = self | other
return self
def __ror__(
self,
other: Union[
-11
View File
@@ -1,11 +0,0 @@
from langgraph.utils import StrEnum
class ReservedChannels(StrEnum):
"""Channels managed by the framework."""
is_last_step = "is_last_step"
"""A channel that is True if the current step is the last step, False otherwise."""
AllReservedChannels = {channel.value for channel in ReservedChannels}
+21 -17
View File
@@ -1,10 +1,8 @@
from typing import Any, Mapping, Optional, Sequence, Union
from typing import Any, Mapping, Optional, Sequence, Type, Union
from langgraph.channels.base import BaseChannel
from langgraph.channels.last_value import LastValue
from langgraph.constants import INTERRUPT
from langgraph.pregel.read import ChannelInvoke
from langgraph.pregel.reserved import ReservedChannels
def validate_graph(
@@ -15,13 +13,14 @@ def validate_graph(
stream_channels: Optional[Union[str, Sequence[str]]],
interrupt_after_nodes: Sequence[str],
interrupt_before_nodes: Sequence[str],
default_channel_cls: Type[BaseChannel],
) -> None:
subscribed_channels = set[str]()
for name, node in nodes.items():
if name == INTERRUPT:
raise ValueError(f"Node name {INTERRUPT} is reserved")
if isinstance(node, ChannelInvoke):
subscribed_channels.update(node.channels.values())
subscribed_channels.update(node.triggers)
else:
raise TypeError(
f"Invalid node type {type(node)}, expected Channel.subscribe_to()"
@@ -29,11 +28,11 @@ def validate_graph(
for chan in subscribed_channels:
if chan not in channels:
channels[chan] = LastValue(Any) # type: ignore[arg-type]
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
if isinstance(input_channels, str):
if input_channels not in channels:
channels[input_channels] = LastValue(Any) # type: ignore[arg-type]
channels[input_channels] = default_channel_cls(Any) # type: ignore[arg-type]
if input_channels not in subscribed_channels:
raise ValueError(
f"Input channel {input_channels} is not subscribed to by any node"
@@ -41,27 +40,32 @@ def validate_graph(
else:
for chan in input_channels:
if chan not in channels:
channels[chan] = LastValue(Any) # type: ignore[arg-type]
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
if all(chan not in subscribed_channels for chan in input_channels):
raise ValueError(
f"None of the input channels {input_channels} are subscribed to by any node"
)
all_output_channels = set[str]()
if isinstance(output_channels, str):
if output_channels not in channels:
channels[output_channels] = LastValue(Any) # type: ignore[arg-type]
all_output_channels.add(output_channels)
else:
for chan in output_channels:
if chan not in channels:
channels[chan] = LastValue(Any) # type: ignore[arg-type]
all_output_channels.update(output_channels)
if isinstance(stream_channels, str):
all_output_channels.add(stream_channels)
elif stream_channels is not None:
all_output_channels.update(stream_channels)
for chan in ReservedChannels:
for chan in all_output_channels:
if chan not in channels:
channels[chan] = LastValue(Any) # type: ignore[arg-type]
channels[chan] = default_channel_cls(Any) # type: ignore[arg-type]
validate_keys(stream_channels, channels)
validate_keys(interrupt_after_nodes, channels)
validate_keys(interrupt_before_nodes, channels)
for node in interrupt_after_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
for node in interrupt_before_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
def validate_keys(
+19 -2
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import asyncio
from typing import Any, Callable, NamedTuple, Optional, Sequence, Union
from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar, Union
from langchain_core.runnables import (
Runnable,
@@ -13,6 +13,7 @@ from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_SEND
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
R = TypeVar("R", bound=Runnable)
SKIP_WRITE = object()
@@ -36,13 +37,17 @@ class ChannelWrite(RunnablePassthrough):
class Config:
arbitrary_types_allowed = True
def __init__(self, *, channels: Sequence[ChannelWriteEntry]):
def __init__(self, channels: Sequence[ChannelWriteEntry]):
super().__init__(func=self._write, afunc=self._awrite, channels=channels)
self.name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.channels)}>"
def __repr_args__(self) -> Any:
return [("channels", self.channels)]
@property
def is_channel_writer(self) -> bool:
return True
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
return [
@@ -99,6 +104,18 @@ class ChannelWrite(RunnablePassthrough):
write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE])
@staticmethod
def is_writer(runnable: Runnable) -> bool:
return (
isinstance(runnable, ChannelWrite)
or getattr(runnable, "_is_channel_writer", False) is True
)
@staticmethod
def register_writer(runnable: R) -> R:
object.__setattr__(runnable, "_is_channel_writer", True)
return runnable
def _mk_future(val: Any) -> asyncio.Future:
fut = asyncio.Future()
+162 -162
View File
@@ -50,20 +50,7 @@
}
},
{
"id": "left_condition",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "RunnableLambda"
}
},
{
"id": "should_start",
"id": "__start___should_start",
"type": "runnable",
"data": {
"id": [
@@ -74,6 +61,19 @@
],
"name": "should_start"
}
},
{
"id": "left_left_condition",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "left_condition"
}
}
],
"edges": [
@@ -81,38 +81,38 @@
"source": "right",
"target": "__end__"
},
{
"source": "left",
"target": "left_condition"
},
{
"source": "left_condition",
"target": "left",
"data": "left"
},
{
"source": "left_condition",
"target": "right",
"data": "right"
},
{
"source": "left_condition",
"target": "__end__",
"data": "__end__"
},
{
"source": "__start__",
"target": "should_start"
"target": "__start___should_start"
},
{
"source": "should_start",
"source": "__start___should_start",
"target": "left",
"data": "go-left"
},
{
"source": "should_start",
"source": "__start___should_start",
"target": "right",
"data": "go-right"
},
{
"source": "left",
"target": "left_left_condition"
},
{
"source": "left_left_condition",
"target": "left",
"data": "left"
},
{
"source": "left_left_condition",
"target": "right",
"data": "right"
},
{
"source": "left_left_condition",
"target": "__end__",
"data": "__end__"
}
]
}
@@ -120,39 +120,39 @@
# ---
# name: test_conditional_entrypoint_graph.3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+--------------+
| should_start |
+--------------+
*** **
* **
** **
+------+ *
| left | *
+------+ *
* *
* *
* *
+----------------+ *
| left_condition | *
+----------------+* *
* ***** *
* *** *
* *** *
** +-------+
* | right |
*** +-------+
* ***
*** *
* **
+---------+
| __end__ |
+---------+
+-----------+
| __start__ |
+-----------+
*
*
*
+------------------------+
| __start___should_start |
+------------------------+
*** **
* **
** **
+------+ **
| left | *
+------+ *
* *
* *
* *
+---------------------+ *
| left_left_condition | *
+---------------------+ *
* ***** *
* **** *
* *** *
** +-------+
* | right |
*** +-------+
* ***
*** *
* **
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_conditional_entrypoint_graph_state
@@ -234,20 +234,7 @@
}
},
{
"id": "left_condition",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "RunnableLambda"
}
},
{
"id": "should_start",
"id": "__start___should_start",
"type": "runnable",
"data": {
"id": [
@@ -258,6 +245,19 @@
],
"name": "should_start"
}
},
{
"id": "left_left_condition",
"type": "runnable",
"data": {
"id": [
"langchain_core",
"runnables",
"base",
"RunnableLambda"
],
"name": "left_condition"
}
}
],
"edges": [
@@ -265,38 +265,38 @@
"source": "right",
"target": "__end__"
},
{
"source": "left",
"target": "left_condition"
},
{
"source": "left_condition",
"target": "left",
"data": "left"
},
{
"source": "left_condition",
"target": "right",
"data": "right"
},
{
"source": "left_condition",
"target": "__end__",
"data": "__end__"
},
{
"source": "__start__",
"target": "should_start"
"target": "__start___should_start"
},
{
"source": "should_start",
"source": "__start___should_start",
"target": "left",
"data": "go-left"
},
{
"source": "should_start",
"source": "__start___should_start",
"target": "right",
"data": "go-right"
},
{
"source": "left",
"target": "left_left_condition"
},
{
"source": "left_left_condition",
"target": "left",
"data": "left"
},
{
"source": "left_left_condition",
"target": "right",
"data": "right"
},
{
"source": "left_left_condition",
"target": "__end__",
"data": "__end__"
}
]
}
@@ -304,39 +304,39 @@
# ---
# name: test_conditional_entrypoint_graph_state.3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+--------------+
| should_start |
+--------------+
*** **
* **
** **
+------+ *
| left | *
+------+ *
* *
* *
* *
+----------------+ *
| left_condition | *
+----------------+* *
* ***** *
* *** *
* *** *
** +-------+
* | right |
*** +-------+
* ***
*** *
* **
+---------+
| __end__ |
+---------+
+-----------+
| __start__ |
+-----------+
*
*
*
+------------------------+
| __start___should_start |
+------------------------+
*** **
* **
** **
+------+ **
| left | *
+------+ *
* *
* *
* *
+---------------------+ *
| left_left_condition | *
+---------------------+ *
* ***** *
* **** *
* *** *
** +-------+
* | right |
*** +-------+
* ***
*** *
* **
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_conditional_graph
@@ -398,6 +398,10 @@
}
],
"edges": [
{
"source": "__start__",
"target": "agent"
},
{
"source": "tools",
"target": "agent"
@@ -415,10 +419,6 @@
"source": "agent_should_continue",
"target": "__end__",
"data": "exit"
},
{
"source": "__start__",
"target": "agent"
}
]
}
@@ -669,6 +669,10 @@
"source": 7,
"target": 3
},
{
"source": "__start__",
"target": 2
},
{
"source": "tools",
"target": 2
@@ -686,10 +690,6 @@
"source": "agent_should_continue",
"target": "__end__",
"data": "exit"
},
{
"source": "__start__",
"target": 2
}
]
}
@@ -1019,6 +1019,10 @@
}
],
"edges": [
{
"source": "__start__",
"target": "agent"
},
{
"source": "tools",
"target": "agent"
@@ -1036,10 +1040,6 @@
"source": "agent_should_continue",
"target": "__end__",
"data": "exit"
},
{
"source": "__start__",
"target": "agent"
}
]
}
@@ -1842,6 +1842,10 @@
"source": "action",
"target": "agent"
},
{
"source": "__start__",
"target": "agent"
},
{
"source": "agent",
"target": "agent_should_continue"
@@ -1855,10 +1859,6 @@
"source": "agent_should_continue",
"target": "__end__",
"data": "end"
},
{
"source": "__start__",
"target": "agent"
}
]
}
@@ -2085,6 +2085,10 @@
"source": "action",
"target": "agent"
},
{
"source": "__start__",
"target": "agent"
},
{
"source": "agent",
"target": "agent_should_continue"
@@ -2098,10 +2102,6 @@
"source": "agent_should_continue",
"target": "__end__",
"data": "end"
},
{
"source": "__start__",
"target": "agent"
}
]
}
@@ -2328,6 +2328,10 @@
"source": "action",
"target": "agent"
},
{
"source": "__start__",
"target": "agent"
},
{
"source": "agent",
"target": "agent_should_continue"
@@ -2341,10 +2345,6 @@
"source": "agent_should_continue",
"target": "__end__",
"data": "end"
},
{
"source": "__start__",
"target": "agent"
}
]
}
+6
View File
@@ -0,0 +1,6 @@
class AnyStr(str):
def __init__(self) -> None:
super().__init__()
def __eq__(self, other: object) -> bool:
return isinstance(other, str)
+189 -147
View File
@@ -26,7 +26,7 @@ from langgraph.prebuilt.chat_agent_executor import (
)
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from langgraph.pregel.reserved import ReservedChannels
from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
@@ -64,7 +64,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
assert app.invoke(2, output_keys=["output"]) == {"output": 3}
assert repr(app), "does not raise recursion error"
assert gapp.invoke(2) == 3
assert gapp.invoke(2, debug=True) == 3
def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None:
@@ -103,23 +103,6 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
chain = (
Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
| add_one
| Channel.write_to("output")
)
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert app.invoke(2) == {"input": 3, "is_last_step": False}
assert app.invoke(2, {"recursion_limit": 1}) == {"input": 3, "is_last_step": True}
def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
@@ -214,7 +197,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
app = Pregel(
nodes={"one": one, "two": two},
checkpointer=memory,
interrupt_after_nodes=["inbox"],
interrupt_after_nodes=["one"],
)
# start execution, stop at inbox
@@ -249,7 +232,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
assert snapshot.next == ("two",)
# update the state, resume
app.update_state({"configurable": {"thread_id": 2}}, {"inbox": 25})
app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26
# no pending tasks
@@ -507,9 +490,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
state = app.get_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 7
assert state.next == ()
state = app.get_state(thread_2)
assert state is not None
assert state.values.get("total") == 5
assert state.next == ()
# list all checkpoints for thread 1
thread_1_history = [c for c in app.get_state_history(thread_1)]
@@ -534,9 +519,7 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
== thread_1_history[1].config["configurable"]["thread_ts"]
)
thread_1_next_config = app.update_state(
thread_1_history[1].config, {"total": 10}
)
thread_1_next_config = app.update_state(thread_1_history[1].config, 10)
# update creates a new checkpoint
assert (
thread_1_next_config["configurable"]["thread_ts"]
@@ -927,9 +910,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert (
@@ -942,14 +924,12 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
app_w_interrupt.update_state(
config,
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
)
@@ -963,9 +943,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
),
"input": "what is weather in sf",
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1010,6 +989,26 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
app_w_interrupt.update_state(
config,
{
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
@@ -1026,12 +1025,12 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
}
},
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == []
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
@@ -1061,23 +1060,20 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
app_w_interrupt.update_state(
config,
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
)
@@ -1091,9 +1087,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
),
"input": "what is weather in sf",
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1138,6 +1133,26 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
app_w_interrupt.update_state(
config,
{
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
@@ -1154,12 +1169,12 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
}
},
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == []
# test re-invoke to continue with interrupt_before
app_w_interrupt = workflow.compile(
@@ -1189,9 +1204,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1286,6 +1300,50 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
]
def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
def left(data: str) -> str:
return data + "->left"
def right(data: str) -> str:
return data + "->right"
def should_start(data: str) -> str:
# Logic to decide where to start
if len(data) > 10:
return "go-right"
else:
return "go-left"
# Define a new graph
workflow = Graph()
workflow.add_node("left", left)
workflow.add_node("right", right)
workflow.set_conditional_entry_point(
should_start, {"go-left": "left", "go-right": "right"}
)
workflow.add_conditional_edges("left", lambda data: END)
workflow.add_edge("right", END)
app = workflow.compile()
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert (
app.invoke("what is weather in sf", debug=True)
== "what is weather in sf->right"
)
assert [*app.stream("what is weather in sf")] == [
{"right": "what is weather in sf->right"},
]
def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
@@ -1479,7 +1537,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1504,7 +1562,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1544,7 +1602,27 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
},
)
assert [c for c in app_w_interrupt.stream(None, config)] == []
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
# test state get/update methods with interrupt_before
@@ -1577,7 +1655,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1602,7 +1680,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -1642,49 +1720,28 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
},
)
assert [c for c in app_w_interrupt.stream(None, config)] == []
def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
def left(data: str) -> str:
return data + "->left"
def right(data: str) -> str:
return data + "->right"
def should_start(data: str) -> str:
# Logic to decide where to start
if len(data) > 10:
return "go-right"
else:
return "go-left"
# Define a new graph
workflow = Graph()
workflow.add_node("left", left)
workflow.add_node("right", right)
workflow.set_conditional_entry_point(
should_start, {"go-left": "left", "go-right": "right"}
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
workflow.add_conditional_edges("left", lambda data: END)
workflow.add_edge("right", END)
app = workflow.compile()
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.invoke("what is weather in sf") == "what is weather in sf->right"
assert [*app.stream("what is weather in sf")] == [
{"right": "what is weather in sf->right"},
]
def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict, total=False):
@@ -2200,7 +2257,7 @@ def test_message_graph(
assert app.invoke(HumanMessage(content="what is weather in sf")) == [
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000002", # adds missing ids
id="00000000-0000-4000-8000-000000000005", # adds missing ids
),
AIMessage(
content="",
@@ -2212,7 +2269,7 @@ def test_message_graph(
FunctionMessage(
content="result for query",
name="search_api",
id="00000000-0000-4000-8000-000000000015",
id="00000000-0000-4000-8000-000000000019",
),
AIMessage(
content="",
@@ -2224,7 +2281,7 @@ def test_message_graph(
FunctionMessage(
content="result for another",
name="search_api",
id="00000000-0000-4000-8000-000000000028",
id="00000000-0000-4000-8000-000000000033",
),
AIMessage(content="answer", id="ai3"),
]
@@ -2234,7 +2291,7 @@ def test_message_graph(
"__start__": [
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000037",
id="00000000-0000-4000-8000-000000000045",
)
]
},
@@ -2251,7 +2308,7 @@ def test_message_graph(
"action": FunctionMessage(
content="result for query",
name="search_api",
id="00000000-0000-4000-8000-000000000050",
id="00000000-0000-4000-8000-000000000059",
)
},
{
@@ -2267,7 +2324,7 @@ def test_message_graph(
"action": FunctionMessage(
content="result for another",
name="search_api",
id="00000000-0000-4000-8000-000000000063",
id="00000000-0000-4000-8000-000000000073",
)
},
{"agent": AIMessage(content="answer", id="ai3")},
@@ -2284,12 +2341,7 @@ def test_message_graph(
HumanMessage(content="what is weather in sf"), config
)
] == [
{
"__start__": HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
)
},
{"__start__": HumanMessage(content="what is weather in sf", id=AnyStr())},
{
"agent": AIMessage(
content="",
@@ -2303,10 +2355,7 @@ def test_message_graph(
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
),
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
content="",
additional_kwargs={
@@ -2315,7 +2364,7 @@ def test_message_graph(
id="ai1",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2327,10 +2376,7 @@ def test_message_graph(
# message was replaced instead of appended
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
),
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
content="",
additional_kwargs={
@@ -2342,7 +2388,7 @@ def test_message_graph(
id="ai1",
),
],
next=("agent:edges",),
next=("action",),
config=next_config,
)
@@ -2351,7 +2397,7 @@ def test_message_graph(
"action": FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000086",
id=AnyStr(),
)
},
{
@@ -2369,7 +2415,7 @@ def test_message_graph(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2384,7 +2430,7 @@ def test_message_graph(
FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000086",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2394,7 +2440,7 @@ def test_message_graph(
id="ai2",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2408,7 +2454,7 @@ def test_message_graph(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2423,16 +2469,14 @@ def test_message_graph(
FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000086",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
],
next=("agent:edges",),
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == []
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(), interrupt_before=["action"]
)
@@ -2448,7 +2492,7 @@ def test_message_graph(
{
"__start__": HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000096",
id=AnyStr(),
)
},
{
@@ -2466,7 +2510,7 @@ def test_message_graph(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000096",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2476,7 +2520,7 @@ def test_message_graph(
id="ai1",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2490,7 +2534,7 @@ def test_message_graph(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000096",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2503,7 +2547,7 @@ def test_message_graph(
id="ai1",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2512,7 +2556,7 @@ def test_message_graph(
"action": FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000113",
id=AnyStr(),
)
},
{
@@ -2530,7 +2574,7 @@ def test_message_graph(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000096",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2545,7 +2589,7 @@ def test_message_graph(
FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000113",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2555,7 +2599,7 @@ def test_message_graph(
id="ai2",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2569,7 +2613,7 @@ def test_message_graph(
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000096",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2584,16 +2628,14 @@ def test_message_graph(
FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000113",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
],
next=("agent:edges",),
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == []
def test_in_one_fan_out_out_one_graph_state() -> None:
def sorted_add(x: list[str], y: list[str]) -> list[str]:
+137 -109
View File
@@ -31,7 +31,7 @@ from langgraph.prebuilt.chat_agent_executor import (
)
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from langgraph.pregel.reserved import ReservedChannels
from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
@@ -102,28 +102,6 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
async def test_invoke_single_process_in_out_reserved_is_last(
mocker: MockerFixture,
) -> None:
add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
chain = (
Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
| add_one
| Channel.write_to("output")
)
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert await app.ainvoke(2) == {"input": 3, "is_last_step": False}
assert await app.ainvoke(2, {"recursion_limit": 1}) == {
"input": 3,
"is_last_step": True,
}
async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
@@ -227,7 +205,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
app = Pregel(
nodes={"one": one, "two": two},
checkpointer=memory,
interrupt_after_nodes=["inbox"],
interrupt_after_nodes=["one"],
)
# start execution, stop at inbox
@@ -262,7 +240,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
assert snapshot.next == ("two",)
# update the state, resume
await app.aupdate_state({"configurable": {"thread_id": 2}}, {"inbox": 25})
await app.aupdate_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
assert await app.ainvoke(None, {"configurable": {"thread_id": 2}}) == 26
# no pending tasks
@@ -536,9 +514,11 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
state = await app.aget_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 7
assert state.next == ()
state = await app.aget_state(thread_2)
assert state is not None
assert state.values.get("total") == 5
assert state.next == ()
# list all checkpoints for thread 1
thread_1_history = [c async for c in app.aget_state_history(thread_1)]
@@ -561,9 +541,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
"ts"
] == thread_1_history[1].config["configurable"]["thread_ts"]
thread_1_next_config = await app.aupdate_state(
thread_1_history[1].config, {"total": 10}
)
thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10)
# update creates a new checkpoint
assert (
thread_1_next_config["configurable"]["thread_ts"]
@@ -981,23 +959,20 @@ async def test_conditional_graph() -> None:
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
await app_w_interrupt.aupdate_state(
config,
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
)
@@ -1011,9 +986,8 @@ async def test_conditional_graph() -> None:
),
"input": "what is weather in sf",
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1058,6 +1032,26 @@ async def test_conditional_graph() -> None:
await app_w_interrupt.aupdate_state(
config,
{
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
@@ -1074,12 +1068,12 @@ async def test_conditional_graph() -> None:
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
}
},
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == []
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
@@ -1112,23 +1106,20 @@ async def test_conditional_graph() -> None:
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
await app_w_interrupt.aupdate_state(
config,
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
)
@@ -1142,9 +1133,8 @@ async def test_conditional_graph() -> None:
),
"input": "what is weather in sf",
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1189,6 +1179,26 @@ async def test_conditional_graph() -> None:
await app_w_interrupt.aupdate_state(
config,
{
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
@@ -1205,12 +1215,12 @@ async def test_conditional_graph() -> None:
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
}
},
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == []
# test re-invoke to continue with interrupt_before
app_w_interrupt = workflow.compile(
@@ -1243,9 +1253,8 @@ async def test_conditional_graph() -> None:
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1527,11 +1536,13 @@ async def test_conditional_graph_state() -> None:
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1556,7 +1567,7 @@ async def test_conditional_graph_state() -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1596,7 +1607,27 @@ async def test_conditional_graph_state() -> None:
},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == []
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
# test state get/update methods with interrupt_before
@@ -1630,7 +1661,7 @@ async def test_conditional_graph_state() -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1655,7 +1686,7 @@ async def test_conditional_graph_state() -> None:
),
"intermediate_steps": [],
},
next=("agent:edges",),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -1695,7 +1726,27 @@ async def test_conditional_graph_state() -> None:
},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == []
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
async def test_conditional_entrypoint_graph() -> None:
@@ -2092,7 +2143,7 @@ async def test_prebuilt_chat() -> None:
]
async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
async def test_message_graph() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.agents import AgentAction
@@ -2206,9 +2257,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
app = workflow.compile()
assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [
HumanMessage(
content="what is weather in sf", id="00000000-0000-4000-8000-000000000002"
),
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
content="",
additional_kwargs={
@@ -2216,11 +2265,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
},
id="ai1",
),
FunctionMessage(
content="result for query",
name="search_api",
id="00000000-0000-4000-8000-000000000015",
),
FunctionMessage(content="result for query", name="search_api", id=AnyStr()),
AIMessage(
content="",
additional_kwargs={
@@ -2228,25 +2273,14 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
},
id="ai2",
),
FunctionMessage(
content="result for another",
name="search_api",
id="00000000-0000-4000-8000-000000000028",
),
FunctionMessage(content="result for another", name="search_api", id=AnyStr()),
AIMessage(content="answer", id="ai3"),
]
assert [
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
] == [
{
"__start__": [
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000037",
)
]
},
{"__start__": [HumanMessage(content="what is weather in sf", id=AnyStr())]},
{
"agent": AIMessage(
content="",
@@ -2258,9 +2292,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
},
{
"action": FunctionMessage(
content="result for query",
name="search_api",
id="00000000-0000-4000-8000-000000000050",
content="result for query", name="search_api", id=AnyStr()
)
},
{
@@ -2274,9 +2306,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
},
{
"action": FunctionMessage(
content="result for another",
name="search_api",
id="00000000-0000-4000-8000-000000000063",
content="result for another", name="search_api", id=AnyStr()
)
},
{"agent": AIMessage(content="answer", id="ai3")},
@@ -2296,7 +2326,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
{
"__start__": HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
)
},
{
@@ -2314,7 +2344,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2324,7 +2354,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
id="ai1",
),
],
next=("agent:edges",),
next=("action",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
@@ -2338,7 +2368,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2351,7 +2381,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
id="ai1",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2360,7 +2390,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
"action": FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000086",
id=AnyStr(),
)
},
{
@@ -2378,7 +2408,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2393,7 +2423,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000086",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2403,7 +2433,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
id="ai2",
),
],
next=("agent:edges",),
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
@@ -2417,7 +2447,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
values=[
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000072",
id=AnyStr(),
),
AIMessage(
content="",
@@ -2432,16 +2462,14 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None:
FunctionMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000086",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
],
next=("agent:edges",),
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == []
async def test_in_one_fan_out_out_one_graph_state() -> None:
def sorted_add(x: list[str], y: list[str]) -> list[str]: