From 35e3276e34a183bae9113b1ef9e0a52b06bf35af Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Thu, 7 Nov 2024 13:16:34 -0500 Subject: [PATCH] langgraph: add add_sequence to StateGraph (#2352) --- libs/langgraph/langgraph/graph/state.py | 59 ++++++++++- libs/langgraph/tests/test_pregel.py | 130 ++++++++++++++++++++++++ libs/langgraph/tests/test_state.py | 36 ++++++- 3 files changed, 219 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 89ee6ecca..1b1b88cf1 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ff56d7394..71b410ed4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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 diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 5475c250f..0a4a8725a 100644 --- a/libs/langgraph/tests/test_state.py +++ b/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) == "" + + # 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"