diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py index f80519ec3..082fb8a5e 100644 --- a/libs/langgraph/langgraph/graph/branch.py +++ b/libs/langgraph/langgraph/graph/branch.py @@ -142,8 +142,8 @@ class Branch(NamedTuple): ), list( zip_longest( - writer([e for e in self.ends.values() if e != END]), - [str(la) for la, e in self.ends.items() if e != END], + writer([e for e in self.ends.values()]), + [str(la) for la, e in self.ends.items()], ) ) if self.ends @@ -211,6 +211,7 @@ class Branch(NamedTuple): ] else: destinations = cast(Sequence[Union[Send, str]], result) + destinations = [d for d in destinations if d != END] if any(dest is None or dest == START for dest in destinations): raise ValueError("Branch did not return a valid destination") if any(p.node == END for p in destinations if isinstance(p, Send)): diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 9604152d4..c6426efe0 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -845,29 +845,26 @@ class CompiledStateGraph(CompiledGraph): def get_writes( packets: Sequence[Union[str, Send]], ) -> Sequence[Union[ChannelWriteEntry, Send]]: - if filtered := [p for p in packets if p != END]: - writes = [ - ( - ChannelWriteEntry(CHANNEL_BRANCH_TO.format(p), None) - if not isinstance(p, Send) - else p + writes = [ + ( + ChannelWriteEntry(CHANNEL_BRANCH_TO.format(p), None) + if not isinstance(p, Send) + else p + ) + for p in packets + ] + if branch.then and branch.then != END: + writes.append( + ChannelWriteEntry( + f"branch:{start}:{name}::then", + WaitForNames( + frozenset( + p.node if isinstance(p, Send) else p for p in packets + ) + ), ) - for p in filtered - ] - if branch.then and branch.then != END: - writes.append( - ChannelWriteEntry( - f"branch:{start}:{name}::then", - WaitForNames( - frozenset( - p.node if isinstance(p, Send) else p - for p in filtered - ) - ), - ) - ) - return writes - return [] + ) + return writes if with_reader: # get schema @@ -1067,13 +1064,9 @@ 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() - if k != END - ] + return [(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 if e != END] + return [(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/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index a32aaeded..518368f2b 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -110,6 +110,11 @@ def draw_graph( static_seen.add(w) # apply static writes if writes := ChannelWrite.get_static_writes(w): + # END writes are not written, but become edges directly + for w in writes: + if w[0] == END: + edges.add((task.name, w[0], True, w[2])) + writes = [t for t in writes if t[0] != END] conditionals.update( {(task.name, *t[:2]): t[2] for t in writes} ) @@ -175,31 +180,36 @@ def draw_graph( if START not in nodes: graph.add_node(None, START) for task in start_tasks.values(): - graph.add_edge(graph.nodes[START], graph.nodes[task.name]) + add_edge(graph, START, task.name) # add discovered edges for src, dest, is_conditional, label in sorted(edges): - graph.add_edge( - graph.nodes[src], - graph.nodes[dest], + add_edge( + graph, + src, + dest, data=label if label != dest else None, conditional=is_conditional, ) # add end edges - if step_sources: - end = graph.add_node(None, END) - termini = {d for _, d, _, _ in edges}.difference(s for s, _, _, _ in edges) - for src in sorted(termini.union(step_sources)): - graph.add_edge(graph.nodes[src], end, conditional=src not in termini) + termini = {d for _, d, _, _ in edges if d != END}.difference( + s for s, _, _, _ in edges + ) + if termini: + for src in sorted(termini): + add_edge(graph, src, END) + elif len(step_sources) == 1: + for src in sorted(step_sources): + add_edge(graph, src, END, conditional=True) # replace subgraphs for name, subgraph in subgraphs.items(): - subgraph.trim_first_node() - subgraph.trim_last_node() if ( len(subgraph.nodes) > 1 and name in graph.nodes and subgraph.first_node() and subgraph.last_node() ): + subgraph.trim_first_node() + subgraph.trim_last_node() # replace the node with the subgraph graph.nodes.pop(name) first, last = graph.extend(subgraph, prefix=name) @@ -210,3 +220,20 @@ def draw_graph( graph.edges[idx] = edge.copy(target=cast(Node, first).id) return graph + + +def add_edge( + graph: Graph, + source: str, + target: str, + *, + data: Optional[Any] = None, + conditional: bool = False, +) -> None: + """Add an edge to the graph.""" + for edge in graph.edges: + if edge.source == source and edge.target == target: + return + if target not in graph.nodes and target == END: + graph.add_node(None, END) + graph.add_edge(graph.nodes[source], graph.nodes[target], data, conditional) diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr index 35a8de7fc..4a9d05004 100644 --- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr +++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr @@ -86,6 +86,12 @@ "source": "__start__", "target": "agent" }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + }, { "source": "agent", "target": "tools", @@ -95,11 +101,6 @@ { "source": "tools", "target": "agent" - }, - { - "source": "agent", - "target": "__end__", - "conditional": true } ] } @@ -109,9 +110,9 @@ ''' graph TD; __start__ --> agent; + agent -.  exit  .-> __end__; agent -.  continue  .-> tools; tools --> agent; - agent -.-> __end__; ''' # --- @@ -130,9 +131,9 @@ __start__([

__start__

]):::first __end__([

__end__

]):::last __start__ --> agent; + agent -.  exit  .-> __end__; agent -.  continue  .-> tools; tools --> agent; - agent -.-> __end__; classDef default fill:#f2f0ff,line-height:1.2 classDef first fill-opacity:0 classDef last fill:#bfb6fc diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 816273093..d4815ce86 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -67,7 +67,8 @@ }, { "source": "left", - "target": "__end__" + "target": "__end__", + "conditional": true }, { "source": "right", @@ -82,7 +83,7 @@ graph TD; __start__ -.  go-left  .-> left; __start__ -.  go-right  .-> right; - left --> __end__; + left -.-> __end__; right --> __end__; ''' @@ -481,6 +482,17 @@ ''' # --- +# name: test_migration_graph + ''' + graph TD; + B -.  X  .-> C; + B -.  Y  .-> D; + D --> B; + __start__ --> B; + C --> __end__; + + ''' +# --- # name: test_multiple_sinks_subgraphs ''' --- @@ -532,12 +544,14 @@ --- graph TD; __start__([

