WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-11 04:07:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix behaviour of interrupt_before

Browse Source 
- checkpoints are now copied before being mutated
- compiled graph and compiled state graph no longer need to override get_state/update_state
- pregel class now natively supports interrupt before/after
This commit is contained in:
Nuno Campos
2024-02-24 19:43:30 -08:00
parent 82c704e57e
commit 5c2afc9418
13 changed files with 832 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -18,7 +18,7 @@ test:
poetry run pytest
test_watch:
poetry run ptw tests
poetry run ptw .
######################
# LINTING AND FORMATTING
langgraph/channels/base.py
+8 -8
View File
@@ -116,16 +116,16 @@ def create_checkpoint(
checkpoint: Checkpoint, channels: Mapping[str, BaseChannel]
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
checkpoint = Checkpoint(
values: dict[str, Any] = {}
for k, v in channels.items():
try:
values[k] = v.checkpoint()
except EmptyChannelError:
pass
return Checkpoint(
v=1,
ts=datetime.now(timezone.utc).isoformat(),
channel_values=checkpoint["channel_values"],
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
)
for k, v in channels.items():
try:
checkpoint["channel_values"][k] = v.checkpoint()
except EmptyChannelError:
pass
return checkpoint
langgraph/checkpoint/base.py
+11
View File
@@ -1,6 +1,7 @@
import asyncio
from abc import ABC, abstractmethod
from collections import defaultdict
from copy import deepcopy
from datetime import datetime, timezone
from typing import Any, Optional, TypedDict
@@ -33,6 +34,16 @@ def empty_checkpoint() -> Checkpoint:
)
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
return Checkpoint(
v=checkpoint["v"],
ts=checkpoint["ts"],
channel_values=checkpoint["channel_values"].copy(),
channel_versions=checkpoint["channel_versions"].copy(),
versions_seen=deepcopy(checkpoint["versions_seen"]),
)
class CheckpointAt(StrEnum):
END_OF_STEP = "end_of_step"
END_OF_RUN = "end_of_run"
langgraph/constants.py
+1
View File
@@ -1,2 +1,3 @@
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
INTERRUPT = "__interrupt__"
langgraph/graph/graph.py
+5 -20
View File
@@ -178,6 +178,7 @@ class Graph:
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
debug: bool = False,
) -> "CompiledGraph":
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
@@ -226,11 +227,11 @@ class Graph:
input=f"{self.entry_point}:inbox" if self.entry_point else START,
output=END,
hidden=[f"{node}:inbox" for node in self.nodes],
snapshot_channels=list(self.nodes),
checkpointer=checkpointer,
interrupt=(
[f"{node}:inbox" for node in interrupt_before]
+ [node for node in interrupt_after]
),
interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],
interrupt_after_nodes=interrupt_after,
debug=debug,
)
@@ -279,19 +280,3 @@ class CompiledGraph(Pregel):
graph.add_edge(graph.nodes[START], graph.nodes[self.graph.entry_point])
return graph
def get_state(self, config: RunnableConfig) -> StateSnapshot:
snapshot = super().get_state(config)
return StateSnapshot(
values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes},
next=snapshot.next,
)
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
snapshot = await super().aget_state(config)
return StateSnapshot(
values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes},
next=snapshot.next,
)
langgraph/graph/state.py
+5 -40
View File
@@ -40,6 +40,7 @@ class StateGraph(Graph):
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
debug: bool = False,
) -> CompiledGraph:
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
@@ -147,11 +148,11 @@ class StateGraph(Graph):
input=f"{START}:inbox",
output=END,
hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys,
snapshot_channels=state_keys_read,
checkpointer=checkpointer,
interrupt=(
[f"{node}:inbox" for node in interrupt_before]
+ [node for node in interrupt_after]
),
interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],
interrupt_after_nodes=interrupt_after,
debug=debug,
)
@@ -209,39 +210,3 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
class CompiledStateGraph(CompiledGraph):
graph: StateGraph
def get_state(self, config: RunnableConfig) -> StateSnapshot:
snapshot = super(CompiledGraph, self).get_state(config)
return StateSnapshot(
values=snapshot.values.get("__root__")
if "__root__" in self.graph.channels
else {k: v for k, v in snapshot.values.items() if k in self.graph.channels},
next=snapshot.next,
)
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
snapshot = await super(CompiledGraph, self).aget_state(config)
return StateSnapshot(
values=snapshot.values.get("__root__")
if "__root__" in self.graph.channels
else {k: v for k, v in snapshot.values.items() if k in self.graph.channels},
next=snapshot.next,
)
def update_state(
self, config: RunnableConfig, values: Union[Any, dict[str, Any]]
) -> None:
return super(CompiledGraph, self).update_state(
config,
{"__root__": values} if "__root__" in self.graph.channels else values,
)
async def aupdate_state(
self, config: RunnableConfig, values: Union[Any, dict[str, Any]]
) -> None:
return await super(CompiledGraph, self).aupdate_state(
config,
{"__root__": values} if "__root__" in self.graph.channels else values,
)
langgraph/pregel/__init__.py
+218 -78
View File
@@ -57,9 +57,10 @@ from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
copy_checkpoint,
empty_checkpoint,
)
from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
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
from langgraph.pregel.log import logger
@@ -162,7 +163,7 @@ class Channel:
class StateSnapshot(NamedTuple):
values: dict[str, Any]
values: dict[str, Any] | Any
"""Current values of channels"""
next: tuple[str]
"""Nodes to execute in the next step, if any"""
@@ -175,12 +176,19 @@ class Pregel(
channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
# TODO Rename to `output_channels`
output: Union[str, Sequence[str]] = "output"
# TODO Replace with `stream_channels`
hidden: Sequence[str] = Field(default_factory=list)
interrupt: Sequence[str] = Field(default_factory=list)
snapshot_channels: Union[str, Sequence[str]] = Field(default_factory=list)
interrupt_after_nodes: Sequence[str] = Field(default_factory=list)
interrupt_before_nodes: Sequence[str] = Field(default_factory=list)
# TODO Rename to `input_channels`
input: Union[str, Sequence[str]] = "input"
step_timeout: Optional[float] = None
@@ -202,8 +210,12 @@ class Pregel(
values["input"],
values["output"],
values["hidden"],
values["interrupt"],
values["interrupt_after_nodes"],
values["interrupt_before_nodes"],
)
if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]:
if not values["checkpointer"]:
raise ValueError("Interrupts require a checkpointer")
return values
@property
@@ -257,6 +269,14 @@ class Pregel(
**{k: (self.channels[k].ValueType, None) for k in self.output},
)
@property
def snapshot_channels_list(self) -> Sequence[str]:
return (
[self.snapshot_channels]
if isinstance(self.snapshot_channels, str)
else self.snapshot_channels or list(self.channels.keys())
)
def get_state(self, config: RunnableConfig) -> StateSnapshot:
if not self.checkpointer:
raise ValueError("No checkpointer set")
@@ -264,15 +284,19 @@ class Pregel(
checkpoint = self.checkpointer.get(config)
checkpoint = checkpoint or empty_checkpoint()
with ChannelsManager(self.channels, checkpoint) as channels:
next_tasks = _prepare_next_tasks(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, update_seen=False
)
values = {
k: _read_channel(channels, k)
for k in channels
if k in self.snapshot_channels_list
and k not in [k.value for k in ReservedChannels]
}
return StateSnapshot(
{
k: _read_channel(channels, k)
for k in channels
if k not in [k.value for k in ReservedChannels]
},
values[self.snapshot_channels]
if isinstance(self.snapshot_channels, str)
else values,
tuple(name for _, _, name in next_tasks),
)
@@ -283,71 +307,130 @@ class Pregel(
checkpoint = await self.checkpointer.aget(config)
checkpoint = checkpoint or empty_checkpoint()
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
next_tasks = _prepare_next_tasks(
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, update_seen=False
)
values = {
k: _read_channel(channels, k)
for k in channels
if k in self.snapshot_channels_list
and k not in [k.value for k in ReservedChannels]
}
return StateSnapshot(
{
k: _read_channel(channels, k)
for k in channels
if k not in [k.value for k in ReservedChannels]
},
values[self.snapshot_channels]
if isinstance(self.snapshot_channels, str)
else values,
tuple(name for _, _, name in next_tasks),
)
def update_state(self, config: RunnableConfig, values: dict[str, Any]) -> None:
if not self.checkpointer:
raise ValueError("No checkpointer set")
checkpoint = self.checkpointer.get(config)
checkpoint = checkpoint or empty_checkpoint()
with ChannelsManager(self.channels, checkpoint) as channels:
for k, v in values.items():
channels[k].update([v])
checkpoint["channel_versions"][k] += 1
self.checkpointer.put(config, create_checkpoint(checkpoint, channels))
async def aupdate_state(
self, config: RunnableConfig, values: dict[str, Any]
def update_state(
self, config: RunnableConfig, values: dict[str, Any] | Any
) -> None:
if not self.checkpointer:
raise ValueError("No checkpointer set")
values = (
{self.snapshot_channels: values}
if isinstance(self.snapshot_channels, str)
else values
)
checkpoint = self.checkpointer.get(config)
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
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 or self.channels:
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
self.checkpointer.put(config, create_checkpoint(checkpoint, channels))
async def aupdate_state(
self, config: RunnableConfig, values: dict[str, Any] | Any
) -> None:
if not self.checkpointer:
raise ValueError("No checkpointer set")
values = (
{self.snapshot_channels: values}
if isinstance(self.snapshot_channels, str)
else values
)
checkpoint = await self.checkpointer.aget(config)
checkpoint = checkpoint or empty_checkpoint()
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
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.snapshot_channels or self.channels:
version = checkpoint["channel_versions"][k]
checkpoint["versions_seen"][INTERRUPT][k] = version
await self.checkpointer.aput(
config, create_checkpoint(checkpoint, channels)
)
def _defaults(
self,
debug: Optional[bool] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
) -> tuple[
bool,
Union[str, Sequence[str]],
Union[str, Sequence[str]],
Optional[Sequence[str]],
Optional[Sequence[str]],
]:
debug = debug if debug is not None else self.debug
if output_keys is None:
output_keys = [chan for chan in self.channels if chan not in self.hidden]
else:
validate_keys(output_keys, self.channels)
if input_keys is None:
input_keys = self.input
else:
validate_keys(input_keys, self.channels)
interrupt_before_nodes = interrupt_before_nodes or self.interrupt_before_nodes
interrupt_after_nodes = interrupt_after_nodes or self.interrupt_after_nodes
return (
debug,
input_keys,
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
)
def _transform(
self,
input: Iterator[Union[dict[str, Any], Any]],
run_manager: CallbackManagerForChainRun,
config: RunnableConfig,
*,
debug: Optional[bool] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt: Optional[Sequence[str]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
) -> Iterator[Union[dict[str, Any], Any]]:
try:
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
# assign defaults
if output_keys is None:
output_keys = [
chan for chan in self.channels if chan not in self.hidden
]
else:
validate_keys(output_keys, self.channels)
if input_keys is None:
input_keys = self.input
else:
validate_keys(input_keys, self.channels)
interrupt = interrupt or self.interrupt
(
debug,
input_keys,
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
) = self._defaults(
debug,
input_keys,
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
)
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
@@ -362,7 +445,7 @@ class Pregel(
w for c in input for w in map_input(input_keys, c)
):
# discard any unfinished tasks from previous checkpoint
_prepare_next_tasks(checkpoint, processes, channels)
checkpoint, _ = _prepare_next_tasks(checkpoint, processes, channels)
# apply input writes
_apply_writes(
checkpoint,
@@ -380,7 +463,9 @@ class Pregel(
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
for step in range(config["recursion_limit"] + 1):
next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
checkpoint, next_tasks = _prepare_next_tasks(
checkpoint, processes, channels
)
# if no more tasks, we're done
if not next_tasks:
@@ -392,7 +477,7 @@ class Pregel(
"by setting the `recursion_limit` config key."
)
if self.debug:
if debug:
print_step_start(step, next_tasks)
# collect all writes to channels, without applying them yet
@@ -430,15 +515,15 @@ class Pregel(
timeout=self.step_timeout,
)
# interrupt on failure or timeout
_interrupt_or_proceed(done, inflight, step)
# panic on failure or timeout
_panic_or_proceed(done, inflight, step)
# apply writes to channels
_apply_writes(
checkpoint, channels, pending_writes, config, step + 1
)
if self.debug:
if debug:
print_checkpoint(step, channels)
# yield current value and checkpoint view
@@ -449,22 +534,37 @@ class Pregel(
# if view was updated, apply writes to channels
_apply_writes_from_view(checkpoint, channels, step_output)
# with previous step's checkpoint
if do_interrupt_before := _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
if self.checkpointer is not None and (
self.checkpointer.at == CheckpointAt.END_OF_STEP
or interrupt_before_nodes
):
checkpoint = create_checkpoint(checkpoint, channels)
self.checkpointer.put(config, checkpoint)
# interrupt if any channel written to is in interrupt list
if any(chan for chan, _ in pending_writes if chan in interrupt):
# with this step's checkpoint,
if _should_interrupt(
checkpoint,
interrupt_after_nodes,
self.snapshot_channels_list,
pending_writes,
):
break
# save end of run checkpoint
if (
self.checkpointer is not None
and self.checkpointer.at == CheckpointAt.END_OF_RUN
and not do_interrupt_before
):
checkpoint = create_checkpoint(checkpoint, channels)
self.checkpointer.put(config, checkpoint)
@@ -482,9 +582,11 @@ class Pregel(
run_manager: AsyncCallbackManagerForChainRun,
config: RunnableConfig,
*,
debug: Optional[bool] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt: Optional[Sequence[str]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
try:
if config["recursion_limit"] < 1:
@@ -499,17 +601,19 @@ class Pregel(
None,
)
# assign defaults
if output_keys is None:
output_keys = [
chan for chan in self.channels if chan not in self.hidden
]
else:
validate_keys(output_keys, self.channels)
if input_keys is None:
input_keys = self.input
else:
validate_keys(input_keys, self.channels)
interrupt = interrupt or self.interrupt
(
debug,
input_keys,
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
) = self._defaults(
debug,
input_keys,
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
)
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
@@ -524,7 +628,7 @@ class Pregel(
[w async for c in input for w in map_input(input_keys, c)]
):
# discard any unfinished tasks from previous checkpoint
_prepare_next_tasks(checkpoint, processes, channels)
checkpoint, _ = _prepare_next_tasks(checkpoint, processes, channels)
# apply input writes
_apply_writes(
checkpoint,
@@ -542,7 +646,9 @@ class Pregel(
# channels are guaranteed to be immutable for the duration of the step,
# channel updates being applied only at the transition between steps
for step in range(config["recursion_limit"] + 1):
next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
checkpoint, next_tasks = _prepare_next_tasks(
checkpoint, processes, channels
)
# if no more tasks, we're done
if not next_tasks:
@@ -554,7 +660,7 @@ class Pregel(
"by setting the `recursion_limit` config key."
)
if self.debug:
if debug:
print_step_start(step, next_tasks)
# collect all writes to channels, without applying them yet
@@ -599,15 +705,15 @@ class Pregel(
timeout=self.step_timeout,
)
# interrupt on failure or timeout
_interrupt_or_proceed(done, inflight, step)
# panic on failure or timeout
_panic_or_proceed(done, inflight, step)
# apply writes to channels
_apply_writes(
checkpoint, channels, pending_writes, config, step + 1
)
if self.debug:
if debug:
print_checkpoint(step, channels)
# yield current value and checkpoint view
@@ -618,6 +724,15 @@ class Pregel(
# if view was updated, apply writes to channels
_apply_writes_from_view(checkpoint, channels, step_output)
# with previous step's checkpoint
if do_interrupt_before := _should_interrupt(
checkpoint,
interrupt_before_nodes,
self.snapshot_channels_list,
pending_writes,
):
break
# save end of step checkpoint
if (
self.checkpointer is not None
@@ -626,14 +741,20 @@ class Pregel(
checkpoint = create_checkpoint(checkpoint, channels)
await self.checkpointer.aput(config, checkpoint)
# interrupt if any channel written to is in interrupt list
if any(chan for chan, _ in pending_writes if chan in interrupt):
# with this step's checkpoint
if _should_interrupt(
checkpoint,
interrupt_after_nodes,
self.snapshot_channels_list,
pending_writes,
):
break
# save end of run checkpoint
if (
self.checkpointer is not None
and self.checkpointer.at == CheckpointAt.END_OF_RUN
and not do_interrupt_before
):
checkpoint = create_checkpoint(checkpoint, channels)
await self.checkpointer.aput(config, checkpoint)
@@ -762,7 +883,7 @@ class Pregel(
yield chunk
def _interrupt_or_proceed(
def _panic_or_proceed(
done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
step: int,
@@ -786,6 +907,24 @@ def _interrupt_or_proceed(
raise TimeoutError(f"Timed out at step {step}")
def _should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Sequence[str],
snapshot_channels: Sequence[str],
pending_writes: Sequence[tuple[str, Any]],
) -> bool:
return (
# interrupt if any of snapshopt_channels has been updated since last interrupt
any(
checkpoint["channel_versions"][chan]
> checkpoint["versions_seen"][INTERRUPT][chan]
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)
)
def _read_channel(
channels: Mapping[str, BaseChannel], chan: str, catch: bool = True
) -> Any:
@@ -840,6 +979,7 @@ def _apply_writes(
def _apply_writes_from_view(
checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], values: dict[str, Any]
) -> None:
# Apply writes to channels
for chan, value in values.items():
if value == _read_channel(channels, chan):
continue
@@ -857,7 +997,8 @@ def _prepare_next_tasks(
processes: Mapping[str, Union[ChannelInvoke, ChannelBatch]],
channels: Mapping[str, BaseChannel],
update_seen: bool = True,
) -> list[tuple[Runnable, Any, str]]:
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]:
checkpoint = copy_checkpoint(checkpoint) if update_seen else checkpoint
tasks: list[tuple[Runnable, Any, str]] = []
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
@@ -916,8 +1057,7 @@ def _prepare_next_tasks(
tasks.append((proc, val, name))
if update_seen:
seen[proc.channel] = checkpoint["channel_versions"][proc.channel]
return tasks
return checkpoint, tasks
async def _aconsume(iterator: AsyncIterator[Any]) -> None:
langgraph/pregel/validate.py
+8 -3
View File
@@ -2,6 +2,7 @@ from typing import Any, Mapping, Sequence, Union
from langgraph.channels.base import BaseChannel
from langgraph.channels.last_value import LastValue
from langgraph.constants import INTERRUPT
from langgraph.pregel.read import ChannelBatch, ChannelInvoke
from langgraph.pregel.reserved import ReservedChannels
@@ -12,10 +13,13 @@ def validate_graph(
input: Union[str, Sequence[str]],
output: Union[str, Sequence[str]],
hidden: Sequence[str],
interrupt: Sequence[str],
interrupt_after: Sequence[str],
interrupt_before: Sequence[str],
) -> None:
subscribed_channels = set[str]()
for node in nodes.values():
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())
elif isinstance(node, ChannelBatch):
@@ -56,7 +60,8 @@ def validate_graph(
channels[chan] = LastValue(Any) # type: ignore[arg-type]
validate_keys(hidden, channels)
validate_keys(interrupt, channels)
validate_keys(interrupt_after, channels)
validate_keys(interrupt_before, channels)
def validate_keys(
poetry.lock
Generated
+5 -6
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand.
[[package]]
name = "aiohttp"
@@ -2520,13 +2520,13 @@ dev = ["pre-commit", "pytest-asyncio", "tox"]
[[package]]
name = "pytest-watcher"
version = "0.3.5"
version = "0.4.1"
description = "Automatically rerun your tests on file modifications"
optional = false
python-versions = ">=3.7.0,<4.0.0"
files = [
{file = "pytest_watcher-0.3.5-py3-none-any.whl", hash = "sha256:af00ca52c7be22dc34c0fd3d7ffef99057207a73b05dc5161fe3b2fe91f58130"},
{file = "pytest_watcher-0.3.5.tar.gz", hash = "sha256:8896152460ba2b1a8200c12117c6611008ec96c8b2d811f0a05ab8a82b043ff8"},
{file = "pytest_watcher-0.4.1-py3-none-any.whl", hash = "sha256:29435669cb0124fb32d6de649fe9b1350f6dac94176313fff559ee4c2a66fd6e"},
{file = "pytest_watcher-0.4.1.tar.gz", hash = "sha256:5a793c4c883e3a55ab2abbfa3a8cd6fa6495b3767d5f6644052cc5f3236f511a"},
]
[package.dependencies]
@@ -2635,7 +2635,6 @@ files = [
{file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
{file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
{file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
@@ -3760,4 +3759,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "0e7777d77d3b34acbfdead224a2b5c5e65ecbf890c57c29bf43f5ebff09d4c0d"
content-hash = "2d35e923bf3902e0e11a305f58d17b0efc3fbb444dff8d6cb92e070a993115c9"
pyproject.toml
+1 -1
View File
@@ -23,7 +23,7 @@ pytest-asyncio = "^0.20.3"
pytest-mock = "^3.10.0"
syrupy = "^4.0.2"
httpx = "^0.26.0"
pytest-watcher = "^0.3.4"
pytest-watcher = "^0.4.1"
langchain = "^0.1.0"
aiosqlite = "^0.19.0"
grandalf = "^0.8"
tests/memory_assert.py
+19
View File
@@ -0,0 +1,19 @@
from langchain_core.pydantic_v1 import Field
from langgraph.checkpoint.base import Checkpoint, CheckpointAt, copy_checkpoint
from langgraph.checkpoint.memory import MemorySaver
class MemorySaverAssertImmutable(MemorySaver):
storage_for_copies: dict[str, Checkpoint] = Field(default_factory=dict)
at = CheckpointAt.END_OF_STEP
def put(self, config: dict, checkpoint: dict) -> None:
# assert checkpoint hasn't been modified since last written
thread_id = config["configurable"]["thread_id"]
if saved := super().get(config):
assert self.storage_for_copies[thread_id] == saved
self.storage_for_copies[thread_id] = copy_checkpoint(checkpoint)
# call super to write checkpoint
super().put(config, checkpoint)
tests/test_pregel.py
+273 -8
View File
@@ -16,7 +16,6 @@ from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, Graph
from langgraph.graph.message import MessageGraph
@@ -28,6 +27,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.memory_assert import MemorySaverAssertImmutable
def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -254,9 +254,11 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
memory = MemorySaver()
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"]
nodes={"one": one, "two": two},
checkpointer=memory,
interrupt_after_nodes=["inbox"],
)
# start execution, stop at inbox
@@ -447,7 +449,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
| raise_if_above_10
)
memory = MemorySaver()
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one},
@@ -899,10 +901,10 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
},
]
# test state get/update methods
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaver(), interrupt_after=["agent"]
checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"]
)
config = {"configurable": {"thread_id": "1"}}
@@ -1044,6 +1046,152 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
}
]
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"]
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
)
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",
},
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
"tools": None,
},
next=("agent:edges",),
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"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": {
"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": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
app_w_interrupt.update_state(
config,
{
"agent": {
"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 [c for c in app_w_interrupt.stream(None, config)] == [
{
"__end__": {
"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",
)
],
}
}
]
def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
from langchain.llms.fake import FakeStreamingListLLM
@@ -1235,8 +1383,10 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
},
]
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaver(), interrupt_after=["agent"]
checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"]
)
config = {"configurable": {"thread_id": "1"}}
@@ -1345,6 +1495,121 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
}
]
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
debug=True,
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("agent:edges",),
)
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"intermediate_steps": [],
},
next=("agent:edges",),
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"__end__": {
"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",
)
],
}
}
]
def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
def left(data: str) -> str:
@@ -2022,7 +2287,7 @@ def test_message_graph(snapshot: SnapshotAssertion) -> None:
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaver(), interrupt_after=["agent"]
checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"]
)
config = {"configurable": {"thread_id": "1"}}
tests/test_pregel_async.py
+277 -8
View File
@@ -23,7 +23,6 @@ from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph
from langgraph.prebuilt.chat_agent_executor import (
@@ -33,6 +32,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.memory_assert import MemorySaverAssertImmutable
async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -261,9 +261,11 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
memory = MemorySaver()
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"]
nodes={"one": one, "two": two},
checkpointer=memory,
interrupt_after_nodes=["inbox"],
)
# start execution, stop at inbox
@@ -461,7 +463,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
| raise_if_above_10
)
memory = MemorySaver()
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one},
@@ -944,10 +946,10 @@ async def test_conditional_graph() -> None:
# Check that agent (one of the nodes) has its output streamed to the logs
assert "/logs/agent/streamed_output/-" in patch_paths
# test state get/update methods
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaver(), interrupt_after=["agent"]
checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"]
)
config = {"configurable": {"thread_id": "1"}}
@@ -1092,6 +1094,155 @@ async def test_conditional_graph() -> None:
}
]
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"]
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
"tools": None,
},
next=("agent:edges",),
)
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",
},
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
"tools": None,
},
next=("agent:edges",),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"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": {
"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": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent": {
"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 [c async for c in app_w_interrupt.astream(None, config)] == [
{
"__end__": {
"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",
)
],
}
}
]
async def test_conditional_graph_state() -> None:
from langchain.llms.fake import FakeStreamingListLLM
@@ -1278,8 +1429,10 @@ async def test_conditional_graph_state() -> None:
},
]
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaver(), interrupt_after=["agent"]
checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"]
)
config = {"configurable": {"thread_id": "1"}}
@@ -1391,6 +1544,122 @@ async def test_conditional_graph_state() -> None:
}
]
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"]
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("agent:edges",),
)
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"intermediate_steps": [],
},
next=("agent:edges",),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"__end__": {
"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",
)
],
}
}
]
async def test_conditional_entrypoint_graph() -> None:
async def left(data: str) -> str:
@@ -2051,7 +2320,7 @@ async def test_message_graph() -> None:
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaver(), interrupt_after=["agent"]
checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"]
)
config = {"configurable": {"thread_id": "1"}}