diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 670e85b3d..f8d280b96 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -25,7 +25,7 @@ from langchain_core.load.serializable import Serializable from zoneinfo import ZoneInfo from langgraph.checkpoint.serde.base import SerializerProtocol -from langgraph.checkpoint.serde.types import CommandProtocol, SendProtocol +from langgraph.checkpoint.serde.types import SendProtocol from langgraph.store.base import Item LC_REVIVER = Reviver() @@ -122,11 +122,6 @@ class JsonPlusSerializer(SerializerProtocol): return self._encode_constructor_args( obj.__class__, kwargs={"node": obj.node, "arg": obj.arg} ) - elif isinstance(obj, CommandProtocol): - return self._encode_constructor_args( - obj.__class__, - kwargs={k: getattr(obj, k) for k in obj.__all_slots__}, - ) elif isinstance(obj, (bytes, bytearray)): return self._encode_constructor_args( obj.__class__, method="fromhex", args=(obj.hex(),) @@ -407,17 +402,6 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]: (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)), ), ) - elif isinstance(obj, CommandProtocol): - return msgpack.ExtType( - EXT_CONSTRUCTOR_KW_ARGS, - _msgpack_enc( - ( - obj.__class__.__module__, - obj.__class__.__name__, - {k: getattr(obj, k) for k in obj.__all_slots__}, - ), - ), - ) elif dataclasses.is_dataclass(obj): # doesn't use dataclasses.asdict to avoid deepcopy and recursion return msgpack.ExtType( diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index e258735bc..1df967b5f 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -4,7 +4,6 @@ from typing import ( Protocol, Sequence, TypeVar, - Union, runtime_checkable, ) @@ -51,11 +50,3 @@ class SendProtocol(Protocol): def __repr__(self) -> str: ... def __eq__(self, value: object) -> bool: ... - - -@runtime_checkable -class CommandProtocol(Protocol): - # Mirrors langgraph.types.Command - update: Optional[dict[str, Any]] - send: Union[Any, Sequence[Any]] - __all_slots__: set[str] diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 2e3d13120..2450b42b1 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -70,7 +70,7 @@ class NodeInterrupt(GraphInterrupt): """Raised by a node to interrupt execution.""" def __init__(self, value: Any) -> None: - super().__init__([Interrupt(value)]) + super().__init__([Interrupt(value=value)]) class GraphDelegate(Exception): diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index c5b0cd958..c581d2259 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1,3 +1,4 @@ +import dataclasses import inspect import logging import typing @@ -49,7 +50,7 @@ from langgraph.managed.base import ( from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, Command, N, RetryPolicy +from langgraph.types import _DC_KWARGS, All, Checkpointer, Command, N, RetryPolicy from langgraph.utils.fields import get_field_default from langgraph.utils.pydantic import create_model from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable @@ -78,21 +79,20 @@ def _get_node_name(node: RunnableLike) -> str: raise TypeError(f"Unsupported node type: {type(node)}") +@dataclasses.dataclass(**_DC_KWARGS) class GraphCommand(Generic[N], Command[N]): """One or more commands to update a StateGraph's state and go to, or send messages to nodes.""" - __slots__ = ("goto",) + goto: Union[str, Sequence[str]] = () - def __init__( - self, - *, - update: Optional[dict[str, Any]] = None, - send: Union[Send, Sequence[Send]] = (), - resume: Optional[Union[Any, dict[str, Any]]] = None, - goto: Union[str, Sequence[str]] = (), - ) -> None: - super().__init__(update=update, send=send, resume=resume) - self.goto = goto + def __repr__(self) -> str: + # get all non-None values + contents = ", ".join( + f"{key}={value!r}" + for key, value in dataclasses.asdict(self).items() + if value + ) + return f"Command({contents})" class StateNodeSpec(NamedTuple): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index b910616a2..104412d8e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,5 +1,6 @@ +import dataclasses +import sys from collections import deque -from dataclasses import dataclass from typing import ( TYPE_CHECKING, Any, @@ -47,6 +48,11 @@ StreamWriter = Callable[[Any], None] Always injected into nodes if requested as a keyword argument, but it's a no-op when not using stream_mode="custom".""" +if sys.version_info >= (3, 10): + _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True} +else: + _DC_KWARGS = {"frozen": True} + def default_retry_on(exc: Exception) -> bool: import httpx @@ -104,9 +110,11 @@ class CachePolicy(NamedTuple): pass -@dataclass +@dataclasses.dataclass(**_DC_KWARGS) class Interrupt: value: Any + resumable: bool = False + ns: Optional[Sequence[str]] = None when: Literal["during"] = "during" @@ -227,53 +235,23 @@ class Send: N = TypeVar("N", bound=Hashable) +@dataclasses.dataclass(**_DC_KWARGS) class Command(Generic[N]): """One or more commands to update the graph's state and send messages to nodes.""" - __slots__ = ("update", "send", "resume") - - def __init__( - self, - *, - update: Optional[dict[str, Any]] = None, - send: Union[Send, Sequence[Send]] = (), - resume: Optional[Union[Any, dict[str, Any]]] = None, - ) -> None: - self.update = update - self.send = send - self.resume = resume - - @property - def __all_slots__(self) -> set[str]: - # get all slots from mro - slots = set() - for cls in type(self).__mro__: - if ss := getattr(cls, "__slots__", ()): - if isinstance(ss, str): - slots.add(ss) - else: - slots.update(ss) - return slots + update: Optional[dict[str, Any]] = None + send: Union[Send, Sequence[Send]] = () + resume: Optional[Union[Any, dict[str, Any]]] = None def __repr__(self) -> str: # get all non-None values contents = ", ".join( f"{key}={value!r}" - for key in self.__all_slots__ - if (value := getattr(self, key)) + for key, value in dataclasses.asdict(self).items() + if value ) return f"Command({contents})" - def __eq__(self, value: Any) -> bool: - return type(value) is type(self) and all( - getattr(self, key) == getattr(value, key) for key in self.__all_slots__ - ) - - def copy(self, **kwargs: Any) -> Self: - for slot in self.__all_slots__: - kwargs.setdefault(slot, getattr(self, slot)) - return self.__class__(**kwargs) - StreamChunk = tuple[tuple[str, ...], str, Any] @@ -318,12 +296,25 @@ class LoopProtocol: def interrupt(value: Any) -> Any: - from langgraph.constants import CONFIG_KEY_RESUME_VALUE, MISSING - from langgraph.errors import NodeInterrupt + from langgraph.constants import ( + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_RESUME_VALUE, + MISSING, + NS_SEP, + ) + from langgraph.errors import GraphInterrupt from langgraph.utils.config import get_configurable conf = get_configurable() if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING: return resume else: - raise NodeInterrupt(value) + raise GraphInterrupt( + ( + Interrupt( + value=value, + resumable=True, + ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP), + ), + ) + ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 81d45f4ee..e29476c19 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8,6 +8,7 @@ import warnings from collections import Counter from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from dataclasses import replace from random import randrange from typing import ( Annotated, @@ -1834,9 +1835,8 @@ def test_send_sequences() -> None: if isinstance(state, list) else ["|".join((self.name, str(state)))] ) - if isinstance(state, GraphCommand): - state.update = update - return state + if isinstance(state, Command): + return replace(state, update=update) else: return update @@ -1918,7 +1918,7 @@ def test_send_dedupe_on_resume( else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): - return state.copy(update=update) + return replace(state, update=update) else: return update @@ -8409,7 +8409,15 @@ def test_dynamic_interrupt( assert [ c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] # resume with answer assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ @@ -8447,7 +8455,13 @@ def test_dynamic_interrupt( AnyStr(), "tool_two", (PULL, "tool_two"), - interrupts=(Interrupt("Just because..."),), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), ), ), config=tool_two.checkpointer.get_tuple(thread1).config, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index df85ed650..1469c018e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6,6 +6,7 @@ import sys import uuid from collections import Counter from contextlib import asynccontextmanager, contextmanager +from dataclasses import replace from time import perf_counter from typing import ( Annotated, @@ -318,7 +319,15 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: {"my_key": "value ⛰️", "market": "DE"}, thread2 ) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] # resume with answer assert [ @@ -336,7 +345,15 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: {"my_key": "value ⛰️", "market": "DE"}, thread1 ) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ { @@ -363,7 +380,13 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: AnyStr(), "tool_two", (PULL, "tool_two"), - interrupts=(Interrupt("Just because..."),), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), ), ), config=tup.config, @@ -2111,7 +2134,7 @@ async def test_send_sequences(checkpointer_name: str) -> None: else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): - return state.copy(update=update) + return replace(state, update=update) else: return update @@ -2215,7 +2238,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): - return state.copy(update=update) + return replace(state, update=update) else: return update