From 6f888b6206e340f91e79354c4381a0efade3b058 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 14 May 2024 17:51:25 -0700
Subject: [PATCH] Add single arg overload for add_node, which takes
 function/runnable's name as the node name

---
 langgraph/graph/graph.py | 26 +++++++++++++++++------
 langgraph/graph/state.py | 46 ++++++++++++++++++++++++++++++++++------
 tests/test_pregel.py     | 14 ++++++------
 3 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py
index 44f3a29f3..6595238c2 100644
--- a/langgraph/graph/graph.py
+++ b/langgraph/graph/graph.py
@@ -14,6 +14,7 @@ from typing import (
     get_args,
     get_origin,
     get_type_hints,
+    overload,
 )
 
 from langchain_core.runnables import Runnable
@@ -106,18 +107,31 @@ class Graph:
     def _all_edges(self) -> set[tuple[str, str]]:
         return self.edges
 
-    def add_node(self, key: str, action: RunnableLike) -> None:
+    @overload
+    def add_node(self, node: RunnableLike) -> None:
+        ...
+
+    @overload
+    def add_node(self, node: str, action: RunnableLike) -> None:
+        ...
+
+    def add_node(
+        self, node: Union[str, RunnableLike], action: Optional[RunnableLike] = None
+    ) -> None:
         if self.compiled:
             logger.warning(
                 "Adding a node to a graph that has already been compiled. This will "
                 "not be reflected in the compiled graph."
             )
-        if key in self.nodes:
-            raise ValueError(f"Node `{key}` already present.")
-        if key == END or key == START:
-            raise ValueError(f"Node `{key}` is reserved.")
+        if not isinstance(node, str):
+            action = node
+            node = getattr(action, "name", action.__name__)
+        if node in self.nodes:
+            raise ValueError(f"Node `{node}` already present.")
+        if node == END or node == START:
+            raise ValueError(f"Node `{node}` is reserved.")
 
-        self.nodes[key] = coerce_to_runnable(action, name=key, trace=False)
+        self.nodes[node] = coerce_to_runnable(action, name=node, trace=False)
 
     def add_edge(self, start_key: str, end_key: str) -> None:
         if self.compiled:
diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py
index 2fe123c0e..99b6c0018 100644
--- a/langgraph/graph/state.py
+++ b/langgraph/graph/state.py
@@ -1,7 +1,16 @@
 import logging
 from functools import partial
 from inspect import signature
-from typing import Any, Optional, Sequence, Type, Union, get_origin, get_type_hints
+from typing import (
+    Any,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+    get_origin,
+    get_type_hints,
+    overload,
+)
 
 from langchain_core.pydantic_v1 import BaseModel
 from langchain_core.runnables import Runnable, RunnableConfig
@@ -94,11 +103,28 @@ class StateGraph(Graph):
             (start, end) for starts, end in self.waiting_edges for start in starts
         }
 
-    def add_node(self, key: str, action: RunnableLike) -> None:
+    @overload
+    def add_node(self, node: RunnableLike) -> None:
+        """Adds a new node to the state graph.
+        Will take the name of the function/runnable as the node name.
+
+        Args:
+            node (RunnableLike): The function or runnable this node will run.
+
+        Raises:
+            ValueError: If the key is already being used as a state key.
+
+        Returns:
+            None
+        """
+        ...
+
+    @overload
+    def add_node(self, node: str, action: RunnableLike) -> None:
         """Adds a new node to the state graph.
 
         Args:
-            key (str): The key of the node.
+            node (str): The key of the node.
             action (RunnableLike): The action associated with the node.
 
         Raises:
@@ -107,9 +133,17 @@ class StateGraph(Graph):
         Returns:
             None
         """
-        if key in self.channels:
-            raise ValueError(f"'{key}' is already being used as a state key")
-        return super().add_node(key, action)
+        ...
+
+    def add_node(
+        self, node: Union[str, RunnableLike], action: Optional[RunnableLike] = None
+    ) -> None:
+        if not isinstance(node, str):
+            action = node
+            node = getattr(action, "name", action.__name__)
+        if node in self.channels:
+            raise ValueError(f"'{node}' is already being used as a state key")
+        return super().add_node(node, action)
 
     def add_edge(self, start_key: Union[str, list[str]], end_key: str) -> None:
         """Adds a directed edge from the start node to the end node.
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index b98000b89..f3ec83798 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -4683,6 +4683,9 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
         answer: str
         docs: Annotated[list[str], sorted_add]
 
+    workflow = StateGraph(State)
+
+    @workflow.add_node
     def rewrite_query(data: State) -> State:
         return {"query": f'query: {data["query"]}'}
 
@@ -4698,13 +4701,10 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
     def qa(data: State) -> State:
         return {"answer": ",".join(data["docs"])}
 
-    workflow = StateGraph(State)
-
-    workflow.add_node("rewrite_query", rewrite_query)
-    workflow.add_node("analyzer_one", analyzer_one)
-    workflow.add_node("retriever_one", retriever_one)
-    workflow.add_node("retriever_two", retriever_two)
-    workflow.add_node("qa", qa)
+    workflow.add_node(analyzer_one)
+    workflow.add_node(retriever_one)
+    workflow.add_node(retriever_two)
+    workflow.add_node(qa)
 
     workflow.set_entry_point("rewrite_query")
     workflow.add_edge("rewrite_query", "analyzer_one")