diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py index ce0d21189..2492551e2 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. when current tasks should be cleared + current_tasks={}, ) diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 6b147f3e3..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, @@ -25,6 +24,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. @@ -53,6 +53,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.""" @@ -74,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. @@ -84,6 +88,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: @@ -93,8 +99,9 @@ def empty_checkpoint() -> Checkpoint: ts=datetime.now(timezone.utc).isoformat(), channel_values={}, channel_versions={}, - versions_seen=defaultdict(dict), + versions_seen={}, pending_sends=[], + current_tasks={}, ) @@ -105,11 +112,9 @@ 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(), ) @@ -118,7 +123,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/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 d0d1c461d..7defcdc70 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,8 +59,6 @@ from langgraph.channels.manager import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, - Checkpoint, - CheckpointMetadata, copy_checkpoint, empty_checkpoint, ) @@ -73,43 +66,37 @@ 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, + local_read, + prepare_next_tasks, +) 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, BackgroundExecutor 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 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 ( All, PregelExecutableTask, - PregelTaskDescription, StateSnapshot, + StreamMode, ) from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -196,16 +183,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 +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, @@ -410,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, @@ -452,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, @@ -498,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, @@ -583,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 + apply_writes( + checkpoint, channels, [task], self.checkpointer.get_next_version ) step = saved.metadata.get("step", -2) + 1 if saved else -1 @@ -675,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 + apply_writes( + checkpoint, channels, [task], self.checkpointer.get_next_version ) step = saved.metadata.get("step", -2) + 1 if saved else -1 @@ -711,7 +688,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 +705,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 +716,6 @@ class Pregel( return ( debug, stream_mode, - input_keys, output_keys, interrupt_before, interrupt_after, @@ -757,7 +728,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, @@ -774,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. @@ -848,214 +817,61 @@ 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 SyncPregelLoop( + 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 - - if debug: - print_step_tasks(step, next_tasks) - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_tasks(step, next_tasks), + 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(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(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 +900,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 +914,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, ), ) else: @@ -1108,74 +924,33 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, 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 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 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 - ), - ) - ), - } - ) - - # after execution, check if we should interrupt - if _should_interrupt( - checkpoint, - interrupt_after, - self.stream_channels_list, - 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" + 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(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 +962,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 +978,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. @@ -1277,7 +1050,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"): @@ -1288,228 +1060,70 @@ 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 = ( - 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]: - 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(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 - 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, - self.stream_channels_list, - 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() @@ -1518,7 +1132,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: @@ -1527,13 +1141,11 @@ 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) + # 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( @@ -1547,7 +1159,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 @@ -1556,76 +1170,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, - self.stream_channels_list, - 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 @@ -1637,7 +1211,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 +1223,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 +1242,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 +1263,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 +1275,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 +1295,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 +1334,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..e93d6e14a --- /dev/null +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -0,0 +1,427 @@ +import json +from collections import defaultdict, deque +from functools import partial +from typing import ( + Any, + Callable, + Iterator, + Literal, + Mapping, + NamedTuple, + Optional, + Protocol, + 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, + RESERVED, + 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 + + +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]], + tasks: list[PregelExecutableTask], +) -> bool: + version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) + null_version = version_type() + seen = checkpoint["versions_seen"].get(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], + task: WritesProtocol, + config: RunnableConfig, + select: Union[list[str], str], + fresh: bool = False, +) -> Union[dict[str, Any], Any]: + if fresh: + 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}, + new_checkpoint, + config, + ) as channels: + all_channels = {**channels, **context_channels} + apply_writes(new_checkpoint, all_channels, [task], 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], + 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() + + # Group writes by channel + 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"]: + max_version = max(checkpoint["channel_versions"].values()) + else: + max_version = None + + # Apply writes to channels + updated_channels: set[str] = set() + 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], + manager: Literal[None] = None, +) -> 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], + manager: Union[None, ParentRunManager, AsyncParentRunManager], +) -> 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, + manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, +) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]: + 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, + PregelTaskWrites(packet.node, writes, triggers), + config, + ), + }, + ), + triggers, + proc.retry_policy, + task_id, + ) + ) + else: + tasks.append(PregelTaskDescription(packet.node, packet.arg)) + # 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 tasks + for name, proc in processes.items(): + seen = checkpoint["versions_seen"].get(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) + ): + try: + val = next(_proc_input(step, name, proc, managed, channels)) + except StopIteration: + continue + + 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, + PregelTaskWrites(name, writes, triggers), + config, + ), + }, + ), + triggers, + proc.retry_policy, + task_id, + ) + ) + else: + tasks.append(PregelTaskDescription(name, val)) + return 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/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 new file mode 100644 index 000000000..5d02080b0 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -0,0 +1,423 @@ +import asyncio +from collections import deque +from contextlib import AsyncExitStack, ExitStack +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + AsyncContextManager, + Callable, + ContextManager, + List, + Literal, + Mapping, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + 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 ( + AsyncChannelsManager, + ChannelsManager, + create_checkpoint, +) +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + Checkpoint, + CheckpointMetadata, + CheckpointTuple, + PendingWrite, + copy_checkpoint, + empty_checkpoint, +) +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, + should_interrupt, +) +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 + +if TYPE_CHECKING: + from langgraph.pregel import Pregel + + +V = TypeVar("V") +INPUT_DONE = object() + + +class PregelLoop: + input: Optional[Any] + config: RunnableConfig + checkpointer: Optional[BaseCheckpointSaver] + checkpointer_get_next_version: Callable[[Optional[V]], V] + checkpointer_put_writes: Optional[ + Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any] + ] + checkpointer_put: Optional[ + Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any] + ] + graph: "Pregel" + + submit: Submit + channels: Mapping[str, BaseChannel] + managed: ManagedValueMapping + checkpoint: Checkpoint + checkpoint_config: RunnableConfig + checkpoint_metadata: CheckpointMetadata + checkpoint_pending_writes: Optional[List[PendingWrite]] + + step: int + status: Literal[ + "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" + ] + 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 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, + self.tasks, + 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 + self.tasks = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + 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)) + + # 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(self.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 + discard_tasks = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + ) + # apply input writes + apply_writes( + self.checkpoint, + self.channels, + 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["versions_seen"].setdefault(INTERRUPT, {}) + 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], + *, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", + ) -> None: + self.stream = deque() + self.stack = ExitStack() + 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.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 + ) 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 = copy_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" + self.step = self.checkpoint_metadata["step"] + 1 + + 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) + + +class AsyncPregelLoop(PregelLoop, AsyncContextManager): + def __init__( + self, + input: Optional[Any], + *, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", + ) -> None: + 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 = copy_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 + ) + ) + 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/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/poetry.lock b/libs/langgraph/poetry.lock index f068e1405..3bee7a407 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" @@ -1746,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] @@ -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] 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 ''' diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 264f30db6..a0a8436e5 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -22,11 +22,14 @@ 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, create_react_agent, ) +from tests.any_str import AnyStr +from tests.memory_assert import MemorySaverAssertImmutable class FakeToolCallingModel(BaseChatModel): @@ -56,14 +59,117 @@ 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}) + 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": { + "__input__": {}, + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + }, + "pending_sends": [], + "current_tasks": {}, + } + 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": { + "__input__": {}, + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + }, + "pending_sends": [], + "current_tasks": {}, + } + 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 df6ba34ca..2646a64fd 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -520,10 +520,8 @@ 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}) + app.invoke(2, {"recursion_limit": 1}, debug=1) graph = Graph() graph.add_node("add_one", add_one) @@ -535,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, @@ -6168,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",), @@ -6877,7 +6887,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> }, ) - assert [c for c in app_w_interrupt.stream(None, config)] == [ + 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 4c166cdc5..0e927e518 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}) @@ -4737,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",),