diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 10908eb87..670e85b3d 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 ControlProtocol, SendProtocol +from langgraph.checkpoint.serde.types import CommandProtocol, SendProtocol from langgraph.store.base import Item LC_REVIVER = Reviver() @@ -122,6 +122,11 @@ 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(),) @@ -402,18 +407,14 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]: (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)), ), ) - elif isinstance(obj, ControlProtocol): + elif isinstance(obj, CommandProtocol): 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, - }, + {k: getattr(obj, k) for k in obj.__all_slots__}, ), ), ) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 862cbe83f..154b1450b 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -52,10 +52,8 @@ class SendProtocol(Protocol): @runtime_checkable -class ControlProtocol(Protocol): - # Mirrors langgraph.constants.Control - update_state: Optional[dict[str, Any]] - trigger: Union[str, Sequence[str]] +class CommandProtocol(Protocol): + # Mirrors langgraph.types.Command + update: Optional[dict[str, Any]] send: Union[Any, Sequence[Any]] - - def __repr__(self) -> str: ... + __all_slots__: set[str] diff --git a/libs/langgraph/langgraph/graph/__init__.py b/libs/langgraph/langgraph/graph/__init__.py index c81ad9903..241106a3a 100644 --- a/libs/langgraph/langgraph/graph/__init__.py +++ b/libs/langgraph/langgraph/graph/__init__.py @@ -1,12 +1,13 @@ from langgraph.graph.graph import END, START, Graph from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.graph.state import StateGraph +from langgraph.graph.state import GraphCommand, StateGraph __all__ = [ "END", "START", "Graph", "StateGraph", + "GraphCommand", "MessageGraph", "add_messages", "MessagesState", diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index f5338a024..2cae9fc7e 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -8,11 +8,13 @@ from types import FunctionType from typing import ( Any, Callable, + Generic, Literal, NamedTuple, Optional, Sequence, Type, + TypeVar, Union, cast, get_args, @@ -48,13 +50,15 @@ 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, Control, RetryPolicy +from langgraph.types import All, Checkpointer, Command, 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 logger = logging.getLogger(__name__) +N = TypeVar("N") + def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None: if isinstance(schema, type): @@ -77,6 +81,22 @@ def _get_node_name(node: RunnableLike) -> str: raise TypeError(f"Unsupported node type: {type(node)}") +class GraphCommand(Command, Generic[N]): + """One or more commands to update a StateGraph's state and go to, or send messages to nodes.""" + + __slots__ = ("goto",) + + def __init__( + self, + *, + update: Optional[dict[str, Any]] = None, + goto: Union[str, Sequence[str]] = (), + send: Union[Send, Sequence[Send]] = (), + ) -> None: + super().__init__(update=update, send=send) + self.goto = goto + + class StateNodeSpec(NamedTuple): runnable: Runnable metadata: Optional[dict[str, Any]] @@ -369,7 +389,7 @@ class StateGraph(Graph): input = input_hint if ( (rtn := hints.get("return")) - and get_origin(rtn) is Control + and get_origin(rtn) is GraphCommand and (rargs := get_args(rtn)) and get_origin(rargs[0]) is Literal and (vals := get_args(rargs[0])) @@ -604,8 +624,8 @@ class CompiledStateGraph(CompiledGraph): ] def _get_root(input: Any) -> Any: - if isinstance(input, Control): - return input.state + if isinstance(input, Command): + return input.update else: return input @@ -618,8 +638,8 @@ class CompiledStateGraph(CompiledGraph): f"Expected node {key} to update at least one of {output_keys}, got {input}" ) return input.get(key, SKIP_WRITE) - elif isinstance(input, Control): - return _get_state_key(input.state, key=key) + elif isinstance(input, Command): + return _get_state_key(input.update, key=key) elif get_type_hints(type(input)): value = getattr(input, key, SKIP_WRITE) return value if value is not None else SKIP_WRITE @@ -799,7 +819,7 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: def _control_branch(value: Any) -> Sequence[Union[str, Send]]: if isinstance(value, Send): return [value] - if not isinstance(value, Control): + if not isinstance(value, GraphCommand): return EMPTY_SEQ rtn: list[Union[str, Send]] = [] if isinstance(value.goto, str): @@ -816,7 +836,7 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]: async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: if isinstance(value, Send): return [value] - if not isinstance(value, Control): + if not isinstance(value, GraphCommand): return EMPTY_SEQ rtn: list[Union[str, Send]] = [] if isinstance(value.goto, str): diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 84f8af8b6..b491383df 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -478,11 +478,6 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer or None, manager=None, ) - print( - saved.checkpoint["versions_seen"], - saved.checkpoint["pending_sends"], - # next_tasks, - ) # get the subgraphs subgraphs = dict(self.get_subgraphs()) parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 3b7bc12a3..af71294ae 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -197,7 +197,6 @@ def apply_writes( # sort tasks on path tasks = sorted(tasks, key=lambda t: t.path) - print("versions_seen", checkpoint["versions_seen"], [task.name for task in tasks]) # update seen versions for task in tasks: checkpoint["versions_seen"].setdefault(task.name, {}).update( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 88d2d06f2..0a3c22f34 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -4,13 +4,11 @@ from typing import ( TYPE_CHECKING, Any, Callable, - Generic, Literal, NamedTuple, Optional, Sequence, Type, - TypeVar, Union, cast, ) @@ -223,32 +221,45 @@ class Send: ) -N = TypeVar("N") +class Command: + """One or more commands to update the graph's state and send messages to nodes.""" - -class Control(Generic[N]): - """A control object to update the graph's state, trigger nodes, and send messages.""" - - __slots__ = ("state", "goto", "send") + __slots__ = ("update", "send") def __init__( self, *, - state: Optional[dict[str, Any]] = None, - goto: Union[str, Sequence[str]] = (), + update: Optional[dict[str, Any]] = None, send: Union[Send, Sequence[Send]] = (), ) -> None: - self.state = state - self.goto = goto + self.update = update self.send = send + @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 + def __repr__(self) -> str: + # get all non-None values contents = ", ".join( f"{key}={value!r}" - for key in self.__slots__ + for key in self.__all_slots__ if (value := getattr(self, key)) ) - return f"Control({contents})" + return f"Command({contents})" + + def __eq__(self, value): + return type(value) is type(self) and all( + getattr(self, key) == getattr(value, key) for key in self.__all_slots__ + ) StreamChunk = tuple[tuple[str, ...], str, Any] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 9d49e8551..20b12beb2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -54,12 +54,10 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt -from langgraph.graph import END, Graph -from langgraph.graph.graph import START +from langgraph.graph import END, Graph, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.graph.state import StateGraph from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, @@ -74,7 +72,7 @@ from langgraph.pregel import ( from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Control, Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import Interrupt, PregelTask, Send, StreamWriter from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, @@ -1808,16 +1806,16 @@ def test_send_sequences() -> None: if isinstance(state, list) # or isinstance(state, Control) else ["|".join((self.name, str(state)))] ) - if isinstance(state, Control): - state.state = update + if isinstance(state, GraphCommand): + state.update = update return state else: return update def send_for_fun(state): return [ - Send("2", Control(send=Send("2", 3))), - Send("2", Control(send=Send("2", 4))), + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("2", 4))), "3.1", ] @@ -1837,8 +1835,8 @@ def test_send_sequences() -> None: "0", "1", "3.1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='2', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", "3", "2|3", "2|4", @@ -2311,9 +2309,9 @@ def test_send_react_interrupt_control( tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], ) - def agent(state) -> Control[Literal["foo"]]: - return Control( - state={"messages": ai_message}, + def agent(state) -> GraphCommand[Literal["foo"]]: + return GraphCommand( + update={"messages": ai_message}, send=[Send(call["name"], call) for call in ai_message.tool_calls], ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 364be8bfa..6fa690463 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -51,10 +51,9 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt -from langgraph.graph import END, Graph, StateGraph -from langgraph.graph.graph import START +from langgraph.graph import END, Graph, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor @@ -63,7 +62,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Control, Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import Interrupt, PregelTask, Send, StreamWriter from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, @@ -2035,7 +2034,8 @@ async def test_concurrent_emit_sends() -> None: ] -async def test_send_sequences() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_sequences(checkpointer_name: str) -> None: class Node: def __init__(self, name: str): self.name = name @@ -2047,16 +2047,16 @@ async def test_send_sequences() -> None: if isinstance(state, list) # or isinstance(state, Control) else ["|".join((self.name, str(state)))] ) - if isinstance(state, Control): - state.state = update + if isinstance(state, GraphCommand): + state.update = update return state else: return update async def send_for_fun(state): return [ - Send("2", Control(send=Send("2", 3))), - Send("2", Control(send=Send("2", 4))), + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("2", 4))), "3.1", ] @@ -2076,14 +2076,29 @@ async def test_send_sequences() -> None: "0", "1", "3.1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='2', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", "3", "2|3", "2|4", "3", ] + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1) == [ + "0", + "1", + "3.1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_react_interrupt(checkpointer_name: str) -> None: @@ -2543,9 +2558,9 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], ) - async def agent(state) -> Control[Literal["foo"]]: - return Control( - state={"messages": ai_message}, + async def agent(state) -> GraphCommand[Literal["foo"]]: + return GraphCommand( + update={"messages": ai_message}, send=[Send(call["name"], call) for call in ai_message.tool_calls], ) @@ -2843,13 +2858,13 @@ async def test_max_concurrency(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_max_concurrency_control(checkpointer_name: str) -> None: - async def node1(state) -> Control[Literal["2"]]: - return Control(state=["1"], send=[Send("2", idx) for idx in range(100)]) + async def node1(state) -> GraphCommand[Literal["2"]]: + return GraphCommand(update=["1"], send=[Send("2", idx) for idx in range(100)]) node2_currently = 0 node2_max_currently = 0 - async def node2(state) -> Control[Literal["3"]]: + async def node2(state) -> GraphCommand[Literal["3"]]: nonlocal node2_currently, node2_max_currently node2_currently += 1 if node2_currently > node2_max_currently: @@ -2857,7 +2872,7 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None: await asyncio.sleep(0.1) node2_currently -= 1 - return Control(state=[state], goto="3") + return GraphCommand(update=[state], goto="3") async def node3(state) -> Literal["3"]: return ["3"]