From aa4fea48ddf8fc58c9cb301872dd063d57a4a8a5 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 5 Dec 2024 13:47:38 -0800 Subject: [PATCH] lib: Add support for returning multiple commands from a node --- libs/langgraph/langgraph/graph/state.py | 147 +++++++++++++++-------- libs/langgraph/langgraph/pregel/io.py | 6 +- libs/langgraph/langgraph/pregel/write.py | 71 ++++++----- libs/langgraph/langgraph/types.py | 13 +- libs/langgraph/tests/test_pregel.py | 14 +-- 5 files changed, 152 insertions(+), 99 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index e63f25111..0462008c8 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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 diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index b2596d3ad..b3d6845b5 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/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) diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 3af0fe5e9..baee23ab2 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/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: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3eb4536aa..645bfb67b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/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__" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f69d36ed3..2515270a1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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