Merge pull request #362 from langchain-ai/nc/29apr/support-cond-edge-list-literal

For cond edges support specifying list of possible destinations as a list or typing annotation

langgraph/checkpoint/aiosqlite.py

@@ -75,8 +75,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
             async with self.conn.execute(
                 "SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
                 (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["thread_ts"],
+                    str(config["configurable"]["thread_id"]),
+                    str(config["configurable"]["thread_ts"]),
                 ),
             ) as cursor:
                 if value := await cursor.fetchone():
@@ -95,7 +95,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         else:
             async with self.conn.execute(
                 "SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
-                (config["configurable"]["thread_id"],),
+                (str(config["configurable"]["thread_id"]),),
             ) as cursor:
                 if value := await cursor.fetchone():
                     return CheckpointTuple(
@@ -120,7 +120,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         await self.setup()
         async with self.conn.execute(
             "SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
-            (config["configurable"]["thread_id"],),
+            (str(config["configurable"]["thread_id"]),),
         ) as cursor:
             async for thread_id, thread_ts, parent_ts, value in cursor:
                 yield CheckpointTuple(
@@ -138,7 +138,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         async with self.conn.execute(
             "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
             (
-                config["configurable"]["thread_id"],
+                str(config["configurable"]["thread_id"]),
                 checkpoint["ts"],
                 config["configurable"].get("thread_ts"),
                 self.serde.dumps(checkpoint),

langgraph/checkpoint/sqlite.py

@@ -93,8 +93,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                 cur.execute(
                     "SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
                     (
-                        config["configurable"]["thread_id"],
-                        config["configurable"]["thread_ts"],
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["thread_ts"]),
                     ),
                 )
                 if value := cur.fetchone():
@@ -113,7 +113,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
             else:
                 cur.execute(
                     "SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
-                    (config["configurable"]["thread_id"],),
+                    (str(config["configurable"]["thread_id"]),),
                 )
                 if value := cur.fetchone():
                     return CheckpointTuple(
@@ -138,7 +138,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
         with self.cursor(transaction=False) as cur:
             cur.execute(
                 "SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
-                (config["configurable"]["thread_id"],),
+                (str(config["configurable"]["thread_id"]),),
             )
             for thread_id, thread_ts, parent_ts, value in cur:
                 yield CheckpointTuple(
@@ -159,7 +159,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
             cur.execute(
                 "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
                 (
-                    config["configurable"]["thread_id"],
+                    str(config["configurable"]["thread_id"]),
                     checkpoint["ts"],
                     config["configurable"].get("thread_ts"),
                     self.serde.dumps(checkpoint),

langgraph/graph/graph.py

@@ -1,15 +1,19 @@
 import logging
-from collections import Counter, defaultdict
+from collections import defaultdict
 from typing import (
     Any,
     Awaitable,
     Callable,
     Dict,
+    Literal,
     NamedTuple,
     Optional,
     Sequence,
     Union,
     cast,
+    get_args,
+    get_origin,
+    get_type_hints,
 )
 
 from langchain_core.runnables import Runnable
@@ -142,7 +146,7 @@ class Graph:
             Callable[..., Awaitable[Union[str, list[str]]]],
             Runnable[Any, Union[str, list[str]]],
         ],
-        path_map: Optional[dict[str, str]] = None,
+        path_map: Optional[Union[dict[str, str], list[str]]] = None,
         then: Optional[str] = None,
     ) -> None:
         """Add a conditional edge from the starting node to any number of destination nodes.
@@ -166,6 +170,14 @@ class Graph:
                 "Adding an edge to a graph that has already been compiled. This will "
                 "not be reflected in the compiled graph."
             )
+        # coerce path_map to a dictionary
+        if isinstance(path_map, dict):
+            pass
+        elif isinstance(path_map, list):
+            path_map = {name: name for name in path_map}
+        elif rtn_type := get_type_hints(path).get("return"):
+            if get_origin(rtn_type) is Literal:
+                path_map = {name: name for name in get_args(rtn_type)}
         # find a name for the condition
         path = coerce_to_runnable(path, name=None, trace=True)
         name = path.name or "condition"
@@ -378,7 +390,6 @@ class CompiledGraph(Pregel):
         config: Optional[RunnableConfig] = None,
         *,
         xray: Union[int, bool] = False,
-        add_condition_nodes: bool = True,
     ) -> DrawableGraph:
         """Returns a drawable representation of the computation graph."""
         graph = DrawableGraph()
@@ -415,38 +426,26 @@ class CompiledGraph(Pregel):
                 end_nodes[key] = n
         for start, end in sorted(self.graph._all_edges):
             graph.add_edge(start_nodes[start], end_nodes[end])
-        branches_by_name = Counter(
-            name for _, branches in self.graph.branches.items() for name in branches
-        )
         for start, branches in self.graph.branches.items():
             default_ends = {
                 **{k: k for k in self.graph.nodes if k != start},
                 END: END,
             }
-            for name, branch in branches.items():
+            for _, branch in branches.items():
                 if branch.ends is not None:
                     ends = branch.ends
                 elif branch.then is not None:
                     ends = {k: k for k in default_ends if k not in (END, branch.then)}
                 else:
                     ends = default_ends
-
-                if add_condition_nodes is True:
-                    cond = graph.add_node(
-                        branch.path,
-                        f"{start}_{name}" if branches_by_name[name] > 1 else name,
+                for label, end in ends.items():
+                    graph.add_edge(
+                        start_nodes[start],
+                        end_nodes[end],
+                        label if label != end else None,
+                        conditional=True,
                     )
-                    graph.add_edge(start_nodes[start], cond)
-                    for label, end in ends.items():
-                        graph.add_edge(cond, end_nodes[end], label, conditional=True)
-                        if branch.then is not None:
-                            graph.add_edge(start_nodes[end], end_nodes[branch.then])
-                else:
-                    for label, end in ends.items():
-                        graph.add_edge(
-                            start_nodes[start], end_nodes[end], label, conditional=True
-                        )
-                        if branch.then is not None:
-                            graph.add_edge(start_nodes[end], end_nodes[branch.then])
+                    if branch.then is not None:
+                        graph.add_edge(start_nodes[end], end_nodes[branch.then])
 
         return graph

tests/__snapshots__/test_pregel_async.ambr

@@ -10,12 +10,12 @@
               +---------------+            
               | rewrite_query |            
               +---------------+            
-               ***         ***             
-              *               *            
-            **                 **          
-  +--------------+          +---------+    
-  | analyzer_one |          | decider |    
-  +--------------+          +---------+    
+               ***         ...             
+              *               .            
+            **                 ...         
+  +--------------+                .        
+  | analyzer_one |                .        
+  +--------------+                .        
           *                       .        
           *                       .        
           *                       .        
@@ -47,12 +47,12 @@
               +---------------+            
               | rewrite_query |            
               +---------------+            
-               ***         ***             
-              *               *            
-            **                 **          
-  +--------------+          +---------+    
-  | analyzer_one |          | decider |    
-  +--------------+          +---------+    
+               ***         ...             
+              *               .            
+            **                 ...         
+  +--------------+                .        
+  | analyzer_one |                .        
+  +--------------+                .        
           *                       .        
           *                       .        
           *                       .        
@@ -84,12 +84,12 @@
               +---------------+            
               | rewrite_query |            
               +---------------+            
-               ***         ***             
-              *               *            
-            **                 **          
-  +--------------+         +-----------+   
-  | analyzer_one |         | condition |   
-  +--------------+         +-----------+   
+               ***         ...             
+              *               .            
+            **                 ...         
+  +--------------+                .        
+  | analyzer_one |                .        
+  +--------------+                .        
           *                       .        
           *                       .        
           *                       .        
@@ -121,12 +121,12 @@
               +---------------+            
               | rewrite_query |            
               +---------------+            
-               ***         ***             
-              *               *            
-            **                 **          
-  +--------------+         +-----------+   
-  | analyzer_one |         | condition |   
-  +--------------+         +-----------+   
+               ***         ...             
+              *               .            
+            **                 ...         
+  +--------------+                .        
+  | analyzer_one |                .        
+  +--------------+                .        
           *                       .        
           *                       .        
           *                       .        

tests/test_pregel.py

@@ -4,7 +4,7 @@ import time
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Annotated, Any, Generator, Optional, TypedDict, Union
+from typing import Annotated, Any, Generator, Literal, Optional, TypedDict, Union
 
 import pytest
 from langchain_core.runnables import RunnableLambda, RunnablePassthrough
@@ -882,19 +882,9 @@ def test_conditional_graph(
     app = workflow.compile()
 
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
-    assert (
-        app.get_graph(add_condition_nodes=False).draw_mermaid(with_styles=False)
-        == snapshot
-    )
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
     assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
-    assert app.get_graph(xray=True).draw_ascii() == snapshot
-    assert (
-        app.get_graph(xray=True, add_condition_nodes=False).draw_mermaid(
-            with_styles=False
-        )
-        == snapshot
-    )
+    assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"input": "what is weather in sf"}) == {
         "input": "what is weather in sf",
@@ -1470,7 +1460,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
     assert app.get_input_schema().schema_json() == snapshot
     assert app.get_output_schema().schema_json() == snapshot
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert (
         app.invoke("what is weather in sf", debug=True)
@@ -1570,7 +1560,7 @@ def test_conditional_graph_state(
     assert app.get_input_schema().schema_json() == snapshot
     assert app.get_output_schema().schema_json() == snapshot
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"input": "what is weather in sf"}) == {
         "input": "what is weather in sf",
@@ -1922,7 +1912,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
     assert app.get_input_schema().schema_json() == snapshot
     assert app.get_output_schema().schema_json() == snapshot
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"input": "what is weather in sf"}) == {
         "input": "what is weather in sf",
@@ -1988,7 +1978,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
     assert app.get_input_schema().schema_json() == snapshot
     assert app.get_output_schema().schema_json() == snapshot
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke(
         {"messages": [HumanMessage(content="what is weather in sf")]}
@@ -2251,7 +2241,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
     assert app.get_input_schema().schema_json() == snapshot
     assert app.get_output_schema().schema_json() == snapshot
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke(
         {"messages": [HumanMessage(content="what is weather in sf")]}
@@ -2466,7 +2456,7 @@ def test_message_graph(
     assert app.get_input_schema().schema_json() == snapshot
     assert app.get_output_schema().schema_json() == snapshot
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke(HumanMessage(content="what is weather in sf")) == [
         HumanMessage(
@@ -3031,10 +3021,7 @@ def test_start_branch_then(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
     tool_two = tool_two_graph.compile()
-    assert tool_two.get_graph().draw_ascii() == snapshot
-    assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot
     assert tool_two.get_graph().draw_mermaid() == snapshot
-    assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot
 
     assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
         "my_key": "value slow",
@@ -3115,6 +3102,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
     invalid_graph.add_conditional_edges(
         source="prepare",
         path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
+        path_map=["tool_two_slow", "tool_two_fast"],
     )
     invalid_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
     invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
@@ -3136,10 +3124,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
     tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
     tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
     tool_two = tool_two_graph.compile()
-    assert tool_two.get_graph().draw_ascii() == snapshot
-    assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot
+    assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
     assert tool_two.get_graph().draw_mermaid() == snapshot
-    assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot
 
     assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
         "my_key": "value prepared slow finished",
@@ -3311,7 +3297,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(
 
     app = workflow.compile()
 
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: what is weather in sf",
@@ -3387,6 +3373,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
     def qa(data: State) -> State:
         return {"answer": ",".join(data["docs"])}
 
+    def rewrite_query_then(data: State) -> Literal["retriever_two"]:
+        return "retriever_two"
+
     workflow = StateGraph(State)
 
     workflow.add_node("rewrite_query", rewrite_query)
@@ -3398,15 +3387,13 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
     workflow.set_entry_point("rewrite_query")
     workflow.add_edge("rewrite_query", "analyzer_one")
     workflow.add_edge("analyzer_one", "retriever_one")
-    workflow.add_conditional_edges(
-        "rewrite_query", lambda _: "retriever_two", {"retriever_two": "retriever_two"}
-    )
+    workflow.add_conditional_edges("rewrite_query", rewrite_query_then)
     workflow.add_edge(["retriever_one", "retriever_two"], "qa")
     workflow.set_finish_point("qa")
 
     app = workflow.compile()
 
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
         "query": "analyzed: query: what is weather in sf",
@@ -3507,7 +3494,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
 
     app = workflow.compile()
 
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     with pytest.raises(ValidationError):
         app.invoke({"query": {}})
@@ -3844,7 +3831,7 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
 
     app = graph.compile()
 
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
     assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
 
 
@@ -3873,7 +3860,6 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
     app = graph.compile()
 
     assert app.get_graph(xray=True).to_json() == snapshot
-    assert app.get_graph().draw_ascii() == snapshot
     assert app.get_graph(xray=True).draw_mermaid() == snapshot
 
 
@@ -3911,7 +3897,7 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
 
     app = graph.compile()
 
-    assert app.get_graph().draw_ascii() == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
     assert app.get_graph(xray=True).draw_mermaid() == snapshot
     assert app.invoke(
         {"my_key": "my value", "never_called": never_called}, debug=True