Merge pull request #2410 from langchain-ai/nc/13nov/command-dataclass

Nc/13nov/command dataclass

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(

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]

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):

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):

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),
+                ),
+            )
+        )

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,

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