From 3edf7e4d7372f0a426cc1f3c5fd24bd3d176a0d3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Apr 2024 10:18:56 -0700 Subject: [PATCH] Fix --- tests/__snapshots__/test_pregel.ambr | 144 +++++++-------------------- tests/test_pregel.py | 12 ++- 2 files changed, 46 insertions(+), 110 deletions(-) diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index b68acaf76..eca8025a3 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -387,21 +387,12 @@ # --- # name: test_conditional_graph[end_of_run].2 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+ - * . - ** .. - * . - +-------+ +---------+ - | tools | | __end__ | - +-------+ +---------+ + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + ''' # --- # name: test_conditional_graph[end_of_run].3 @@ -611,45 +602,18 @@ # --- # name: test_conditional_graph[end_of_run].5 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +------------------------------+ - | ParallelInput | - +------------------------------+ - ***** * ***** - ****** * **** - *** * ***** - +----------------+ * *** - | PromptTemplate | * * - +----------------+ * * - * * * - * * * - * * * - +----------------------+ * * - | FakeStreamingListLLM | * * - +----------------------+ * * - * * * - * * * - * * * - +----------------------+ +-------------+ * - | Lambda(agent_parser) | | Passthrough | * - +----------------------+ +-------------+ * - *** *** * - ** ** * - ** ** * - +-------------------------------+ * - | ParallelOutput | * - +-------------------------------+..... * - . ......... * - . ........ * - . ..... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> Lambda_agent_parser_; + Parallel_agent_outcome_Input --> PromptTemplate; + Lambda_agent_parser_ --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -. continue .-> tools; + Parallel_agent_outcome_Output -. exit .-> __end__; + ''' # --- # name: test_conditional_graph[end_of_step] @@ -762,21 +726,12 @@ # --- # name: test_conditional_graph[end_of_step].2 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | agent | - +-------+ - . . - .. .. - . . - +-------+ +---------+ - | tools | | __end__ | - +-------+ +---------+ + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + ''' # --- # name: test_conditional_graph[end_of_step].3 @@ -986,45 +941,18 @@ # --- # name: test_conditional_graph[end_of_step].5 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +------------------------------+ - | ParallelInput | - +------------------------------+ - ***** * ***** - ****** * **** - *** * ***** - +----------------+ * *** - | PromptTemplate | * * - +----------------+ * * - * * * - * * * - * * * - +----------------------+ * * - | FakeStreamingListLLM | * * - +----------------------+ * * - * * * - * * * - * * * - +----------------------+ +-------------+ * - | Lambda(agent_parser) | | Passthrough | * - +----------------------+ +-------------+ * - *** *** * - ** ** * - ** ** * - +-------------------------------+ * - | ParallelOutput | * - +-------------------------------+..... * - . ......... * - . ........ * - . ..... * - +---------+ +-------+ - | __end__ | | tools | - +---------+ +-------+ + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> Lambda_agent_parser_; + Parallel_agent_outcome_Input --> PromptTemplate; + Lambda_agent_parser_ --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -. continue .-> tools; + Parallel_agent_outcome_Output -. exit .-> __end__; + ''' # --- # name: test_conditional_graph_state[end_of_run] diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 3b0c40c4f..25accdf50 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -796,10 +796,18 @@ 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_ascii() == snapshot + assert ( + app.get_graph(add_condition_nodes=False).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_ascii() == snapshot + assert ( + app.get_graph(xray=True, add_condition_nodes=False).draw_mermaid( + with_styles=False + ) + == snapshot + ) assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf",