diff --git a/libs/langgraph/langgraph/pregel/_draw.py b/libs/langgraph/langgraph/pregel/_draw.py index b8ae73389..af6d03427 100644 --- a/libs/langgraph/langgraph/pregel/_draw.py +++ b/libs/langgraph/langgraph/pregel/_draw.py @@ -215,10 +215,11 @@ def draw_graph( termini = {d for _, d, _, _ in edges if d != END}.difference( s for s, _, _, _ in edges ) + end_edge_exists = any(d == END for _, d, _, _ in edges) if termini: for src in sorted(termini): add_edge(graph, src, END) - elif len(step_sources) == 1: + elif len(step_sources) == 1 and not end_edge_exists: for src in sorted(step_sources): add_edge(graph, src, END, conditional=True) # replace subgraphs diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 5167d8924..b7dfbcd0e 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -795,6 +795,99 @@ ''' # --- +# name: test_get_graph_nonterminal_last_step_source + ''' + { + "edges": [ + { + "source": "__start__", + "target": "human" + }, + { + "conditional": true, + "source": "chatbot", + "target": "human" + }, + { + "conditional": true, + "source": "chatbot", + "target": "tools" + }, + { + "conditional": true, + "source": "human", + "target": "__end__" + }, + { + "conditional": true, + "source": "human", + "target": "chatbot" + }, + { + "source": "tools", + "target": "chatbot" + } + ], + "nodes": [ + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "__start__" + }, + "id": "__start__", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "chatbot" + }, + "id": "chatbot", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "id": "tools", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "human" + }, + "id": "human", + "type": "runnable" + }, + { + "id": "__end__" + } + ] + } + ''' +# --- # name: test_repeat_condition ''' graph TD; diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index cbf6d03b6..7cc67bd35 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8343,6 +8343,46 @@ def test_subgraph_streaming_sync() -> None: assert result["num_chunks"] == 9 +def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: list[str] + + def chatbot_node(state: State) -> State: + return {"messages": state["messages"] + ["chatbot"]} + + def tools_node(state: State) -> State: + return {"messages": state["messages"] + ["tools"]} + + def human_node(state: State) -> State: + return {"messages": state["messages"] + ["human"]} + + def tools_condition(_: State) -> str: + return "tools" + + def end_condition(_: State) -> str: + return "chatbot" + + workflow = StateGraph(State) + workflow.add_node("chatbot", chatbot_node) + workflow.add_node("tools", tools_node) + workflow.add_node("human", human_node) + + workflow.add_edge(START, "human") + workflow.add_edge("tools", "chatbot") + + workflow.add_conditional_edges( + "chatbot", tools_condition, {"tools": "tools", "human": "human"} + ) + workflow.add_conditional_edges( + "human", end_condition, {"chatbot": "chatbot", END: END} + ) + + app = workflow.compile() + graph = app.get_graph() + graph_json = graph.to_json() + + assert json.dumps(graph_json, indent=2, sort_keys=True) == snapshot + def test_null_resume_disallowed_with_multiple_interrupts( sync_checkpointer: BaseCheckpointSaver, ) -> None: