diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py index 078f71e7c..aa8f507b6 100644 --- a/libs/langgraph/langgraph/managed/base.py +++ b/libs/langgraph/langgraph/managed/base.py @@ -1,6 +1,4 @@ from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Iterator -from contextlib import asynccontextmanager, contextmanager from inspect import isclass from typing import ( Any, @@ -8,48 +6,18 @@ from typing import ( TypeVar, ) -from typing_extensions import Self, TypeGuard +from typing_extensions import TypeGuard -from langgraph.types import LoopProtocol +from langgraph.types import PregelScratchpad V = TypeVar("V") U = TypeVar("U") class ManagedValue(ABC, Generic[V]): - def __init__(self, loop: LoopProtocol) -> None: - self.loop = loop - - @classmethod - @contextmanager - def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]: - try: - value = cls(loop, **kwargs) - yield value - finally: - # because managed value and Pregel have reference to each other - # let's make sure to break the reference on exit - try: - del value - except UnboundLocalError: - pass - - @classmethod - @asynccontextmanager - async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]: - try: - value = cls(loop, **kwargs) - yield value - finally: - # because managed value and Pregel have reference to each other - # let's make sure to break the reference on exit - try: - del value - except UnboundLocalError: - pass - + @staticmethod @abstractmethod - def __call__(self) -> V: ... + def get(scratchpad: PregelScratchpad) -> V: ... ManagedValueSpec = type[ManagedValue] @@ -59,4 +27,4 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]: return isclass(value) and issubclass(value, ManagedValue) -ManagedValueMapping = dict[str, ManagedValue] +ManagedValueMapping = dict[str, ManagedValueSpec] diff --git a/libs/langgraph/langgraph/managed/is_last_step.py b/libs/langgraph/langgraph/managed/is_last_step.py index 9f25a8121..ccfaea038 100644 --- a/libs/langgraph/langgraph/managed/is_last_step.py +++ b/libs/langgraph/langgraph/managed/is_last_step.py @@ -1,19 +1,22 @@ from typing import Annotated from langgraph.managed.base import ManagedValue +from langgraph.types import PregelScratchpad class IsLastStepManager(ManagedValue[bool]): - def __call__(self) -> bool: - return self.loop.step == self.loop.stop - 1 + @staticmethod + def get(scratchpad: PregelScratchpad) -> bool: + return scratchpad.step == scratchpad.stop - 1 IsLastStep = Annotated[bool, IsLastStepManager] class RemainingStepsManager(ManagedValue[int]): - def __call__(self) -> int: - return self.loop.stop - self.loop.step + @staticmethod + def get(scratchpad: PregelScratchpad) -> int: + return scratchpad.stop - scratchpad.step RemainingSteps = Annotated[int, RemainingStepsManager] diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index b8de81f53..4550c56d0 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -88,12 +88,15 @@ from langgraph.pregel.algo import ( prepare_next_tasks, ) from langgraph.pregel.call import identifier -from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint +from langgraph.pregel.checkpoint import ( + channels_from_checkpoint, + create_checkpoint, + empty_checkpoint, +) from langgraph.pregel.debug import tasks_w_writes from langgraph.pregel.draw import draw_graph from langgraph.pregel.io import map_input, read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop -from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.messages import StreamMessagesHandler from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.read import PregelNode @@ -108,7 +111,6 @@ from langgraph.types import ( CachePolicy, Checkpointer, Interrupt, - LoopProtocol, StateSnapshot, StateUpdate, StreamChunk, @@ -894,104 +896,102 @@ class Pregel(PregelProtocol): # migrate checkpoint if needed self._migrate_checkpoint(saved.checkpoint) - with ChannelsManager( + step = saved.metadata.get("step", -1) + 1 + stop = step + 2 + channels, managed = channels_from_checkpoint( self.channels, saved.checkpoint, - LoopProtocol( - config=saved.config, - step=saved.metadata.get("step", -1) + 1, - stop=saved.metadata.get("step", -1) + 2, + ) + # tasks for this checkpoint + next_tasks = prepare_next_tasks( + saved.checkpoint, + saved.pending_writes or [], + self.nodes, + channels, + managed, + saved.config, + step, + stop, + for_execution=True, + store=self.store, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None ), - ) as (channels, managed): - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - saved.checkpoint, - saved.pending_writes or [], - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), - manager=None, - ) - # get the subgraphs - subgraphs = dict(self.get_subgraphs()) - parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") - task_states: dict[str, RunnableConfig | StateSnapshot] = {} - for task in next_tasks.values(): - if task.name not in subgraphs: - continue - # assemble checkpoint_ns for this task - task_ns = f"{task.name}{NS_END}{task.id}" - if parent_ns: - task_ns = f"{parent_ns}{NS_SEP}{task_ns}" - if not recurse: - # set config as signal that subgraph checkpoints exist - config = { - CONF: { - "thread_id": saved.config[CONF]["thread_id"], - CONFIG_KEY_CHECKPOINT_NS: task_ns, - } + manager=None, + ) + # get the subgraphs + subgraphs = dict(self.get_subgraphs()) + parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") + task_states: dict[str, RunnableConfig | StateSnapshot] = {} + for task in next_tasks.values(): + if task.name not in subgraphs: + continue + # assemble checkpoint_ns for this task + task_ns = f"{task.name}{NS_END}{task.id}" + if parent_ns: + task_ns = f"{parent_ns}{NS_SEP}{task_ns}" + if not recurse: + # set config as signal that subgraph checkpoints exist + config = { + CONF: { + "thread_id": saved.config[CONF]["thread_id"], + CONFIG_KEY_CHECKPOINT_NS: task_ns, } - task_states[task.id] = config - else: - # get the state of the subgraph - config = { - CONF: { - CONFIG_KEY_CHECKPOINTER: recurse, - "thread_id": saved.config[CONF]["thread_id"], - CONFIG_KEY_CHECKPOINT_NS: task_ns, - } + } + task_states[task.id] = config + else: + # get the state of the subgraph + config = { + CONF: { + CONFIG_KEY_CHECKPOINTER: recurse, + "thread_id": saved.config[CONF]["thread_id"], + CONFIG_KEY_CHECKPOINT_NS: task_ns, } - task_states[task.id] = subgraphs[task.name].get_state( - config, subgraphs=True - ) - # apply pending writes - if null_writes := [ - w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - self.trigger_to_nodes, + } + task_states[task.id] = subgraphs[task.name].get_state( + config, subgraphs=True ) - if apply_pending_writes and saved.pending_writes: - for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes( - saved.checkpoint, channels, tasks, None, self.trigger_to_nodes - ) - tasks_with_writes = tasks_w_writes( - next_tasks.values(), - saved.pending_writes, - task_states, - self.stream_channels_asis, - ) - # assemble the state snapshot - return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks.values() if not t.writes), - patch_checkpoint_map(saved.config, saved.metadata), - saved.metadata, - saved.checkpoint["ts"], - patch_checkpoint_map(saved.parent_config, saved.metadata), - tasks_with_writes, - tuple([i for task in tasks_with_writes for i in task.interrupts]), + # apply pending writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + self.trigger_to_nodes, ) + if apply_pending_writes and saved.pending_writes: + for tid, k, v in saved.pending_writes: + if k in (ERROR, INTERRUPT, SCHEDULED): + continue + if tid not in next_tasks: + continue + next_tasks[tid].writes.append((k, v)) + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes( + saved.checkpoint, channels, tasks, None, self.trigger_to_nodes + ) + tasks_with_writes = tasks_w_writes( + next_tasks.values(), + saved.pending_writes, + task_states, + self.stream_channels_asis, + ) + # assemble the state snapshot + return StateSnapshot( + read_channels(channels, self.stream_channels_asis), + tuple(t.name for t in next_tasks.values() if not t.writes), + patch_checkpoint_map(saved.config, saved.metadata), + saved.metadata, + saved.checkpoint["ts"], + patch_checkpoint_map(saved.parent_config, saved.metadata), + tasks_with_writes, + tuple([i for task in tasks_with_writes for i in task.interrupts]), + ) async def _aprepare_state_snapshot( self, @@ -1015,108 +1015,103 @@ class Pregel(PregelProtocol): # migrate checkpoint if needed self._migrate_checkpoint(saved.checkpoint) - async with AsyncChannelsManager( + step = saved.metadata.get("step", -1) + 1 + stop = step + 2 + channels, managed = channels_from_checkpoint( self.channels, saved.checkpoint, - LoopProtocol( - config=saved.config, - step=saved.metadata.get("step", -1) + 1, - stop=saved.metadata.get("step", -1) + 2, - ), - ) as ( + ) + # tasks for this checkpoint + next_tasks = prepare_next_tasks( + saved.checkpoint, + saved.pending_writes or [], + self.nodes, channels, managed, - ): - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - saved.checkpoint, - saved.pending_writes or [], - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), - manager=None, - ) - # get the subgraphs - subgraphs = {n: g async for n, g in self.aget_subgraphs()} - parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") - task_states: dict[str, RunnableConfig | StateSnapshot] = {} - for task in next_tasks.values(): - if task.name not in subgraphs: - continue - # assemble checkpoint_ns for this task - task_ns = f"{task.name}{NS_END}{task.id}" - if parent_ns: - task_ns = f"{parent_ns}{NS_SEP}{task_ns}" - if not recurse: - # set config as signal that subgraph checkpoints exist - config = { - CONF: { - "thread_id": saved.config[CONF]["thread_id"], - CONFIG_KEY_CHECKPOINT_NS: task_ns, - } + saved.config, + step, + stop, + for_execution=True, + store=self.store, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), + manager=None, + ) + # get the subgraphs + subgraphs = {n: g async for n, g in self.aget_subgraphs()} + parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") + task_states: dict[str, RunnableConfig | StateSnapshot] = {} + for task in next_tasks.values(): + if task.name not in subgraphs: + continue + # assemble checkpoint_ns for this task + task_ns = f"{task.name}{NS_END}{task.id}" + if parent_ns: + task_ns = f"{parent_ns}{NS_SEP}{task_ns}" + if not recurse: + # set config as signal that subgraph checkpoints exist + config = { + CONF: { + "thread_id": saved.config[CONF]["thread_id"], + CONFIG_KEY_CHECKPOINT_NS: task_ns, } - task_states[task.id] = config - else: - # get the state of the subgraph - config = { - CONF: { - CONFIG_KEY_CHECKPOINTER: recurse, - "thread_id": saved.config[CONF]["thread_id"], - CONFIG_KEY_CHECKPOINT_NS: task_ns, - } + } + task_states[task.id] = config + else: + # get the state of the subgraph + config = { + CONF: { + CONFIG_KEY_CHECKPOINTER: recurse, + "thread_id": saved.config[CONF]["thread_id"], + CONFIG_KEY_CHECKPOINT_NS: task_ns, } - task_states[task.id] = await subgraphs[task.name].aget_state( - config, subgraphs=True - ) - # apply pending writes - if null_writes := [ - w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - self.trigger_to_nodes, + } + task_states[task.id] = await subgraphs[task.name].aget_state( + config, subgraphs=True + ) + # apply pending writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + self.trigger_to_nodes, + ) + if apply_pending_writes and saved.pending_writes: + for tid, k, v in saved.pending_writes: + if k in (ERROR, INTERRUPT, SCHEDULED): + continue + if tid not in next_tasks: + continue + next_tasks[tid].writes.append((k, v)) + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes( + saved.checkpoint, channels, tasks, None, self.trigger_to_nodes ) - if apply_pending_writes and saved.pending_writes: - for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes( - saved.checkpoint, channels, tasks, None, self.trigger_to_nodes - ) - tasks_with_writes = tasks_w_writes( - next_tasks.values(), - saved.pending_writes, - task_states, - self.stream_channels_asis, - ) - # assemble the state snapshot - return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks.values() if not t.writes), - patch_checkpoint_map(saved.config, saved.metadata), - saved.metadata, - saved.checkpoint["ts"], - patch_checkpoint_map(saved.parent_config, saved.metadata), - tasks_with_writes, - tuple([i for task in tasks_with_writes for i in task.interrupts]), - ) + tasks_with_writes = tasks_w_writes( + next_tasks.values(), + saved.pending_writes, + task_states, + self.stream_channels_asis, + ) + # assemble the state snapshot + return StateSnapshot( + read_channels(channels, self.stream_channels_asis), + tuple(t.name for t in next_tasks.values() if not t.writes), + patch_checkpoint_map(saved.config, saved.metadata), + saved.metadata, + saved.checkpoint["ts"], + patch_checkpoint_map(saved.parent_config, saved.metadata), + tasks_with_writes, + tuple([i for task in tasks_with_writes for i in task.interrupts]), + ) def get_state( self, config: RunnableConfig, *, subgraphs: bool = False @@ -1383,211 +1378,35 @@ class Pregel(PregelProtocol): if saved: checkpoint_config = patch_configurable(config, saved.config[CONF]) checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} - with ChannelsManager( + channels, managed = channels_from_checkpoint( self.channels, checkpoint, - LoopProtocol(config=config, step=step + 1, stop=step + 2), - ) as (channels, managed): - values, as_node = updates[0] + ) + values, as_node = updates[0] - # no values as END, just clear all tasks - if values is None and as_node == END: - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot apply multiple updates when clearing state" - ) - - if saved is not None: - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - checkpoint, - saved.pending_writes or [], - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, - manager=None, - ) - # apply null writes - if null_writes := [ - w[1:] - for w in saved.pending_writes or [] - if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - self.trigger_to_nodes, - ) - # apply writes from tasks that already ran - for tid, k, v in saved.pending_writes or []: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - # clear all current tasks - apply_writes( - checkpoint, - channels, - next_tasks.values(), - None, - self.trigger_to_nodes, - ) - # save checkpoint - next_config = checkpointer.put( - checkpoint_config, - create_checkpoint(checkpoint, None, step), - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, empty checkpoint - if values is None and as_node is None: - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot create empty checkpoint with multiple updates" - ) - - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = checkpointer.put( - checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None + # no values as END, just clear all tasks + if values is None and as_node == END: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when clearing state" ) - # act as an input - if as_node == INPUT: - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot apply multiple updates when updating as input" - ) - - if input_writes := deque(map_input(self.input_channels, values)): - apply_writes( - checkpoint, - channels, - [PregelTaskWrites((), INPUT, input_writes, [])], - checkpointer.get_next_version, - self.trigger_to_nodes, - ) - - # apply input write to channels - next_step = ( - step + 1 - if saved and saved.metadata.get("step") is not None - else -1 - ) - next_config = checkpointer.put( - checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), - { - **checkpoint_metadata, - "source": "input", - "step": next_step, - "writes": dict(input_writes), - }, - get_new_channel_versions( - checkpoint_previous_versions, - checkpoint["channel_versions"], - ), - ) - - # store the writes - checkpointer.put_writes( - next_config, - input_writes, - str(uuid5(UUID(checkpoint["id"]), INPUT)), - ) - - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - else: - raise InvalidUpdateError( - f"Received no input writes for {self.input_channels}" - ) - - # no values, copy checkpoint - if values is None and as_node == "__copy__": - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot copy checkpoint with multiple updates" - ) - - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = checkpointer.put( - saved.parent_config or saved.config - if saved - else checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "fork", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # apply pending writes, if not on specific checkpoint - if ( - CONFIG_KEY_CHECKPOINT_ID not in config[CONF] - and saved is not None - and saved.pending_writes - ): + if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, - saved.pending_writes, + saved.pending_writes or [], self.nodes, channels, managed, saved.config, - saved.metadata.get("step", -1) + 1, + step + 1, + step + 3, for_execution=True, store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # apply null writes @@ -1603,17 +1422,184 @@ class Pregel(PregelProtocol): None, self.trigger_to_nodes, ) - # apply writes - for tid, k, v in saved.pending_writes: + # apply writes from tasks that already ran + for tid, k, v in saved.pending_writes or []: if k in (ERROR, INTERRUPT, SCHEDULED): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes( - checkpoint, channels, tasks, None, self.trigger_to_nodes - ) + # clear all current tasks + apply_writes( + checkpoint, + channels, + next_tasks.values(), + None, + self.trigger_to_nodes, + ) + # save checkpoint + next_config = checkpointer.put( + checkpoint_config, + create_checkpoint(checkpoint, None, step), + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # no values, empty checkpoint + if values is None and as_node is None: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot create empty checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = checkpointer.put( + checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + + # act as an input + if as_node == INPUT: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when updating as input" + ) + + if input_writes := deque(map_input(self.input_channels, values)): + apply_writes( + checkpoint, + channels, + [PregelTaskWrites((), INPUT, input_writes, [])], + checkpointer.get_next_version, + self.trigger_to_nodes, + ) + + # apply input write to channels + next_step = ( + step + 1 + if saved and saved.metadata.get("step") is not None + else -1 + ) + next_config = checkpointer.put( + checkpoint_config, + create_checkpoint(checkpoint, channels, next_step), + { + **checkpoint_metadata, + "source": "input", + "step": next_step, + "writes": dict(input_writes), + }, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), + ) + + # store the writes + checkpointer.put_writes( + next_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) + + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + else: + raise InvalidUpdateError( + f"Received no input writes for {self.input_channels}" + ) + + # no values, copy checkpoint + if values is None and as_node == "__copy__": + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot copy checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = checkpointer.put( + saved.parent_config or saved.config if saved else checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "fork", + "step": step + 1, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # apply pending writes, if not on specific checkpoint + if ( + CONFIG_KEY_CHECKPOINT_ID not in config[CONF] + and saved is not None + and saved.pending_writes + ): + # tasks for this checkpoint + next_tasks = prepare_next_tasks( + checkpoint, + saved.pending_writes, + self.nodes, + channels, + managed, + saved.config, + step + 1, + step + 3, + for_execution=True, + store=self.store, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), + manager=None, + ) + # apply null writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + self.trigger_to_nodes, + ) + # apply writes + for tid, k, v in saved.pending_writes: + if k in (ERROR, INTERRUPT, SCHEDULED): + continue + if tid not in next_tasks: + continue + next_tasks[tid].writes.append((k, v)) + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes( + checkpoint, channels, tasks, None, self.trigger_to_nodes + ) valid_updates: list[tuple[str, dict[str, Any] | None]] = [] if len(updates) == 1: values, as_node = updates[0] @@ -1809,212 +1795,33 @@ class Pregel(PregelProtocol): if saved: checkpoint_config = patch_configurable(config, saved.config[CONF]) checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} - async with AsyncChannelsManager( + channels, managed = channels_from_checkpoint( self.channels, checkpoint, - LoopProtocol(config=config, step=step + 1, stop=step + 2), - ) as ( - channels, - managed, - ): - values, as_node = updates[0] - # no values, just clear all tasks - if values is None and as_node == END: - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot apply multiple updates when clearing state" - ) - if saved is not None: - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - checkpoint, - saved.pending_writes or [], - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, - manager=None, - ) - # apply null writes - if null_writes := [ - w[1:] - for w in saved.pending_writes or [] - if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - self.trigger_to_nodes, - ) - # apply writes from tasks that already ran - for tid, k, v in saved.pending_writes or []: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - # clear all current tasks - apply_writes( - checkpoint, - channels, - next_tasks.values(), - None, - self.trigger_to_nodes, - ) - # save checkpoint - next_config = await checkpointer.aput( - checkpoint_config, - create_checkpoint(checkpoint, None, step), - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, - {}, + ) + values, as_node = updates[0] + # no values, just clear all tasks + if values is None and as_node == END: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when clearing state" ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, empty checkpoint - if values is None and as_node is None: - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot create empty checkpoint with multiple updates" - ) - - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = await checkpointer.aput( - checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - - # act as an input - if as_node == INPUT: - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot apply multiple updates when updating as input" - ) - - if input_writes := deque(map_input(self.input_channels, values)): - apply_writes( - checkpoint, - channels, - [PregelTaskWrites((), INPUT, input_writes, [])], - checkpointer.get_next_version, - self.trigger_to_nodes, - ) - - # apply input write to channels - next_step = ( - step + 1 - if saved and saved.metadata.get("step") is not None - else -1 - ) - next_config = await checkpointer.aput( - checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), - { - **checkpoint_metadata, - "source": "input", - "step": next_step, - "writes": dict(input_writes), - }, - get_new_channel_versions( - checkpoint_previous_versions, - checkpoint["channel_versions"], - ), - ) - - # store the writes - await checkpointer.aput_writes( - next_config, - input_writes, - str(uuid5(UUID(checkpoint["id"]), INPUT)), - ) - - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - else: - raise InvalidUpdateError( - f"Received no input writes for {self.input_channels}" - ) - - # no values, copy checkpoint - if values is None and as_node == "__copy__": - if len(updates) > 1: - raise InvalidUpdateError( - "Cannot copy checkpoint with multiple updates" - ) - - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = await checkpointer.aput( - saved.parent_config or saved.config - if saved - else checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "fork", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # apply pending writes, if not on specific checkpoint - if ( - CONFIG_KEY_CHECKPOINT_ID not in config[CONF] - and saved is not None - and saved.pending_writes - ): + if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, - saved.pending_writes, + saved.pending_writes or [], self.nodes, channels, managed, saved.config, - saved.metadata.get("step", -1) + 1, + step + 1, + step + 3, for_execution=True, store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # apply null writes @@ -2030,16 +1837,183 @@ class Pregel(PregelProtocol): None, self.trigger_to_nodes, ) - for tid, k, v in saved.pending_writes: + # apply writes from tasks that already ran + for tid, k, v in saved.pending_writes or []: if k in (ERROR, INTERRUPT, SCHEDULED): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes( - checkpoint, channels, tasks, None, self.trigger_to_nodes - ) + # clear all current tasks + apply_writes( + checkpoint, + channels, + next_tasks.values(), + None, + self.trigger_to_nodes, + ) + # save checkpoint + next_config = await checkpointer.aput( + checkpoint_config, + create_checkpoint(checkpoint, None, step), + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # no values, empty checkpoint + if values is None and as_node is None: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot create empty checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = await checkpointer.aput( + checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + + # act as an input + if as_node == INPUT: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when updating as input" + ) + + if input_writes := deque(map_input(self.input_channels, values)): + apply_writes( + checkpoint, + channels, + [PregelTaskWrites((), INPUT, input_writes, [])], + checkpointer.get_next_version, + self.trigger_to_nodes, + ) + + # apply input write to channels + next_step = ( + step + 1 + if saved and saved.metadata.get("step") is not None + else -1 + ) + next_config = await checkpointer.aput( + checkpoint_config, + create_checkpoint(checkpoint, channels, next_step), + { + **checkpoint_metadata, + "source": "input", + "step": next_step, + "writes": dict(input_writes), + }, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), + ) + + # store the writes + await checkpointer.aput_writes( + next_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) + + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + else: + raise InvalidUpdateError( + f"Received no input writes for {self.input_channels}" + ) + + # no values, copy checkpoint + if values is None and as_node == "__copy__": + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot copy checkpoint with multiple updates" + ) + + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = await checkpointer.aput( + saved.parent_config or saved.config if saved else checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "fork", + "step": step + 1, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + # apply pending writes, if not on specific checkpoint + if ( + CONFIG_KEY_CHECKPOINT_ID not in config[CONF] + and saved is not None + and saved.pending_writes + ): + # tasks for this checkpoint + next_tasks = prepare_next_tasks( + checkpoint, + saved.pending_writes, + self.nodes, + channels, + managed, + saved.config, + step + 1, + step + 3, + for_execution=True, + store=self.store, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), + manager=None, + ) + # apply null writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + self.trigger_to_nodes, + ) + for tid, k, v in saved.pending_writes: + if k in (ERROR, INTERRUPT, SCHEDULED): + continue + if tid not in next_tasks: + continue + next_tasks[tid].writes.append((k, v)) + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes( + checkpoint, channels, tasks, None, self.trigger_to_nodes + ) valid_updates: list[tuple[str, dict[str, Any] | None]] = [] if len(updates) == 1: values, as_node = updates[0] diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index f8bd96ce8..471fcfb99 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -364,6 +364,7 @@ def prepare_next_tasks( managed: ManagedValueMapping, config: RunnableConfig, step: int, + stop: int, *, for_execution: Literal[False], store: Literal[None] = None, @@ -385,6 +386,7 @@ def prepare_next_tasks( managed: ManagedValueMapping, config: RunnableConfig, step: int, + stop: int, *, for_execution: Literal[True], store: Optional[BaseStore], @@ -405,6 +407,7 @@ def prepare_next_tasks( managed: ManagedValueMapping, config: RunnableConfig, step: int, + stop: int, *, for_execution: bool, store: Optional[BaseStore] = None, @@ -459,6 +462,7 @@ def prepare_next_tasks( managed=managed, config=config, step=step, + stop=stop, for_execution=for_execution, store=store, checkpointer=checkpointer, @@ -504,6 +508,7 @@ def prepare_next_tasks( managed=managed, config=config, step=step, + stop=stop, for_execution=for_execution, store=store, checkpointer=checkpointer, @@ -532,6 +537,7 @@ def prepare_single_task( managed: ManagedValueMapping, config: RunnableConfig, step: int, + stop: int, for_execution: bool, store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, @@ -632,6 +638,8 @@ def prepare_single_task( task_id, xxh3_128_hexdigest(task_checkpoint_ns.encode()), config[CONF].get(CONFIG_KEY_RESUME_MAP), + step, + stop, ), }, ), @@ -752,6 +760,8 @@ def prepare_single_task( task_id, xxh3_128_hexdigest(task_checkpoint_ns.encode()), config[CONF].get(CONFIG_KEY_RESUME_MAP), + step, + stop, ), CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( PREVIOUS, None @@ -785,23 +795,6 @@ def prepare_single_task( proc, ): triggers = tuple(sorted(proc.triggers)) - try: - val = _proc_input( - proc, - managed, - channels, - for_execution=for_execution, - input_cache=input_cache, - ) - if val is MISSING: - return - except Exception as exc: - if SUPPORTS_EXC_NOTES: - exc.add_note( - f"Before task with name '{name}' and path '{task_path[:3]}'" - ) - raise - # create task id checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name task_id = task_id_func( @@ -813,6 +806,35 @@ def prepare_single_task( *triggers, ) task_checkpoint_ns = f"{checkpoint_ns}{NS_END}{task_id}" + # create scratchpad + scratchpad = _scratchpad( + config[CONF].get(CONFIG_KEY_SCRATCHPAD), + pending_writes, + task_id, + xxh3_128_hexdigest(task_checkpoint_ns.encode()), + config[CONF].get(CONFIG_KEY_RESUME_MAP), + step, + stop, + ) + # create task input + try: + val = _proc_input( + proc, + managed, + channels, + for_execution=for_execution, + input_cache=input_cache, + scratchpad=scratchpad, + ) + if val is MISSING: + return + except Exception as exc: + if SUPPORTS_EXC_NOTES: + exc.add_note( + f"Before task with name '{name}' and path '{task_path[:3]}'" + ) + raise + metadata = { "langgraph_step": step, "langgraph_node": name, @@ -888,13 +910,7 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_SCRATCHPAD: _scratchpad( - config[CONF].get(CONFIG_KEY_SCRATCHPAD), - pending_writes, - task_id, - xxh3_128_hexdigest(task_checkpoint_ns.encode()), - config[CONF].get(CONFIG_KEY_RESUME_MAP), - ), + CONFIG_KEY_SCRATCHPAD: scratchpad, CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( PREVIOUS, None ), @@ -947,6 +963,8 @@ def _scratchpad( task_id: str, namespace_hash: str, resume_map: Optional[dict[str, Any]], + step: int, + stop: int, ) -> PregelScratchpad: if len(pending_writes) > 0: # find global resume value @@ -994,6 +1012,8 @@ def _scratchpad( # using itertools.count as an atomic counter (+= 1 is not thread-safe) return PregelScratchpad( + step=step, + stop=stop, # call call_counter=LazyAtomicCounter(), # interrupt @@ -1011,6 +1031,7 @@ def _proc_input( channels: Mapping[str, BaseChannel], *, for_execution: bool, + scratchpad: PregelScratchpad, input_cache: Optional[dict[INPUT_CACHE_KEY_TYPE, Any]], ) -> Any: """Prepare input for a PULL task, based on the process's channels and triggers.""" @@ -1026,7 +1047,7 @@ def _proc_input( if channels[chan].is_available(): val[k] = channels[chan].get() else: - val[k] = managed[k]() + val[k] = managed[k].get(scratchpad) elif isinstance(proc.channels, list): for chan in proc.channels: if chan in channels: @@ -1034,7 +1055,7 @@ def _proc_input( val = channels[chan].get() break else: - val = managed[chan]() + val = managed[chan].get(scratchpad) break else: return MISSING diff --git a/libs/langgraph/langgraph/pregel/checkpoint.py b/libs/langgraph/langgraph/pregel/checkpoint.py index fec604345..47ef04404 100644 --- a/libs/langgraph/langgraph/pregel/checkpoint.py +++ b/libs/langgraph/langgraph/pregel/checkpoint.py @@ -1,11 +1,12 @@ from collections.abc import Mapping from datetime import datetime, timezone -from typing import Optional +from typing import Optional, Union from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint from langgraph.checkpoint.base.id import uuid6 from langgraph.constants import MISSING +from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec LATEST_VERSION = 3 @@ -50,3 +51,24 @@ def create_checkpoint( versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), ) + + +def channels_from_checkpoint( + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + checkpoint: Checkpoint, +) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]: + """Get channels from a checkpoint.""" + channel_specs: dict[str, BaseChannel] = {} + managed_specs: dict[str, ManagedValueSpec] = {} + for k, v in specs.items(): + if isinstance(v, BaseChannel): + channel_specs[k] = v + else: + managed_specs[k] = v + return ( + { + k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) + for k, v in channel_specs.items() + }, + managed_specs, + ) diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index 2518130a5..2e7783f21 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -15,12 +15,11 @@ from langgraph.pregel.algo import ( increment, prepare_next_tasks, ) -from langgraph.pregel.checkpoint import empty_checkpoint +from langgraph.pregel.checkpoint import channels_from_checkpoint, empty_checkpoint from langgraph.pregel.io import map_input -from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite -from langgraph.types import All, Checkpointer, LoopProtocol +from langgraph.types import All, Checkpointer def draw_graph( @@ -56,31 +55,102 @@ def draw_graph( if isinstance(checkpointer, BaseCheckpointSaver) else increment ) - with ChannelsManager( + channels, managed = channels_from_checkpoint( specs, checkpoint, - LoopProtocol(step=step, stop=-1, config=config), - ) as (channels, managed): - static_seen: set[Any] = set() - sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {} - step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {} - # remove node mappers - nodes = { - k: v.copy(update={"mapper": None}) if v.mapper is not None else v - for k, v in nodes.items() + ) + static_seen: set[Any] = set() + sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {} + step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {} + # remove node mappers + nodes = { + k: v.copy(update={"mapper": None}) if v.mapper is not None else v + for k, v in nodes.items() + } + # apply input writes + input_writes = list(map_input(input_channels, {})) + updated_channels = apply_writes( + checkpoint, + channels, + [ + PregelTaskWrites((), INPUT, input_writes, []), + ], + get_next_version, + trigger_to_nodes, + ) + # prepare first tasks + tasks = prepare_next_tasks( + checkpoint, + [], + nodes, + channels, + managed, + config, + step, + -1, + for_execution=True, + store=None, + checkpointer=None, + manager=None, + trigger_to_nodes=trigger_to_nodes, + updated_channels=updated_channels, + ) + start_tasks = tasks + # run the pregel loop + for step in range(step, limit): + if not tasks: + break + conditionals: dict[tuple[str, str, Any], Optional[str]] = {} + # run task writers + for task in tasks.values(): + for w in task.writers: + # apply regular writes + if isinstance(w, ChannelWrite): + empty_input = ( + cast(BaseChannel, specs["__root__"]).ValueType() + if "__root__" in specs + else None + ) + w.invoke(empty_input, task.config) + # apply conditional writes declared for static analysis, only once + if w not in static_seen: + static_seen.add(w) + # apply static writes + if writes := ChannelWrite.get_static_writes(w): + # END writes are not written, but become edges directly + for t in writes: + if t[0] == END: + edges.add((task.name, t[0], True, t[2])) + writes = [t for t in writes if t[0] != END] + conditionals.update( + {(task.name, t[0], t[1] or None): t[2] for t in writes} + ) + task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes]) + # collect sources + step_sources = { + task.name: { + ( + w[0], + (task.name, w[0], w[1] or None) in conditionals, + conditionals.get((task.name, w[0], w[1] or None)), + ) + for w in task.writes + } + for task in tasks.values() } - # apply input writes - input_writes = list(map_input(input_channels, {})) - updated_channels = apply_writes( - checkpoint, - channels, - [ - PregelTaskWrites((), INPUT, input_writes, []), - ], - get_next_version, - trigger_to_nodes, + sources.update(step_sources) + # invert triggers + trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = ( + defaultdict(set) ) - # prepare first tasks + for src, triggers in sources.items(): + for trigger, cond, label in triggers: + trigger_to_sources[trigger].add((src, cond, label)) + # apply writes + updated_channels = apply_writes( + checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes + ) + # prepare next tasks tasks = prepare_next_tasks( checkpoint, [], @@ -89,6 +159,7 @@ def draw_graph( managed, config, step, + limit, for_execution=True, store=None, checkpointer=None, @@ -96,149 +167,78 @@ def draw_graph( trigger_to_nodes=trigger_to_nodes, updated_channels=updated_channels, ) - start_tasks = tasks - # run the pregel loop - for step in range(step, limit): - if not tasks: - break - conditionals: dict[tuple[str, str, Any], Optional[str]] = {} - # run task writers - for task in tasks.values(): - for w in task.writers: - # apply regular writes - if isinstance(w, ChannelWrite): - empty_input = ( - cast(BaseChannel, specs["__root__"]).ValueType() - if "__root__" in specs - else None - ) - w.invoke(empty_input, task.config) - # apply conditional writes declared for static analysis, only once - if w not in static_seen: - static_seen.add(w) - # apply static writes - if writes := ChannelWrite.get_static_writes(w): - # END writes are not written, but become edges directly - for t in writes: - if t[0] == END: - edges.add((task.name, t[0], True, t[2])) - writes = [t for t in writes if t[0] != END] - conditionals.update( - {(task.name, t[0], t[1] or None): t[2] for t in writes} - ) - task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes]) - # collect sources - step_sources = { - task.name: { - ( - w[0], - (task.name, w[0], w[1] or None) in conditionals, - conditionals.get((task.name, w[0], w[1] or None)), - ) - for w in task.writes - } - for task in tasks.values() - } - sources.update(step_sources) - # invert triggers - trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = ( - defaultdict(set) - ) - for src, triggers in sources.items(): - for trigger, cond, label in triggers: - trigger_to_sources[trigger].add((src, cond, label)) - # apply writes - updated_channels = apply_writes( - checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes - ) - # prepare next tasks - tasks = prepare_next_tasks( - checkpoint, - [], - nodes, - channels, - managed, - config, - step, - for_execution=True, - store=None, - checkpointer=None, - manager=None, - trigger_to_nodes=trigger_to_nodes, - updated_channels=updated_channels, - ) - # collect edges - for task in tasks.values(): - added = False - for trigger in task.triggers: - for src, cond, label in sorted(trigger_to_sources[trigger]): - edges.add((src, task.name, cond, label)) - # if the edge is from this step, skip adding the implicit edges - if (trigger, cond, label) in step_sources.get(src, set()): - added = True - else: - sources[src].discard((trigger, cond, label)) - # if no edges from this step, add implicit edges from all previous tasks - if not added: - for src in step_sources: - edges.add((src, task.name, True, None)) - # assemble the graph - graph = Graph() - # add nodes - for name, node in nodes.items(): - metadata = dict(node.metadata or {}) - if name in interrupt_before_nodes and name in interrupt_after_nodes: - metadata["__interrupt"] = "before,after" - elif name in interrupt_before_nodes: - metadata["__interrupt"] = "before" - elif name in interrupt_after_nodes: - metadata["__interrupt"] = "after" - graph.add_node(node.bound, name, metadata=metadata or None) - # add start node - if START not in nodes: - graph.add_node(None, START) - for task in start_tasks.values(): - add_edge(graph, START, task.name) - # add discovered edges - for src, dest, is_conditional, label in sorted(edges): - add_edge( - graph, - src, - dest, - data=label if label != dest else None, - conditional=is_conditional, - ) - # add end edges - termini = {d for _, d, _, _ in edges if d != END}.difference( - s for s, _, _, _ in edges + # collect edges + for task in tasks.values(): + added = False + for trigger in task.triggers: + for src, cond, label in sorted(trigger_to_sources[trigger]): + edges.add((src, task.name, cond, label)) + # if the edge is from this step, skip adding the implicit edges + if (trigger, cond, label) in step_sources.get(src, set()): + added = True + else: + sources[src].discard((trigger, cond, label)) + # if no edges from this step, add implicit edges from all previous tasks + if not added: + for src in step_sources: + edges.add((src, task.name, True, None)) + # assemble the graph + graph = Graph() + # add nodes + for name, node in nodes.items(): + metadata = dict(node.metadata or {}) + if name in interrupt_before_nodes and name in interrupt_after_nodes: + metadata["__interrupt"] = "before,after" + elif name in interrupt_before_nodes: + metadata["__interrupt"] = "before" + elif name in interrupt_after_nodes: + metadata["__interrupt"] = "after" + graph.add_node(node.bound, name, metadata=metadata or None) + # add start node + if START not in nodes: + graph.add_node(None, START) + for task in start_tasks.values(): + add_edge(graph, START, task.name) + # add discovered edges + for src, dest, is_conditional, label in sorted(edges): + add_edge( + graph, + src, + dest, + data=label if label != dest else None, + conditional=is_conditional, ) - if termini: - for src in sorted(termini): - add_edge(graph, src, END) - elif len(step_sources) == 1: - for src in sorted(step_sources): - add_edge(graph, src, END, conditional=True) - # replace subgraphs - for name, subgraph in subgraphs.items(): - if ( - len(subgraph.nodes) > 1 - and name in graph.nodes - and subgraph.first_node() - and subgraph.last_node() - ): - subgraph.trim_first_node() - subgraph.trim_last_node() - # replace the node with the subgraph - graph.nodes.pop(name) - first, last = graph.extend(subgraph, prefix=name) - for idx, edge in enumerate(graph.edges): - if edge.source == name: - edge = edge.copy(source=cast(Node, last).id) - if edge.target == name: - edge = edge.copy(target=cast(Node, first).id) - graph.edges[idx] = edge + # add end edges + termini = {d for _, d, _, _ in edges if d != END}.difference( + s for s, _, _, _ in edges + ) + if termini: + for src in sorted(termini): + add_edge(graph, src, END) + elif len(step_sources) == 1: + for src in sorted(step_sources): + add_edge(graph, src, END, conditional=True) + # replace subgraphs + for name, subgraph in subgraphs.items(): + if ( + len(subgraph.nodes) > 1 + and name in graph.nodes + and subgraph.first_node() + and subgraph.last_node() + ): + subgraph.trim_first_node() + subgraph.trim_last_node() + # replace the node with the subgraph + graph.nodes.pop(name) + first, last = graph.extend(subgraph, prefix=name) + for idx, edge in enumerate(graph.edges): + if edge.source == name: + edge = edge.copy(source=cast(Node, last).id) + if edge.target == name: + edge = edge.copy(target=cast(Node, first).id) + graph.edges[idx] = edge - return graph + return graph def add_edge( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 88eb06886..a07778d8e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -89,7 +89,11 @@ from langgraph.pregel.algo import ( should_interrupt, task_path_str, ) -from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint +from langgraph.pregel.checkpoint import ( + channels_from_checkpoint, + create_checkpoint, + empty_checkpoint, +) from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results, @@ -111,7 +115,6 @@ from langgraph.pregel.io import ( read_channels, single, ) -from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.pregel.utils import get_new_channel_versions, is_xxh3_128_hexdigest from langgraph.store.base import BaseStore @@ -119,7 +122,6 @@ from langgraph.types import ( All, CachePolicy, Command, - LoopProtocol, PregelExecutableTask, PregelScratchpad, RetryPolicy, @@ -146,7 +148,13 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: return StreamProtocol(__call__, {mode for s in streams for mode in s.modes}) -class PregelLoop(LoopProtocol): +class PregelLoop: + config: RunnableConfig + store: Optional["BaseStore"] + stream: Optional[StreamProtocol] + step: int + stop: int + input: Optional[Any] input_model: Optional[type[BaseModel]] cache: Optional[BaseCache[WritesT]] @@ -226,13 +234,11 @@ class PregelLoop(LoopProtocol): cache_policy: Optional[CachePolicy] = None, checkpoint_during: bool = True, ) -> None: - super().__init__( - step=0, - stop=0, - config=config, - stream=stream, - store=store, - ) + self.stream = stream + self.config = config + self.store = store + self.step = 0 + self.stop = 0 self.input = input self.input_model = input_model self.checkpointer = checkpointer @@ -423,6 +429,7 @@ class PregelLoop(LoopProtocol): managed=self.managed, config=task.config, step=self.step, + stop=self.stop, for_execution=True, store=self.store, checkpointer=self.checkpointer, @@ -554,6 +561,7 @@ class PregelLoop(LoopProtocol): self.managed, self.config, self.step, + self.stop, for_execution=True, manager=self.manager, store=self.store, @@ -738,6 +746,7 @@ class PregelLoop(LoopProtocol): self.managed, self.config, self.step, + self.stop, for_execution=True, store=None, checkpointer=None, @@ -1146,8 +1155,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) - self.channels, self.managed = self.stack.enter_context( - ChannelsManager(self.specs, self.checkpoint, self) + self.channels, self.managed = channels_from_checkpoint( + self.specs, self.checkpoint ) self.stack.push(self._suppress_interrupt) self.status = "pending" @@ -1341,8 +1350,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): self.submit = await self.stack.enter_async_context( AsyncBackgroundExecutor(self.config) ) - self.channels, self.managed = await self.stack.enter_async_context( - AsyncChannelsManager(self.specs, self.checkpoint, self) + self.channels, self.managed = channels_from_checkpoint( + self.specs, self.checkpoint ) self.stack.push(self._suppress_interrupt) self.status = "pending" diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py deleted file mode 100644 index bdb583974..000000000 --- a/libs/langgraph/langgraph/pregel/manager.py +++ /dev/null @@ -1,76 +0,0 @@ -import asyncio -from collections.abc import AsyncIterator, Iterator, Mapping -from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager -from typing import Union - -from langgraph.channels.base import BaseChannel -from langgraph.checkpoint.base import Checkpoint -from langgraph.constants import MISSING -from langgraph.managed.base import ( - ManagedValueMapping, - ManagedValueSpec, -) -from langgraph.types import LoopProtocol - - -@contextmanager -def ChannelsManager( - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], - checkpoint: Checkpoint, - loop: LoopProtocol, -) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: - """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" - channel_specs: dict[str, BaseChannel] = {} - managed_specs: dict[str, ManagedValueSpec] = {} - for k, v in specs.items(): - if isinstance(v, BaseChannel): - channel_specs[k] = v - else: - managed_specs[k] = v - with ExitStack() as stack: - yield ( - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) - for k, v in channel_specs.items() - }, - ManagedValueMapping( - { - key: stack.enter_context(value.enter(loop)) - for key, value in managed_specs.items() - } - ), - ) - - -@asynccontextmanager -async def AsyncChannelsManager( - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], - checkpoint: Checkpoint, - loop: LoopProtocol, -) -> AsyncIterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: - """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" - channel_specs: dict[str, BaseChannel] = {} - managed_specs: dict[str, ManagedValueSpec] = {} - for k, v in specs.items(): - if isinstance(v, BaseChannel): - channel_specs[k] = v - else: - managed_specs[k] = v - async with AsyncExitStack() as stack: - # managed: create enter tasks with reference to spec, await them - if tasks := { - asyncio.create_task(stack.enter_async_context(value.aenter(loop))): key - for key, value in managed_specs.items() - }: - done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED) - else: - done = set() - yield ( - # channels: enter each channel with checkpoint - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) - for k, v in channel_specs.items() - }, - # managed: build mapping from spec to result - ManagedValueMapping({tasks[task]: task.result() for task in done}), - ) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 06bc20cfd..cc92e032c 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -27,7 +27,6 @@ from langgraph.utils.fields import get_update_as_tuples if TYPE_CHECKING: from langgraph.pregel.protocol import PregelProtocol - from langgraph.store.base import BaseStore try: @@ -380,31 +379,10 @@ class StreamProtocol: self.modes = modes -class LoopProtocol: - config: RunnableConfig - store: Optional["BaseStore"] - stream: Optional[StreamProtocol] - step: int - stop: int - - def __init__( - self, - *, - step: int, - stop: int, - config: RunnableConfig, - store: Optional["BaseStore"] = None, - stream: Optional[StreamProtocol] = None, - ) -> None: - self.stream = stream - self.config = config - self.store = store - self.step = step - self.stop = stop - - @dataclasses.dataclass(**_DC_KWARGS) class PregelScratchpad: + step: int + stop: int # call call_counter: Callable[[], int] # interrupt diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index a3a588f6b..dfe892e1d 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -1,46 +1,48 @@ from langgraph.checkpoint.base import empty_checkpoint from langgraph.constants import PULL, PUSH from langgraph.pregel.algo import prepare_next_tasks, task_path_str -from langgraph.pregel.manager import ChannelsManager +from langgraph.pregel.checkpoint import channels_from_checkpoint def test_prepare_next_tasks() -> None: config = {} processes = {} checkpoint = empty_checkpoint() + channels, managed = channels_from_checkpoint({}, checkpoint) - with ChannelsManager({}, checkpoint, config) as (channels, managed): - assert ( - prepare_next_tasks( - checkpoint, - {}, - processes, - channels, - managed, - config, - 0, - for_execution=False, - ) - == {} + assert ( + prepare_next_tasks( + checkpoint, + {}, + processes, + channels, + managed, + config, + 0, + -1, + for_execution=False, ) - assert ( - prepare_next_tasks( - checkpoint, - {}, - processes, - channels, - managed, - config, - 0, - for_execution=True, - checkpointer=None, - store=None, - manager=None, - ) - == {} + == {} + ) + assert ( + prepare_next_tasks( + checkpoint, + {}, + processes, + channels, + managed, + config, + 0, + -1, + for_execution=True, + checkpointer=None, + store=None, + manager=None, ) + == {} + ) - # TODO: add more tests + # TODO: add more tests def test_tuple_str() -> None: diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 9d0af1c8b..2725fdf79 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -22,12 +22,12 @@ from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel from langgraph.pregel.algo import checkpoint_null_version, prepare_single_task +from langgraph.pregel.checkpoint import channels_from_checkpoint from langgraph.pregel.executor import ( AsyncBackgroundExecutor, BackgroundExecutor, Submit, ) -from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.runner import PregelRunner from langgraph.scheduler.kafka.retry import aretry, retry from langgraph.scheduler.kafka.types import ( @@ -41,7 +41,7 @@ from langgraph.scheduler.kafka.types import ( Sendable, Topics, ) -from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy +from langgraph.types import PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable, recast_checkpoint_ns @@ -184,19 +184,10 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): raise RuntimeError("Checkpoint not found") if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() - async with ( - AsyncChannelsManager( - graph.channels, - saved.checkpoint, - LoopProtocol( - config=msg["config"], - store=self.graph.store, - step=saved.metadata["step"] + 1, - stop=saved.metadata["step"] + 2, - ), - ) as (channels, managed), - AsyncBackgroundExecutor(msg["config"]) as submit, - ): + async with AsyncBackgroundExecutor(msg["config"]) as submit: + channels, managed = channels_from_checkpoint( + graph.channels, saved.checkpoint + ) if task := await asyncio.to_thread( prepare_single_task, msg["task"]["path"], @@ -208,6 +199,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): managed=managed, config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}), step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, for_execution=True, checkpointer=self.graph.checkpointer, store=self.graph.store, @@ -404,19 +396,10 @@ class KafkaExecutor(AbstractContextManager): raise RuntimeError("Checkpoint not found") if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() - with ( - ChannelsManager( - graph.channels, - saved.checkpoint, - LoopProtocol( - config=msg["config"], - store=self.graph.store, - step=saved.metadata["step"] + 1, - stop=saved.metadata["step"] + 2, - ), - ) as (channels, managed), - BackgroundExecutor({}) as submit, - ): + with BackgroundExecutor({}) as submit: + channels, managed = channels_from_checkpoint( + graph.channels, saved.checkpoint + ) if task := prepare_single_task( msg["task"]["path"], msg["task"]["id"], @@ -427,6 +410,7 @@ class KafkaExecutor(AbstractContextManager): managed=managed, config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}), step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, for_execution=True, checkpointer=self.graph.checkpointer, checkpoint_id_bytes=binascii.unhexlify( diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 4febf76f0..36eca328b 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics_async from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -199,6 +199,8 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -272,6 +274,8 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -375,6 +379,8 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -488,6 +494,8 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -556,6 +564,8 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -680,6 +690,8 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index 1071e0524..af6176d22 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -15,7 +15,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.default_sync import DefaultProducer from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -198,6 +198,8 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -271,6 +273,8 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -374,6 +378,8 @@ def test_subgraph_w_interrupt( "__pregel_task_id": history[0].tasks[0].id, "__pregel_previous": None, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -486,6 +492,8 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -554,6 +562,8 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None, @@ -678,6 +688,8 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "step": AnyInt(), + "stop": AnyInt(), "subgraph_counter": None, "call_counter": None, "interrupt_counter": None,