From 84a0eca935b4e0b1d1eff62b1e6710fad18f4b02 Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Tue, 11 Feb 2025 11:01:45 -0800 Subject: [PATCH] langgraph: allow passing destinations to .add_node (#3384) --- libs/langgraph/langgraph/graph/graph.py | 7 ++- libs/langgraph/langgraph/graph/state.py | 17 +++++-- libs/langgraph/tests/test_pregel.py | 66 +++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 65bc908b7..4bda590c3 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -47,7 +47,7 @@ logger = logging.getLogger(__name__) class NodeSpec(NamedTuple): runnable: Runnable metadata: Optional[dict[str, Any]] = None - ends: Optional[tuple[str, ...]] = EMPTY_SEQ + ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ class Branch(NamedTuple): @@ -625,7 +625,10 @@ class CompiledGraph(Pregel): if branch.then is not None: add_edge(end, branch.then) for key, n in self.builder.nodes.items(): - if n.ends: + if isinstance(n.ends, dict): + for end, label in n.ends.items(): + add_edge(key, end, label, conditional=True) + elif isinstance(n.ends, tuple): for end in n.ends: add_edge(key, end, conditional=True) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 4d9f118a4..8c93c0c1d 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -90,7 +90,7 @@ class StateNodeSpec(NamedTuple): metadata: Optional[dict[str, Any]] input: Type[Any] retry_policy: Optional[RetryPolicy] - ends: Optional[tuple[str, ...]] = EMPTY_SEQ + ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ class StateGraph(Graph): @@ -230,6 +230,7 @@ class StateGraph(Graph): metadata: Optional[dict[str, Any]] = None, input: Optional[Type[Any]] = None, retry: Optional[RetryPolicy] = None, + destinations: Optional[Union[dict[str, str], tuple[str]]] = None, ) -> Self: """Adds a new node to the state graph. Will take the name of the function/runnable as the node name. @@ -254,6 +255,7 @@ class StateGraph(Graph): metadata: Optional[dict[str, Any]] = None, input: Optional[Type[Any]] = None, retry: Optional[RetryPolicy] = None, + destinations: Optional[Union[dict[str, str], tuple[str]]] = None, ) -> Self: """Adds a new node to the state graph. @@ -277,6 +279,7 @@ class StateGraph(Graph): metadata: Optional[dict[str, Any]] = None, input: Optional[Type[Any]] = None, retry: Optional[RetryPolicy] = None, + destinations: Optional[Union[dict[str, str], tuple[str]]] = None, ) -> Self: """Adds a new node to the state graph. @@ -288,7 +291,11 @@ class StateGraph(Graph): metadata (Optional[dict[str, Any]]): The metadata associated with the node. (default: None) input (Optional[Type[Any]]): The input schema for the node. (default: the graph's input schema) retry (Optional[RetryPolicy]): The policy for retrying the node. (default: None) - + destinations (Optional[Union[dict[str, str], tuple[str]]]): Destinations that indicate where a node can route to. + This is useful for edgeless graphs with nodes that return `Command` objects. + If a dict is provided, the keys will be used as the target node names and the values will be used as the labels for the edges. + If a tuple is provided, the values will be used as the target node names. + NOTE: this is only used for graph rendering and doesn't have any effect on the graph execution. Raises: ValueError: If the key is already being used as a state key. @@ -357,7 +364,7 @@ class StateGraph(Graph): f"'{character}' is a reserved character and is not allowed in the node names." ) - ends = EMPTY_SEQ + ends: Union[tuple[str, ...], dict[str, str]] = EMPTY_SEQ try: if ( isfunction(action) @@ -401,6 +408,10 @@ class StateGraph(Graph): ends = vals except (TypeError, StopIteration): pass + + if destinations is not None: + ends = destinations + if input is not None: self._add_schema(input) self.nodes[cast(str, node)] = StateNodeSpec( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 248bb7be2..ccd4fae94 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -35,6 +35,7 @@ from langchain_core.runnables import ( RunnableLambda, RunnablePassthrough, ) +from langchain_core.runnables.graph import Edge from langsmith import traceable from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -6416,3 +6417,68 @@ def test_tags_stream_mode_messages() -> None: }, ) ] + + +def test_node_destinations() -> None: + class State(TypedDict): + foo: Annotated[str, operator.add] + + def node_a(state: State): + value = state["foo"] + if value == "a": + goto = "node_b" + else: + goto = "node_c" + + return Command( + update={"foo": value}, + goto=goto, + graph=Command.PARENT, + ) + + subgraph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile() + + # test calling subgraph inside a node function + def call_subgraph(state: State): + return subgraph.invoke(state) + + def node_b(state: State): + return {"foo": "b"} + + def node_c(state: State): + return {"foo": "c"} + + for subgraph_node in (subgraph, call_subgraph): + # destinations w/ tuples + builder = StateGraph(State) + builder.add_edge(START, "child") + builder.add_node("child", subgraph_node, destinations=("node_b", "node_c")) + builder.add_node(node_b) + builder.add_node(node_c) + compiled_graph = builder.compile() + assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"} + + graph = compiled_graph.get_graph() + assert [ + Edge(source="__start__", target="child", data=None, conditional=False), + Edge(source="child", target="node_b", data=None, conditional=True), + Edge(source="child", target="node_c", data=None, conditional=True), + ] == graph.edges + + # destinations w/ dicts + builder = StateGraph(State) + builder.add_edge(START, "child") + builder.add_node( + "child", subgraph_node, destinations={"node_b": "foo", "node_c": "bar"} + ) + builder.add_node(node_b) + builder.add_node(node_c) + compiled_graph = builder.compile() + assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"} + + graph = compiled_graph.get_graph() + assert [ + Edge(source="__start__", target="child", data=None, conditional=False), + Edge(source="child", target="node_b", data="foo", conditional=True), + Edge(source="child", target="node_c", data="bar", conditional=True), + ] == graph.edges