__start__

]):::first - inner(inner) side(side) __end__([

__end__

]):::last - __start__ --> inner; - inner --> side; + __start__ --> inner_up; + inner_up --> side; side --> __end__; + subgraph inner + inner_up(up) + end classDef default fill:#f2f0ff,line-height:1.2 classDef first fill-opacity:0 classDef last fill:#bfb6fc @@ -550,8 +564,48 @@ dict({ 'conditional': True, 'source': '__start__', + 'target': 'tool_one', + }), + dict({ + 'conditional': True, + 'source': '__start__', + 'target': 'tool_three', + }), + dict({ + 'conditional': True, + 'source': '__start__', + 'target': 'tool_two:__start__', + }), + dict({ + 'source': 'tool_one', 'target': '__end__', }), + dict({ + 'source': 'tool_three', + 'target': '__end__', + }), + dict({ + 'source': 'tool_two:__end__', + 'target': '__end__', + }), + dict({ + 'conditional': True, + 'source': 'tool_two:__start__', + 'target': 'tool_two:tool_two_fast', + }), + dict({ + 'conditional': True, + 'source': 'tool_two:__start__', + 'target': 'tool_two:tool_two_slow', + }), + dict({ + 'source': 'tool_two:tool_two_fast', + 'target': 'tool_two:__end__', + }), + dict({ + 'source': 'tool_two:tool_two_slow', + 'target': 'tool_two:__end__', + }), ]), 'nodes': list([ dict({ @@ -580,19 +634,6 @@ 'id': 'tool_one', 'type': 'runnable', }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'graph', - 'state', - 'CompiledStateGraph', - ]), - 'name': 'tool_two', - }), - 'id': 'tool_two', - 'type': 'runnable', - }), dict({ 'data': dict({ 'id': list([ @@ -609,6 +650,48 @@ dict({ 'id': '__end__', }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnablePassthrough', + ]), + 'name': 'tool_two:__start__', + }), + 'id': 'tool_two:__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tool_two:tool_two_slow', + }), + 'id': 'tool_two:tool_two_slow', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tool_two:tool_two_fast', + }), + 'id': 'tool_two:tool_two_fast', + 'type': 'runnable', + }), + dict({ + 'id': 'tool_two:__end__', + }), ]), }) # --- @@ -620,12 +703,26 @@ curve: linear --- graph TD; - __start__(

