diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index e88b6db99..177a78f80 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -847,7 +847,9 @@ class CompiledStateGraph(CompiledGraph): ) -> Sequence[Union[ChannelWriteEntry, Send]]: writes = [ ( - ChannelWriteEntry(CHANNEL_BRANCH_TO.format(p), None) + ChannelWriteEntry( + p if p == END else CHANNEL_BRANCH_TO.format(p), None + ) if not isinstance(p, Send) else p ) @@ -1067,9 +1069,14 @@ def _control_static( ends: Union[tuple[str, ...], dict[str, str]], ) -> Sequence[tuple[str, Any, Optional[str]]]: if isinstance(ends, dict): - return [(CHANNEL_BRANCH_TO.format(k), None, label) for k, label in ends.items()] + return [ + (k if k == END else CHANNEL_BRANCH_TO.format(k), None, label) + for k, label in ends.items() + ] else: - return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends] + return [ + (e if e == END else CHANNEL_BRANCH_TO.format(e), None, None) for e in ends + ] def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]: diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr index 4a9d05004..cdd896b8d 100644 --- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr +++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr @@ -198,6 +198,12 @@ "source": "__start__", "target": "agent" }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + }, { "source": "agent", "target": "tools", @@ -207,11 +213,6 @@ { "source": "tools", "target": "agent" - }, - { - "source": "agent", - "target": "__end__", - "conditional": true } ] } @@ -221,9 +222,9 @@ ''' graph TD; __start__ --> agent; + agent -.  exit  .-> __end__; agent -.  continue  .-> tools; tools --> agent; - agent -.-> __end__; ''' # --- @@ -284,6 +285,12 @@ "source": "__start__", "target": "agent" }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + }, { "source": "agent", "target": "tools", @@ -293,11 +300,6 @@ { "source": "tools", "target": "agent" - }, - { - "source": "agent", - "target": "__end__", - "conditional": true } ] } @@ -307,9 +309,9 @@ ''' graph TD; __start__ --> agent; + agent -.  end  .-> __end__; agent -.  continue  .-> tools; tools --> agent; - agent -.-> __end__; ''' # --- @@ -371,6 +373,11 @@ "source": "__start__", "target": "agent" }, + { + "source": "agent", + "target": "__end__", + "conditional": true + }, { "source": "agent", "target": "tools", @@ -379,11 +386,6 @@ { "source": "tools", "target": "agent" - }, - { - "source": "agent", - "target": "__end__", - "conditional": true } ] } @@ -393,9 +395,9 @@ ''' graph TD; __start__ --> agent; + agent -.-> __end__; agent -.-> tools; tools --> agent; - agent -.-> __end__; ''' # --- diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index d4815ce86..69303056d 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -156,7 +156,8 @@ }, { "source": "left", - "target": "__end__" + "target": "__end__", + "conditional": true }, { "source": "right", @@ -171,7 +172,7 @@ graph TD; __start__ -.  go-left  .-> left; __start__ -.  go-right  .-> right; - left --> __end__; + left -.-> __end__; right --> __end__; ''' @@ -736,8 +737,10 @@ Call_Tool -.-> Researcher; Chart_Generator -.  call_tool  .-> Call_Tool; Chart_Generator -.  continue  .-> Researcher; + Chart_Generator -.  end  .-> __end__; Researcher -.  call_tool  .-> Call_Tool; Researcher -.  continue  .-> Chart_Generator; + Researcher -.  end  .-> __end__; __start__ --> Researcher; Researcher -.  redo  .-> Researcher; @@ -777,26 +780,26 @@ gp_one(gp_one) __end__([

__end__

]):::last __start__ --> gp_one; + gp_one -.  1  .-> __end__; gp_one -.  0  .-> gp_two___start__; gp_two___end__ --> gp_one; - gp_one -.-> __end__; subgraph gp_two gp_two___start__(

__start__

) gp_two_p_one(p_one) gp_two___end__(

__end__

) gp_two___start__ --> gp_two_p_one; + gp_two_p_one -.  1  .-> gp_two___end__; gp_two_p_one -.  0  .-> gp_two_p_two___start__; gp_two_p_two___end__ --> gp_two_p_one; - gp_two_p_one -.-> gp_two___end__; subgraph p_two gp_two_p_two___start__(

__start__

) gp_two_p_two_c_one(c_one) gp_two_p_two_c_two(c_two) gp_two_p_two___end__(

__end__

) gp_two_p_two___start__ --> gp_two_p_two_c_one; + gp_two_p_two_c_one -.  1  .-> gp_two_p_two___end__; gp_two_p_two_c_one -.  0  .-> gp_two_p_two_c_two; gp_two_p_two_c_two --> gp_two_p_two_c_one; - gp_two_p_two_c_one -.-> gp_two_p_two___end__; end end classDef default fill:#f2f0ff,line-height:1.2 @@ -817,18 +820,18 @@ p_one(p_one) __end__([

