lib: Add support for returning multiple commands from a node

libs/langgraph/langgraph/graph/state.py

@@ -51,7 +51,11 @@ from langgraph.managed.base import (
     is_writable_managed_value,
 )
 from langgraph.pregel.read import ChannelRead, PregelNode
-from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
+from langgraph.pregel.write import (
+    ChannelWrite,
+    ChannelWriteEntry,
+    ChannelWriteTupleEntry,
+)
 from langgraph.store.base import BaseStore
 from langgraph.types import All, Checkpointer, Command, RetryPolicy
 from langgraph.utils.fields import get_field_default
@@ -608,33 +612,53 @@ class CompiledStateGraph(CompiledGraph):
                 if is_writable_managed_value(v)
             ]
 
-        def _get_root(input: Any) -> Any:
-            if isinstance(input, Command):
-                if input.graph == Command.PARENT:
-                    return SKIP_WRITE
-                return input.update
-            else:
-                return input
-
-        # to avoid name collision below
-        node_key = key
-
-        def _get_state_key(input: Union[None, dict, Any], *, key: str) -> Any:
-            if input is None:
-                return SKIP_WRITE
-            elif isinstance(input, dict):
-                if all(k not in output_keys for k in input):
-                    raise InvalidUpdateError(
-                        f"Expected node {node_key} to update at least one of {output_keys}, got {input}"
-                    )
-                return input.get(key, SKIP_WRITE)
+        def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
+            if (
+                isinstance(input, (list, tuple))
+                and input
+                and all(isinstance(i, Command) for i in input)
+            ):
+                updates: list[tuple[str, Any]] = []
+                for i in input:
+                    if i.graph == Command.PARENT:
+                        continue
+                    updates.extend(i._update_as_tuples())
+                return updates
             elif isinstance(input, Command):
                 if input.graph == Command.PARENT:
-                    return SKIP_WRITE
-                return _get_state_key(input.update, key=key)
+                    return ()
+                return input._update_as_tuples()
+            elif input is not None:
+                return [("__root__", input)]
+
+        def _get_updates(
+            input: Union[None, dict, Any],
+        ) -> Optional[Sequence[tuple[str, Any]]]:
+            if input is None:
+                return None
+            elif isinstance(input, dict):
+                return [(k, v) for k, v in input.items() if k in output_keys]
+            elif isinstance(input, Command):
+                if input.graph == Command.PARENT:
+                    return None
+                return input._update_as_tuples()
+            elif (
+                isinstance(input, (list, tuple))
+                and input
+                and all(isinstance(i, Command) for i in input)
+            ):
+                updates: list[tuple[str, Any]] = []
+                for i in input:
+                    if i.graph == Command.PARENT:
+                        continue
+                    updates.extend(i._update_as_tuples())
+                return updates
             elif get_type_hints(type(input)):
-                value = getattr(input, key, SKIP_WRITE)
-                return value if value is not None else SKIP_WRITE
+                return [
+                    (k, getattr(input, k))
+                    for k in output_keys
+                    if getattr(input, k, None) is not None
+                ]
             else:
                 msg = create_error_message(
                     message=f"Expected dict, got {input}",
@@ -643,14 +667,11 @@ class CompiledStateGraph(CompiledGraph):
                 raise InvalidUpdateError(msg)
 
         # state updaters
-        write_entries = (
-            [ChannelWriteEntry("__root__", skip_none=True, mapper=_get_root)]
-            if output_keys == ["__root__"]
-            else [
-                ChannelWriteEntry(key, mapper=partial(_get_state_key, key=key))
-                for key in output_keys
-            ]
-        )
+        write_entries = [
+            ChannelWriteTupleEntry(
+                mapper=_get_root if output_keys == ["__root__"] else _get_updates
+            )
+        ]
 
         # add node and output channel
         if key == START:
@@ -811,34 +832,54 @@ 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, Command):
-        return EMPTY_SEQ
-    if value.graph == Command.PARENT:
-        raise ParentCommand(value)
-    rtn: list[Union[str, Send]] = []
-    if isinstance(value.goto, Send):
-        rtn.append(value.goto)
-    elif isinstance(value.goto, str):
-        rtn.append(value.goto)
+    commands: list[Command] = []
+    if isinstance(value, Command):
+        commands.append(value)
+    elif (
+        isinstance(value, (list, tuple))
+        and value
+        and all(isinstance(i, Command) for i in value)
+    ):
+        commands.extend(value)
     else:
