Merge pull request #2347 from langchain-ai/nc/5nov/control-serializable

lib: Make Control object serializable

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 SendProtocol
+from langgraph.checkpoint.serde.types import ControlProtocol, SendProtocol
 from langgraph.store.base import Item
 
 LC_REVIVER = Reviver()
@@ -402,6 +402,21 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]:
                 (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)),
             ),
         )
+    elif isinstance(obj, ControlProtocol):
+        return msgpack.ExtType(
+            EXT_CONSTRUCTOR_KW_ARGS,
+            _msgpack_enc(
+                (
+                    obj.__class__.__module__,
+                    obj.__class__.__name__,
+                    {
+                        "update_state": obj.update_state,
+                        "trigger": obj.trigger,
+                        "send": obj.send,
+                    },
+                ),
+            ),
+        )
     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,6 +4,7 @@ from typing import (
     Protocol,
     Sequence,
     TypeVar,
+    Union,
     runtime_checkable,
 )
 
@@ -48,3 +49,13 @@ class SendProtocol(Protocol):
     def __repr__(self) -> str: ...
 
     def __eq__(self, value: object) -> bool: ...
+
+
+@runtime_checkable
+class ControlProtocol(Protocol):
+    # Mirrors langgraph.constants.Control
+    update_state: Optional[dict[str, Any]]
+    trigger: Union[str, Sequence[str]]
+    send: Union[Any, Sequence[Any]]
+
+    def __repr__(self) -> str: ...

libs/langgraph/langgraph/types.py

@@ -227,6 +227,10 @@ N = TypeVar("N")
 
 
 class Control(Generic[N]):
+    """A control object to update the graph's state, trigger nodes, and send messages."""
+
+    __slots__ = ("update_state", "trigger", "send")
+
     def __init__(
         self,
         *,
@@ -240,7 +244,9 @@ class Control(Generic[N]):
 
     def __repr__(self) -> str:
         contents = ", ".join(
-            f"{key}={value!r}" for key, value in self.__dict__.items() if value
+            f"{key}={value!r}"
+            for key in self.__slots__
+            if (value := getattr(self, key))
         )
         return f"Control({contents})"