Make Command a 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/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,5 @@
 from collections import deque
-from dataclasses import dataclass
+import dataclasses
 import sys
 from typing import (
     TYPE_CHECKING,
@@ -49,9 +49,9 @@ 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}
+    _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
 else:
-    _DC_KWARGS = {}
+    _DC_KWARGS = {"frozen": True}
 
 
 def default_retry_on(exc: Exception) -> bool:
@@ -110,7 +110,7 @@ class CachePolicy(NamedTuple):
     pass
 
 
-@dataclass(**_DC_KWARGS)
+@dataclasses.dataclass(**_DC_KWARGS)
 class Interrupt:
     value: Any
     resumable: bool = False
@@ -235,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]
 

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
 

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,
@@ -2133,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
 
@@ -2237,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