From 644a6c3b6355d628efd2c51f616ece4bc21c1a80 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Thu, 24 Apr 2025 13:26:01 -0700 Subject: [PATCH] use resume instead of resume_map and deprecate old mapping task_id -> resume logic --- libs/langgraph/langgraph/pregel/algo.py | 12 +----------- libs/langgraph/langgraph/pregel/io.py | 23 ++--------------------- libs/langgraph/langgraph/pregel/loop.py | 13 ++++++++----- libs/langgraph/langgraph/pregel/utils.py | 6 ++++++ libs/langgraph/langgraph/types.py | 7 +++++-- 5 files changed, 22 insertions(+), 39 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 029cf73be..b855ccd7e 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -908,17 +908,7 @@ def _scratchpad( if not isinstance(task_resume_write, list): task_resume_write = [task_resume_write] else: - # find task-specific resume value - for w in pending_writes: - if w[0] == task_id and w[1] == RESUME: - task_resume_write = w[2] - if not isinstance(task_resume_write, list): - task_resume_write = [task_resume_write] - break - else: - task_resume_write = [] - # clear var - del w + task_resume_write = [] else: null_resume_write = None task_resume_write = [] diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 026051a4d..45dba331d 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -24,15 +24,6 @@ from langgraph.pregel.log import logger from langgraph.types import Command, PregelExecutableTask, Send -def is_task_id(task_id: str) -> bool: - """Check if a string is a valid task id.""" - try: - UUID(task_id) - except Exception: - return False - return True - - def read_channel( channels: Mapping[str, BaseChannel], chan: str, @@ -66,9 +57,7 @@ def read_channels( return values -def map_command( - cmd: Command, pending_writes: list[PendingWrite] -) -> Iterator[tuple[str, str, Any]]: +def map_command(cmd: Command) -> Iterator[tuple[str, str, Any]]: """Map input chunk to a sequence of pending writes in the form (channel, value).""" if cmd.graph == Command.PARENT: raise InvalidUpdateError("There is no parent graph") @@ -87,15 +76,7 @@ def map_command( f"In Command.goto, expected Send/str, got {type(send).__name__}" ) if cmd.resume is not None: - if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume): - for tid, resume in cmd.resume.items(): - existing: list[Any] = next( - (w[2] for w in pending_writes if w[0] == tid and w[1] == RESUME), [] - ) - existing.append(resume) - yield (tid, RESUME, existing) - else: - yield (NULL_TASK_ID, RESUME, cmd.resume) + yield (NULL_TASK_ID, RESUME, cmd.resume) if cmd.update: for k, v in cmd._update_as_tuples(): yield (NULL_TASK_ID, k, v) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 978cbd0df..e71972fb0 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -113,7 +113,7 @@ from langgraph.pregel.io import ( ) from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.read import PregelNode -from langgraph.pregel.utils import get_new_channel_versions +from langgraph.pregel.utils import get_new_channel_versions, is_xxh3_128_hexdigest from langgraph.store.base import BaseStore from langgraph.types import ( All, @@ -650,17 +650,20 @@ class PregelLoop(LoopProtocol): # map command to writes if isinstance(self.input, Command): - if self.input.resume_map: - self.config[CONF][CONFIG_KEY_RESUME_MAP] = self.input.resume_map + if (resume := self.input.resume) is not None: + if isinstance(resume, dict) and all( + is_xxh3_128_hexdigest(k) for k in resume + ): + self.config[CONF][CONFIG_KEY_RESUME_MAP] = self.input.resume if self.input.resume is not None and not self.checkpointer: raise RuntimeError( "Cannot use Command(resume=...) without checkpointer" ) writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list) # group writes by task ID - for tid, c, v in map_command(self.input, self.checkpoint_pending_writes): + for tid, c, v in map_command(self.input): writes[tid].append((c, v)) - if not writes and not self.input.resume_map: + if not writes: raise EmptyInputError("Received empty Command input") # save writes for tid, ws in writes.items(): diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index fd62e345f..da0076f6b 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,5 +1,6 @@ import ast import inspect +import re import textwrap from typing import Any, Callable, Optional @@ -207,3 +208,8 @@ class NonLocals(ast.NodeVisitor): parent = parent.value if isinstance(parent, ast.Name): self.loads.add(parent.id + "." + attr_expr) + + +def is_xxh3_128_hexdigest(value: str) -> bool: + """Check if the given string matches the format of xxh3_128_hexdigest.""" + return bool(re.fullmatch(r"[0-9a-f]{32}", value)) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4c088ebe8..8062481f9 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -294,6 +294,10 @@ class Command(Generic[N], ToolOutputMixin): - Command.PARENT: closest parent graph update: update to apply to the graph's state. resume: value to resume execution with. To be used together with [`interrupt()`][langgraph.types.interrupt]. + Can be one of the following: + + - mapping of interrupt ids to resume values + - a single value with which to resume the next interrupt goto: can be one of the following: - name of the node to navigate to next (any node that belongs to the specified `graph`) @@ -305,9 +309,8 @@ class Command(Generic[N], ToolOutputMixin): graph: Optional[str] = None update: Optional[Any] = None - resume: Optional[Union[Any, dict[str, Any]]] = None + resume: Optional[Union[dict[str, Any], Any]] = None goto: Union[Send, Sequence[Union[Send, str]], str] = () - resume_map: Optional[dict[str, Any]] = None def __repr__(self) -> str: # get all non-None values