-        rtn.extend(value.goto)
+        return EMPTY_SEQ
+    rtn: list[Union[str, Send]] = []
+    for command in commands:
+        if command.graph == Command.PARENT:
+            raise ParentCommand(command)
+        if isinstance(command.goto, Send):
+            rtn.append(command.goto)
+        elif isinstance(command.goto, str):
+            rtn.append(command.goto)
+        else:
+            rtn.extend(command.goto)
     return rtn
 
 
 async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
     if isinstance(value, Send):
         return [value]
-    if not isinstance(value, Command):
-        return EMPTY_SEQ
-    if value.graph == Command.PARENT:
-        raise ParentCommand(value)
-    rtn: list[Union[str, Send]] = []
-    if isinstance(value.goto, Send):
-        rtn.append(value.goto)
-    elif isinstance(value.goto, str):
-        rtn.append(value.goto)
+    commands: list[Command] = []
+    if isinstance(value, Command):
+        commands.append(value)
+    elif (
+        isinstance(value, (list, tuple))
+        and value
+        and all(isinstance(i, Command) for i in value)
+    ):
+        commands.extend(value)
     else:
-        rtn.extend(value.goto)
+        return EMPTY_SEQ
+    rtn: list[Union[str, Send]] = []
+    for command in commands:
+        if command.graph == Command.PARENT:
+            raise ParentCommand(command)
+        if isinstance(command.goto, Send):
+            rtn.append(command.goto)
+        elif isinstance(command.goto, str):
+            rtn.append(command.goto)
+        else:
+            rtn.extend(command.goto)
     return rtn
 
 

libs/langgraph/langgraph/pregel/io.py

