diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 2e6042bbc..23c86d83a 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -390,7 +390,6 @@ class CompiledGraph(Pregel): config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False, - add_condition_nodes: bool = True, ) -> DrawableGraph: """Returns a drawable representation of the computation graph.""" graph = DrawableGraph() @@ -406,7 +405,6 @@ class CompiledGraph(Pregel): subgraph = ( node.get_graph( config=config, - add_condition_nodes=add_condition_nodes, xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray, ) if isinstance(node, CompiledGraph) @@ -428,46 +426,26 @@ class CompiledGraph(Pregel): end_nodes[key] = n for start, end in sorted(self.graph._all_edges): graph.add_edge(start_nodes[start], end_nodes[end]) - branches_by_name = Counter( - name for _, branches in self.graph.branches.items() for name in branches - ) for start, branches in self.graph.branches.items(): default_ends = { **{k: k for k in self.graph.nodes if k != start}, END: END, } - for name, branch in branches.items(): + for _, branch in branches.items(): if branch.ends is not None: ends = branch.ends elif branch.then is not None: ends = {k: k for k in default_ends if k not in (END, branch.then)} else: ends = default_ends - - if add_condition_nodes is True: - cond = graph.add_node( - branch.path, - f"{start}_{name}" if branches_by_name[name] > 1 else name, + for label, end in ends.items(): + graph.add_edge( + start_nodes[start], + end_nodes[end], + label if label != end else None, + conditional=True, ) - graph.add_edge(start_nodes[start], cond) - for label, end in ends.items(): - graph.add_edge( - cond, - end_nodes[end], - label if label != end else None, - conditional=True, - ) - if branch.then is not None: - graph.add_edge(start_nodes[end], end_nodes[branch.then]) - else: - for label, end in ends.items(): - graph.add_edge( - start_nodes[start], - end_nodes[end], - label if label != end else None, - conditional=True, - ) - if branch.then is not None: - graph.add_edge(start_nodes[end], end_nodes[branch.then]) + if branch.then is not None: + graph.add_edge(start_nodes[end], end_nodes[branch.then]) return graph diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index adc8c6368..41542464d 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1,42 +1,5 @@ # serializer version: 1 # name: test_branch_then[end_of_run] - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------+ - | prepare | - +---------+ - * - * - * - +-----------+ - | condition | - +-----------+ - .. .. - .. .. - .. .. - +---------------+ +---------------+ - | tool_two_slow | | tool_two_fast | - +---------------+ +---------------+ - ** ** - ** ** - ** ** - +--------+ - | finish | - +--------+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_branch_then[end_of_run].1 ''' +-----------+ | __start__ | @@ -67,31 +30,7 @@ +---------+ ''' # --- -# name: test_branch_then[end_of_run].2 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - prepare([prepare]):::otherclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - finish([finish]):::otherclass; - condition([condition]):::otherclass; - __start__ --> prepare; - finish --> __end__; - prepare --> condition; - condition -.-> tool_two_slow; - tool_two_slow --> finish; - condition -.-> tool_two_fast; - tool_two_fast --> finish; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_branch_then[end_of_run].3 +# name: test_branch_then[end_of_run].1 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; @@ -114,43 +53,6 @@ ''' # --- # name: test_branch_then[end_of_step] - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------+ - | prepare | - +---------+ - * - * - * - +-----------+ - | condition | - +-----------+ - .. .. - .. .. - .. .. - +---------------+ +---------------+ - | tool_two_slow | | tool_two_fast | - +---------------+ +---------------+ - ** ** - ** ** - ** ** - +--------+ - | finish | - +--------+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_branch_then[end_of_step].1 ''' +-----------+ | __start__ | @@ -181,31 +83,7 @@ +---------+ ''' # --- -# name: test_branch_then[end_of_step].2 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - prepare([prepare]):::otherclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - finish([finish]):::otherclass; - condition([condition]):::otherclass; - __start__ --> prepare; - finish --> __end__; - prepare --> condition; - condition -.-> tool_two_slow; - tool_two_slow --> finish; - condition -.-> tool_two_fast; - tool_two_fast --> finish; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_branch_then[end_of_step].3 +# name: test_branch_then[end_of_step].1 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; @@ -270,30 +148,6 @@ ], "name": "right" } - }, - { - "id": "should_start", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_start" - } - }, - { - "id": "condition", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "RunnableCallable" - } } ], "edges": [ @@ -303,26 +157,18 @@ }, { "source": "__start__", - "target": "should_start" - }, - { - "source": "should_start", "target": "left", "data": "go-left", "conditional": true }, { - "source": "should_start", + "source": "__start__", "target": "right", "data": "go-right", "conditional": true }, { "source": "left", - "target": "condition" - }, - { - "source": "condition", "target": "__end__", "conditional": true } @@ -332,33 +178,21 @@ # --- # name: test_conditional_entrypoint_graph.3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +--------------+ - | should_start | - +--------------+ - . .. - .. . - . .. - +------+ . - | left | . - +------+ . - * . - * . - * . - +-----------+ +-------+ - | condition | | right | - +-----------+ +-------+ - . * - .. ** - . * - +---------+ - | __end__ | - +---------+ + +-----------+ + | __start__ | + +-----------+ + .. . + . .. + . . + +-------+ +------+ + | right | | left | + +-------+ +------+ + ** . + * .. + * . + +---------+ + | __end__ | + +---------+ ''' # --- # name: test_conditional_entrypoint_graph_state @@ -404,30 +238,6 @@ ], "name": "right" } - }, - { - "id": "should_start", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_start" - } - }, - { - "id": "condition", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "RunnableCallable" - } } ], "edges": [ @@ -437,26 +247,18 @@ }, { "source": "__start__", - "target": "should_start" - }, - { - "source": "should_start", "target": "left", "data": "go-left", "conditional": true }, { - "source": "should_start", + "source": "__start__", "target": "right", "data": "go-right", "conditional": true }, { "source": "left", - "target": "condition" - }, - { - "source": "condition", "target": "__end__", "conditional": true } @@ -466,33 +268,21 @@ # --- # name: test_conditional_entrypoint_graph_state.3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +--------------+ - | should_start | - +--------------+ - . .. - .. . - . .. - +------+ . - | left | . - +------+ . - * . - * . - * . - +-----------+ +-------+ - | condition | | right | - +-----------+ +-------+ - . * - .. ** - . * - +---------+ - | __end__ | - +---------+ + +-----------+ + | __start__ | + +-----------+ + .. . + . .. + . . + +-------+ +------+ + | right | | left | + +-------+ +------+ + ** . + * .. + * . + +---------+ + | __end__ | + +---------+ ''' # --- # name: test_conditional_graph[end_of_run] @@ -533,18 +323,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -558,16 +336,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "tools", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -578,27 +352,21 @@ # --- # name: test_conditional_graph[end_of_run].1 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * . + ** .. + * . + +-------+ +---------+ + | tools | | __end__ | + +-------+ +---------+ ''' # --- # name: test_conditional_graph[end_of_run].2 @@ -698,18 +466,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -747,16 +503,12 @@ }, { "source": 3, - "target": "should_continue" - }, - { - "source": "should_continue", "target": "tools", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": 3, "target": "__end__", "data": "exit", "conditional": true @@ -767,51 +519,45 @@ # --- # name: test_conditional_graph[end_of_run].4 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +------------------------------+ - | ParallelInput | - +------------------------------+ - ***** * ***** - ****** * **** - *** * ***** - +----------------+ * *** - | PromptTemplate | * * - +----------------+ * * - * * * - * * * - * * * - +----------------------+ * * - | FakeStreamingListLLM | * * - +----------------------+ * * - * * * - * * * - * * * - +----------------------+ +-------------+ * - | Lambda(agent_parser) | | Passthrough | * - +----------------------+ +-------------+ * - *** *** * - ** ** * - ** ** * - +-------------------------------+ * - | ParallelOutput | * - +-------------------------------+ * - * * - * * - * * - +-----------------+ * - | should_continue |... * - +-----------------+ ......... * - . ........ * - . ......... * - . ..... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +------------------------------+ + | ParallelInput | + +------------------------------+ + ***** * ***** + ****** * **** + *** * ***** + +----------------+ * *** + | PromptTemplate | * * + +----------------+ * * + * * * + * * * + * * * + +----------------------+ * * + | FakeStreamingListLLM | * * + +----------------------+ * * + * * * + * * * + * * * + +----------------------+ +-------------+ * + | Lambda(agent_parser) | | Passthrough | * + +----------------------+ +-------------+ * + *** *** * + ** ** * + ** ** * + +-------------------------------+ * + | ParallelOutput | * + +-------------------------------+..... * + . ......... * + . ........ * + . ..... * + +---------+ +-------+ + | __end__ | | tools | + +---------+ +-------+ ''' # --- # name: test_conditional_graph[end_of_run].5 @@ -868,18 +614,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -893,16 +627,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "tools", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -913,27 +643,21 @@ # --- # name: test_conditional_graph[end_of_step].1 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * . + ** .. + * . + +-------+ +---------+ + | tools | | __end__ | + +-------+ +---------+ ''' # --- # name: test_conditional_graph[end_of_step].2 @@ -1033,18 +757,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -1082,16 +794,12 @@ }, { "source": 3, - "target": "should_continue" - }, - { - "source": "should_continue", "target": "tools", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": 3, "target": "__end__", "data": "exit", "conditional": true @@ -1102,51 +810,45 @@ # --- # name: test_conditional_graph[end_of_step].4 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +------------------------------+ - | ParallelInput | - +------------------------------+ - ***** * ***** - ****** * **** - *** * ***** - +----------------+ * *** - | PromptTemplate | * * - +----------------+ * * - * * * - * * * - * * * - +----------------------+ * * - | FakeStreamingListLLM | * * - +----------------------+ * * - * * * - * * * - * * * - +----------------------+ +-------------+ * - | Lambda(agent_parser) | | Passthrough | * - +----------------------+ +-------------+ * - *** *** * - ** ** * - ** ** * - +-------------------------------+ * - | ParallelOutput | * - +-------------------------------+ * - * * - * * - * * - +-----------------+ * - | should_continue |... * - +-----------------+ ......... * - . ........ * - . ......... * - . ..... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +------------------------------+ + | ParallelInput | + +------------------------------+ + ***** * ***** + ****** * **** + *** * ***** + +----------------+ * *** + | PromptTemplate | * * + +----------------+ * * + * * * + * * * + * * * + +----------------------+ * * + | FakeStreamingListLLM | * * + +----------------------+ * * + * * * + * * * + * * * + +----------------------+ +-------------+ * + | Lambda(agent_parser) | | Passthrough | * + +----------------------+ +-------------+ * + *** *** * + ** ** * + ** ** * + +-------------------------------+ * + | ParallelOutput | * + +-------------------------------+..... * + . ......... * + . ........ * + . ..... * + +---------+ +-------+ + | __end__ | | tools | + +---------+ +-------+ ''' # --- # name: test_conditional_graph[end_of_step].5 @@ -1209,18 +911,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -1234,16 +924,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "tools", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -1254,27 +940,21 @@ # --- # name: test_conditional_graph_state[end_of_run].3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + . . + .. .. + . . + +-------+ +---------+ + | tools | | __end__ | + +-------+ +---------+ ''' # --- # name: test_conditional_graph_state[end_of_step] @@ -1321,18 +1001,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -1346,16 +1014,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "tools", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -1366,27 +1030,21 @@ # --- # name: test_conditional_graph_state[end_of_step].3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * . + ** .. + * . + +-------+ +---------+ + | tools | | __end__ | + +-------+ +---------+ ''' # --- # name: test_in_one_fan_out_state_graph_waiting_edge[end_of_run] @@ -1474,12 +1132,12 @@ +---------------+ | rewrite_query | +---------------+ - *** *** - * * - ** ** - +--------------+ +---------+ - | analyzer_one | | decider | - +--------------+ +---------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . * . * . * . @@ -1511,12 +1169,12 @@ +---------------+ | rewrite_query | +---------------+ - *** *** - * * - ** ** - +--------------+ +---------+ - | analyzer_one | | decider | - +--------------+ +---------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . * . * . * . @@ -1539,76 +1197,76 @@ # --- # name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_run] ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------+ - | rewrite_query | - +---------------+ - *** *** - * * - ** ** - +--------------+ +--------------------+ - | analyzer_one | | rewrite_query_then | - +--------------+ +--------------------+ - * . - * . - * . - +---------------+ +---------------+ - | retriever_one | | retriever_two | - +---------------+ +---------------+ - *** *** - * * - ** ** - +----+ - | qa | - +----+ - * - * - * - +---------+ - | __end__ | - +---------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ ''' # --- # name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_step] ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------+ - | rewrite_query | - +---------------+ - *** *** - * * - ** ** - +--------------+ +--------------------+ - | analyzer_one | | rewrite_query_then | - +--------------+ +--------------------+ - * . - * . - * . - +---------------+ +---------------+ - | retriever_one | | retriever_two | - +---------------+ +---------------+ - *** *** - * * - ** ** - +----+ - | qa | - +----+ - * - * - * - +---------+ - | __end__ | - +---------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ ''' # --- # name: test_message_graph[end_of_run] @@ -1655,18 +1313,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -1680,16 +1326,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "action", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "end", "conditional": true @@ -1700,27 +1342,21 @@ # --- # name: test_message_graph[end_of_run].3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +--------+ - | __end__ | | action | - +---------+ +--------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * .. + ** .. + * . + +--------+ +---------+ + | action | | __end__ | + +--------+ +---------+ ''' # --- # name: test_message_graph[end_of_step] @@ -1767,18 +1403,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -1792,16 +1416,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "action", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "end", "conditional": true @@ -1812,27 +1432,21 @@ # --- # name: test_message_graph[end_of_step].3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +--------+ - | __end__ | | action | - +---------+ +--------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * .. + ** .. + * . + +--------+ +---------+ + | action | | __end__ | + +--------+ +---------+ ''' # --- # name: test_nested_graph @@ -1882,7 +1496,7 @@ 'edges': list([ dict({ 'conditional': True, - 'source': 'tool_two:condition', + 'source': 'tool_two:__start__', 'target': 'tool_two:tool_two_slow', }), dict({ @@ -1891,20 +1505,16 @@ }), dict({ 'conditional': True, - 'source': 'tool_two:condition', + 'source': 'tool_two:__start__', 'target': 'tool_two:tool_two_fast', }), dict({ 'source': 'tool_two:tool_two_fast', 'target': 'tool_two:__end__', }), - dict({ - 'source': '__start__', - 'target': 'condition', - }), dict({ 'conditional': True, - 'source': 'condition', + 'source': '__start__', 'target': 'tool_one', }), dict({ @@ -1913,8 +1523,8 @@ }), dict({ 'conditional': True, - 'source': 'condition', - 'target': 'tool_two:condition', + 'source': '__start__', + 'target': 'tool_two:__start__', }), dict({ 'source': 'tool_two:__end__', @@ -1922,7 +1532,7 @@ }), dict({ 'conditional': True, - 'source': 'condition', + 'source': '__start__', 'target': 'tool_three', }), dict({ @@ -1953,6 +1563,11 @@ 'id': 'tool_one', 'type': 'runnable', }), + dict({ + 'data': 'tool_two:__start__', + 'id': 'tool_two:__start__', + 'type': 'schema', + }), dict({ 'data': 'tool_two:__end__', 'id': 'tool_two:__end__', @@ -1982,18 +1597,6 @@ 'id': 'tool_two:tool_two_fast', 'type': 'runnable', }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'RunnableCallable', - ]), - 'name': 'RunnableCallable', - }), - 'id': 'tool_two:condition', - 'type': 'runnable', - }), dict({ 'data': dict({ 'id': list([ @@ -2006,18 +1609,6 @@ 'id': 'tool_three', 'type': 'runnable', }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'RunnableCallable', - ]), - 'name': 'RunnableCallable', - }), - 'id': 'condition', - 'type': 'runnable', - }), ]), }) # --- @@ -2025,12 +1616,6 @@ ''' +-----------+ | __start__ | - +-----------+ - * - * - * - +-----------+ - | condition | ..+-----------+... ..... . ..... ... . ... @@ -2047,38 +1632,6 @@ ''' # --- # name: test_nested_graph_xray.2 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - tool_one([tool_one]):::otherclass; - tool_two___end__([__end__]):::otherclass; - tool_two_tool_two_slow([tool_two_slow]):::otherclass; - tool_two_tool_two_fast([tool_two_fast]):::otherclass; - tool_two_condition([condition]):::otherclass; - tool_three([tool_three]):::otherclass; - condition([condition]):::otherclass; - subgraph tool_two - tool_two_condition -.-> tool_two_tool_two_slow; - tool_two_tool_two_slow --> tool_two___end__; - tool_two_condition -.-> tool_two_tool_two_fast; - tool_two_tool_two_fast --> tool_two___end__; - end - __start__ --> condition; - condition -.-> tool_one; - tool_one --> __end__; - condition -.-> tool_two_condition; - tool_two___end__ --> __end__; - condition -.-> tool_three; - tool_three --> __end__; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_nested_graph_xray.3 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; @@ -2153,18 +1706,6 @@ ], "name": "call_tool" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -2178,16 +1719,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "action", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "end", "conditional": true @@ -2198,27 +1735,21 @@ # --- # name: test_prebuilt_chat.3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +--------+ - | __end__ | | action | - +---------+ +--------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * .. + ** .. + * . + +--------+ +---------+ + | action | | __end__ | + +--------+ +---------+ ''' # --- # name: test_prebuilt_tool_chat @@ -2266,18 +1797,6 @@ ], "name": "tools" } - }, - { - "id": "should_continue", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "should_continue" - } } ], "edges": [ @@ -2291,16 +1810,12 @@ }, { "source": "agent", - "target": "should_continue" - }, - { - "source": "should_continue", "target": "action", "data": "continue", "conditional": true }, { - "source": "should_continue", + "source": "agent", "target": "__end__", "data": "end", "conditional": true @@ -2311,44 +1826,35 @@ # --- # name: test_prebuilt_tool_chat.3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+* - *** *** - * * - ** *** - +-----------------+ * - | should_continue | * - +-----------------+. * - . ..... * - . ... * - . ... * - +---------+ +--------+ - | __end__ | | action | - +---------+ +--------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + . .. + .. .. + . . + +--------+ +---------+ + | action | | __end__ | + +--------+ +---------+ ''' # --- # name: test_repeat_condition ''' graph TD; __start__ --> Researcher; - Researcher --> Researcher_router; - Researcher_router -. continue .-> Chart_Generator; - Researcher_router -. call_tool .-> Call_Tool; - Researcher_router -. end .-> __end__; - Chart_Generator --> Chart_Generator_router; - Chart_Generator_router -. continue .-> Researcher; - Chart_Generator_router -. call_tool .-> Call_Tool; - Chart_Generator_router -. end .-> __end__; - Call_Tool --> condition; - condition -.-> Researcher; - condition -.-> Chart_Generator; + Researcher -. continue .-> Chart_Generator; + Researcher -. call_tool .-> Call_Tool; + Researcher -. end .-> __end__; + Chart_Generator -. continue .-> Researcher; + Chart_Generator -. call_tool .-> Call_Tool; + Chart_Generator -. end .-> __end__; + Call_Tool -.-> Researcher; + Call_Tool -.-> Chart_Generator; ''' # --- @@ -2387,12 +1893,6 @@ ''' +-----------+ | __start__ | - +-----------+ - * - * - * - +-----------+ - | condition | +-----------+ .. .. .. .. @@ -2409,45 +1909,6 @@ ''' # --- # name: test_start_branch_then[end_of_run].1 - ''' - +-----------+ - | __start__ | - +-----------+ - .. .. - .. .. - .. .. - +---------------+ +---------------+ - | tool_two_slow | | tool_two_fast | - +---------------+ +---------------+ - ** ** - ** ** - ** ** - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_start_branch_then[end_of_run].2 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - condition([condition]):::otherclass; - __start__ --> condition; - condition -.-> tool_two_slow; - tool_two_slow --> __end__; - condition -.-> tool_two_fast; - tool_two_fast --> __end__; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_start_branch_then[end_of_run].3 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; @@ -2469,12 +1930,6 @@ ''' +-----------+ | __start__ | - +-----------+ - * - * - * - +-----------+ - | condition | +-----------+ .. .. .. .. @@ -2491,45 +1946,6 @@ ''' # --- # name: test_start_branch_then[end_of_step].1 - ''' - +-----------+ - | __start__ | - +-----------+ - .. .. - .. .. - .. .. - +---------------+ +---------------+ - | tool_two_slow | | tool_two_fast | - +---------------+ +---------------+ - ** ** - ** ** - ** ** - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_start_branch_then[end_of_step].2 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - condition([condition]):::otherclass; - __start__ --> condition; - condition -.-> tool_two_slow; - tool_two_slow --> __end__; - condition -.-> tool_two_fast; - tool_two_fast --> __end__; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_start_branch_then[end_of_step].3 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; diff --git a/tests/__snapshots__/test_pregel_async.ambr b/tests/__snapshots__/test_pregel_async.ambr index 709726945..2d68a82a5 100644 --- a/tests/__snapshots__/test_pregel_async.ambr +++ b/tests/__snapshots__/test_pregel_async.ambr @@ -10,12 +10,12 @@ +---------------+ | rewrite_query | +---------------+ - *** *** - * * - ** ** - +--------------+ +---------+ - | analyzer_one | | decider | - +--------------+ +---------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . * . * . * . @@ -47,12 +47,12 @@ +---------------+ | rewrite_query | +---------------+ - *** *** - * * - ** ** - +--------------+ +---------+ - | analyzer_one | | decider | - +--------------+ +---------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . * . * . * . @@ -84,12 +84,12 @@ +---------------+ | rewrite_query | +---------------+ - *** *** - * * - ** ** - +--------------+ +-----------+ - | analyzer_one | | condition | - +--------------+ +-----------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . * . * . * . @@ -121,12 +121,12 @@ +---------------+ | rewrite_query | +---------------+ - *** *** - * * - ** ** - +--------------+ +-----------+ - | analyzer_one | | condition | - +--------------+ +-----------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . * . * . * . diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 8e8222697..900683242 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -883,18 +883,10 @@ def test_conditional_graph( assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot assert app.get_graph().draw_ascii() == snapshot - assert ( - app.get_graph(add_condition_nodes=False).draw_mermaid(with_styles=False) - == snapshot - ) + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot assert app.get_graph(xray=True).draw_ascii() == snapshot - assert ( - app.get_graph(xray=True, add_condition_nodes=False).draw_mermaid( - with_styles=False - ) - == snapshot - ) + assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", @@ -3032,9 +3024,7 @@ def test_start_branch_then( ) tool_two = tool_two_graph.compile() assert tool_two.get_graph().draw_ascii() == snapshot - assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot assert tool_two.get_graph().draw_mermaid() == snapshot - assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { "my_key": "value slow", @@ -3138,9 +3128,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) tool_two = tool_two_graph.compile() assert tool_two.get_graph().draw_ascii() == snapshot - assert tool_two.get_graph(add_condition_nodes=False).draw_ascii() == snapshot assert tool_two.get_graph().draw_mermaid() == snapshot - assert tool_two.get_graph(add_condition_nodes=False).draw_mermaid() == snapshot assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == { "my_key": "value prepared slow finished", @@ -3877,9 +3865,6 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None: assert app.get_graph(xray=True).to_json() == snapshot assert app.get_graph().draw_ascii() == snapshot assert app.get_graph(xray=True).draw_mermaid() == snapshot - assert ( - app.get_graph(xray=True, add_condition_nodes=False).draw_mermaid() == snapshot - ) def test_nested_graph(snapshot: SnapshotAssertion) -> None: