diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py index ce0d21189..524611e5a 100644 --- a/libs/langgraph/langgraph/channels/manager.py +++ b/libs/langgraph/langgraph/channels/manager.py @@ -61,4 +61,6 @@ def create_checkpoint( channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), + # checkpoints are saved only at the end of a step, ie. no current tasks + current_tasks={}, ) diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 2f73b52a6..8cb278f32 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -54,6 +54,10 @@ class CheckpointMetadata(TypedDict, total=False): """ +class TaskInfo(TypedDict): + status: Literal["scheduled", "success", "error"] + + class Checkpoint(TypedDict): """State snapshot at a given point in time.""" @@ -85,6 +89,8 @@ class Checkpoint(TypedDict): pending_sends: List[Send] """List of packets sent to nodes but not yet processed. Cleared by the next checkpoint.""" + current_tasks: Dict[str, TaskInfo] + """Map from task ID to task info.""" def empty_checkpoint() -> Checkpoint: @@ -96,6 +102,7 @@ def empty_checkpoint() -> Checkpoint: channel_versions={}, versions_seen=defaultdict(dict), pending_sends=[], + current_tasks={}, ) @@ -111,6 +118,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: {k: v.copy() for k, v in checkpoint["versions_seen"].items()}, ), pending_sends=checkpoint.get("pending_sends", []).copy(), + current_tasks=checkpoint.get("current_tasks", {}).copy(), ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index b61d05982..41b517e0a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -59,7 +59,6 @@ from langgraph.channels.manager import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, - CheckpointMetadata, copy_checkpoint, empty_checkpoint, ) @@ -77,28 +76,20 @@ from langgraph.managed.base import ( ) from langgraph.pregel.algo import ( apply_writes, - increment, local_read, prepare_next_tasks, - should_interrupt, ) from langgraph.pregel.debug import ( - map_debug_checkpoint, map_debug_task_results, - map_debug_tasks, print_step_checkpoint, print_step_tasks, print_step_writes, ) -from langgraph.pregel.executor import AsyncBackgroundExecutor from langgraph.pregel.io import ( - map_input, map_output_updates, - map_output_values, read_channels, - single, ) -from langgraph.pregel.loop import PregelLoop +from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry from langgraph.pregel.types import ( @@ -839,8 +830,7 @@ class Pregel( debug=debug, ) - # create channels from checkpoint - with PregelLoop( + with SyncPregelLoop( input, config=config, checkpointer=self.checkpointer, graph=self ) as loop: # Similarly to Bulk Synchronous Parallel / Pregel model @@ -852,6 +842,7 @@ class Pregel( output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, + manager=run_manager, ): # debug flag if self.debug: @@ -871,19 +862,16 @@ class Pregel( # debug flag if debug: print_step_tasks(loop.step, loop.tasks) - # TODO move to tick() ? - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_tasks(loop.step, loop.tasks), - ) # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - loop.submit(run_with_retry, task, self.retry_policy): task + loop.submit( + run_with_retry, + task, + self.retry_policy, + ): task for task in loop.tasks if not task.writes } @@ -1063,7 +1051,6 @@ class Pregel( None, ) try: - loop = asyncio.get_event_loop() if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -1085,211 +1072,59 @@ class Pregel( interrupt_after=interrupt_after, debug=debug, ) - # copy nodes to ignore mutations during execution - processes = {**self.nodes} - # get checkpoint from saver, or create an empty one - saved = ( - await self.checkpointer.aget_tuple(config) - if self.checkpointer - else None - ) - checkpoint = saved.checkpoint if saved else empty_checkpoint() - - # merge configurable fields with previous checkpoint config - checkpoint_config = config - if saved: - checkpoint_config = { - **config, - **saved.config, - "configurable": { - **config.get("configurable", {}), - **saved.config["configurable"], - }, - } - - start = saved.metadata.get("step", -2) + 1 if saved else -1 - # create channels from checkpoint - async with AsyncBackgroundExecutor() as submit, AsyncChannelsManager( - self.channels, checkpoint, config - ) as channels, AsyncManagedValuesManager( - self.managed_values_dict, config, self - ) as managed: - - def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None: - if self.checkpointer is not None: - submit( - self.checkpointer.aput_writes, - { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - }, - writes, - task_id, - ) - - def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: - print(metadata) - nonlocal checkpoint, checkpoint_config, channels - - if self.checkpointer is None: - return - if debug: - print_step_checkpoint( - metadata["step"], channels, self.stream_channels_list - ) - - # create new checkpoint - checkpoint = create_checkpoint( - checkpoint, channels, metadata["step"] - ) - # save it, without blocking - submit( - self.checkpointer.aput, - checkpoint_config, - copy_checkpoint(checkpoint), - metadata, - ) - - # update checkpoint config - checkpoint_config = { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - } - # yield debug checkpoint event - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - metadata["step"], - checkpoint_config, - channels, - self.stream_channels_asis, - metadata, - ), - ) - - # map inputs to channel updates - if input_writes := deque(map_input(self.input_channels, input)): - # discard any unfinished tasks from previous checkpoint - checkpoint, _ = prepare_next_tasks( - checkpoint, - processes, - channels, - managed, - config, - -1, - for_execution=True, - get_next_version=( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - # apply input writes - apply_writes( - checkpoint, - channels, - input_writes, - ( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - # save input checkpoint - for chunk in put_checkpoint( - {"source": "input", "step": start, "writes": input} - ): - yield chunk - # increment start to 0 - start += 1 - else: - # no input is taken as signal to proceed past previous interrupt - checkpoint = copy_checkpoint(checkpoint) - for k in channels: - if k in checkpoint["channel_versions"]: - version = checkpoint["channel_versions"][k] - checkpoint["versions_seen"][INTERRUPT][k] = version - + async with AsyncPregelLoop( + input, config=config, checkpointer=self.checkpointer, graph=self + ) as loop: + aioloop = asyncio.get_event_loop() # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates - # channel updates from step N are only visible in step N+1, + # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, - # channel updates being applied only at the transition between steps - stop = start + config["recursion_limit"] + 1 - for step in range(start, stop): - next_checkpoint, next_tasks = prepare_next_tasks( - checkpoint, - processes, - channels, - managed, - config, - step, - for_execution=True, - manager=run_manager, - get_next_version=( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - - # assign pending writes to tasks - if saved and saved.pending_writes: - for task in next_tasks: - task.writes.extend( - (c, v) - for tid, c, v in saved.pending_writes - if tid == task.id - ) - - # if no more tasks, we're done - if not next_tasks: - if step == start: - raise ValueError("No tasks to run in graph.") - else: - break - - # before execution, check if we should interrupt - if should_interrupt(checkpoint, interrupt_before, next_tasks): - break - else: - checkpoint = next_checkpoint - + # with channel updates applied only at the transition between steps + while loop.tick( + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + manager=run_manager, + ): + # debug flag + if self.debug: + print_step_checkpoint( + loop.checkpoint_metadata, + loop.channels, + self.stream_channels_list, + ) + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # debug flag if debug: - print_step_tasks(step, next_tasks) - if "debug" in stream_modes: - for chunk in _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_tasks(step, next_tasks), - ): - yield chunk + print_step_tasks(loop.step, loop.tasks) # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - submit( + loop.submit( arun_with_retry, task, self.retry_policy, - do_stream, + stream=do_stream, __name__=task.name, __cancel_on_exit__=True, ): task - for task in next_tasks + for task in loop.tasks if not task.writes } end_time = ( - self.step_timeout + loop.time() if self.step_timeout else None + self.step_timeout + aioloop.time() + if self.step_timeout + else None ) if not futures: done, inflight = set(), set() @@ -1298,7 +1133,7 @@ class Pregel( futures, return_when=asyncio.FIRST_COMPLETED, timeout=( - max(0, end_time - loop.time()) if end_time else None + max(0, end_time - aioloop.time()) if end_time else None ), ) if not done: @@ -1307,13 +1142,12 @@ class Pregel( task = futures.pop(fut) if fut.exception() is not None: # we got an exception, break out of while loop - # exception will be handle in panic_or_proceed + # exception will be handled in panic_or_proceed futures.clear() else: - # save task writes to checkpointer, unless this - # is the single or last task in this step - if futures: - put_writes(task.id, task.writes) + print(loop.step, task.name, stream_modes) + # save task writes to checkpointer + loop.put_writes(task.id, task.writes) # yield updates output for the finished task if "updates" in stream_modes: for chunk in _with_mode( @@ -1327,7 +1161,9 @@ class Pregel( "debug", isinstance(stream_mode, list), map_debug_task_results( - step, [task], self.stream_channels_list + loop.step, + [task], + self.stream_channels_list, ), ): yield chunk @@ -1336,71 +1172,36 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, step, asyncio.TimeoutError) + _panic_or_proceed(done, inflight, loop.step, asyncio.TimeoutError) # don't keep futures around in memory longer than needed del done, inflight, futures - - # combine pending writes from all tasks - pending_writes = deque[tuple[str, Any]]() - for task in next_tasks: - pending_writes.extend(task.writes) - + # debug flag if debug: print_step_writes( - step, pending_writes, self.stream_channels_list + loop.step, + [w for t in loop.tasks for w in t.writes], + self.stream_channels_list, ) - - # apply writes to channels - apply_writes( - checkpoint, - channels, - pending_writes, - ( - self.checkpointer.get_next_version - if self.checkpointer - else increment - ), - ) - - # yield current values - if "values" in stream_modes: - for chunk in _with_mode( - "values", - isinstance(stream_mode, list), - map_output_values(output_keys, pending_writes, channels), - ): - yield chunk - - # save end of step checkpoint - for chunk in put_checkpoint( - { - "source": "loop", - "step": step, - "writes": ( - single(map_output_updates(output_keys, next_tasks)) - if self.stream_mode == "updates" - else single( - map_output_values( - output_keys, pending_writes, channels - ) - ) - ), - } - ): - yield chunk - - # after execution, check if we should interrupt - if should_interrupt(checkpoint, interrupt_after, next_tasks): - break - else: + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # handle exit + if loop.status == "out_of_steps": raise GraphRecursionError( - f"Recursion limit of {config['recursion_limit']} reached" - "without hitting a stop condition. You can increase the limit" - "by setting the `recursion_limit` config key." + f"Recursion limit of {config['recursion_limit']} reached " + "without hitting a stop condition. You can increase the " + "limit by setting the `recursion_limit` config key." ) # set final channel values as run output - await run_manager.on_chain_end(read_channels(channels, output_keys)) + await run_manager.on_chain_end( + read_channels(loop.channels, output_keys) + ) except BaseException as e: await asyncio.shield(run_manager.on_chain_error(e)) raise diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 4d43aab60..352b8cf51 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -6,6 +6,7 @@ from contextvars import copy_context from types import TracebackType from typing import ( AsyncContextManager, + Awaitable, Callable, Iterator, Optional, @@ -78,7 +79,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): def submit( self, - fn: Callable[P, T], + fn: Callable[P, Awaitable[T]], *args: P.args, __name__: Optional[str] = None, __cancel_on_exit__: bool = False, @@ -101,7 +102,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): else: self.tasks.pop(task) - async def __aenter__(self) -> Submit: + async def __aenter__(self) -> "submit": return self.submit async def exit(self) -> None: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 003204b4e..b8c3fdc4e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,9 +1,11 @@ +import asyncio from collections import deque -from contextlib import ExitStack +from contextlib import AsyncExitStack, ExitStack from types import TracebackType from typing import ( TYPE_CHECKING, Any, + AsyncContextManager, Callable, ContextManager, List, @@ -17,11 +19,16 @@ from typing import ( Union, ) +from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager from langchain_core.runnables import RunnableConfig from typing_extensions import Self from langgraph.channels.base import BaseChannel -from langgraph.channels.manager import ChannelsManager, create_checkpoint +from langgraph.channels.manager import ( + AsyncChannelsManager, + ChannelsManager, + create_checkpoint, +) from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -32,15 +39,23 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.constants import INTERRUPT -from langgraph.managed.base import ManagedValueMapping, ManagedValuesManager +from langgraph.managed.base import ( + AsyncManagedValuesManager, + ManagedValueMapping, + ManagedValuesManager, +) from langgraph.pregel.algo import ( apply_writes, increment, prepare_next_tasks, should_interrupt, ) -from langgraph.pregel.debug import map_debug_checkpoint -from langgraph.pregel.executor import BackgroundExecutor, Submit +from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks +from langgraph.pregel.executor import ( + AsyncBackgroundExecutor, + BackgroundExecutor, + Submit, +) from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single from langgraph.pregel.types import PregelExecutableTask @@ -52,11 +67,17 @@ V = TypeVar("V") INPUT_DONE = object() -class PregelLoop(ContextManager): +class PregelLoop: input: Optional[Any] config: RunnableConfig checkpointer: Optional[BaseCheckpointSaver] - get_next_version: Callable[[Optional[V]], V] + checkpointer_get_next_version: Callable[[Optional[V]], V] + checkpointer_put_writes: Optional[ + Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], ...] + ] + checkpointer_put: Optional[ + Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], ...] + ] graph: "Pregel" submit: Submit @@ -73,6 +94,203 @@ class PregelLoop(ContextManager): tasks: Sequence[PregelExecutableTask] stream: deque[Tuple[str, Any]] + # public + + def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: + """Mark tasks as scheduled, to be used by queue-based executors.""" + raise NotImplementedError + + def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + """Put writes for a task, to be read by the next tick.""" + self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) + if self.checkpointer_put_writes is not None: + self.submit( + self.checkpointer_put_writes, + { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + }, + writes, + task_id, + ) + + def tick( + self, + *, + output_keys: Union[str, Sequence[str]] = None, + interrupt_after: Optional[Sequence[str]] = None, + interrupt_before: Optional[Sequence[str]] = None, + manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, + ) -> bool: + """Execute a single iteration of the Pregel loop. + Returns True if more iterations are needed.""" + + if self.status != "pending": + raise RuntimeError("Cannot tick when status is no longer 'pending'") + + if self.input is not INPUT_DONE: + self._first() + elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( + self.tasks + ): + # assign writes to tasks, apply them in order + grouped: dict[str, list[tuple[str, Any]]] = {} + for tid, k, v in self.checkpoint_pending_writes: + grouped.setdefault(tid, []).append((k, v)) + writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] + # all tasks have finished + apply_writes( + self.checkpoint, + self.channels, + writes, + self.checkpointer_get_next_version, + ) + # produce values output + self.stream.extend( + ("values", v) + for v in map_output_values(output_keys, writes, self.channels) + ) + # clear pending writes + self.checkpoint_pending_writes.clear() + # save checkpoint + self._put_checkpoint( + { + "source": "loop", + "writes": single( + map_output_updates(output_keys, self.tasks) + if self.graph.stream_mode == "updates" + else map_output_values(output_keys, writes, self.channels) + ), + } + ) + # after execution, check if we should interrupt + if should_interrupt(self.checkpoint, interrupt_after, self.tasks): + self.status = "interrupt_after" + return False + else: + return False + + # check if iteration limit is reached + if self.step > self.config["recursion_limit"]: + self.status = "out_of_steps" + return False + + # prepare next tasks + prev_checkpoint = self.checkpoint + self.checkpoint, self.tasks = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + get_next_version=self.checkpointer_get_next_version, + manager=manager, + ) + + # if no more tasks, we're done + if not self.tasks: + self.status = "done" + return False + + # if there are pending writes from a previous loop, apply them + if self.checkpoint_pending_writes: + for tid, k, v in self.checkpoint_pending_writes: + if task := next((t for t in self.tasks if t.id == tid), None): + task.writes.append((k, v)) + + # before execution, check if we should interrupt + if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): + self.status = "interrupt_before" + return False + + # produce debug output + self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks)) + + return True + + # private + + def _first(self) -> None: + # map inputs to channel updates + if input_writes := deque(map_input(self.graph.input_channels, self.input)): + # discard any unfinished tasks from previous checkpoint + self.checkpoint, _ = prepare_next_tasks( + self.checkpoint, + self.graph.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + get_next_version=self.checkpointer_get_next_version, + ) + # apply input writes + apply_writes( + self.checkpoint, + self.channels, + input_writes, + self.checkpointer_get_next_version, + ) + # save input checkpoint + self._put_checkpoint({"source": "input", "writes": self.input}) + else: + # no input is taken as signal to proceed past previous interrupt + self.checkpoint = copy_checkpoint(self.checkpoint) + for k in self.channels: + if k in self.checkpoint["channel_versions"]: + version = self.checkpoint["channel_versions"][k] + self.checkpoint["versions_seen"][INTERRUPT][k] = version + # done with input + self.input = INPUT_DONE + + def _put_checkpoint( + self, + metadata: CheckpointMetadata, + ) -> None: + # assign step + metadata["step"] = self.step + # bail if no checkpointer + if self.checkpointer_put is not None: + # create new checkpoint + self.checkpoint_metadata = metadata + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels, self.step + ) + # save it, without blocking + self.submit( + self.checkpointer_put, + self.checkpoint_config, + copy_checkpoint(self.checkpoint), + self.checkpoint_metadata, + ) + self.checkpoint_config = { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + } + # produce debug output + self.stream.extend( + ("debug", v) + for v in map_debug_checkpoint( + self.step, + self.checkpoint_config, + self.channels, + self.graph.stream_channels_asis, + self.checkpoint_metadata, + ) + ) + # increment step + self.step += 1 + + +class SyncPregelLoop(PregelLoop, ContextManager): def __init__( self, input: Optional[Any], @@ -86,13 +304,17 @@ class PregelLoop(ContextManager): self.input = input self.config = config self.checkpointer = checkpointer - self.get_next_version = ( + self.checkpointer_get_next_version = ( checkpointer.get_next_version if checkpointer else increment ) + self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None + self.checkpointer_put = checkpointer.put if checkpointer else None self.graph = graph # TODO if managed values no longer needs graph we can replace with # managed_specs, channel_specs + # context manager + def __enter__(self) -> Self: saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None @@ -132,183 +354,73 @@ class PregelLoop(ContextManager): del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) - def tick( + +class AsyncPregelLoop(PregelLoop, AsyncContextManager): + def __init__( self, + input: Optional[Any], *, - output_keys: Union[str, Sequence[str]] = None, - interrupt_after: Optional[Sequence[str]] = None, - interrupt_before: Optional[Sequence[str]] = None, - ) -> bool: - if self.status != "pending": - raise RuntimeError(f"Cannot tick when status is {self.status}") - if self.input is not INPUT_DONE: - self.first() - elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( - self.tasks - ): - # assign writes to tasks, apply them in order - grouped: dict[str, list[tuple[str, Any]]] = {} - for tid, k, v in self.checkpoint_pending_writes: - grouped.setdefault(tid, []).append((k, v)) - writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] - # all tasks have finished - apply_writes( - self.checkpoint, - self.channels, - writes, - self.get_next_version, - ) - # produce values output - self.stream.extend( - ("values", v) - for v in map_output_values(output_keys, writes, self.channels) - ) - # clear pending writes - self.checkpoint_pending_writes.clear() - # save checkpoint - self.put_checkpoint( - { - "source": "loop", - "writes": single( - map_output_updates(output_keys, self.tasks) - if self.graph.stream_mode == "updates" - else map_output_values(output_keys, writes, self.channels) - ), - } - ) - # after execution, check if we should interrupt - if should_interrupt(self.checkpoint, interrupt_after, self.tasks): - self.status = "interrupt_after" - return False - else: - return False - - # check if iteration limit is reached - if self.step > self.config["recursion_limit"]: - self.status = "out_of_steps" - return False - - # prepare next tasks - prev_checkpoint = self.checkpoint - self.checkpoint, self.tasks = prepare_next_tasks( - self.checkpoint, - self.graph.nodes, - self.channels, - self.managed, - self.config, - self.step, - for_execution=True, - get_next_version=self.get_next_version, - ) - - # if no more tasks, we're done - if not self.tasks: - self.status = "done" - return False - - # TODO how to make this work for both - # - online case: we should schedule remaining tasks - # - offline case: we should just bail, as other tasks were scheduled before - # if there are pending writes from a previous loop, apply them - if self.checkpoint_pending_writes: - for tid, k, v in self.checkpoint_pending_writes: - if task := next((t for t in self.tasks if t.id == tid), None): - task.writes.append((k, v)) - - # before execution, check if we should interrupt - if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): - self.status = "interrupt_before" - return False - - return True - - def first(self) -> None: - # map inputs to channel updates - if input_writes := deque(map_input(self.graph.input_channels, self.input)): - # discard any unfinished tasks from previous checkpoint - self.checkpoint, _ = prepare_next_tasks( - self.checkpoint, - self.graph.nodes, - self.channels, - self.managed, - self.config, - self.step, - for_execution=True, - get_next_version=self.get_next_version, - # TODO missing run_manager - ) - # apply input writes - apply_writes( - self.checkpoint, - self.channels, - input_writes, - self.get_next_version, - ) - # save input checkpoint - self.put_checkpoint({"source": "input", "writes": self.input}) - else: - # no input is taken as signal to proceed past previous interrupt - self.checkpoint = copy_checkpoint(self.checkpoint) - for k in self.channels: - if k in self.checkpoint["channel_versions"]: - version = self.checkpoint["channel_versions"][k] - self.checkpoint["versions_seen"][INTERRUPT][k] = version - # done with input - self.input = INPUT_DONE - - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: - self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) - if self.checkpointer is not None: - self.submit( - self.checkpointer.put_writes, - { - **self.checkpoint_config, - "configurable": { - **self.checkpoint_config["configurable"], - "thread_ts": self.checkpoint["id"], - }, - }, - writes, - task_id, - ) - - def put_checkpoint( - self, - metadata: CheckpointMetadata, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", ) -> None: - # assign step - metadata["step"] = self.step - # bail if no checkpointer - if self.checkpointer is not None: - # create new checkpoint - self.checkpoint_metadata = metadata - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.step + self.stream = deque() + self.stack = AsyncExitStack() + self.input = input + self.config = config + self.checkpointer = checkpointer + self.checkpointer_get_next_version = ( + checkpointer.get_next_version if checkpointer else increment + ) + self.checkpointer_put_writes = ( + checkpointer.aput_writes if checkpointer else None + ) + self.checkpointer_put = checkpointer.aput if checkpointer else None + self.graph = graph + # TODO if managed values no longer needs graph we can replace with + # managed_specs, channel_specs + + # context manager + + async def __aenter__(self) -> Self: + saved = ( + await self.checkpointer.aget_tuple(self.config) + if self.checkpointer + else None + ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) + self.checkpoint_config = { + **self.config, + **saved.config, + "configurable": { + **self.config.get("configurable", {}), + **saved.config.get("configurable", {}), + }, + } + self.checkpoint = saved.checkpoint + self.checkpoint_metadata = saved.metadata + self.checkpoint_pending_writes = saved.pending_writes + + self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor()) + self.channels = await self.stack.enter_async_context( + AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config) + ) + self.managed = await self.stack.enter_async_context( + AsyncManagedValuesManager( + self.graph.managed_values_dict, self.config, self.graph ) - # save it, without blocking - self.submit( - self.checkpointer.put, - self.checkpoint_config, - copy_checkpoint(self.checkpoint), - self.checkpoint_metadata, - ) - self.checkpoint_config = { - **self.checkpoint_config, - "configurable": { - **self.checkpoint_config["configurable"], - "thread_ts": self.checkpoint["id"], - }, - } - # produce debug output - self.stream.extend( - ("debug", v) - for v in map_debug_checkpoint( - self.step, - self.checkpoint_config, - self.channels, - self.graph.stream_channels_asis, - self.checkpoint_metadata, - ) - ) - # increment step - self.step += 1 + ) + self.status = "pending" + self.step = self.checkpoint_metadata["step"] + 1 + + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + del self.graph + return await asyncio.shield( + self.stack.__aexit__(exc_type, exc_value, traceback) + ) diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 7b56eba46..2e3c2098b 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,8 +1,6 @@ from collections import defaultdict from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union -from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.checkpoint.sqlite import SqliteSaver import pytest from langchain_core.callbacks import ( CallbackManagerForLLMRun, @@ -25,6 +23,7 @@ from langchain_core.tools import BaseTool from langchain_core.tools import tool as dec_tool from pydantic import BaseModel as BaseModelV2 +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.prebuilt import ( ToolNode, ValidationNode, @@ -110,6 +109,7 @@ def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]): }, ), "pending_sends": [], + "current_tasks": {}, } assert saved.metadata == { "source": "loop", @@ -168,6 +168,7 @@ async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]): }, ), "pending_sends": [], + "current_tasks": {}, } assert saved.metadata == { "source": "loop",