@@ -95,11 +95,7 @@ def map_command(
         else:
             yield (NULL_TASK_ID, RESUME, cmd.resume)
     if cmd.update:
-        if not isinstance(cmd.update, dict):
-            raise TypeError(
-                f"Expected cmd.update to be a dict mapping channel names to update values, got {type(cmd.update).__name__}"
-            )
-        for k, v in cmd.update.items():
+        for k, v in cmd._update_as_tuples():
             yield (NULL_TASK_ID, k, v)
 
 

libs/langgraph/langgraph/pregel/write.py

@@ -36,31 +36,40 @@ class ChannelWriteEntry(NamedTuple):
     """Function to transform the value before writing."""
 
 
+class ChannelWriteTupleEntry(NamedTuple):
+    mapper: Callable[[Any], Sequence[tuple[str, Any]]]
+    """Function to extract tuples from value."""
+    value: Any = PASSTHROUGH
+    """Value to write, or PASSTHROUGH to use the input."""
+
+
 class ChannelWrite(RunnableCallable):
-    """Implements th logic for sending writes to CONFIG_KEY_SEND.
+    """Implements the logic for sending writes to CONFIG_KEY_SEND.
     Can be used as a runnable or as a static method to call imperatively."""
 
-    writes: list[Union[ChannelWriteEntry, Send]]
+    writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]]
     """Sequence of write entries or Send objects to write."""
     require_at_least_one_of: Optional[Sequence[str]]
     """If defined, at least one of these channels must be written to."""
 
     def __init__(
         self,
-        writes: Sequence[Union[ChannelWriteEntry, Send]],
+        writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
         *,
         tags: Optional[Sequence[str]] = None,
         require_at_least_one_of: Optional[Sequence[str]] = None,
     ):
         super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags)
-        self.writes = cast(list[Union[ChannelWriteEntry, Send]], writes)
+        self.writes = cast(
+            list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes
+        )
         self.require_at_least_one_of = require_at_least_one_of
 
     def get_name(
         self, suffix: Optional[str] = None, *, name: Optional[str] = None
     ) -> str:
         if not name:
-            name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else w.node for w in self.writes)}>"
+            name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else '...' if isinstance(w, ChannelWriteTupleEntry) else w.node for w in self.writes)}>"
         return super().get_name(suffix, name=name)
 
     @property
@@ -79,6 +88,8 @@ class ChannelWrite(RunnableCallable):
         writes = [
             ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
             if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
+            else ChannelWriteTupleEntry(write.mapper, input)
+            if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH
             else write
             for write in self.writes
         ]
@@ -93,6 +104,8 @@ class ChannelWrite(RunnableCallable):
         writes = [
             ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
             if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
+            else ChannelWriteTupleEntry(write.mapper, input)
+            if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH
             else write
             for write in self.writes
         ]
@@ -106,7 +119,7 @@ class ChannelWrite(RunnableCallable):
     @staticmethod
     def do_write(
         config: RunnableConfig,
-        writes: Sequence[Union[ChannelWriteEntry, Send]],
+        writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
         require_at_least_one_of: Optional[Sequence[str]] = None,
     ) -> None:
         # validate
@@ -118,32 +131,36 @@ class ChannelWrite(RunnableCallable):
                     )
                 if w.value is PASSTHROUGH:
                     raise InvalidUpdateError("PASSTHROUGH value must be replaced")
-        # split packets and entries
-        sends = [
-            (PUSH if FF_SEND_V2 else TASKS, packet)
-            for packet in writes
-            if isinstance(packet, Send)
-        ]
-        entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
-        # process entries into values
-        values = [
-            write.mapper(write.value) if write.mapper is not None else write.value
-            for write in entries
-        ]
-        values = [
-            (write.channel, val)
-            for val, write in zip(values, entries)
-            if not write.skip_none or val is not None
-        ]
-        # filter out SKIP_WRITE values
-        filtered = [(chan, val) for chan, val in values if val is not SKIP_WRITE]
+            if isinstance(w, ChannelWriteTupleEntry):
+                if w.value is PASSTHROUGH:
+                    raise InvalidUpdateError("PASSTHROUGH value must be replaced")
+        # assemble writes
+        tuples: list[tuple[str, Any]] = []
+        print(writes)
+        for w in writes:
+            if isinstance(w, Send):
+                tuples.append((PUSH if FF_SEND_V2 else TASKS, w))
+            elif isinstance(w, ChannelWriteTupleEntry):
+                if ww := w.mapper(w.value):
+                    tuples.extend(ww)
+            elif isinstance(w, ChannelWriteEntry):
+                value = w.mapper(w.value) if w.mapper is not None else w.value
+                if value is SKIP_WRITE:
+                    continue
+                if w.skip_none and value is None:
+                    continue
+                tuples.append((w.channel, value))
+            else:
+                raise ValueError(f"Invalid write entry: {w}")
+        print(tuples, require_at_least_one_of)
+        # assert required channels
         if require_at_least_one_of is not None:
-            if not {chan for chan, _ in filtered} & set(require_at_least_one_of):
+            if not {chan for chan, _ in tuples} & set(require_at_least_one_of):
                 raise InvalidUpdateError(
                     f"Must write to at least one of {require_at_least_one_of}"
                 )
         write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
-        write(sends + filtered)
+        write(tuples)
 
     @staticmethod
     def is_writer(runnable: Runnable) -> bool:

libs/langgraph/langgraph/types.py

@@ -263,7 +263,7 @@ class Command(Generic[N]):
     """
 
     graph: Optional[str] = None
-    update: Optional[dict[str, Any]] = None
+    update: Union[dict[str, Any], Sequence[tuple[str, Any]]] = ()
     resume: Optional[Union[Any, dict[str, Any]]] = None
     goto: Union[Send, Sequence[Union[Send, str]], str] = ()
 
@@ -276,6 +276,17 @@ class Command(Generic[N]):
         )
         return f"Command({contents})"
 
+    def _update_as_tuples(self) -> Sequence[tuple[str, Any]]:
+        if isinstance(self.update, dict):
+            return list(self.update.items())
+        elif isinstance(self.update, (list, tuple)) and all(
+            isinstance(t, tuple) and len(t) == 2 and isinstance(t[0], str)
+            for t in self.update
+        ):
+            return self.update
+        else:
+            return [("__root__", self.update)]
+
     PARENT: ClassVar[Literal["__parent__"]] = "__parent__"
 
 

libs/langgraph/tests/test_pregel.py

@@ -220,18 +220,6 @@ def test_graph_validation() -> None:
     class State(TypedDict):
         hello: str
 
-    def node_a(state: State) -> State:
-        # typo
-        return {"hell": "world"}
-
-    builder = StateGraph(State)
-    builder.add_node("a", node_a)
-    builder.set_entry_point("a")
-    builder.set_finish_point("a")
-    graph = builder.compile()
-    with pytest.raises(InvalidUpdateError):
-        graph.invoke({"hello": "there"})
-
     graph = StateGraph(State)
     graph.add_node("start", lambda x: x)
     graph.add_edge("__start__", "start")
@@ -1919,7 +1907,7 @@ def test_send_sequences() -> None:
                 else ["|".join((self.name, str(state)))]
             )
             if isinstance(state, Command):
-                return replace(state, update=update)
+                return [state, Command(update=update)]
             else:
                 return update