From 2bf66b883499febc384d19d804a9c9c14d1595b9 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 1 Apr 2024 12:29:00 -0700
Subject: [PATCH] Add some more docstrings around multi edges

- prompted by some github issues
---
 langgraph/graph/graph.py |  5 +++-
 langgraph/graph/state.py |  8 ++++++
 tests/test_pregel.py     | 59 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py
index ea89ac3c9..37ec6be78 100644
--- a/langgraph/graph/graph.py
+++ b/langgraph/graph/graph.py
@@ -151,7 +151,10 @@ class Graph:
         if not self.support_multiple_edges and start_key in set(
             start for start, _ in self.edges
         ):
-            raise ValueError(f"Already found path for {start_key}")
+            raise ValueError(
+                f"Already found path for node '{start_key}'.\n"
+                "For multiple edges, use StateGraph with an annotated state key."
+            )
 
         self.edges.add((start_key, end_key))
 
diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py
index 164ecaac4..814bc5c46 100644
--- a/langgraph/graph/state.py
+++ b/langgraph/graph/state.py
@@ -21,6 +21,14 @@ logger = logging.getLogger(__name__)
 
 
 class StateGraph(Graph):
+    """A graph whose nodes communicate by reading and writing to a shared state.
+    The signature of each node is State -> Partial<State>.
+
+    Each state key can optionally be annotated with a reducer function that
+    will be used to aggregate the values of that key received from multiple nodes.
+    The signature of a reducer function is (Value, Value) -> Value.
+    """
+
     def __init__(self, schema: Type[Any]) -> None:
         super().__init__()
         self.schema = schema
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 977ec74c1..85952f468 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -2996,3 +2996,62 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         },
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
+
+
+def test_simple_multi_edge() -> None:
+    class State(TypedDict):
+        my_key: Annotated[str, operator.add]
+
+    def up(state: State):
+        pass
+
+    def side(state: State):
+        pass
+
+    def down(state: State):
+        pass
+
+    graph = StateGraph(State)
+
+    graph.add_node("up", up)
+    graph.add_node("side", side)
+    graph.add_node("down", down)
+
+    graph.set_entry_point("up")
+    graph.add_edge("up", "side")
+    graph.add_edge("up", "down")
+    graph.add_edge(["up", "side"], "down")
+    graph.set_finish_point("down")
+
+    app = graph.compile()
+
+    assert app.get_graph().draw_ascii() == (
+        """    +-----------+  
+    | __start__ |  
+    +-----------+  
+           *       
+           *       
+           *       
+        +----+     
+        | up |     
+        +----+     
+       **     **   
+      *         *  
+     *           * 
++------+          *
+| side |         * 
++------+        *  
+       **     **   
+         *   *     
+          * *      
+       +------+    
+       | down |    
+       +------+    
+           *       
+           *       
+           *       
+      +---------+  
+      | __end__ |  
+      +---------+  """
+    )
+    assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}