diff --git a/Makefile b/Makefile index c50d63a07..418703a74 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,7 @@ test: poetry run pytest test_watch: - poetry run ptw tests + poetry run ptw . ###################### # LINTING AND FORMATTING diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index b30bee9b2..71680c8ea 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -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 diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 699dc9dd8..99b0e294a 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -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" diff --git a/langgraph/constants.py b/langgraph/constants.py index 6e5f9c37f..4bf8b1335 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,2 +1,3 @@ CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" +INTERRUPT = "__interrupt__" diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 56e1afa33..53331d2af 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -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, - ) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index bfc459a1c..b1f048074 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -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, - ) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index fa061d619..76df9f591 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -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: diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 8283d093c..498d96aa2 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -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( diff --git a/poetry.lock b/poetry.lock index 0c5dc41ef..46889d252 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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" diff --git a/pyproject.toml b/pyproject.toml index 9e2d94a0b..53d60b19d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/tests/memory_assert.py b/tests/memory_assert.py new file mode 100644 index 000000000..4362cf939 --- /dev/null +++ b/tests/memory_assert.py @@ -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) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6ae5f5ee9..b943c4506 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -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"}} diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 6db086683..fe716c953 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -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"}}