__end__

]):::last __start__ --> p_one; + p_one -.  1  .-> __end__; p_one -.  0  .-> p_two___start__; p_two___end__ --> p_one; - p_one -.-> __end__; subgraph p_two p_two___start__(

__start__

) p_two_c_one(c_one) p_two_c_two(c_two) p_two___end__(

__end__

) p_two___start__ --> p_two_c_one; + p_two_c_one -.  1  .-> p_two___end__; p_two_c_one -.  0  .-> p_two_c_two; p_two_c_two --> p_two_c_one; - p_two_c_one -.-> p_two___end__; end classDef default fill:#f2f0ff,line-height:1.2 classDef first fill-opacity:0 @@ -843,6 +846,11 @@ 'source': '__start__', 'target': 'ask_question', }), + dict({ + 'conditional': True, + 'source': 'answer_question', + 'target': '__end__', + }), dict({ 'conditional': True, 'source': 'answer_question', @@ -852,11 +860,6 @@ 'source': 'ask_question', 'target': 'answer_question', }), - dict({ - 'conditional': True, - 'source': 'answer_question', - 'target': '__end__', - }), ]), 'nodes': list([ dict({ @@ -1008,6 +1011,11 @@ 'source': 'conduct_interview:__start__', 'target': 'conduct_interview:ask_question', }), + dict({ + 'conditional': True, + 'source': 'conduct_interview:answer_question', + 'target': 'conduct_interview:__end__', + }), dict({ 'conditional': True, 'source': 'conduct_interview:answer_question', @@ -1017,11 +1025,6 @@ 'source': 'conduct_interview:ask_question', 'target': 'conduct_interview:answer_question', }), - dict({ - 'conditional': True, - 'source': 'conduct_interview:answer_question', - 'target': 'conduct_interview:__end__', - }), ]), 'nodes': list([ dict({ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e5ca5248d..55c04c853 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8703,3 +8703,88 @@ async def test_batch_update_as_input(checkpointer_name: str) -> None: ] assert new_history == history + + +async def test_draw_invalid(): + from langchain_core.messages import BaseMessage + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + workflow = StateGraph(AgentState) + + async def call_model(state: AgentState) -> AgentState: + return state + + async def call_tool(state: AgentState) -> AgentState: + return state + + async def do_nothing(state: AgentState) -> AgentState: + return state + + def should_continue(state): + messages = state["messages"] + last_message = messages[-1] + if last_message.content.startswith("end"): + return END + else: + return [Send("tool", last_message), Send("nothing", last_message)] + + workflow.add_node("agent", call_model) + workflow.add_node("tool", call_tool) + workflow.add_node("nothing", do_nothing) + workflow.set_entry_point("agent") + workflow.add_conditional_edges( + "agent", + should_continue, + path_map=["tool", "nothing", END], + ) + workflow.add_edge("tool", "agent") + + graph = workflow.compile() + + assert graph.get_graph().to_json() == { + "nodes": [ + { + "id": "__start__", + "type": "runnable", + "data": { + "id": ["langchain", "schema", "runnable", "RunnablePassthrough"], + "name": "__start__", + }, + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "runnable", "RunnableCallable"], + "name": "agent", + }, + }, + { + "id": "tool", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "runnable", "RunnableCallable"], + "name": "tool", + }, + }, + { + "id": "nothing", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "runnable", "RunnableCallable"], + "name": "nothing", + }, + }, + {"id": "__end__"}, + ], + "edges": [ + {"source": "__start__", "target": "agent"}, + {"source": "agent", "target": "__end__", "conditional": True}, + {"source": "agent", "target": "nothing", "conditional": True}, + {"source": "agent", "target": "tool", "conditional": True}, + {"source": "tool", "target": "agent"}, + {"source": "nothing", "target": "__end__"}, + ], + }