langgraph: allow passing destinations to .add_node (#3384)

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)
 

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(

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