From 56238036d760a49ee0608e35b7f666e2f6659e1e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 09:29:12 -0700 Subject: [PATCH 01/15] WIP: Split out loop logic from Pregel class --- libs/langgraph/langgraph/checkpoint/base.py | 3 +- libs/langgraph/langgraph/pregel/__init__.py | 737 ++------------------ libs/langgraph/langgraph/pregel/algo.py | 413 +++++++++++ libs/langgraph/langgraph/pregel/debug.py | 2 +- libs/langgraph/langgraph/pregel/loop.py | 282 ++++++++ libs/langgraph/langgraph/pregel/types.py | 9 + libs/langgraph/tests/test_prebuilt.py | 2 +- libs/langgraph/tests/test_pregel.py | 2 - libs/langgraph/tests/test_pregel_async.py | 2 - 9 files changed, 785 insertions(+), 667 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/algo.py create mode 100644 libs/langgraph/langgraph/pregel/loop.py diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 6b147f3e3..2f73b52a6 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -25,6 +25,7 @@ from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer V = TypeVar("V", int, float, str) +PendingWrite = Tuple[str, str, Any] # Marked as total=False to allow for future expansion. @@ -118,7 +119,7 @@ class CheckpointTuple(NamedTuple): checkpoint: Checkpoint metadata: CheckpointMetadata parent_config: Optional[RunnableConfig] = None - pending_writes: Optional[List[Tuple[str, str, Any]]] = None + pending_writes: Optional[List[PendingWrite]] = None CheckpointThreadId = ConfigurableFieldSpec( diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d0d1c461d..5b1d5b9f1 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2,9 +2,8 @@ from __future__ import annotations import asyncio import concurrent.futures -import json import time -from collections import defaultdict, deque +from collections import deque from functools import partial from inspect import signature from typing import ( @@ -14,7 +13,6 @@ from typing import ( Callable, Dict, Iterator, - Literal, Mapping, Optional, Sequence, @@ -26,7 +24,6 @@ from typing import ( ) from uuid import UUID, uuid5 -from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager from langchain_core.globals import get_debug from langchain_core.load.dump import dumpd from langchain_core.pydantic_v1 import BaseModel, Field, root_validator @@ -41,7 +38,6 @@ from langchain_core.runnables.config import ( ensure_config, get_async_callback_manager_for_config, get_callback_manager_for_config, - merge_configs, patch_config, ) from langchain_core.runnables.utils import ( @@ -54,7 +50,6 @@ from typing_extensions import Self from langgraph.channels.base import ( BaseChannel, - EmptyChannelError, ) from langgraph.channels.context import Context from langgraph.channels.manager import ( @@ -64,7 +59,6 @@ from langgraph.channels.manager import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, - Checkpoint, CheckpointMetadata, copy_checkpoint, empty_checkpoint, @@ -73,18 +67,21 @@ from langgraph.constants import ( CONFIG_KEY_READ, CONFIG_KEY_SEND, INTERRUPT, - TAG_HIDDEN, - TASKS, - Send, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( AsyncManagedValuesManager, - ManagedValueMapping, ManagedValuesManager, ManagedValueSpec, is_managed_value, ) +from langgraph.pregel.algo import ( + apply_writes, + increment, + local_read, + prepare_next_tasks, + should_interrupt, +) from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results, @@ -93,23 +90,22 @@ from langgraph.pregel.debug import ( print_step_tasks, print_step_writes, ) -from langgraph.pregel.executor import AsyncBackgroundExecutor, BackgroundExecutor +from langgraph.pregel.executor import AsyncBackgroundExecutor from langgraph.pregel.io import ( map_input, map_output_updates, map_output_values, - read_channel, read_channels, single, ) -from langgraph.pregel.log import logger +from langgraph.pregel.loop import PregelLoop from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry from langgraph.pregel.types import ( All, PregelExecutableTask, - PregelTaskDescription, StateSnapshot, + StreamMode, ) from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -196,16 +192,6 @@ class Channel: ) -StreamMode = Literal["values", "updates", "debug"] -"""How the stream method should emit outputs. - -- 'values': Emit all values of the state for each step. -- 'updates': Emit only the node name(s) and updates - that were returned by the node(s) **after** each step. -- 'debug': Emit debug events for each step. -""" - - class Pregel( RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] ): @@ -378,7 +364,7 @@ class Pregel( ) as channels, ManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = _prepare_next_tasks( + _, next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -410,7 +396,7 @@ class Pregel( ) as channels, AsyncManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = _prepare_next_tasks( + _, next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -452,7 +438,7 @@ class Pregel( ) as channels, ManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = _prepare_next_tasks( + _, next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -498,7 +484,7 @@ class Pregel( ) as channels, AsyncManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = _prepare_next_tasks( + _, next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -583,13 +569,13 @@ class Pregel( # deque.extend is thread-safe CONFIG_KEY_SEND: task.writes.extend, CONFIG_KEY_READ: partial( - _local_read, checkpoint, channels, task.writes, config + local_read, checkpoint, channels, task.writes, config ), }, ), ) # apply to checkpoint and save - _apply_writes( + apply_writes( checkpoint, channels, task.writes, self.checkpointer.get_next_version ) step = saved.metadata.get("step", -2) + 1 if saved else -1 @@ -675,13 +661,13 @@ class Pregel( # deque.extend is thread-safe CONFIG_KEY_SEND: task.writes.extend, CONFIG_KEY_READ: partial( - _local_read, checkpoint, channels, task.writes, config + local_read, checkpoint, channels, task.writes, config ), }, ), ) # apply to checkpoint and save - _apply_writes( + apply_writes( checkpoint, channels, task.writes, self.checkpointer.get_next_version ) step = saved.metadata.get("step", -2) + 1 if saved else -1 @@ -711,7 +697,6 @@ class Pregel( config: Optional[RunnableConfig] = None, *, stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, @@ -729,10 +714,6 @@ class Pregel( output_keys = self.stream_channels_asis else: validate_keys(output_keys, self.channels) - if input_keys is None: - input_keys = self.input_channels - else: - validate_keys(input_keys, self.channels) interrupt_before = interrupt_before or self.interrupt_before_nodes interrupt_after = interrupt_after or self.interrupt_after_nodes stream_mode = stream_mode if stream_mode is not None else self.stream_mode @@ -744,7 +725,6 @@ class Pregel( return ( debug, stream_mode, - input_keys, output_keys, interrupt_before, interrupt_after, @@ -757,7 +737,6 @@ class Pregel( *, stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: Optional[bool] = None, @@ -848,214 +827,68 @@ class Pregel( ( debug, stream_modes, - input_keys, output_keys, interrupt_before, interrupt_after, ) = self._defaults( config, stream_mode=stream_mode, - input_keys=input_keys, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, debug=debug, ) - # copy nodes to ignore mutations during execution - processes = {**self.nodes} - # get checkpoint from saver, or create an empty one - saved = self.checkpointer.get_tuple(config) if self.checkpointer else None - checkpoint = saved.checkpoint if saved else empty_checkpoint() - # merge configurable fields with previous checkpoint config - checkpoint_config = config - if saved: - checkpoint_config = { - **config, - **saved.config, - "configurable": { - **config.get("configurable", {}), - **saved.config["configurable"], - }, - } - - start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint - with BackgroundExecutor(config) as submit, ChannelsManager( - self.channels, checkpoint, config - ) as channels, ManagedValuesManager( - self.managed_values_dict, config, self - ) as managed: - - def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None: - if self.checkpointer is not None: - submit( - self.checkpointer.put_writes, - { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - }, - writes, - task_id, - ) - - def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: - nonlocal checkpoint, checkpoint_config, channels - - if self.checkpointer is None: - return - if debug: - print_step_checkpoint( - metadata["step"], channels, self.stream_channels_list - ) - - # create new checkpoint - checkpoint = create_checkpoint( - checkpoint, channels, metadata["step"] - ) - # save it, without blocking - submit( - self.checkpointer.put, - checkpoint_config, - copy_checkpoint(checkpoint), - metadata, - ) - # update checkpoint config - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - } - # yield debug checkpoint event - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - metadata["step"], - checkpoint_config, - channels, - self.stream_channels_asis, - metadata, - ), - ) - - # map inputs to channel updates - if input_writes := deque(map_input(input_keys, input)): - # discard any unfinished tasks from previous checkpoint - checkpoint, _ = _prepare_next_tasks( - checkpoint, - processes, - channels, - managed, - config, - -1, - for_execution=True, - get_next_version=( - self.checkpointer.get_next_version - if self.checkpointer - else _increment - ), - ) - # apply input writes - _apply_writes( - checkpoint, - channels, - input_writes, - ( - self.checkpointer.get_next_version - if self.checkpointer - else _increment - ), - ) - # save input checkpoint - yield from put_checkpoint( - { - "source": "input", - "step": start, - "writes": input, - } - ) - # increment start to 0 - start += 1 - else: - # no input is taken as signal to proceed past previous interrupt - checkpoint = copy_checkpoint(checkpoint) - for k in channels: - if k in checkpoint["channel_versions"]: - version = checkpoint["channel_versions"][k] - checkpoint["versions_seen"][INTERRUPT][k] = version - + with PregelLoop( + input, config=config, checkpointer=self.checkpointer, graph=self + ) as loop: # 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 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps - stop = start + config["recursion_limit"] + 1 - for step in range(start, stop): - next_checkpoint, next_tasks = _prepare_next_tasks( - checkpoint, - processes, - channels, - managed, - config, - step, - for_execution=True, - manager=run_manager, - get_next_version=( - self.checkpointer.get_next_version - if self.checkpointer - else _increment - ), - ) - - # assign pending writes to tasks - if saved and saved.pending_writes: - for task in next_tasks: - task.writes.extend( - (c, v) - for tid, c, v in saved.pending_writes - if tid == task.id - ) - - # if no more tasks, we're done - if not next_tasks: - if step == start: - raise ValueError("No tasks to run in graph.") - else: - break - - # before execution, check if we should interrupt - if _should_interrupt( - checkpoint, - interrupt_before, - self.stream_channels_list, - next_tasks, - ): - break - else: - checkpoint = next_checkpoint + while loop.tick( + interrupt_before=interrupt_before, interrupt_after=interrupt_after + ): + # print debug output + if self.debug: + print_step_checkpoint( + loop.checkpoint_metadata, + loop.channels, + self.stream_channels_list, + ) + # emit debug output + if self.checkpointer and "debug" in stream_modes: + yield from _with_mode( + "debug", + isinstance(stream_mode, list), + map_debug_checkpoint( + loop.checkpoint_metadata["step"], + loop.config, + loop.channels, + self.stream_channels_asis, + loop.checkpoint_metadata, + ), + ) if debug: - print_step_tasks(step, next_tasks) + print_step_tasks(loop.checkpoint_metadata["step"], loop.tasks) if "debug" in stream_modes: yield from _with_mode( "debug", isinstance(stream_mode, list), - map_debug_tasks(step, next_tasks), + map_debug_tasks( + loop.checkpoint_metadata["step"], loop.tasks + ), ) # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - submit(run_with_retry, task, self.retry_policy): task - for task in next_tasks + loop.submit(run_with_retry, task, self.retry_policy): task + for task in loop.tasks if not task.writes } end_time = ( @@ -1084,10 +917,8 @@ class Pregel( # exception will be handled in panic_or_proceed futures.clear() else: - # save task writes to checkpointer, unless this - # is the single or last task in this step - if futures: - put_writes(task.id, task.writes) + # save task writes to checkpointer + loop.put_writes(task.id, task.writes) # yield updates output for the finished task if "updates" in stream_modes: yield from _with_mode( @@ -1100,7 +931,9 @@ class Pregel( "debug", isinstance(stream_mode, list), map_debug_task_results( - step, [task], self.stream_channels_list + loop.checkpoint_metadata["step"], + [task], + self.stream_channels_list, ), ) else: @@ -1108,66 +941,32 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, step) + _panic_or_proceed(done, inflight, loop.checkpoint_metadata["step"]) # don't keep futures around in memory longer than needed del done, inflight, futures # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for task in next_tasks: + for task in loop.tasks: pending_writes.extend(task.writes) if debug: print_step_writes( - step, pending_writes, self.stream_channels_list + loop.checkpoint_metadata["step"], + pending_writes, + self.stream_channels_list, ) - # apply writes to channels - _apply_writes( - checkpoint, - channels, - pending_writes, - ( - self.checkpointer.get_next_version - if self.checkpointer - else _increment - ), - ) - # yield values output if "values" in stream_modes: yield from _with_mode( "values", isinstance(stream_mode, list), - map_output_values(output_keys, pending_writes, channels), - ) - - # save end of step checkpoint - yield from put_checkpoint( - { - "source": "loop", - "step": step, - "writes": ( - single(map_output_updates(output_keys, next_tasks)) - if self.stream_mode == "updates" - else single( - map_output_values( - output_keys, pending_writes, channels - ), - ) + map_output_values( + output_keys, pending_writes, loop.channels ), - } - ) - - # after execution, check if we should interrupt - if _should_interrupt( - checkpoint, - interrupt_after, - self.stream_channels_list, - next_tasks, - ): - break - else: + ) + if loop.status == "out_of_steps": raise GraphRecursionError( f"Recursion limit of {config['recursion_limit']} reached" "without hitting a stop condition. You can increase the " @@ -1175,7 +974,7 @@ class Pregel( ) # set final channel values as run output - run_manager.on_chain_end(read_channels(channels, output_keys)) + run_manager.on_chain_end(read_channels(loop.channels, output_keys)) except BaseException as e: run_manager.on_chain_error(e) raise @@ -1187,7 +986,6 @@ class Pregel( *, stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: Optional[bool] = None, @@ -1204,7 +1002,6 @@ class Pregel( Output is a dict with the node name as key and the updated values as value. debug: Emit debug events for each step. output_keys: The keys to stream, defaults to all non-context channels. - input_keys: The keys to use from the input, defaults to all input channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. debug: Whether to print debug information during execution, defaults to False. @@ -1288,14 +1085,12 @@ class Pregel( ( debug, stream_modes, - input_keys, output_keys, interrupt_before, interrupt_after, ) = self._defaults( config, stream_mode=stream_mode, - input_keys=input_keys, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, @@ -1391,9 +1186,9 @@ class Pregel( ) # map inputs to channel updates - if input_writes := deque(map_input(input_keys, input)): + if input_writes := deque(map_input(self.input_channels, input)): # discard any unfinished tasks from previous checkpoint - checkpoint, _ = _prepare_next_tasks( + checkpoint, _ = prepare_next_tasks( checkpoint, processes, channels, @@ -1404,18 +1199,18 @@ class Pregel( get_next_version=( self.checkpointer.get_next_version if self.checkpointer - else _increment + else increment ), ) # apply input writes - _apply_writes( + apply_writes( checkpoint, channels, input_writes, ( self.checkpointer.get_next_version if self.checkpointer - else _increment + else increment ), ) # save input checkpoint @@ -1440,7 +1235,7 @@ class Pregel( # channel updates being applied only at the transition between steps stop = start + config["recursion_limit"] + 1 for step in range(start, stop): - next_checkpoint, next_tasks = _prepare_next_tasks( + next_checkpoint, next_tasks = prepare_next_tasks( checkpoint, processes, channels, @@ -1452,7 +1247,7 @@ class Pregel( get_next_version=( self.checkpointer.get_next_version if self.checkpointer - else _increment + else increment ), ) @@ -1473,7 +1268,7 @@ class Pregel( break # before execution, check if we should interrupt - if _should_interrupt( + if should_interrupt( checkpoint, interrupt_before, self.stream_channels_list, @@ -1571,14 +1366,14 @@ class Pregel( ) # apply writes to channels - _apply_writes( + apply_writes( checkpoint, channels, pending_writes, ( self.checkpointer.get_next_version if self.checkpointer - else _increment + else increment ), ) @@ -1610,7 +1405,7 @@ class Pregel( yield chunk # after execution, check if we should interrupt - if _should_interrupt( + if should_interrupt( checkpoint, interrupt_after, self.stream_channels_list, @@ -1637,7 +1432,6 @@ class Pregel( *, stream_mode: StreamMode = "values", output_keys: Optional[Union[str, Sequence[str]]] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: Optional[bool] = None, @@ -1650,7 +1444,6 @@ class Pregel( config: Optional. The configuration for the graph run. stream_mode: Optional[str]. The stream mode for the graph run. Default is "values". output_keys: Optional. The output keys to retrieve from the graph run. - input_keys: Optional. The input keys to provide for the graph run. interrupt_before: Optional. The nodes to interrupt the graph run before. interrupt_after: Optional. The nodes to interrupt the graph run after. debug: Optional. Enable debug mode for the graph run. @@ -1670,7 +1463,6 @@ class Pregel( config, stream_mode=stream_mode, output_keys=output_keys, - input_keys=input_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, debug=debug, @@ -1692,7 +1484,6 @@ class Pregel( *, stream_mode: StreamMode = "values", output_keys: Optional[Union[str, Sequence[str]]] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: Optional[bool] = None, @@ -1705,7 +1496,6 @@ class Pregel( config: Optional. The configuration for the computation. stream_mode: Optional. The stream mode for the computation. Default is "values". output_keys: Optional. The output keys to include in the result. Default is None. - input_keys: Optional. The input keys to include in the result. Default is None. interrupt_before: Optional. The nodes to interrupt before. Default is None. interrupt_after: Optional. The nodes to interrupt after. Default is None. debug: Optional. Whether to enable debug mode. Default is None. @@ -1726,7 +1516,6 @@ class Pregel( config, stream_mode=stream_mode, output_keys=output_keys, - input_keys=input_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, debug=debug, @@ -1766,378 +1555,6 @@ def _panic_or_proceed( raise timeout_exc_cls(f"Timed out at step {step}") -def _should_interrupt( - checkpoint: Checkpoint, - interrupt_nodes: Union[All, Sequence[str]], - snapshot_channels: Sequence[str], - tasks: list[PregelExecutableTask], -) -> bool: - version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) - null_version = version_type() - # defaultdicts are mutated on access :( so we need to copy - seen = checkpoint["versions_seen"].copy()[INTERRUPT] - return ( - # interrupt if any channel has been updated since last interrupt - any( - version > seen.get(chan, null_version) - for chan, version in checkpoint["channel_versions"].items() - ) - # and any triggered node is in interrupt_nodes list - and any( - task.name - for task in tasks - if ( - (not task.config or TAG_HIDDEN not in task.config.get("tags")) - if interrupt_nodes == "*" - else task.name in interrupt_nodes - ) - ) - ) - - -def _local_read( - checkpoint: Checkpoint, - channels: Mapping[str, BaseChannel], - writes: Sequence[tuple[str, Any]], - config: RunnableConfig, - select: Union[list[str], str], - fresh: bool = False, -) -> Union[dict[str, Any], Any]: - if fresh: - checkpoint = create_checkpoint(checkpoint, channels, -1) - context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)} - with ChannelsManager( - {k: v for k, v in channels.items() if k not in context_channels}, - checkpoint, - config, - ) as channels: - all_channels = {**channels, **context_channels} - _apply_writes(copy_checkpoint(checkpoint), all_channels, writes, None) - return read_channels(all_channels, select) - else: - return read_channels(channels, select) - - -def _local_write( - commit: Callable[[Sequence[tuple[str, Any]]], None], - processes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - writes: Sequence[tuple[str, Any]], -) -> None: - for chan, value in writes: - if chan == TASKS: - if not isinstance(value, Send): - raise InvalidUpdateError( - f"Invalid packet type, expected Packet, got {value}" - ) - if value.node not in processes: - raise InvalidUpdateError(f"Invalid node name {value.node} in packet") - elif chan not in channels: - logger.warning(f"Skipping write for channel '{chan}' which has no readers") - commit(writes) - - -def _increment(current: Optional[int], channel: BaseChannel) -> int: - return current + 1 if current is not None else 1 - - -def _apply_writes( - checkpoint: Checkpoint, - channels: Mapping[str, BaseChannel], - pending_writes: Sequence[tuple[str, Any]], - get_next_version: Optional[Callable[[int, BaseChannel], int]], -) -> None: - if checkpoint["pending_sends"]: - checkpoint["pending_sends"].clear() - - pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) - # Group writes by channel - for chan, val in pending_writes: - if chan == TASKS: - checkpoint["pending_sends"].append(val) - else: - pending_writes_by_channel[chan].append(val) - - # Find the highest version of all channels - if checkpoint["channel_versions"]: - max_version = max(checkpoint["channel_versions"].values()) - else: - max_version = None - - updated_channels: set[str] = set() - # Apply writes to channels - for chan, vals in pending_writes_by_channel.items(): - if chan in channels: - try: - updated = channels[chan].update(vals) - except InvalidUpdateError as e: - raise InvalidUpdateError( - f"Invalid update for channel {chan} with values {vals}" - ) from e - if updated and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, channels[chan] - ) - updated_channels.add(chan) - # Channels that weren't updated in this step are notified of a new step - for chan in channels: - if chan not in updated_channels: - if channels[chan].update([]) and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, channels[chan] - ) - - -@overload -def _prepare_next_tasks( - checkpoint: Checkpoint, - processes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, - step: int, - for_execution: Literal[False], - get_next_version: Literal[None] = None, - manager: Literal[None] = None, -) -> tuple[Checkpoint, list[PregelTaskDescription]]: - ... - - -@overload -def _prepare_next_tasks( - checkpoint: Checkpoint, - processes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, - step: int, - for_execution: Literal[True], - get_next_version: Callable[[int, BaseChannel], int], - manager: Union[None, ParentRunManager, AsyncParentRunManager], -) -> tuple[Checkpoint, list[PregelExecutableTask]]: - ... - - -def _prepare_next_tasks( - checkpoint: Checkpoint, - processes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, - step: int, - *, - for_execution: bool, - get_next_version: Union[None, Callable[[int, BaseChannel], int]] = None, - manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, -) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: - checkpoint = copy_checkpoint(checkpoint) - tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] - # Consume pending packets - for packet in checkpoint["pending_sends"]: - if not isinstance(packet, Send): - logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends") - continue - if for_execution: - proc = processes[packet.node] - if node := proc.get_node(): - triggers = [TASKS] - metadata = { - "langgraph_step": step, - "langgraph_node": packet.node, - "langgraph_triggers": triggers, - "langgraph_task_idx": len(tasks), - } - task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) - writes = deque() - tasks.append( - PregelExecutableTask( - packet.node, - packet.arg, - node, - writes, - patch_config( - merge_configs( - config, - processes[packet.node].config, - {"metadata": metadata}, - ), - run_name=packet.node, - callbacks=( - manager.get_child(f"graph:step:{step}") - if manager - else None - ), - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - _local_write, writes.extend, processes, channels - ), - CONFIG_KEY_READ: partial( - _local_read, checkpoint, channels, writes, config - ), - }, - ), - triggers, - proc.retry_policy, - task_id, - ) - ) - else: - tasks.append(PregelTaskDescription(packet.node, packet.arg)) - if for_execution: - checkpoint["pending_sends"].clear() - # Collect channels to consume - channels_to_consume = set() - # Check if any processes should be run in next step - # If so, prepare the values to be passed to them - version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) - null_version = version_type() - if null_version is None: - return checkpoint, tasks - for name, proc in processes.items(): - seen = checkpoint["versions_seen"][name] - # If any of the channels read by this process were updated - if triggers := sorted( - chan - for chan in proc.triggers - if not isinstance( - read_channel(channels, chan, return_exception=True), EmptyChannelError - ) - and checkpoint["channel_versions"].get(chan, null_version) - > seen.get(chan, null_version) - ): - channels_to_consume.update(triggers) - try: - val = next(_proc_input(step, name, proc, managed, channels)) - except StopIteration: - continue - - # update seen versions - if for_execution: - seen.update( - { - chan: checkpoint["channel_versions"][chan] - for chan in proc.triggers - if chan in checkpoint["channel_versions"] - } - ) - - if for_execution: - if node := proc.get_node(): - metadata = { - "langgraph_step": step, - "langgraph_node": name, - "langgraph_triggers": triggers, - "langgraph_task_idx": len(tasks), - } - task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) - writes = deque() - tasks.append( - PregelExecutableTask( - name, - val, - node, - writes, - patch_config( - merge_configs( - config, - proc.config, - {"metadata": metadata}, - ), - run_name=name, - callbacks=( - manager.get_child(f"graph:step:{step}") - if manager - else None - ), - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - _local_write, writes.extend, processes, channels - ), - CONFIG_KEY_READ: partial( - _local_read, - checkpoint, - channels, - writes, - config, - ), - }, - ), - triggers, - proc.retry_policy, - task_id, - ) - ) - else: - tasks.append(PregelTaskDescription(name, val)) - # Find the highest version of all channels - if checkpoint["channel_versions"]: - max_version = max(checkpoint["channel_versions"].values()) - else: - max_version = None - # Consume all channels that were read - if for_execution: - for chan in channels_to_consume: - if channels[chan].consume(): - checkpoint["channel_versions"][chan] = get_next_version( - max_version, channels[chan] - ) - return checkpoint, tasks - - -def _proc_input( - step: int, - name: str, - proc: PregelNode, - managed: ManagedValueMapping, - channels: Mapping[str, BaseChannel], -) -> Iterator[Any]: - # If all trigger channels subscribed by this process are not empty - # then invoke the process with the values of all non-empty channels - if isinstance(proc.channels, dict): - try: - val: dict = { - k: read_channel( - channels, - chan, - catch=chan not in proc.triggers, - ) - for k, chan in proc.channels.items() - if isinstance(chan, str) - } - - managed_values = {} - for key, chan in proc.channels.items(): - if is_managed_value(chan): - managed_values[key] = managed[key]( - step, PregelTaskDescription(name, val) - ) - - val.update(managed_values) - except EmptyChannelError: - return - elif isinstance(proc.channels, list): - for chan in proc.channels: - try: - val = read_channel(channels, chan, catch=False) - break - except EmptyChannelError: - pass - else: - return - 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) - - yield val - - def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]: if on: for chunk in iter: diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py new file mode 100644 index 000000000..256e62fd8 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -0,0 +1,413 @@ +import json +from collections import defaultdict, deque +from functools import partial +from typing import ( + Any, + Callable, + Iterator, + Literal, + Mapping, + Optional, + Sequence, + Union, + overload, +) +from uuid import UUID, uuid5 + +from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager +from langchain_core.runnables.config import ( + RunnableConfig, + merge_configs, + patch_config, +) + +from langgraph.channels.base import BaseChannel +from langgraph.channels.context import Context +from langgraph.channels.manager import ChannelsManager, create_checkpoint +from langgraph.checkpoint.base import Checkpoint, copy_checkpoint +from langgraph.constants import ( + CONFIG_KEY_READ, + CONFIG_KEY_SEND, + INTERRUPT, + TAG_HIDDEN, + TASKS, + Send, +) +from langgraph.errors import EmptyChannelError, InvalidUpdateError +from langgraph.managed.base import ManagedValueMapping, is_managed_value +from langgraph.pregel.io import read_channel, read_channels +from langgraph.pregel.log import logger +from langgraph.pregel.read import PregelNode +from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription + + +def should_interrupt( + checkpoint: Checkpoint, + interrupt_nodes: Union[All, Sequence[str]], + snapshot_channels: Sequence[str], + tasks: list[PregelExecutableTask], +) -> bool: + version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) + null_version = version_type() + # defaultdicts are mutated on access :( so we need to copy + seen = checkpoint["versions_seen"].copy()[INTERRUPT] + return ( + # interrupt if any channel has been updated since last interrupt + any( + version > seen.get(chan, null_version) + for chan, version in checkpoint["channel_versions"].items() + ) + # and any triggered node is in interrupt_nodes list + and any( + task.name + for task in tasks + if ( + (not task.config or TAG_HIDDEN not in task.config.get("tags")) + if interrupt_nodes == "*" + else task.name in interrupt_nodes + ) + ) + ) + + +def local_read( + checkpoint: Checkpoint, + channels: Mapping[str, BaseChannel], + writes: Sequence[tuple[str, Any]], + config: RunnableConfig, + select: Union[list[str], str], + fresh: bool = False, +) -> Union[dict[str, Any], Any]: + if fresh: + checkpoint = create_checkpoint(checkpoint, channels, -1) + context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)} + with ChannelsManager( + {k: v for k, v in channels.items() if k not in context_channels}, + checkpoint, + config, + ) as channels: + all_channels = {**channels, **context_channels} + apply_writes(copy_checkpoint(checkpoint), all_channels, writes, None) + return read_channels(all_channels, select) + else: + return read_channels(channels, select) + + +def local_write( + commit: Callable[[Sequence[tuple[str, Any]]], None], + processes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + writes: Sequence[tuple[str, Any]], +) -> None: + for chan, value in writes: + if chan == TASKS: + if not isinstance(value, Send): + raise InvalidUpdateError( + f"Invalid packet type, expected Packet, got {value}" + ) + if value.node not in processes: + raise InvalidUpdateError(f"Invalid node name {value.node} in packet") + elif chan not in channels: + logger.warning(f"Skipping write for channel '{chan}' which has no readers") + commit(writes) + + +def increment(current: Optional[int], channel: BaseChannel) -> int: + return current + 1 if current is not None else 1 + + +def apply_writes( + checkpoint: Checkpoint, + channels: Mapping[str, BaseChannel], + pending_writes: Sequence[tuple[str, Any]], + get_next_version: Optional[Callable[[int, BaseChannel], int]], +) -> None: + if checkpoint["pending_sends"]: + checkpoint["pending_sends"].clear() + + pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) + # Group writes by channel + for chan, val in pending_writes: + if chan == TASKS: + checkpoint["pending_sends"].append(val) + else: + pending_writes_by_channel[chan].append(val) + + # Find the highest version of all channels + if checkpoint["channel_versions"]: + max_version = max(checkpoint["channel_versions"].values()) + else: + max_version = None + + updated_channels: set[str] = set() + # Apply writes to channels + for chan, vals in pending_writes_by_channel.items(): + if chan in channels: + try: + updated = channels[chan].update(vals) + except InvalidUpdateError as e: + raise InvalidUpdateError( + f"Invalid update for channel {chan} with values {vals}" + ) from e + if updated and get_next_version is not None: + checkpoint["channel_versions"][chan] = get_next_version( + max_version, channels[chan] + ) + updated_channels.add(chan) + # Channels that weren't updated in this step are notified of a new step + for chan in channels: + if chan not in updated_channels: + if channels[chan].update([]) and get_next_version is not None: + checkpoint["channel_versions"][chan] = get_next_version( + max_version, channels[chan] + ) + + +@overload +def prepare_next_tasks( + checkpoint: Checkpoint, + processes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + managed: ManagedValueMapping, + config: RunnableConfig, + step: int, + for_execution: Literal[False], + get_next_version: Literal[None] = None, + manager: Literal[None] = None, +) -> tuple[Checkpoint, list[PregelTaskDescription]]: + ... + + +@overload +def prepare_next_tasks( + checkpoint: Checkpoint, + processes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + managed: ManagedValueMapping, + config: RunnableConfig, + step: int, + for_execution: Literal[True], + get_next_version: Callable[[int, BaseChannel], int], + manager: Union[None, ParentRunManager, AsyncParentRunManager], +) -> tuple[Checkpoint, list[PregelExecutableTask]]: + ... + + +def prepare_next_tasks( + checkpoint: Checkpoint, + processes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + managed: ManagedValueMapping, + config: RunnableConfig, + step: int, + *, + for_execution: bool, + get_next_version: Union[None, Callable[[int, BaseChannel], int]] = None, + manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, +) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: + checkpoint = copy_checkpoint(checkpoint) + tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] + # Consume pending packets + for packet in checkpoint["pending_sends"]: + if not isinstance(packet, Send): + logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends") + continue + if for_execution: + proc = processes[packet.node] + if node := proc.get_node(): + triggers = [TASKS] + metadata = { + "langgraph_step": step, + "langgraph_node": packet.node, + "langgraph_triggers": triggers, + "langgraph_task_idx": len(tasks), + } + task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) + writes = deque() + tasks.append( + PregelExecutableTask( + packet.node, + packet.arg, + node, + writes, + patch_config( + merge_configs( + config, + processes[packet.node].config, + {"metadata": metadata}, + ), + run_name=packet.node, + callbacks=( + manager.get_child(f"graph:step:{step}") + if manager + else None + ), + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, writes.extend, processes, channels + ), + CONFIG_KEY_READ: partial( + local_read, checkpoint, channels, writes, config + ), + }, + ), + triggers, + proc.retry_policy, + task_id, + ) + ) + else: + tasks.append(PregelTaskDescription(packet.node, packet.arg)) + if for_execution: + checkpoint["pending_sends"].clear() + # Collect channels to consume + channels_to_consume = set() + # Check if any processes should be run in next step + # If so, prepare the values to be passed to them + version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) + null_version = version_type() + if null_version is None: + return checkpoint, tasks + for name, proc in processes.items(): + seen = checkpoint["versions_seen"][name] + # If any of the channels read by this process were updated + if triggers := sorted( + chan + for chan in proc.triggers + if not isinstance( + read_channel(channels, chan, return_exception=True), EmptyChannelError + ) + and checkpoint["channel_versions"].get(chan, null_version) + > seen.get(chan, null_version) + ): + channels_to_consume.update(triggers) + try: + val = next(_proc_input(step, name, proc, managed, channels)) + except StopIteration: + continue + + # update seen versions + if for_execution: + seen.update( + { + chan: checkpoint["channel_versions"][chan] + for chan in proc.triggers + if chan in checkpoint["channel_versions"] + } + ) + + if for_execution: + if node := proc.get_node(): + metadata = { + "langgraph_step": step, + "langgraph_node": name, + "langgraph_triggers": triggers, + "langgraph_task_idx": len(tasks), + } + task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) + writes = deque() + tasks.append( + PregelExecutableTask( + name, + val, + node, + writes, + patch_config( + merge_configs( + config, + proc.config, + {"metadata": metadata}, + ), + run_name=name, + callbacks=( + manager.get_child(f"graph:step:{step}") + if manager + else None + ), + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, writes.extend, processes, channels + ), + CONFIG_KEY_READ: partial( + local_read, + checkpoint, + channels, + writes, + config, + ), + }, + ), + triggers, + proc.retry_policy, + task_id, + ) + ) + else: + tasks.append(PregelTaskDescription(name, val)) + # Find the highest version of all channels + if checkpoint["channel_versions"]: + max_version = max(checkpoint["channel_versions"].values()) + else: + max_version = None + # Consume all channels that were read + if for_execution: + for chan in channels_to_consume: + if channels[chan].consume(): + checkpoint["channel_versions"][chan] = get_next_version( + max_version, channels[chan] + ) + return checkpoint, tasks + + +def _proc_input( + step: int, + name: str, + proc: PregelNode, + managed: ManagedValueMapping, + channels: Mapping[str, BaseChannel], +) -> Iterator[Any]: + # If all trigger channels subscribed by this process are not empty + # then invoke the process with the values of all non-empty channels + if isinstance(proc.channels, dict): + try: + val: dict = { + k: read_channel( + channels, + chan, + catch=chan not in proc.triggers, + ) + for k, chan in proc.channels.items() + if isinstance(chan, str) + } + + managed_values = {} + for key, chan in proc.channels.items(): + if is_managed_value(chan): + managed_values[key] = managed[key]( + step, PregelTaskDescription(name, val) + ) + + val.update(managed_values) + except EmptyChannelError: + return + elif isinstance(proc.channels, list): + for chan in proc.channels: + try: + val = read_channel(channels, chan, catch=False) + break + except EmptyChannelError: + pass + else: + return + 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) + + yield val diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 389c59e81..5a317bb26 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -110,7 +110,7 @@ def map_debug_task_results( def map_debug_checkpoint( - step: int, + step: int, # TODO remove arg, it's in metadata config: RunnableConfig, channels: Mapping[str, BaseChannel], stream_channels: Union[str, Sequence[str]], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py new file mode 100644 index 000000000..b601b6446 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -0,0 +1,282 @@ +from collections import deque +from contextlib import ExitStack +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + ContextManager, + List, + Literal, + Mapping, + Optional, + Self, + Sequence, + Type, +) + +from langchain_core.runnables import RunnableConfig + +from langgraph.channels.base import BaseChannel +from langgraph.channels.manager import ChannelsManager, create_checkpoint +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + Checkpoint, + CheckpointMetadata, + CheckpointTuple, + PendingWrite, + copy_checkpoint, + empty_checkpoint, +) +from langgraph.constants import INTERRUPT +from langgraph.managed.base import ManagedValueMapping, ManagedValuesManager +from langgraph.pregel.algo import ( + apply_writes, + increment, + prepare_next_tasks, + should_interrupt, +) +from langgraph.pregel.executor import BackgroundExecutor, Submit +from langgraph.pregel.io import map_input +from langgraph.pregel.types import PregelExecutableTask + +if TYPE_CHECKING: + from langgraph.pregel import Pregel + + +INPUT_DONE = object() + + +class PregelLoop(ContextManager): + config: RunnableConfig + checkpoint: Checkpoint + checkpoint_metadata: CheckpointMetadata + checkpoint_pending_writes: Optional[List[PendingWrite]] + + submit: Submit + channels: Mapping[str, BaseChannel] + managed: ManagedValueMapping + + status: Literal[ + "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" + ] + tasks: Sequence[PregelExecutableTask] + + def __init__( + self, + input: Optional[Any], + *, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", + ) -> None: + self.stack = ExitStack() + self.input = input + self.config = config + self.checkpointer = checkpointer + self.get_next_version = ( + checkpointer.get_next_version if checkpointer else increment + ) + self.graph = graph + # TODO if managed values no longer needs graph we can replace with + # managed_specs, channel_specs + + def __enter__(self) -> Self: + saved = ( + self.checkpointer.get_tuple(self.config) if self.checkpointer else None + ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) + self.config = { + **self.config, + **saved.config, + "configurable": { + **self.config.get("configurable", {}), + **saved.config.get("configurable", {}), + }, + } + self.checkpoint = saved.checkpoint + self.checkpoint_metadata = saved.metadata + self.checkpoint_pending_writes = saved.pending_writes + + self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) + self.channels = self.stack.enter_context( + ChannelsManager(self.graph.channels, self.checkpoint, self.config) + ) + self.managed = self.stack.enter_context( + ManagedValuesManager( + self.graph.managed_values_dict, self.config, self.graph + ) + ) + self.status = "pending" + + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + del self.graph + return self.stack.__exit__(exc_type, exc_value, traceback) + + def tick( + self, + *, + interrupt_after: Optional[Sequence[str]] = None, + interrupt_before: Optional[Sequence[str]] = None, + ) -> bool: + print("ticking", self.status, self.checkpoint_metadata["step"]) + if self.status != "pending": + raise RuntimeError(f"Cannot tick when status is {self.status}") + if self.input is not INPUT_DONE: + self.first() + elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( + self.tasks + ): + # all tasks have finished + apply_writes( + self.checkpoint, + self.channels, + ((k, v) for _, k, v in self.checkpoint_pending_writes), + self.get_next_version, + ) + # clear pending writes + self.checkpoint_pending_writes.clear() + # save checkpoint + self.put_checkpoint({"source": "loop", "writes": None}) # TODO + # after execution, check if we should interrupt + if should_interrupt( + self.checkpoint, + interrupt_after, + self.graph.stream_channels_list, + self.tasks, + ): + self.status = "interrupt_after" + return False + else: + return False + + # check if iteration limit is reached + if self.checkpoint_metadata["step"] >= self.config["recursion_limit"]: + self.status = "out_of_steps" + return False + + # prepare next tasks + self.checkpoint, self.tasks = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.checkpoint_metadata["step"], + for_execution=True, + get_next_version=self.get_next_version, + ) + + # if no more tasks, we're done + if not self.tasks: + self.status = "done" + return False + + # TODO how to make this work for both + # - online case: we should schedule remaining tasks + # - offline case: we should just bail, as other tasks were scheduled before + # assign pending writes to tasks + # if self.checkpoint_pending_writes: + # # if there are pending writes from a previous loop, apply them + # for tid, k, v in self.checkpoint_pending_writes: + # if task := next((t for t in self.tasks if t.id == tid), None): + # task.writes.append((k, v)) + # if + + # before execution, check if we should interrupt + if should_interrupt( + self.checkpoint, + interrupt_before, + self.graph.stream_channels_list, + self.tasks, + ): + self.status = "interrupt_before" + return False + + return True + + def first(self) -> None: + # map inputs to channel updates + if input_writes := deque(map_input(self.graph.input_channels, self.input)): + # discard any unfinished tasks from previous checkpoint + self.checkpoint, _ = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + -1, + for_execution=True, + get_next_version=self.get_next_version, + ) + # apply input writes + apply_writes( + self.checkpoint, + self.channels, + input_writes, + self.get_next_version, + ) + # save input checkpoint + self.put_checkpoint({"source": "input", "writes": self.input}) + else: + # no input is taken as signal to proceed past previous interrupt + self.checkpoint = copy_checkpoint(self.checkpoint) + for k in self.channels: + if k in self.checkpoint["channel_versions"]: + version = self.checkpoint["channel_versions"][k] + self.checkpoint["versions_seen"][INTERRUPT][k] = version + # done with input + self.input = INPUT_DONE + + def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) + if self.checkpointer is not None: + self.submit( + self.checkpointer.put_writes, + { + **self.config, + "configurable": { + **self.config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + }, + writes, + task_id, + ) + + def put_checkpoint( + self, + metadata: CheckpointMetadata, + ) -> None: + # increment step + self.checkpoint_metadata = { + **metadata, + "step": self.checkpoint_metadata["step"] + 1, + } + # bail if no checkpointer + if self.checkpointer is None: + return + # create new checkpoint + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels, self.checkpoint_metadata["step"] + ) + # save it, without blocking + self.submit( + self.checkpointer.put, + self.config, + copy_checkpoint(self.checkpoint), + self.checkpoint_metadata, + ) + # update checkpoint config + self.config = { + **self.config, + "configurable": { + **self.config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + } diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 9a1afb596..19d3c3301 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -88,3 +88,12 @@ class StateSnapshot(NamedTuple): All = Literal["*"] + +StreamMode = Literal["values", "updates", "debug"] +"""How the stream method should emit outputs. + +- 'values': Emit all values of the state for each step. +- 'updates': Emit only the node name(s) and updates + that were returned by the node(s) **after** each step. +- 'debug': Emit debug events for each step. +""" diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 264f30db6..6053a4e5b 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -60,7 +60,7 @@ def test_no_modifier(): model = FakeToolCallingModel() agent = create_react_agent(model, []) inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}) + response = agent.invoke({"messages": inputs}, debug=True) expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} assert response == expected_response diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index df6ba34ca..805794bdf 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -520,8 +520,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert app.invoke(2) == 4 - assert app.invoke(2, input_keys="inbox") == 3 - with pytest.raises(GraphRecursionError): app.invoke(2, {"recursion_limit": 1}) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4c166cdc5..1cf5c35d4 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -623,8 +623,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert await app.ainvoke(2) == 4 - assert await app.ainvoke(2, input_keys="inbox") == 3 - with pytest.raises(GraphRecursionError): await app.ainvoke(2, {"recursion_limit": 1}) From 60f7a7d593d88bc55d0694dda654745a180a81d0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 13:23:17 -0700 Subject: [PATCH 02/15] Lint --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index b601b6446..f217e201a 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -9,12 +9,12 @@ from typing import ( Literal, Mapping, Optional, - Self, Sequence, Type, ) from langchain_core.runnables import RunnableConfig +from typing_extensions import Self from langgraph.channels.base import BaseChannel from langgraph.channels.manager import ChannelsManager, create_checkpoint From b3ee7288396d58ea7d18c74ef793c2c7b8b2d136 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 15:48:28 -0700 Subject: [PATCH 03/15] Pregel.stream passing all tests --- libs/langgraph/langgraph/pregel/__init__.py | 88 +++++------- libs/langgraph/langgraph/pregel/algo.py | 1 - libs/langgraph/langgraph/pregel/debug.py | 4 + libs/langgraph/langgraph/pregel/loop.py | 149 ++++++++++++-------- libs/langgraph/tests/test_prebuilt.py | 118 +++++++++++++++- libs/langgraph/tests/test_pregel.py | 19 ++- libs/langgraph/tests/test_pregel_async.py | 12 ++ 7 files changed, 272 insertions(+), 119 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 5b1d5b9f1..b61d05982 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -849,38 +849,34 @@ class Pregel( # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick( - interrupt_before=interrupt_before, interrupt_after=interrupt_after + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, ): - # print debug output + # debug flag if self.debug: print_step_checkpoint( loop.checkpoint_metadata, loop.channels, self.stream_channels_list, ) - # emit debug output - if self.checkpointer and "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - loop.checkpoint_metadata["step"], - loop.config, - loop.channels, - self.stream_channels_asis, - loop.checkpoint_metadata, - ), - ) - + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # debug flag if debug: - print_step_tasks(loop.checkpoint_metadata["step"], loop.tasks) + print_step_tasks(loop.step, loop.tasks) + # TODO move to tick() ? if "debug" in stream_modes: yield from _with_mode( "debug", isinstance(stream_mode, list), - map_debug_tasks( - loop.checkpoint_metadata["step"], loop.tasks - ), + map_debug_tasks(loop.step, loop.tasks), ) # execute tasks, and wait for one to fail or all to finish. @@ -931,7 +927,7 @@ class Pregel( "debug", isinstance(stream_mode, list), map_debug_task_results( - loop.checkpoint_metadata["step"], + loop.step, [task], self.stream_channels_list, ), @@ -941,38 +937,31 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, loop.checkpoint_metadata["step"]) + _panic_or_proceed(done, inflight, loop.step) # don't keep futures around in memory longer than needed del done, inflight, futures - - # combine pending writes from all tasks - pending_writes = deque[tuple[str, Any]]() - for task in loop.tasks: - pending_writes.extend(task.writes) - + # debug flag if debug: print_step_writes( - loop.checkpoint_metadata["step"], - pending_writes, + loop.step, + [w for t in loop.tasks for w in t.writes], self.stream_channels_list, ) - - # yield values output - if "values" in stream_modes: - yield from _with_mode( - "values", - isinstance(stream_mode, list), - map_output_values( - output_keys, pending_writes, loop.channels - ), - ) + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # handle exit if loop.status == "out_of_steps": raise GraphRecursionError( - f"Recursion limit of {config['recursion_limit']} reached" + f"Recursion limit of {config['recursion_limit']} reached " "without hitting a stop condition. You can increase the " "limit by setting the `recursion_limit` config key." ) - # set final channel values as run output run_manager.on_chain_end(read_channels(loop.channels, output_keys)) except BaseException as e: @@ -1142,6 +1131,7 @@ class Pregel( ) def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: + print(metadata) nonlocal checkpoint, checkpoint_config, channels if self.checkpointer is None: @@ -1268,12 +1258,7 @@ class Pregel( break # before execution, check if we should interrupt - if should_interrupt( - checkpoint, - interrupt_before, - self.stream_channels_list, - next_tasks, - ): + if should_interrupt(checkpoint, interrupt_before, next_tasks): break else: checkpoint = next_checkpoint @@ -1405,12 +1390,7 @@ class Pregel( yield chunk # after execution, check if we should interrupt - if should_interrupt( - checkpoint, - interrupt_after, - self.stream_channels_list, - next_tasks, - ): + if should_interrupt(checkpoint, interrupt_after, next_tasks): break else: raise GraphRecursionError( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 256e62fd8..46950b9f8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -44,7 +44,6 @@ from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescript def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], - snapshot_channels: Sequence[str], tasks: list[PregelExecutableTask], ) -> bool: version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 5a317bb26..33b7f440e 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -70,6 +70,8 @@ def map_debug_tasks( if config is not None and TAG_HIDDEN in config.get("tags", []): continue + print("map_debug_tasks", json.dumps((name, step, config["metadata"]))) + yield { "type": "task", "timestamp": ts, @@ -95,6 +97,8 @@ def map_debug_task_results( if config is not None and TAG_HIDDEN in config.get("tags", []): continue + print("map_debug_tasks_r", json.dumps((name, step, config["metadata"]))) + yield { "type": "task_result", "timestamp": ts, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f217e201a..d618ffa24 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -4,13 +4,17 @@ from types import TracebackType from typing import ( TYPE_CHECKING, Any, + Callable, ContextManager, List, Literal, Mapping, Optional, Sequence, + Tuple, Type, + TypeVar, + Union, ) from langchain_core.runnables import RunnableConfig @@ -35,31 +39,39 @@ from langgraph.pregel.algo import ( prepare_next_tasks, should_interrupt, ) +from langgraph.pregel.debug import map_debug_checkpoint from langgraph.pregel.executor import BackgroundExecutor, Submit -from langgraph.pregel.io import map_input +from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single from langgraph.pregel.types import PregelExecutableTask if TYPE_CHECKING: from langgraph.pregel import Pregel +V = TypeVar("V") INPUT_DONE = object() class PregelLoop(ContextManager): + input: Optional[Any] config: RunnableConfig - checkpoint: Checkpoint - checkpoint_metadata: CheckpointMetadata - checkpoint_pending_writes: Optional[List[PendingWrite]] + checkpointer: Optional[BaseCheckpointSaver] + get_next_version: Callable[[Optional[V]], V] + graph: "Pregel" submit: Submit channels: Mapping[str, BaseChannel] managed: ManagedValueMapping + checkpoint: Checkpoint + checkpoint_config: RunnableConfig + checkpoint_metadata: CheckpointMetadata + checkpoint_pending_writes: Optional[List[PendingWrite]] status: Literal[ "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" ] tasks: Sequence[PregelExecutableTask] + stream: deque[Tuple[str, Any]] def __init__( self, @@ -69,6 +81,7 @@ class PregelLoop(ContextManager): checkpointer: Optional[BaseCheckpointSaver], graph: "Pregel", ) -> None: + self.stream = deque() self.stack = ExitStack() self.input = input self.config = config @@ -84,7 +97,7 @@ class PregelLoop(ContextManager): saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) - self.config = { + self.checkpoint_config = { **self.config, **saved.config, "configurable": { @@ -106,6 +119,7 @@ class PregelLoop(ContextManager): ) ) self.status = "pending" + self.step = self.checkpoint_metadata["step"] + 1 return self @@ -121,10 +135,11 @@ class PregelLoop(ContextManager): def tick( self, *, + output_keys: Union[str, Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, interrupt_before: Optional[Sequence[str]] = None, ) -> bool: - print("ticking", self.status, self.checkpoint_metadata["step"]) + print("tick", self.status, self.step) if self.status != "pending": raise RuntimeError(f"Cannot tick when status is {self.status}") if self.input is not INPUT_DONE: @@ -132,42 +147,57 @@ class PregelLoop(ContextManager): elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( self.tasks ): + # assign writes to tasks, apply them in order + grouped: dict[str, list[tuple[str, Any]]] = {} + for tid, k, v in self.checkpoint_pending_writes: + grouped.setdefault(tid, []).append((k, v)) + writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] # all tasks have finished apply_writes( self.checkpoint, self.channels, - ((k, v) for _, k, v in self.checkpoint_pending_writes), + writes, self.get_next_version, ) + # produce values output + self.stream.extend( + ("values", v) + for v in map_output_values(output_keys, writes, self.channels) + ) # clear pending writes self.checkpoint_pending_writes.clear() # save checkpoint - self.put_checkpoint({"source": "loop", "writes": None}) # TODO + self.put_checkpoint( + { + "source": "loop", + "writes": single( + map_output_updates(output_keys, self.tasks) + if self.graph.stream_mode == "updates" + else map_output_values(output_keys, writes, self.channels) + ), + } + ) # after execution, check if we should interrupt - if should_interrupt( - self.checkpoint, - interrupt_after, - self.graph.stream_channels_list, - self.tasks, - ): + if should_interrupt(self.checkpoint, interrupt_after, self.tasks): self.status = "interrupt_after" return False else: return False # check if iteration limit is reached - if self.checkpoint_metadata["step"] >= self.config["recursion_limit"]: + if self.step > self.config["recursion_limit"]: self.status = "out_of_steps" return False # prepare next tasks + prev_checkpoint = self.checkpoint self.checkpoint, self.tasks = prepare_next_tasks( self.checkpoint, self.graph.nodes, self.channels, self.managed, self.config, - self.checkpoint_metadata["step"], + self.step, for_execution=True, get_next_version=self.get_next_version, ) @@ -180,21 +210,14 @@ class PregelLoop(ContextManager): # TODO how to make this work for both # - online case: we should schedule remaining tasks # - offline case: we should just bail, as other tasks were scheduled before - # assign pending writes to tasks - # if self.checkpoint_pending_writes: - # # if there are pending writes from a previous loop, apply them - # for tid, k, v in self.checkpoint_pending_writes: - # if task := next((t for t in self.tasks if t.id == tid), None): - # task.writes.append((k, v)) - # if + # if there are pending writes from a previous loop, apply them + if self.checkpoint_pending_writes: + for tid, k, v in self.checkpoint_pending_writes: + if task := next((t for t in self.tasks if t.id == tid), None): + task.writes.append((k, v)) # before execution, check if we should interrupt - if should_interrupt( - self.checkpoint, - interrupt_before, - self.graph.stream_channels_list, - self.tasks, - ): + if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" return False @@ -210,9 +233,10 @@ class PregelLoop(ContextManager): self.channels, self.managed, self.config, - -1, + self.step, for_execution=True, get_next_version=self.get_next_version, + # TODO missing run_manager ) # apply input writes apply_writes( @@ -239,9 +263,9 @@ class PregelLoop(ContextManager): self.submit( self.checkpointer.put_writes, { - **self.config, + **self.checkpoint_config, "configurable": { - **self.config["configurable"], + **self.checkpoint_config["configurable"], "thread_ts": self.checkpoint["id"], }, }, @@ -253,30 +277,39 @@ class PregelLoop(ContextManager): self, metadata: CheckpointMetadata, ) -> None: - # increment step - self.checkpoint_metadata = { - **metadata, - "step": self.checkpoint_metadata["step"] + 1, - } + # assign step + metadata["step"] = self.step # bail if no checkpointer - if self.checkpointer is None: - return - # create new checkpoint - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.checkpoint_metadata["step"] - ) - # save it, without blocking - self.submit( - self.checkpointer.put, - self.config, - copy_checkpoint(self.checkpoint), - self.checkpoint_metadata, - ) - # update checkpoint config - self.config = { - **self.config, - "configurable": { - **self.config["configurable"], - "thread_ts": self.checkpoint["id"], - }, - } + if self.checkpointer is not None: + # create new checkpoint + self.checkpoint_metadata = metadata + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels, self.step + ) + # save it, without blocking + self.submit( + self.checkpointer.put, + self.checkpoint_config, + copy_checkpoint(self.checkpoint), + self.checkpoint_metadata, + ) + self.checkpoint_config = { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + } + # produce debug output + self.stream.extend( + ("debug", v) + for v in map_debug_checkpoint( + self.step, + self.checkpoint_config, + self.channels, + self.graph.stream_channels_asis, + self.checkpoint_metadata, + ) + ) + # increment step + self.step += 1 diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 6053a4e5b..7b56eba46 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,5 +1,8 @@ +from collections import defaultdict from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.sqlite import SqliteSaver import pytest from langchain_core.callbacks import ( CallbackManagerForLLMRun, @@ -27,6 +30,8 @@ from langgraph.prebuilt import ( ValidationNode, create_react_agent, ) +from tests.any_str import AnyStr +from tests.memory_assert import MemorySaverAssertImmutable class FakeToolCallingModel(BaseChatModel): @@ -56,14 +61,121 @@ class FakeToolCallingModel(BaseChatModel): return self -def test_no_modifier(): +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + None, + ], + ids=[ + "memory", + "none", + ], +) +def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]): model = FakeToolCallingModel() - agent = create_react_agent(model, []) + agent = create_react_agent(model, [], checkpointer=checkpointer) inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}, debug=True) + thread = {"configurable": {"thread_id": "123"}} + response = agent.invoke({"messages": inputs}, thread, debug=True) expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} assert response == expected_response + if checkpointer: + saved = checkpointer.get_tuple(thread) + assert saved is not None + assert saved.checkpoint == { + "v": 1, + "ts": AnyStr(), + "id": AnyStr(), + "channel_values": { + "messages": [ + HumanMessage(content="hi?", id=AnyStr()), + AIMessage(content="hi?", id="0"), + ], + "agent": "agent", + }, + "channel_versions": { + "__start__": 2, + "messages": 3, + "start:agent": 3, + "agent": 3, + }, + "versions_seen": defaultdict( + dict, + { + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + "tools": {}, + }, + ), + "pending_sends": [], + } + assert saved.metadata == { + "source": "loop", + "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, + "step": 1, + } + assert saved.pending_writes == [] + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + None, + ], + ids=[ + "memory", + "none", + ], +) +async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]): + model = FakeToolCallingModel() + agent = create_react_agent(model, [], checkpointer=checkpointer) + inputs = [HumanMessage("hi?")] + thread = {"configurable": {"thread_id": "123"}} + response = await agent.ainvoke({"messages": inputs}, thread, debug=True) + expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} + assert response == expected_response + + if checkpointer: + saved = await checkpointer.aget_tuple(thread) + assert saved is not None + assert saved.checkpoint == { + "v": 1, + "ts": AnyStr(), + "id": AnyStr(), + "channel_values": { + "messages": [ + HumanMessage(content="hi?", id=AnyStr()), + AIMessage(content="hi?", id="0"), + ], + "agent": "agent", + }, + "channel_versions": { + "__start__": 2, + "messages": 3, + "start:agent": 3, + "agent": 3, + }, + "versions_seen": defaultdict( + dict, + { + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + "tools": {}, + }, + ), + "pending_sends": [], + } + assert saved.metadata == { + "source": "loop", + "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, + "step": 1, + } + assert saved.pending_writes == [] + def test_passing_two_modifiers(): model = FakeToolCallingModel() diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 805794bdf..9b32fc149 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -521,7 +521,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert app.invoke(2) == 4 with pytest.raises(GraphRecursionError): - app.invoke(2, {"recursion_limit": 1}) + app.invoke(2, {"recursion_limit": 1}, debug=1) graph = Graph() graph.add_node("add_one", add_one) @@ -533,7 +533,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert gapp.invoke(2) == 4 - for step, values in enumerate(gapp.stream(2), start=1): + for step, values in enumerate(gapp.stream(2, debug=1), start=1): if step == 1: assert values == { "add_one": 3, @@ -6166,6 +6166,18 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value ⛰️", "market": "DE", } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, next=("tool_two_slow",), @@ -6875,7 +6887,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> }, ) - assert [c for c in app_w_interrupt.stream(None, config)] == [ + print("yo") + assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1cf5c35d4..0e927e518 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4735,6 +4735,18 @@ async def test_start_branch_then() -> None: "my_key": "value", "market": "DE", } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + ] assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value", "market": "DE"}, next=("tool_two_slow",), From 06f83710c8bce3437dde99f704202469a10f9f78 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 15:49:50 -0700 Subject: [PATCH 04/15] Lint --- libs/langgraph/langgraph/pregel/debug.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 33b7f440e..5a317bb26 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -70,8 +70,6 @@ def map_debug_tasks( if config is not None and TAG_HIDDEN in config.get("tags", []): continue - print("map_debug_tasks", json.dumps((name, step, config["metadata"]))) - yield { "type": "task", "timestamp": ts, @@ -97,8 +95,6 @@ def map_debug_task_results( if config is not None and TAG_HIDDEN in config.get("tags", []): continue - print("map_debug_tasks_r", json.dumps((name, step, config["metadata"]))) - yield { "type": "task_result", "timestamp": ts, From 73b4b6ec5b38526d5e993cdbe05bcfc4bb81ff9e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 15:50:45 -0700 Subject: [PATCH 05/15] Lint --- libs/langgraph/langgraph/pregel/debug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 5a317bb26..389c59e81 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -110,7 +110,7 @@ def map_debug_task_results( def map_debug_checkpoint( - step: int, # TODO remove arg, it's in metadata + step: int, config: RunnableConfig, channels: Mapping[str, BaseChannel], stream_channels: Union[str, Sequence[str]], From 909bf4433fc508ab4ffb918c3df502b854ccae4d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 15:51:37 -0700 Subject: [PATCH 06/15] Lint --- libs/langgraph/langgraph/pregel/loop.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index d618ffa24..003204b4e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -139,7 +139,6 @@ class PregelLoop(ContextManager): interrupt_after: Optional[Sequence[str]] = None, interrupt_before: Optional[Sequence[str]] = None, ) -> bool: - print("tick", self.status, self.step) if self.status != "pending": raise RuntimeError(f"Cannot tick when status is {self.status}") if self.input is not INPUT_DONE: From f974471f7d031164f9b09896bb8e92432a9a1498 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 15:52:20 -0700 Subject: [PATCH 07/15] Lint --- libs/langgraph/tests/test_pregel.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 9b32fc149..2646a64fd 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6887,7 +6887,6 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> }, ) - print("yo") assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, ] From 7ed5f9e4bcdb713f97e87f26fd8027596a717b29 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 17:28:44 -0700 Subject: [PATCH 08/15] Implement AsyncLoop --- libs/langgraph/langgraph/channels/manager.py | 2 + libs/langgraph/langgraph/checkpoint/base.py | 8 + libs/langgraph/langgraph/pregel/__init__.py | 343 +++---------- libs/langgraph/langgraph/pregel/executor.py | 5 +- libs/langgraph/langgraph/pregel/loop.py | 480 ++++++++++++------- libs/langgraph/tests/test_prebuilt.py | 5 +- 6 files changed, 384 insertions(+), 459 deletions(-) diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py index ce0d21189..524611e5a 100644 --- a/libs/langgraph/langgraph/channels/manager.py +++ b/libs/langgraph/langgraph/channels/manager.py @@ -61,4 +61,6 @@ def create_checkpoint( channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), + # checkpoints are saved only at the end of a step, ie. no current tasks + current_tasks={}, ) diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 2f73b52a6..8cb278f32 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -54,6 +54,10 @@ class CheckpointMetadata(TypedDict, total=False): """ +class TaskInfo(TypedDict): + status: Literal["scheduled", "success", "error"] + + class Checkpoint(TypedDict): """State snapshot at a given point in time.""" @@ -85,6 +89,8 @@ class Checkpoint(TypedDict): pending_sends: List[Send] """List of packets sent to nodes but not yet processed. Cleared by the next checkpoint.""" + current_tasks: Dict[str, TaskInfo] + """Map from task ID to task info.""" def empty_checkpoint() -> Checkpoint: @@ -96,6 +102,7 @@ def empty_checkpoint() -> Checkpoint: channel_versions={}, versions_seen=defaultdict(dict), pending_sends=[], + current_tasks={}, ) @@ -111,6 +118,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: {k: v.copy() for k, v in checkpoint["versions_seen"].items()}, ), pending_sends=checkpoint.get("pending_sends", []).copy(), + current_tasks=checkpoint.get("current_tasks", {}).copy(), ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index b61d05982..41b517e0a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -59,7 +59,6 @@ from langgraph.channels.manager import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, - CheckpointMetadata, copy_checkpoint, empty_checkpoint, ) @@ -77,28 +76,20 @@ from langgraph.managed.base import ( ) from langgraph.pregel.algo import ( apply_writes, - increment, local_read, prepare_next_tasks, - should_interrupt, ) from langgraph.pregel.debug import ( - map_debug_checkpoint, map_debug_task_results, - map_debug_tasks, print_step_checkpoint, print_step_tasks, print_step_writes, ) -from langgraph.pregel.executor import AsyncBackgroundExecutor from langgraph.pregel.io import ( - map_input, map_output_updates, - map_output_values, read_channels, - single, ) -from langgraph.pregel.loop import PregelLoop +from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry from langgraph.pregel.types import ( @@ -839,8 +830,7 @@ class Pregel( debug=debug, ) - # create channels from checkpoint - with PregelLoop( + with SyncPregelLoop( input, config=config, checkpointer=self.checkpointer, graph=self ) as loop: # Similarly to Bulk Synchronous Parallel / Pregel model @@ -852,6 +842,7 @@ class Pregel( output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, + manager=run_manager, ): # debug flag if self.debug: @@ -871,19 +862,16 @@ class Pregel( # debug flag if debug: print_step_tasks(loop.step, loop.tasks) - # TODO move to tick() ? - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_tasks(loop.step, loop.tasks), - ) # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - loop.submit(run_with_retry, task, self.retry_policy): task + loop.submit( + run_with_retry, + task, + self.retry_policy, + ): task for task in loop.tasks if not task.writes } @@ -1063,7 +1051,6 @@ class Pregel( None, ) try: - loop = asyncio.get_event_loop() if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -1085,211 +1072,59 @@ class Pregel( interrupt_after=interrupt_after, debug=debug, ) - # copy nodes to ignore mutations during execution - processes = {**self.nodes} - # get checkpoint from saver, or create an empty one - saved = ( - await self.checkpointer.aget_tuple(config) - if self.checkpointer - else None - ) - checkpoint = saved.checkpoint if saved else empty_checkpoint() - - # merge configurable fields with previous checkpoint config - checkpoint_config = config - if saved: - checkpoint_config = { - **config, - **saved.config, - "configurable": { - **config.get("configurable", {}), - **saved.config["configurable"], - }, - } - - start = saved.metadata.get("step", -2) + 1 if saved else -1 - # create channels from checkpoint - async with AsyncBackgroundExecutor() as submit, AsyncChannelsManager( - self.channels, checkpoint, config - ) as channels, AsyncManagedValuesManager( - self.managed_values_dict, config, self - ) as managed: - - def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None: - if self.checkpointer is not None: - submit( - self.checkpointer.aput_writes, - { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - }, - writes, - task_id, - ) - - def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: - print(metadata) - nonlocal checkpoint, checkpoint_config, channels - - if self.checkpointer is None: - return - if debug: - print_step_checkpoint( - metadata["step"], channels, self.stream_channels_list - ) - - # create new checkpoint - checkpoint = create_checkpoint( - checkpoint, channels, metadata["step"] - ) - # save it, without blocking - submit( - self.checkpointer.aput, - checkpoint_config, - copy_checkpoint(checkpoint), - metadata, - ) - - # update checkpoint config - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - } - # yield debug checkpoint event - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - metadata["step"], - checkpoint_config, - channels, - self.stream_channels_asis, - metadata, - ), - ) - - # map inputs to channel updates - if input_writes := deque(map_input(self.input_channels, input)): - # discard any unfinished tasks from previous checkpoint - checkpoint, _ = prepare_next_tasks( - checkpoint, - processes, - channels, - managed, - config, - -1, - for_execution=True, - get_next_version=( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - # apply input writes - apply_writes( - checkpoint, - channels, - input_writes, - ( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - # save input checkpoint - for chunk in put_checkpoint( - {"source": "input", "step": start, "writes": input} - ): - yield chunk - # increment start to 0 - start += 1 - else: - # no input is taken as signal to proceed past previous interrupt - checkpoint = copy_checkpoint(checkpoint) - for k in channels: - if k in checkpoint["channel_versions"]: - version = checkpoint["channel_versions"][k] - checkpoint["versions_seen"][INTERRUPT][k] = version - + async with AsyncPregelLoop( + input, config=config, checkpointer=self.checkpointer, graph=self + ) as loop: + aioloop = asyncio.get_event_loop() # 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, + # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, - # channel updates being applied only at the transition between steps - stop = start + config["recursion_limit"] + 1 - for step in range(start, stop): - next_checkpoint, next_tasks = prepare_next_tasks( - checkpoint, - processes, - channels, - managed, - config, - step, - for_execution=True, - manager=run_manager, - get_next_version=( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - - # assign pending writes to tasks - if saved and saved.pending_writes: - for task in next_tasks: - task.writes.extend( - (c, v) - for tid, c, v in saved.pending_writes - if tid == task.id - ) - - # if no more tasks, we're done - if not next_tasks: - if step == start: - raise ValueError("No tasks to run in graph.") - else: - break - - # before execution, check if we should interrupt - if should_interrupt(checkpoint, interrupt_before, next_tasks): - break - else: - checkpoint = next_checkpoint - + # with channel updates applied only at the transition between steps + while loop.tick( + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + manager=run_manager, + ): + # debug flag + if self.debug: + print_step_checkpoint( + loop.checkpoint_metadata, + loop.channels, + self.stream_channels_list, + ) + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # debug flag if debug: - print_step_tasks(step, next_tasks) - if "debug" in stream_modes: - for chunk in _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_tasks(step, next_tasks), - ): - yield chunk + print_step_tasks(loop.step, loop.tasks) # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - submit( + loop.submit( arun_with_retry, task, self.retry_policy, - do_stream, + stream=do_stream, __name__=task.name, __cancel_on_exit__=True, ): task - for task in next_tasks + for task in loop.tasks if not task.writes } end_time = ( - self.step_timeout + loop.time() if self.step_timeout else None + self.step_timeout + aioloop.time() + if self.step_timeout + else None ) if not futures: done, inflight = set(), set() @@ -1298,7 +1133,7 @@ class Pregel( futures, return_when=asyncio.FIRST_COMPLETED, timeout=( - max(0, end_time - loop.time()) if end_time else None + max(0, end_time - aioloop.time()) if end_time else None ), ) if not done: @@ -1307,13 +1142,12 @@ class Pregel( task = futures.pop(fut) if fut.exception() is not None: # we got an exception, break out of while loop - # exception will be handle in panic_or_proceed + # exception will be handled in panic_or_proceed futures.clear() else: - # save task writes to checkpointer, unless this - # is the single or last task in this step - if futures: - put_writes(task.id, task.writes) + print(loop.step, task.name, stream_modes) + # save task writes to checkpointer + loop.put_writes(task.id, task.writes) # yield updates output for the finished task if "updates" in stream_modes: for chunk in _with_mode( @@ -1327,7 +1161,9 @@ class Pregel( "debug", isinstance(stream_mode, list), map_debug_task_results( - step, [task], self.stream_channels_list + loop.step, + [task], + self.stream_channels_list, ), ): yield chunk @@ -1336,71 +1172,36 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, step, asyncio.TimeoutError) + _panic_or_proceed(done, inflight, loop.step, asyncio.TimeoutError) # don't keep futures around in memory longer than needed del done, inflight, futures - - # combine pending writes from all tasks - pending_writes = deque[tuple[str, Any]]() - for task in next_tasks: - pending_writes.extend(task.writes) - + # debug flag if debug: print_step_writes( - step, pending_writes, self.stream_channels_list + loop.step, + [w for t in loop.tasks for w in t.writes], + self.stream_channels_list, ) - - # apply writes to channels - apply_writes( - checkpoint, - channels, - pending_writes, - ( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - - # yield current values - if "values" in stream_modes: - for chunk in _with_mode( - "values", - isinstance(stream_mode, list), - map_output_values(output_keys, pending_writes, channels), - ): - yield chunk - - # save end of step checkpoint - for chunk in put_checkpoint( - { - "source": "loop", - "step": step, - "writes": ( - single(map_output_updates(output_keys, next_tasks)) - if self.stream_mode == "updates" - else single( - map_output_values( - output_keys, pending_writes, channels - ) - ) - ), - } - ): - yield chunk - - # after execution, check if we should interrupt - if should_interrupt(checkpoint, interrupt_after, next_tasks): - break - else: + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # handle exit + if loop.status == "out_of_steps": raise GraphRecursionError( - f"Recursion limit of {config['recursion_limit']} reached" - "without hitting a stop condition. You can increase the limit" - "by setting the `recursion_limit` config key." + f"Recursion limit of {config['recursion_limit']} reached " + "without hitting a stop condition. You can increase the " + "limit by setting the `recursion_limit` config key." ) # set final channel values as run output - await run_manager.on_chain_end(read_channels(channels, output_keys)) + await run_manager.on_chain_end( + read_channels(loop.channels, output_keys) + ) except BaseException as e: await asyncio.shield(run_manager.on_chain_error(e)) raise diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 4d43aab60..352b8cf51 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -6,6 +6,7 @@ from contextvars import copy_context from types import TracebackType from typing import ( AsyncContextManager, + Awaitable, Callable, Iterator, Optional, @@ -78,7 +79,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): def submit( self, - fn: Callable[P, T], + fn: Callable[P, Awaitable[T]], *args: P.args, __name__: Optional[str] = None, __cancel_on_exit__: bool = False, @@ -101,7 +102,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): else: self.tasks.pop(task) - async def __aenter__(self) -> Submit: + async def __aenter__(self) -> "submit": return self.submit async def exit(self) -> None: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 003204b4e..b8c3fdc4e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,9 +1,11 @@ +import asyncio from collections import deque -from contextlib import ExitStack +from contextlib import AsyncExitStack, ExitStack from types import TracebackType from typing import ( TYPE_CHECKING, Any, + AsyncContextManager, Callable, ContextManager, List, @@ -17,11 +19,16 @@ from typing import ( Union, ) +from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager from langchain_core.runnables import RunnableConfig from typing_extensions import Self from langgraph.channels.base import BaseChannel -from langgraph.channels.manager import ChannelsManager, create_checkpoint +from langgraph.channels.manager import ( + AsyncChannelsManager, + ChannelsManager, + create_checkpoint, +) from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -32,15 +39,23 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.constants import INTERRUPT -from langgraph.managed.base import ManagedValueMapping, ManagedValuesManager +from langgraph.managed.base import ( + AsyncManagedValuesManager, + ManagedValueMapping, + ManagedValuesManager, +) from langgraph.pregel.algo import ( apply_writes, increment, prepare_next_tasks, should_interrupt, ) -from langgraph.pregel.debug import map_debug_checkpoint -from langgraph.pregel.executor import BackgroundExecutor, Submit +from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks +from langgraph.pregel.executor import ( + AsyncBackgroundExecutor, + BackgroundExecutor, + Submit, +) from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single from langgraph.pregel.types import PregelExecutableTask @@ -52,11 +67,17 @@ V = TypeVar("V") INPUT_DONE = object() -class PregelLoop(ContextManager): +class PregelLoop: input: Optional[Any] config: RunnableConfig checkpointer: Optional[BaseCheckpointSaver] - get_next_version: Callable[[Optional[V]], V] + checkpointer_get_next_version: Callable[[Optional[V]], V] + checkpointer_put_writes: Optional[ + Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], ...] + ] + checkpointer_put: Optional[ + Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], ...] + ] graph: "Pregel" submit: Submit @@ -73,6 +94,203 @@ class PregelLoop(ContextManager): tasks: Sequence[PregelExecutableTask] stream: deque[Tuple[str, Any]] + # public + + def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: + """Mark tasks as scheduled, to be used by queue-based executors.""" + raise NotImplementedError + + def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + """Put writes for a task, to be read by the next tick.""" + self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) + if self.checkpointer_put_writes is not None: + self.submit( + self.checkpointer_put_writes, + { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + }, + writes, + task_id, + ) + + def tick( + self, + *, + output_keys: Union[str, Sequence[str]] = None, + interrupt_after: Optional[Sequence[str]] = None, + interrupt_before: Optional[Sequence[str]] = None, + manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, + ) -> bool: + """Execute a single iteration of the Pregel loop. + Returns True if more iterations are needed.""" + + if self.status != "pending": + raise RuntimeError("Cannot tick when status is no longer 'pending'") + + if self.input is not INPUT_DONE: + self._first() + elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( + self.tasks + ): + # assign writes to tasks, apply them in order + grouped: dict[str, list[tuple[str, Any]]] = {} + for tid, k, v in self.checkpoint_pending_writes: + grouped.setdefault(tid, []).append((k, v)) + writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] + # all tasks have finished + apply_writes( + self.checkpoint, + self.channels, + writes, + self.checkpointer_get_next_version, + ) + # produce values output + self.stream.extend( + ("values", v) + for v in map_output_values(output_keys, writes, self.channels) + ) + # clear pending writes + self.checkpoint_pending_writes.clear() + # save checkpoint + self._put_checkpoint( + { + "source": "loop", + "writes": single( + map_output_updates(output_keys, self.tasks) + if self.graph.stream_mode == "updates" + else map_output_values(output_keys, writes, self.channels) + ), + } + ) + # after execution, check if we should interrupt + if should_interrupt(self.checkpoint, interrupt_after, self.tasks): + self.status = "interrupt_after" + return False + else: + return False + + # check if iteration limit is reached + if self.step > self.config["recursion_limit"]: + self.status = "out_of_steps" + return False + + # prepare next tasks + prev_checkpoint = self.checkpoint + self.checkpoint, self.tasks = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + get_next_version=self.checkpointer_get_next_version, + manager=manager, + ) + + # if no more tasks, we're done + if not self.tasks: + self.status = "done" + return False + + # if there are pending writes from a previous loop, apply them + if self.checkpoint_pending_writes: + for tid, k, v in self.checkpoint_pending_writes: + if task := next((t for t in self.tasks if t.id == tid), None): + task.writes.append((k, v)) + + # before execution, check if we should interrupt + if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): + self.status = "interrupt_before" + return False + + # produce debug output + self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks)) + + return True + + # private + + def _first(self) -> None: + # map inputs to channel updates + if input_writes := deque(map_input(self.graph.input_channels, self.input)): + # discard any unfinished tasks from previous checkpoint + self.checkpoint, _ = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + get_next_version=self.checkpointer_get_next_version, + ) + # apply input writes + apply_writes( + self.checkpoint, + self.channels, + input_writes, + self.checkpointer_get_next_version, + ) + # save input checkpoint + self._put_checkpoint({"source": "input", "writes": self.input}) + else: + # no input is taken as signal to proceed past previous interrupt + self.checkpoint = copy_checkpoint(self.checkpoint) + for k in self.channels: + if k in self.checkpoint["channel_versions"]: + version = self.checkpoint["channel_versions"][k] + self.checkpoint["versions_seen"][INTERRUPT][k] = version + # done with input + self.input = INPUT_DONE + + def _put_checkpoint( + self, + metadata: CheckpointMetadata, + ) -> None: + # assign step + metadata["step"] = self.step + # bail if no checkpointer + if self.checkpointer_put is not None: + # create new checkpoint + self.checkpoint_metadata = metadata + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels, self.step + ) + # save it, without blocking + self.submit( + self.checkpointer_put, + self.checkpoint_config, + copy_checkpoint(self.checkpoint), + self.checkpoint_metadata, + ) + self.checkpoint_config = { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + } + # produce debug output + self.stream.extend( + ("debug", v) + for v in map_debug_checkpoint( + self.step, + self.checkpoint_config, + self.channels, + self.graph.stream_channels_asis, + self.checkpoint_metadata, + ) + ) + # increment step + self.step += 1 + + +class SyncPregelLoop(PregelLoop, ContextManager): def __init__( self, input: Optional[Any], @@ -86,13 +304,17 @@ class PregelLoop(ContextManager): self.input = input self.config = config self.checkpointer = checkpointer - self.get_next_version = ( + self.checkpointer_get_next_version = ( checkpointer.get_next_version if checkpointer else increment ) + self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None + self.checkpointer_put = checkpointer.put if checkpointer else None self.graph = graph # TODO if managed values no longer needs graph we can replace with # managed_specs, channel_specs + # context manager + def __enter__(self) -> Self: saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None @@ -132,183 +354,73 @@ class PregelLoop(ContextManager): del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) - def tick( + +class AsyncPregelLoop(PregelLoop, AsyncContextManager): + def __init__( self, + input: Optional[Any], *, - output_keys: Union[str, Sequence[str]] = None, - interrupt_after: Optional[Sequence[str]] = None, - interrupt_before: Optional[Sequence[str]] = None, - ) -> bool: - if self.status != "pending": - raise RuntimeError(f"Cannot tick when status is {self.status}") - if self.input is not INPUT_DONE: - self.first() - elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( - self.tasks - ): - # assign writes to tasks, apply them in order - grouped: dict[str, list[tuple[str, Any]]] = {} - for tid, k, v in self.checkpoint_pending_writes: - grouped.setdefault(tid, []).append((k, v)) - writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] - # all tasks have finished - apply_writes( - self.checkpoint, - self.channels, - writes, - self.get_next_version, - ) - # produce values output - self.stream.extend( - ("values", v) - for v in map_output_values(output_keys, writes, self.channels) - ) - # clear pending writes - self.checkpoint_pending_writes.clear() - # save checkpoint - self.put_checkpoint( - { - "source": "loop", - "writes": single( - map_output_updates(output_keys, self.tasks) - if self.graph.stream_mode == "updates" - else map_output_values(output_keys, writes, self.channels) - ), - } - ) - # after execution, check if we should interrupt - if should_interrupt(self.checkpoint, interrupt_after, self.tasks): - self.status = "interrupt_after" - return False - else: - return False - - # check if iteration limit is reached - if self.step > self.config["recursion_limit"]: - self.status = "out_of_steps" - return False - - # prepare next tasks - prev_checkpoint = self.checkpoint - self.checkpoint, self.tasks = prepare_next_tasks( - self.checkpoint, - self.graph.nodes, - self.channels, - self.managed, - self.config, - self.step, - for_execution=True, - get_next_version=self.get_next_version, - ) - - # if no more tasks, we're done - if not self.tasks: - self.status = "done" - return False - - # TODO how to make this work for both - # - online case: we should schedule remaining tasks - # - offline case: we should just bail, as other tasks were scheduled before - # if there are pending writes from a previous loop, apply them - if self.checkpoint_pending_writes: - for tid, k, v in self.checkpoint_pending_writes: - if task := next((t for t in self.tasks if t.id == tid), None): - task.writes.append((k, v)) - - # before execution, check if we should interrupt - if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): - self.status = "interrupt_before" - return False - - return True - - def first(self) -> None: - # map inputs to channel updates - if input_writes := deque(map_input(self.graph.input_channels, self.input)): - # discard any unfinished tasks from previous checkpoint - self.checkpoint, _ = prepare_next_tasks( - self.checkpoint, - self.graph.nodes, - self.channels, - self.managed, - self.config, - self.step, - for_execution=True, - get_next_version=self.get_next_version, - # TODO missing run_manager - ) - # apply input writes - apply_writes( - self.checkpoint, - self.channels, - input_writes, - self.get_next_version, - ) - # save input checkpoint - self.put_checkpoint({"source": "input", "writes": self.input}) - else: - # no input is taken as signal to proceed past previous interrupt - self.checkpoint = copy_checkpoint(self.checkpoint) - for k in self.channels: - if k in self.checkpoint["channel_versions"]: - version = self.checkpoint["channel_versions"][k] - self.checkpoint["versions_seen"][INTERRUPT][k] = version - # done with input - self.input = INPUT_DONE - - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: - self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) - if self.checkpointer is not None: - self.submit( - self.checkpointer.put_writes, - { - **self.checkpoint_config, - "configurable": { - **self.checkpoint_config["configurable"], - "thread_ts": self.checkpoint["id"], - }, - }, - writes, - task_id, - ) - - def put_checkpoint( - self, - metadata: CheckpointMetadata, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", ) -> None: - # assign step - metadata["step"] = self.step - # bail if no checkpointer - if self.checkpointer is not None: - # create new checkpoint - self.checkpoint_metadata = metadata - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.step + self.stream = deque() + self.stack = AsyncExitStack() + self.input = input + self.config = config + self.checkpointer = checkpointer + self.checkpointer_get_next_version = ( + checkpointer.get_next_version if checkpointer else increment + ) + self.checkpointer_put_writes = ( + checkpointer.aput_writes if checkpointer else None + ) + self.checkpointer_put = checkpointer.aput if checkpointer else None + self.graph = graph + # TODO if managed values no longer needs graph we can replace with + # managed_specs, channel_specs + + # context manager + + async def __aenter__(self) -> Self: + saved = ( + await self.checkpointer.aget_tuple(self.config) + if self.checkpointer + else None + ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) + self.checkpoint_config = { + **self.config, + **saved.config, + "configurable": { + **self.config.get("configurable", {}), + **saved.config.get("configurable", {}), + }, + } + self.checkpoint = saved.checkpoint + self.checkpoint_metadata = saved.metadata + self.checkpoint_pending_writes = saved.pending_writes + + self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor()) + self.channels = await self.stack.enter_async_context( + AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config) + ) + self.managed = await self.stack.enter_async_context( + AsyncManagedValuesManager( + self.graph.managed_values_dict, self.config, self.graph ) - # save it, without blocking - self.submit( - self.checkpointer.put, - self.checkpoint_config, - copy_checkpoint(self.checkpoint), - self.checkpoint_metadata, - ) - self.checkpoint_config = { - **self.checkpoint_config, - "configurable": { - **self.checkpoint_config["configurable"], - "thread_ts": self.checkpoint["id"], - }, - } - # produce debug output - self.stream.extend( - ("debug", v) - for v in map_debug_checkpoint( - self.step, - self.checkpoint_config, - self.channels, - self.graph.stream_channels_asis, - self.checkpoint_metadata, - ) - ) - # increment step - self.step += 1 + ) + self.status = "pending" + self.step = self.checkpoint_metadata["step"] + 1 + + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + del self.graph + return await asyncio.shield( + self.stack.__aexit__(exc_type, exc_value, traceback) + ) diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 7b56eba46..2e3c2098b 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,8 +1,6 @@ from collections import defaultdict from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union -from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.checkpoint.sqlite import SqliteSaver import pytest from langchain_core.callbacks import ( CallbackManagerForLLMRun, @@ -25,6 +23,7 @@ from langchain_core.tools import BaseTool from langchain_core.tools import tool as dec_tool from pydantic import BaseModel as BaseModelV2 +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.prebuilt import ( ToolNode, ValidationNode, @@ -110,6 +109,7 @@ def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]): }, ), "pending_sends": [], + "current_tasks": {}, } assert saved.metadata == { "source": "loop", @@ -168,6 +168,7 @@ async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]): }, ), "pending_sends": [], + "current_tasks": {}, } assert saved.metadata == { "source": "loop", From 29e860acca3c2cf67a21f8e7cbfb5f3cc3b3c004 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 08:52:08 -0700 Subject: [PATCH 09/15] Update manager.py --- libs/langgraph/langgraph/channels/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py index 524611e5a..2492551e2 100644 --- a/libs/langgraph/langgraph/channels/manager.py +++ b/libs/langgraph/langgraph/channels/manager.py @@ -61,6 +61,6 @@ def create_checkpoint( channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), - # checkpoints are saved only at the end of a step, ie. no current tasks + # checkpoints are saved only at the end of a step, ie. when current tasks should be cleared current_tasks={}, ) From ed5d114087c5c9817e02310f45262e08723cdc00 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 12:19:43 -0700 Subject: [PATCH 10/15] Fix --- libs/langgraph/langgraph/pregel/loop.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index b8c3fdc4e..c49d20697 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -73,10 +73,10 @@ class PregelLoop: checkpointer: Optional[BaseCheckpointSaver] checkpointer_get_next_version: Callable[[Optional[V]], V] checkpointer_put_writes: Optional[ - Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], ...] + Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any] ] checkpointer_put: Optional[ - Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], ...] + Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any] ] graph: "Pregel" From dc703857dd5b4ded719a9ca5c9f76789e69b9344 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 15:43:40 -0700 Subject: [PATCH 11/15] Run tests in parallel --- libs/langgraph/poetry.lock | 37 ++++++++++++++++++++++++++++++++++- libs/langgraph/pyproject.toml | 3 ++- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index f068e1405..b77864484 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -747,6 +747,20 @@ files = [ [package.extras] test = ["pytest (>=6)"] +[[package]] +name = "execnet" +version = "2.1.1" +description = "execnet: rapid multi-Python deployment" +optional = false +python-versions = ">=3.8" +files = [ + {file = "execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc"}, + {file = "execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3"}, +] + +[package.extras] +testing = ["hatch", "pre-commit", "pytest", "tox"] + [[package]] name = "executing" version = "2.0.1" @@ -2784,6 +2798,27 @@ files = [ tomli = {version = ">=2.0.1,<3.0.0", markers = "python_version < \"3.11\""} watchdog = ">=2.0.0" +[[package]] +name = "pytest-xdist" +version = "3.6.1" +description = "pytest xdist plugin for distributed testing, most importantly across multiple CPUs" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pytest_xdist-3.6.1-py3-none-any.whl", hash = "sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7"}, + {file = "pytest_xdist-3.6.1.tar.gz", hash = "sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d"}, +] + +[package.dependencies] +execnet = ">=2.1" +psutil = {version = ">=3.0", optional = true, markers = "extra == \"psutil\""} +pytest = ">=7.0.0" + +[package.extras] +psutil = ["psutil (>=3.0)"] +setproctitle = ["setproctitle"] +testing = ["filelock"] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -4130,4 +4165,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "170eaa0e542a02d5f2fb0d42d1f04c5d010bd3735a44b927b28e6b742c689eb0" +content-hash = "5fb6190a1b01d0cd351ea9a0023c8c8d6acf4fe831101ab87f9fc41308f74b74" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 14c90c401..7b4ab2078 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -31,6 +31,7 @@ langchainhub = "^0.1.14" langchain-openai = ">=0.1.2" langchain-anthropic = ">=0.1.8" dataclasses-json = "^0.6.7" +pytest-xdist = {extras = ["psutil"], version = "^3.6.1"} [tool.poetry.group.dev] optional = true @@ -61,7 +62,7 @@ omit = ["tests/*"] [tool.pytest-watcher] now = true delay = 0.1 -runner_args = ["-x", "--ff", "-vv", "--snapshot-update"] +runner_args = ["-x", "--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"] patterns = ["*.py"] [build-system] From 27cd4e222146ff84240f424ed6a0b822c78ddd1e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 15:44:21 -0700 Subject: [PATCH 12/15] Update core --- libs/langgraph/poetry.lock | 6 +++--- libs/langgraph/tests/__snapshots__/test_pregel.ambr | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index b77864484..3bee7a407 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -1760,13 +1760,13 @@ langchain-core = ">=0.2.2rc1,<0.3" [[package]] name = "langchain-core" -version = "0.2.19" +version = "0.2.22" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langchain_core-0.2.19-py3-none-any.whl", hash = "sha256:5b3cd34395be274c89e822c84f0e03c4da14168c177a83921c5b9414ac7a0651"}, - {file = "langchain_core-0.2.19.tar.gz", hash = "sha256:13043a83e5c9ab58b9f5ce2a56896e7e88b752e8891b2958960a98e71801471e"}, + {file = "langchain_core-0.2.22-py3-none-any.whl", hash = "sha256:7731a86440c0958b3186c003fb9b26b2d5a682a6344bda7bfb9174e2898f8b43"}, + {file = "langchain_core-0.2.22.tar.gz", hash = "sha256:582d6f929a43b830139444e4124123cd415331ad62f25757b1406252958cdcac"}, ] [package.dependencies] diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index bbfb425eb..250acc712 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -509,10 +509,10 @@ ''' # --- # name: test_conditional_state_graph - '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- # name: test_conditional_state_graph.1 - '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- # name: test_conditional_state_graph.2 ''' From 1dbf7a33925e9c26c4d955cd82fa2f07d88eea47 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 16:37:27 -0700 Subject: [PATCH 13/15] Move all checkpoint edits to apply_writes --- libs/langgraph/langgraph/checkpoint/base.py | 10 +- libs/langgraph/langgraph/constants.py | 3 +- libs/langgraph/langgraph/pregel/__init__.py | 17 ++- libs/langgraph/langgraph/pregel/algo.py | 121 +++++++++++--------- libs/langgraph/langgraph/pregel/loop.py | 38 +++--- libs/langgraph/tests/test_prebuilt.py | 27 ++--- 6 files changed, 109 insertions(+), 107 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 8cb278f32..7f1ffab1f 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -1,5 +1,4 @@ from abc import ABC -from collections import defaultdict from datetime import datetime, timezone from typing import ( Any, @@ -79,7 +78,7 @@ class Checkpoint(TypedDict): The keys are channel names and the values are the logical time step at which the channel was last updated. """ - versions_seen: defaultdict[str, dict[str, Union[str, int, float]]] + versions_seen: dict[str, dict[str, Union[str, int, float]]] """Map from node ID to map from channel name to version seen. This keeps track of the versions of the channels that each node has seen. @@ -100,7 +99,7 @@ def empty_checkpoint() -> Checkpoint: ts=datetime.now(timezone.utc).isoformat(), channel_values={}, channel_versions={}, - versions_seen=defaultdict(dict), + versions_seen={}, pending_sends=[], current_tasks={}, ) @@ -113,10 +112,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: id=checkpoint["id"], channel_values=checkpoint["channel_values"].copy(), channel_versions=checkpoint["channel_versions"].copy(), - versions_seen=defaultdict( - dict, - {k: v.copy() for k, v in checkpoint["versions_seen"].items()}, - ), + versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()}, pending_sends=checkpoint.get("pending_sends", []).copy(), current_tasks=checkpoint.get("current_tasks", {}).copy(), ) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 41970ae16..4a50e8833 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,10 +1,11 @@ from typing import Any +INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" INTERRUPT = "__interrupt__" TASKS = "__pregel_tasks" -RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ} +RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT} TAG_HIDDEN = "langsmith:hidden" START = "__start__" diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 41b517e0a..14edeacbd 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -355,7 +355,7 @@ class Pregel( ) as channels, ManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = prepare_next_tasks( + next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -387,7 +387,7 @@ class Pregel( ) as channels, AsyncManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = prepare_next_tasks( + next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -429,7 +429,7 @@ class Pregel( ) as channels, ManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = prepare_next_tasks( + next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -475,7 +475,7 @@ class Pregel( ) as channels, AsyncManagedValuesManager( self.managed_values_dict, ensure_config(config), self ) as managed: - _, next_tasks = prepare_next_tasks( + next_tasks = prepare_next_tasks( checkpoint, self.nodes, channels, @@ -560,14 +560,14 @@ class Pregel( # deque.extend is thread-safe CONFIG_KEY_SEND: task.writes.extend, CONFIG_KEY_READ: partial( - local_read, checkpoint, channels, task.writes, config + local_read, checkpoint, channels, task, config ), }, ), ) # apply to checkpoint and save apply_writes( - checkpoint, channels, task.writes, self.checkpointer.get_next_version + checkpoint, channels, [task], self.checkpointer.get_next_version ) step = saved.metadata.get("step", -2) + 1 if saved else -1 @@ -652,14 +652,14 @@ class Pregel( # deque.extend is thread-safe CONFIG_KEY_SEND: task.writes.extend, CONFIG_KEY_READ: partial( - local_read, checkpoint, channels, task.writes, config + local_read, checkpoint, channels, task, config ), }, ), ) # apply to checkpoint and save apply_writes( - checkpoint, channels, task.writes, self.checkpointer.get_next_version + checkpoint, channels, [task], self.checkpointer.get_next_version ) step = saved.metadata.get("step", -2) + 1 if saved else -1 @@ -1145,7 +1145,6 @@ class Pregel( # exception will be handled in panic_or_proceed futures.clear() else: - print(loop.step, task.name, stream_modes) # save task writes to checkpointer loop.put_writes(task.id, task.writes) # yield updates output for the finished task diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 46950b9f8..e93d6e14a 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -7,7 +7,9 @@ from typing import ( Iterator, Literal, Mapping, + NamedTuple, Optional, + Protocol, Sequence, Union, overload, @@ -29,6 +31,7 @@ from langgraph.constants import ( CONFIG_KEY_READ, CONFIG_KEY_SEND, INTERRUPT, + RESERVED, TAG_HIDDEN, TASKS, Send, @@ -41,6 +44,18 @@ from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription +class WritesProtocol(Protocol): + name: str + writes: Sequence[tuple[str, Any]] + triggers: Sequence[str] + + +class PregelTaskWrites(NamedTuple): + name: str + writes: Sequence[tuple[str, Any]] + triggers: Sequence[str] + + def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], @@ -48,8 +63,7 @@ def should_interrupt( ) -> bool: version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) null_version = version_type() - # defaultdicts are mutated on access :( so we need to copy - seen = checkpoint["versions_seen"].copy()[INTERRUPT] + seen = checkpoint["versions_seen"].get(INTERRUPT, {}) return ( # interrupt if any channel has been updated since last interrupt any( @@ -72,21 +86,21 @@ def should_interrupt( def local_read( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], - writes: Sequence[tuple[str, Any]], + task: WritesProtocol, config: RunnableConfig, select: Union[list[str], str], fresh: bool = False, ) -> Union[dict[str, Any], Any]: if fresh: - checkpoint = create_checkpoint(checkpoint, channels, -1) + new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1) context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)} with ChannelsManager( {k: v for k, v in channels.items() if k not in context_channels}, - checkpoint, + new_checkpoint, config, ) as channels: all_channels = {**channels, **context_channels} - apply_writes(copy_checkpoint(checkpoint), all_channels, writes, None) + apply_writes(new_checkpoint, all_channels, [task], None) return read_channels(all_channels, select) else: return read_channels(channels, select) @@ -118,19 +132,46 @@ def increment(current: Optional[int], channel: BaseChannel) -> int: def apply_writes( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], - pending_writes: Sequence[tuple[str, Any]], + tasks: Sequence[WritesProtocol], get_next_version: Optional[Callable[[int, BaseChannel], int]], ) -> None: + # update seen versions + for task in tasks: + checkpoint["versions_seen"].setdefault(task.name, {}).update( + { + chan: checkpoint["channel_versions"][chan] + for chan in task.triggers + if chan in checkpoint["channel_versions"] + } + ) + + # Find the highest version of all channels + if checkpoint["channel_versions"]: + max_version = max(checkpoint["channel_versions"].values()) + else: + max_version = None + # Consume all channels that were read + for chan in { + chan for task in tasks for chan in task.triggers if chan not in RESERVED + }: + if channels[chan].consume(): + if get_next_version is not None: + checkpoint["channel_versions"][chan] = get_next_version( + max_version, channels[chan] + ) + + # clear pending sends if checkpoint["pending_sends"]: checkpoint["pending_sends"].clear() - pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel - for chan, val in pending_writes: - if chan == TASKS: - checkpoint["pending_sends"].append(val) - else: - pending_writes_by_channel[chan].append(val) + pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) + for task in tasks: + for chan, val in task.writes: + if chan == TASKS: + checkpoint["pending_sends"].append(val) + else: + pending_writes_by_channel[chan].append(val) # Find the highest version of all channels if checkpoint["channel_versions"]: @@ -138,8 +179,8 @@ def apply_writes( else: max_version = None - updated_channels: set[str] = set() # Apply writes to channels + updated_channels: set[str] = set() for chan, vals in pending_writes_by_channel.items(): if chan in channels: try: @@ -153,6 +194,7 @@ def apply_writes( max_version, channels[chan] ) updated_channels.add(chan) + # Channels that weren't updated in this step are notified of a new step for chan in channels: if chan not in updated_channels: @@ -171,9 +213,8 @@ def prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[False], - get_next_version: Literal[None] = None, manager: Literal[None] = None, -) -> tuple[Checkpoint, list[PregelTaskDescription]]: +) -> list[PregelTaskDescription]: ... @@ -186,9 +227,8 @@ def prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[True], - get_next_version: Callable[[int, BaseChannel], int], manager: Union[None, ParentRunManager, AsyncParentRunManager], -) -> tuple[Checkpoint, list[PregelExecutableTask]]: +) -> list[PregelExecutableTask]: ... @@ -201,10 +241,8 @@ def prepare_next_tasks( step: int, *, for_execution: bool, - get_next_version: Union[None, Callable[[int, BaseChannel], int]] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, -) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: - checkpoint = copy_checkpoint(checkpoint) +) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]: tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] # Consume pending packets for packet in checkpoint["pending_sends"]: @@ -247,7 +285,11 @@ def prepare_next_tasks( local_write, writes.extend, processes, channels ), CONFIG_KEY_READ: partial( - local_read, checkpoint, channels, writes, config + local_read, + checkpoint, + channels, + PregelTaskWrites(packet.node, writes, triggers), + config, ), }, ), @@ -258,18 +300,14 @@ def prepare_next_tasks( ) else: tasks.append(PregelTaskDescription(packet.node, packet.arg)) - if for_execution: - checkpoint["pending_sends"].clear() - # Collect channels to consume - channels_to_consume = set() # Check if any processes should be run in next step # If so, prepare the values to be passed to them version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) null_version = version_type() if null_version is None: - return checkpoint, tasks + return tasks for name, proc in processes.items(): - seen = checkpoint["versions_seen"][name] + seen = checkpoint["versions_seen"].get(name, {}) # If any of the channels read by this process were updated if triggers := sorted( chan @@ -280,22 +318,11 @@ def prepare_next_tasks( and checkpoint["channel_versions"].get(chan, null_version) > seen.get(chan, null_version) ): - channels_to_consume.update(triggers) try: val = next(_proc_input(step, name, proc, managed, channels)) except StopIteration: continue - # update seen versions - if for_execution: - seen.update( - { - chan: checkpoint["channel_versions"][chan] - for chan in proc.triggers - if chan in checkpoint["channel_versions"] - } - ) - if for_execution: if node := proc.get_node(): metadata = { @@ -333,7 +360,7 @@ def prepare_next_tasks( local_read, checkpoint, channels, - writes, + PregelTaskWrites(name, writes, triggers), config, ), }, @@ -345,19 +372,7 @@ def prepare_next_tasks( ) else: tasks.append(PregelTaskDescription(name, val)) - # Find the highest version of all channels - if checkpoint["channel_versions"]: - max_version = max(checkpoint["channel_versions"].values()) - else: - max_version = None - # Consume all channels that were read - if for_execution: - for chan in channels_to_consume: - if channels[chan].consume(): - checkpoint["channel_versions"][chan] = get_next_version( - max_version, channels[chan] - ) - return checkpoint, tasks + return tasks def _proc_input( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index c49d20697..62bf6c806 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -38,13 +38,14 @@ from langgraph.checkpoint.base import ( copy_checkpoint, empty_checkpoint, ) -from langgraph.constants import INTERRUPT +from langgraph.constants import INPUT, INTERRUPT from langgraph.managed.base import ( AsyncManagedValuesManager, ManagedValueMapping, ManagedValuesManager, ) from langgraph.pregel.algo import ( + PregelTaskWrites, apply_writes, increment, prepare_next_tasks, @@ -88,6 +89,7 @@ class PregelLoop: checkpoint_metadata: CheckpointMetadata checkpoint_pending_writes: Optional[List[PendingWrite]] + step: int status: Literal[ "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" ] @@ -133,19 +135,13 @@ class PregelLoop: if self.input is not INPUT_DONE: self._first() - elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( - self.tasks - ): - # assign writes to tasks, apply them in order - grouped: dict[str, list[tuple[str, Any]]] = {} - for tid, k, v in self.checkpoint_pending_writes: - grouped.setdefault(tid, []).append((k, v)) - writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] + elif all(task.writes for task in self.tasks): + writes = [w for t in self.tasks for w in t.writes] # all tasks have finished apply_writes( self.checkpoint, self.channels, - writes, + self.tasks, self.checkpointer_get_next_version, ) # produce values output @@ -179,8 +175,7 @@ class PregelLoop: return False # prepare next tasks - prev_checkpoint = self.checkpoint - self.checkpoint, self.tasks = prepare_next_tasks( + self.tasks = prepare_next_tasks( self.checkpoint, self.graph.nodes, self.channels, @@ -188,7 +183,6 @@ class PregelLoop: self.config, self.step, for_execution=True, - get_next_version=self.checkpointer_get_next_version, manager=manager, ) @@ -202,9 +196,14 @@ class PregelLoop: for tid, k, v in self.checkpoint_pending_writes: if task := next((t for t in self.tasks if t.id == tid), None): task.writes.append((k, v)) + # TODO clear checkpoint_pending_writes + + # if all tasks have finished, re-tick + if all(task.writes for task in self.tasks): + return self.tick() # before execution, check if we should interrupt - if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): + if should_interrupt(self.checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" return False @@ -219,7 +218,7 @@ class PregelLoop: # map inputs to channel updates if input_writes := deque(map_input(self.graph.input_channels, self.input)): # discard any unfinished tasks from previous checkpoint - self.checkpoint, _ = prepare_next_tasks( + discard_tasks = prepare_next_tasks( self.checkpoint, self.graph.nodes, self.channels, @@ -227,20 +226,19 @@ class PregelLoop: self.config, self.step, for_execution=True, - get_next_version=self.checkpointer_get_next_version, ) # apply input writes apply_writes( self.checkpoint, self.channels, - input_writes, + discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])], self.checkpointer_get_next_version, ) # save input checkpoint self._put_checkpoint({"source": "input", "writes": self.input}) else: # no input is taken as signal to proceed past previous interrupt - self.checkpoint = copy_checkpoint(self.checkpoint) + self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) for k in self.channels: if k in self.checkpoint["channel_versions"]: version = self.checkpoint["channel_versions"][k] @@ -327,7 +325,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): **saved.config.get("configurable", {}), }, } - self.checkpoint = saved.checkpoint + self.checkpoint = copy_checkpoint(saved.checkpoint) self.checkpoint_metadata = saved.metadata self.checkpoint_pending_writes = saved.pending_writes @@ -396,7 +394,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): **saved.config.get("configurable", {}), }, } - self.checkpoint = saved.checkpoint + self.checkpoint = copy_checkpoint(saved.checkpoint) self.checkpoint_metadata = saved.metadata self.checkpoint_pending_writes = saved.pending_writes diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 2e3c2098b..a0a8436e5 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,4 +1,3 @@ -from collections import defaultdict from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union import pytest @@ -100,14 +99,11 @@ def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]): "start:agent": 3, "agent": 3, }, - "versions_seen": defaultdict( - dict, - { - "__start__": {"__start__": 1}, - "agent": {"start:agent": 2}, - "tools": {}, - }, - ), + "versions_seen": { + "__input__": {}, + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + }, "pending_sends": [], "current_tasks": {}, } @@ -159,14 +155,11 @@ async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]): "start:agent": 3, "agent": 3, }, - "versions_seen": defaultdict( - dict, - { - "__start__": {"__start__": 1}, - "agent": {"start:agent": 2}, - "tools": {}, - }, - ), + "versions_seen": { + "__input__": {}, + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + }, "pending_sends": [], "current_tasks": {}, } From b0e0d269c530c32b01f31ada366d6a27b41849af Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 16:40:50 -0700 Subject: [PATCH 14/15] Lint --- libs/langgraph/langgraph/pregel/loop.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 62bf6c806..5d02080b0 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -196,7 +196,6 @@ class PregelLoop: for tid, k, v in self.checkpoint_pending_writes: if task := next((t for t in self.tasks if t.id == tid), None): task.writes.append((k, v)) - # TODO clear checkpoint_pending_writes # if all tasks have finished, re-tick if all(task.writes for task in self.tasks): From b77ef7d162994146299be5e5717dcc52d4bff929 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 17:13:51 -0700 Subject: [PATCH 15/15] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 14edeacbd..7defcdc70 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -744,7 +744,6 @@ class Pregel( Output is a dict with the node name as key and the updated values as value. debug: Emit debug events for each step. output_keys: The keys to stream, defaults to all non-context channels. - input_keys: The keys to use from the input, defaults to all input channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. debug: Whether to print debug information during execution, defaults to False.