langgraph: add add_sequence to StateGraph (#2352)

libs/langgraph/langgraph/graph/state.py

@@ -68,6 +68,15 @@ def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
     )
 
 
+def _get_node_name(node: RunnableLike) -> str:
+    if isinstance(node, Runnable):
+        return node.get_name()
+    elif callable(node):
+        return getattr(node, "__name__", node.__class__.__name__)
+    else:
+        raise TypeError(f"Unsupported node type: {type(node)}")
+
+
 class StateNodeSpec(NamedTuple):
     runnable: Runnable
     metadata: Optional[dict[str, Any]]
@@ -225,7 +234,7 @@ class StateGraph(Graph):
             ValueError: If the key is already being used as a state key.
 
         Returns:
-            None
+            StateGraph
         """
         ...
 
@@ -249,7 +258,7 @@ class StateGraph(Graph):
             ValueError: If the key is already being used as a state key.
 
         Returns:
-            None
+            StateGraph
         """
         ...
 
@@ -302,7 +311,7 @@ class StateGraph(Graph):
             ```
 
         Returns:
-            None
+            StateGraph
         """
         if not isinstance(node, str):
             action = node
@@ -392,7 +401,7 @@ class StateGraph(Graph):
             ValueError: If the start key is 'END' or if the start key or end key is not present in the graph.
 
         Returns:
-            None
+            StateGraph
         """
         if isinstance(start_key, str):
             return super().add_edge(start_key, end_key)
@@ -415,6 +424,48 @@ class StateGraph(Graph):
         self.waiting_edges.add((tuple(start_key), end_key))
         return self
 
+    def add_sequence(
+        self,
+        nodes: Sequence[Union[RunnableLike, tuple[str, RunnableLike]]],
+    ) -> Self:
+        """Add a sequence of nodes that will be executed in the provided order.
+
+        Args:
+            nodes: A sequence of RunnableLike objects (e.g. a LangChain Runnable or a callable) or (name, RunnableLike) tuples.
+                If no names are provided, the name will be inferred from the node object (e.g. a runnable or a callable name).
+                Each node will be executed in the order provided.
+
+        Raises:
+            ValueError: if the sequence is empty.
+            ValueError: if the sequence contains duplicate node names.
+
+        Returns:
+            StateGraph
+        """
+        if len(nodes) < 1:
+            raise ValueError("Sequence requires at least one node.")
+
+        previous_name: Optional[str] = None
+        for node in nodes:
+            if isinstance(node, tuple) and len(node) == 2:
+                name, node = node
+            else:
+                name = _get_node_name(node)
+
+            if name in self.nodes:
+                raise ValueError(
+                    f"Node names must be unique: node with the name '{name}' already exists. "
+                    "If you need to use two different runnables/callables with the same name (for example, using `lambda`), please provide them as tuples (name, runnable/callable)."
+                )
+
+            self.add_node(name, node)
+            if previous_name is not None:
+                self.add_edge(previous_name, name)
+
+            previous_name = name
+
+        return self
+
     def compile(
         self,
         checkpointer: Checkpointer = None,

libs/langgraph/tests/test_pregel.py

@@ -13467,6 +13467,136 @@ def test_debug_nested_subgraphs():
                 assert stream_task.get("state") == history_task.state
 
 
+def test_add_sequence():
+    class State(TypedDict):
+        foo: Annotated[list[str], operator.add]
+        bar: str
+
+    def step1(state: State):
+        return {"foo": ["step1"], "bar": "baz"}
+
+    def step2(state: State):
+        return {"foo": ["step2"]}
+
+    # test raising if less than 1 steps
+    with pytest.raises(ValueError):
+        StateGraph(State).add_sequence([])
+
+    # test raising if duplicate step names
+    with pytest.raises(ValueError):
+        StateGraph(State).add_sequence([step1, step1])
+
+    with pytest.raises(ValueError):
+        StateGraph(State).add_sequence([("foo", step1), ("foo", step1)])
+
+    # test unnamed steps
+    builder = StateGraph(State)
+    builder.add_sequence([step1, step2])
+    builder.add_edge(START, "step1")
+    graph = builder.compile()
+    result = graph.invoke({"foo": []})
+    assert result == {"foo": ["step1", "step2"], "bar": "baz"}
+    stream_chunks = list(graph.stream({"foo": []}))
+    assert stream_chunks == [
+        {"step1": {"foo": ["step1"], "bar": "baz"}},
+        {"step2": {"foo": ["step2"]}},
+    ]
+
+    # test named steps
+    builder_named_steps = StateGraph(State)
+    builder_named_steps.add_sequence([("meow1", step1), ("meow2", step2)])
+    builder_named_steps.add_edge(START, "meow1")
+    graph_named_steps = builder_named_steps.compile()
+    result = graph_named_steps.invoke({"foo": []})
+    stream_chunks = list(graph_named_steps.stream({"foo": []}))
+    assert result == {"foo": ["step1", "step2"], "bar": "baz"}
+    assert stream_chunks == [
+        {"meow1": {"foo": ["step1"], "bar": "baz"}},
+        {"meow2": {"foo": ["step2"]}},
+    ]
+
+    builder_named_steps = StateGraph(State)
+    builder_named_steps.add_sequence(
+        [
+            ("meow1", lambda state: {"foo": ["foo"]}),
+            ("meow2", lambda state: {"bar": state["foo"][0] + "bar"}),
+        ],
+    )
+    builder_named_steps.add_edge(START, "meow1")
+    graph_named_steps = builder_named_steps.compile()
+    result = graph_named_steps.invoke({"foo": []})
+    stream_chunks = list(graph_named_steps.stream({"foo": []}))
+    # filtered by output schema
+    assert result == {"bar": "foobar", "foo": ["foo"]}
+    assert stream_chunks == [
+        {"meow1": {"foo": ["foo"]}},
+        {"meow2": {"bar": "foobar"}},
+    ]
+
+    # test two sequences
+
+    def a(state: State):
+        return {"foo": ["a"]}
+
+    def b(state: State):
+        return {"foo": ["b"]}
+
+    builder_two_sequences = StateGraph(State)
+    builder_two_sequences.add_sequence([a])
+    builder_two_sequences.add_sequence([b])
+    builder_two_sequences.add_edge(START, "a")
+    builder_two_sequences.add_edge("a", "b")
+    graph_two_sequences = builder_two_sequences.compile()
+
+    result = graph_two_sequences.invoke({"foo": []})
+    assert result == {"foo": ["a", "b"]}
+
+    stream_chunks = list(graph_two_sequences.stream({"foo": []}))
+    assert stream_chunks == [
+        {"a": {"foo": ["a"]}},
+        {"b": {"foo": ["b"]}},
+    ]
+
+    # test mixed nodes and sequences
+
+    def c(state: State):
+        return {"foo": ["c"]}
+
+    def d(state: State):
+        return {"foo": ["d"]}
+
+    def e(state: State):
+        return {"foo": ["e"]}
+
+    def foo(state: State):
+        if state["foo"][0] == "a":
+            return "d"
+        else:
+            return "c"
+
+    builder_complex = StateGraph(State)
+    builder_complex.add_sequence([a, b])
+    builder_complex.add_conditional_edges("b", foo)
+    builder_complex.add_node(c)
+    builder_complex.add_sequence([d, e])
+    builder_complex.add_edge(START, "a")
+    graph_complex = builder_complex.compile()
+
+    result = graph_complex.invoke({"foo": []})
+    assert result == {"foo": ["a", "b", "d", "e"]}
+
+    result = graph_complex.invoke({"foo": ["start"]})
+    assert result == {"foo": ["start", "a", "b", "c"]}
+
+    stream_chunks = list(graph_complex.stream({"foo": []}))
+    assert stream_chunks == [
+        {"a": {"foo": ["a"]}},
+        {"b": {"foo": ["b"]}},
+        {"d": {"foo": ["d"]}},
+        {"e": {"foo": ["e"]}},
+    ]
+
+
 def test_runnable_passthrough_node_graph() -> None:
     class State(TypedDict):
         changeme: str

libs/langgraph/tests/test_state.py

@@ -5,11 +5,11 @@ from typing import Annotated as Annotated2
 from typing import Any, Optional
 
 import pytest
-from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables import RunnableConfig, RunnableLambda
 from pydantic.v1 import BaseModel
 from typing_extensions import Annotated, NotRequired, Required, TypedDict
 
-from langgraph.graph.state import StateGraph, _warn_invalid_state_schema
+from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema
 from langgraph.managed.shared_value import SharedValue
 
 
@@ -286,3 +286,35 @@ def test_raises_invalid_managed():
             match="Invalid managed channels detected in BadOutputState: some_output_channel. Managed channels are not permitted in Input/Output schema.",
         ):
             StateGraph(_state, input=_inp, output=_outp)
+
+
+def test__get_node_name() -> None:
+    # default runnable name
+    assert _get_node_name(RunnableLambda(func=lambda x: x)) == "RunnableLambda"
+    # custom runnable name
+    assert (
+        _get_node_name(RunnableLambda(name="my_runnable", func=lambda x: x))
+        == "my_runnable"
+    )
+
+    # lambda
+    assert _get_node_name(lambda x: x) == "<lambda>"
+
+    # regular function
+    def func(state):
+        return
+
+    assert _get_node_name(func) == "func"
+
+    class MyClass:
+        def __call__(self, state):
+            return
+
+        def class_method(self, state):
+            return
+
+    # callable class
+    assert _get_node_name(MyClass()) == "MyClass"
+
+    # class method
+    assert _get_node_name(MyClass().class_method) == "class_method"