From 85fc26db43f9d9ba79b1105e48178f5b4ddefa15 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 6 Dec 2024 08:03:26 -0800 Subject: [PATCH] lib: Support returning mixed list of commands and state updates --- libs/langgraph/langgraph/graph/state.py | 37 ++++++++++++++----------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 6e37c6150..0c4d0c15c 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -613,21 +613,23 @@ class CompiledStateGraph(CompiledGraph): ] 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 isinstance(input, Command): if input.graph == Command.PARENT: return () return input._update_as_tuples() + elif ( + isinstance(input, (list, tuple)) + and input + and any(isinstance(i, Command) for i in input) + ): + updates: list[tuple[str, Any]] = [] + for i in input: + if isinstance(i, Command): + if i.graph == Command.PARENT: + continue + updates.extend(i._update_as_tuples()) + else: + updates.append(("__root__", i)) elif input is not None: return [("__root__", input)] @@ -645,13 +647,16 @@ class CompiledStateGraph(CompiledGraph): elif ( isinstance(input, (list, tuple)) and input - and all(isinstance(i, Command) for i in input) + and any(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()) + if isinstance(i, Command): + if i.graph == Command.PARENT: + continue + updates.extend(i._update_as_tuples()) + else: + updates.extend(_get_updates(i) or ()) return updates elif get_type_hints(type(input)): return [