From 254a38560e0f036531ac2cde7f55ce93a2cc3ff0 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 20:35:10 -0400 Subject: [PATCH 1/8] langgraph: fix drawing graph with __root__ channel --- libs/langgraph/langgraph/graph/state.py | 2 +- .../tests/__snapshots__/test_pregel.ambr | 55 +++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 15 +++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index f6a29198b..68a75fa8e 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1114,7 +1114,7 @@ def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]: else: updates.append(("__root__", i)) return updates - elif input is not None: + else: return [("__root__", input)] diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 4caad8801..9559651f8 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -396,6 +396,61 @@ ''' # --- +# name: test_get_graph_root_channel + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "__start__" + } + }, + { + "id": "child", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "graph", + "state", + "CompiledStateGraph" + ], + "name": "child" + } + }, + { + "id": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "child" + }, + { + "source": "child", + "target": "__end__" + } + ] + } + ''' +# --- +# name: test_get_graph_root_channel.1 + ''' + graph TD; + __start__ --> child; + child --> __end__; + + ''' +# --- # name: test_get_graph_self_loop ''' { diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b251cf465..461791e00 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8754,3 +8754,18 @@ def test_get_graph_self_loop(snapshot: SnapshotAssertion) -> None: assert json.dumps(self_loop_graph.get_graph().to_json(), indent=2) == snapshot assert self_loop_graph.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None: + child_builder = StateGraph(str) + child_builder.add_node("child_node", lambda x: x) + child_builder.add_edge(START, "child_node") + child_graph = child_builder.compile() + + graph_builder = StateGraph(str) + graph_builder.add_node("child", child_graph) + graph_builder.add_edge(START, "child") + graph = graph_builder.compile() + + assert json.dumps(graph.get_graph().to_json(), indent=2) == snapshot + assert graph.get_graph().draw_mermaid(with_styles=False) == snapshot From efd33d860fdef71bc37f286aeab71fae30a0e607 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 21:03:47 -0400 Subject: [PATCH 2/8] update --- libs/langgraph/langgraph/graph/state.py | 2 +- libs/langgraph/langgraph/pregel/draw.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 68a75fa8e..f6a29198b 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1114,7 +1114,7 @@ def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]: else: updates.append(("__root__", i)) return updates - else: + elif input is not None: return [("__root__", input)] diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index f2d8e2631..edc689890 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -108,7 +108,8 @@ def draw_graph( for w in task.writers: # apply regular writes if isinstance(w, ChannelWrite): - w.invoke(None, task.config) + empty_input = "" if "__root__" in specs else {} + w.invoke(empty_input, task.config) # apply conditional writes declared for static analysis, only once if w not in static_seen: static_seen.add(w) From 60c41ce69e17fff2367eb816a667dd38cf7bb58e Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 21:14:49 -0400 Subject: [PATCH 3/8] update --- libs/langgraph/langgraph/pregel/draw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index edc689890..5b84c30c7 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -108,7 +108,7 @@ def draw_graph( for w in task.writers: # apply regular writes if isinstance(w, ChannelWrite): - empty_input = "" if "__root__" in specs else {} + empty_input = "" if "__root__" in specs else None w.invoke(empty_input, task.config) # apply conditional writes declared for static analysis, only once if w not in static_seen: From 54ddde9d4c73a31cb921420867231542ae7faa72 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 21:30:22 -0400 Subject: [PATCH 4/8] update --- libs/langgraph/langgraph/pregel/draw.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index 5b84c30c7..3ece304f6 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -108,7 +108,11 @@ def draw_graph( for w in task.writers: # apply regular writes if isinstance(w, ChannelWrite): - empty_input = "" if "__root__" in specs else None + empty_input = ( + specs["__root__"].ValueType() + if "__root__" in specs + else None + ) w.invoke(empty_input, task.config) # apply conditional writes declared for static analysis, only once if w not in static_seen: From 4ffae6065f1c7fe4dc00fed21ee9a584cbbb14b3 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 21:35:25 -0400 Subject: [PATCH 5/8] lint + update --- libs/langgraph/langgraph/pregel/draw.py | 4 ++-- libs/langgraph/tests/test_remote_graph.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index 3ece304f6..5421a7bba 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -133,8 +133,8 @@ def draw_graph( task.name: { ( w[0], - (task.name, *w) in conditionals, - conditionals.get((task.name, *w)), + (task.name, w[0], w[1] or None) in conditionals, + conditionals.get((task.name, w[0], w[1] or None)), ) for w in task.writes } diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 1358e6f25..82e5d27f1 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -863,7 +863,9 @@ async def test_ainvoke(): assert result == {"messages": [{"type": "human", "content": "world"}]} -@pytest.mark.skip("Unskip this test to manually test the LangGraph Platform integration") +@pytest.mark.skip( + "Unskip this test to manually test the LangGraph Platform integration" +) @pytest.mark.anyio async def test_langgraph_cloud_integration(): from langgraph_sdk.client import get_client, get_sync_client From c33e64daa6295be3fadd5c8664185fc4268b9f71 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 22:03:40 -0400 Subject: [PATCH 6/8] use list --- libs/langgraph/tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 461791e00..5558daa68 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8757,12 +8757,12 @@ def test_get_graph_self_loop(snapshot: SnapshotAssertion) -> None: def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None: - child_builder = StateGraph(str) + child_builder = StateGraph(list) child_builder.add_node("child_node", lambda x: x) child_builder.add_edge(START, "child_node") child_graph = child_builder.compile() - graph_builder = StateGraph(str) + graph_builder = StateGraph(list) graph_builder.add_node("child", child_graph) graph_builder.add_edge(START, "child") graph = graph_builder.compile() From 5ab2aa79bb4949ec96d593ffdd3f7fe6d1c7cc32 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 May 2025 22:11:09 -0400 Subject: [PATCH 7/8] lint again --- libs/langgraph/langgraph/pregel/draw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index 5421a7bba..8e0df382b 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -109,7 +109,7 @@ def draw_graph( # apply regular writes if isinstance(w, ChannelWrite): empty_input = ( - specs["__root__"].ValueType() + cast(BaseChannel, specs["__root__"]).ValueType() if "__root__" in specs else None ) From 3a1c02ff3364fa68d6659282dc994e5c57fd1833 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 15 May 2025 08:38:25 -0400 Subject: [PATCH 8/8] update for consistency --- libs/langgraph/langgraph/pregel/draw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index 8e0df382b..bcae60d98 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -125,7 +125,7 @@ def draw_graph( edges.add((task.name, t[0], True, t[2])) writes = [t for t in writes if t[0] != END] conditionals.update( - {(task.name, *t[:2]): t[2] for t in writes} + {(task.name, t[0], t[1] or None): t[2] for t in writes} ) task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes]) # collect sources