__start__

) + __start__([

__start__

]):::first tool_one(tool_one) - tool_two(tool_two) tool_three(tool_three) - __end__(

__end__

) - __start__ -.-> __end__; + __end__([

__end__

]):::last + __start__ -.-> tool_one; + __start__ -.-> tool_three; + __start__ -.-> tool_two___start__; + tool_one --> __end__; + tool_three --> __end__; + tool_two___end__ --> __end__; + subgraph tool_two + tool_two___start__(

__start__

) + tool_two_tool_two_slow(tool_two_slow) + tool_two_tool_two_fast(tool_two_fast) + tool_two___end__(

__end__

) + tool_two___start__ -.-> tool_two_tool_two_fast; + tool_two___start__ -.-> tool_two_tool_two_slow; + tool_two_tool_two_fast --> tool_two___end__; + tool_two_tool_two_slow --> tool_two___end__; + end classDef default fill:#f2f0ff,line-height:1.2 classDef first fill-opacity:0 classDef last fill:#bfb6fc @@ -642,9 +739,6 @@ Researcher -.  call_tool  .-> Call_Tool; Researcher -.  continue  .-> Chart_Generator; __start__ --> Researcher; - Call_Tool -.-> __end__; - Chart_Generator -.-> __end__; - Researcher -.-> __end__; Researcher -.  redo  .-> Researcher; ''' @@ -1014,7 +1108,6 @@ dict({ 'id': 'conduct_interview:__end__', }), - ]), }) # --- diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 679a58115..d8259a1e2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3491,6 +3491,7 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None: tool_two_graph.add_node("tool_two_fast", logic) tool_two_graph.set_conditional_entry_point( lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + ["tool_two_slow", "tool_two_fast"], then=END, ) tool_two = tool_two_graph.compile() @@ -3499,7 +3500,9 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None: graph.add_node("tool_one", logic) graph.add_node("tool_two", tool_two) graph.add_node("tool_three", logic) - graph.set_conditional_entry_point(lambda s: "tool_one", then=END) + graph.set_conditional_entry_point( + lambda s: "tool_one", ["tool_one", "tool_two", "tool_three"], then=END + ) app = graph.compile() assert app.get_graph(xray=True).to_json() == snapshot @@ -8120,3 +8123,44 @@ def test_batch_update_as_input( ] assert new_history == history + + +def test_migration_graph(snapshot: SnapshotAssertion) -> None: + from pydantic import BaseModel + + class DummyState(BaseModel): + pass_count: int = 0 + + def increment_pass_count(state: DummyState): + state.pass_count += 1 + return state + + def route_b(state: DummyState): + if state.pass_count == 0: + return "X" + else: + return "Y" + + migration_graph = StateGraph(DummyState) + + migration_graph.add_node("B", increment_pass_count) + migration_graph.add_node("C", increment_pass_count) + migration_graph.add_node("D", increment_pass_count) + + migration_graph.add_edge(START, "B") + + migration_graph.add_conditional_edges( + "B", + route_b, + { + "X": "C", + "Y": "D", + }, + ) + + migration_graph.add_edge("D", "B") + migration_graph.add_edge("C", END) + + app = migration_graph.compile() + + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot