diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index f2d8e2631..bcae60d98 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -108,7 +108,12 @@ def draw_graph( for w in task.writers: # apply regular writes if isinstance(w, ChannelWrite): - w.invoke(None, task.config) + empty_input = ( + cast(BaseChannel, 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: static_seen.add(w) @@ -120,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 @@ -128,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/__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..5558daa68 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(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(list) + 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 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