Update tests

This commit is contained in:
Nuno Campos
2024-04-29 17:46:42 -07:00
parent ca2c93e0a1
commit 7554309abf
3 changed files with 149 additions and 624 deletions
File diff suppressed because one or more lines are too long
+134 -605
View File
@@ -1,33 +1,14 @@
# serializer version: 1
# name: test_branch_then[end_of_run]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------+
| prepare |
+---------+.
.. ..
.. ..
.. ..
+---------------+ +---------------+
| tool_two_slow | | tool_two_fast |
+---------------+ +---------------+
** **
** **
** **
+--------+
| finish |
+--------+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> prepare;
finish --> __end__;
prepare -.-> tool_two_slow;
tool_two_slow --> finish;
prepare -.-> tool_two_fast;
tool_two_fast --> finish;
'''
# ---
# name: test_branch_then[end_of_run].1
@@ -54,33 +35,14 @@
# ---
# name: test_branch_then[end_of_step]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------+
| prepare |
+---------+.
.. ..
.. ..
.. ..
+---------------+ +---------------+
| tool_two_slow | | tool_two_fast |
+---------------+ +---------------+
** **
** **
** **
+--------+
| finish |
+--------+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> prepare;
finish --> __end__;
prepare -.-> tool_two_slow;
tool_two_slow --> finish;
prepare -.-> tool_two_fast;
tool_two_fast --> finish;
'''
# ---
# name: test_branch_then[end_of_step].1
@@ -178,21 +140,12 @@
# ---
# name: test_conditional_entrypoint_graph.3
'''
+-----------+
| __start__ |
+-----------+
.. .
. ..
. .
+-------+ +------+
| right | | left |
+-------+ +------+
** .
* ..
* .
+---------+
| __end__ |
+---------+
graph TD;
right --> __end__;
__start__ -. go-left .-> left;
__start__ -. go-right .-> right;
left -.-> __end__;
'''
# ---
# name: test_conditional_entrypoint_graph_state
@@ -268,21 +221,12 @@
# ---
# name: test_conditional_entrypoint_graph_state.3
'''
+-----------+
| __start__ |
+-----------+
.. .
. ..
. .
+-------+ +------+
| right | | left |
+-------+ +------+
** .
* ..
* .
+---------+
| __end__ |
+---------+
graph TD;
right --> __end__;
__start__ -. go-left .-> left;
__start__ -. go-right .-> right;
left -.-> __end__;
'''
# ---
# name: test_conditional_graph[end_of_run]
@@ -351,25 +295,6 @@
'''
# ---
# name: test_conditional_graph[end_of_run].1
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* .
** ..
* .
+-------+ +---------+
| tools | | __end__ |
+-------+ +---------+
'''
# ---
# name: test_conditional_graph[end_of_run].2
'''
graph TD;
__start__ --> agent;
@@ -379,7 +304,7 @@
'''
# ---
# name: test_conditional_graph[end_of_run].3
# name: test_conditional_graph[end_of_run].2
'''
{
"nodes": [
@@ -517,50 +442,7 @@
}
'''
# ---
# name: test_conditional_graph[end_of_run].4
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+------------------------------+
| Parallel<agent_outcome>Input |
+------------------------------+
***** * *****
****** * ****
*** * *****
+----------------+ * ***
| PromptTemplate | * *
+----------------+ * *
* * *
* * *
* * *
+----------------------+ * *
| FakeStreamingListLLM | * *
+----------------------+ * *
* * *
* * *
* * *
+----------------------+ +-------------+ *
| Lambda(agent_parser) | | Passthrough | *
+----------------------+ +-------------+ *
*** *** *
** ** *
** ** *
+-------------------------------+ *
| Parallel<agent_outcome>Output | *
+-------------------------------+..... *
. ......... *
. ........ *
. ..... *
+---------+ +-------+
| __end__ | | tools |
+---------+ +-------+
'''
# ---
# name: test_conditional_graph[end_of_run].5
# name: test_conditional_graph[end_of_run].3
'''
graph TD;
PromptTemplate --> FakeStreamingListLLM;
@@ -642,25 +524,6 @@
'''
# ---
# name: test_conditional_graph[end_of_step].1
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* .
** ..
* .
+-------+ +---------+
| tools | | __end__ |
+-------+ +---------+
'''
# ---
# name: test_conditional_graph[end_of_step].2
'''
graph TD;
__start__ --> agent;
@@ -670,7 +533,7 @@
'''
# ---
# name: test_conditional_graph[end_of_step].3
# name: test_conditional_graph[end_of_step].2
'''
{
"nodes": [
@@ -808,50 +671,7 @@
}
'''
# ---
# name: test_conditional_graph[end_of_step].4
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+------------------------------+
| Parallel<agent_outcome>Input |
+------------------------------+
***** * *****
****** * ****
*** * *****
+----------------+ * ***
| PromptTemplate | * *
+----------------+ * *
* * *
* * *
* * *
+----------------------+ * *
| FakeStreamingListLLM | * *
+----------------------+ * *
* * *
* * *
* * *
+----------------------+ +-------------+ *
| Lambda(agent_parser) | | Passthrough | *
+----------------------+ +-------------+ *
*** *** *
** ** *
** ** *
+-------------------------------+ *
| Parallel<agent_outcome>Output | *
+-------------------------------+..... *
. ......... *
. ........ *
. ..... *
+---------+ +-------+
| __end__ | | tools |
+---------+ +-------+
'''
# ---
# name: test_conditional_graph[end_of_step].5
# name: test_conditional_graph[end_of_step].3
'''
graph TD;
PromptTemplate --> FakeStreamingListLLM;
@@ -940,21 +760,12 @@
# ---
# name: test_conditional_graph_state[end_of_run].3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
. .
.. ..
. .
+-------+ +---------+
| tools | | __end__ |
+-------+ +---------+
graph TD;
__start__ --> agent;
tools --> agent;
agent -. continue .-> tools;
agent -. exit .-> __end__;
'''
# ---
# name: test_conditional_graph_state[end_of_step]
@@ -1030,243 +841,90 @@
# ---
# name: test_conditional_graph_state[end_of_step].3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* .
** ..
* .
+-------+ +---------+
| tools | | __end__ |
+-------+ +---------+
graph TD;
__start__ --> agent;
tools --> agent;
agent -. continue .-> tools;
agent -. exit .-> __end__;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge[end_of_run]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** ***
+--------------+ *
| analyzer_one | *
+--------------+ *
* *
* *
* *
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query --> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge[end_of_step]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ***
* *
** ***
+--------------+ *
| analyzer_one | *
+--------------+ *
* *
* *
* *
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query --> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_run]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_step]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_run]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_step]
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+---------------+
| rewrite_query |
+---------------+
*** ...
* .
** ...
+--------------+ .
| analyzer_one | .
+--------------+ .
* .
* .
* .
+---------------+ +---------------+
| retriever_one | | retriever_two |
+---------------+ +---------------+
*** ***
* *
** **
+----+
| qa |
+----+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_message_graph[end_of_run]
@@ -1342,21 +1000,12 @@
# ---
# name: test_message_graph[end_of_run].3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* ..
** ..
* .
+--------+ +---------+
| action | | __end__ |
+--------+ +---------+
graph TD;
__start__ --> agent;
action --> agent;
agent -. continue .-> action;
agent -. end .-> __end__;
'''
# ---
# name: test_message_graph[end_of_step]
@@ -1432,46 +1081,21 @@
# ---
# name: test_message_graph[end_of_step].3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* ..
** ..
* .
+--------+ +---------+
| action | | __end__ |
+--------+ +---------+
graph TD;
__start__ --> agent;
action --> agent;
agent -. continue .-> action;
agent -. end .-> __end__;
'''
# ---
# name: test_nested_graph
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| inner |
+-------+
*
*
*
+------+
| side |
+------+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> inner;
inner --> side;
side --> __end__;
'''
# ---
# name: test_nested_graph.1
@@ -1613,25 +1237,6 @@
})
# ---
# name: test_nested_graph_xray.1
'''
+-----------+
| __start__ |
..+-----------+...
..... . .....
... . ...
... . ...
+----------+ +----------+ +------------+
| tool_one |* | tool_two | | tool_three |
+----------+ *** +----------+ **+------------+
***** * *****
*** * ***
*** * ***
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_nested_graph_xray.2
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
@@ -1735,21 +1340,12 @@
# ---
# name: test_prebuilt_chat.3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
* ..
** ..
* .
+--------+ +---------+
| action | | __end__ |
+--------+ +---------+
graph TD;
__start__ --> agent;
action --> agent;
agent -. continue .-> action;
agent -. end .-> __end__;
'''
# ---
# name: test_prebuilt_tool_chat
@@ -1826,21 +1422,12 @@
# ---
# name: test_prebuilt_tool_chat.3
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-------+
| agent |
+-------+
. ..
.. ..
. .
+--------+ +---------+
| action | | __end__ |
+--------+ +---------+
graph TD;
__start__ --> agent;
action --> agent;
agent -. continue .-> action;
agent -. end .-> __end__;
'''
# ---
# name: test_repeat_condition
@@ -1860,55 +1447,16 @@
# ---
# name: test_simple_multi_edge
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+----+
| up |
+----+
** **
* *
* *
+------+ *
| side | *
+------+ *
** **
* *
* *
+------+
| down |
+------+
*
*
*
+---------+
| __end__ |
+---------+
graph TD;
__start__ --> up;
down --> __end__;
side --> down;
up --> down;
up --> side;
'''
# ---
# name: test_start_branch_then[end_of_run]
'''
+-----------+
| __start__ |
+-----------+
.. ..
.. ..
.. ..
+---------------+ +---------------+
| tool_two_slow | | tool_two_fast |
+---------------+ +---------------+
** **
** **
** **
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_start_branch_then[end_of_run].1
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
@@ -1927,25 +1475,6 @@
'''
# ---
# name: test_start_branch_then[end_of_step]
'''
+-----------+
| __start__ |
+-----------+
.. ..
.. ..
.. ..
+---------------+ +---------------+
| tool_two_slow | | tool_two_fast |
+---------------+ +---------------+
** **
** **
** **
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_start_branch_then[end_of_step].1
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
+12 -16
View File
@@ -882,10 +882,8 @@ def test_conditional_graph(
app = workflow.compile()
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == 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).draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"input": "what is weather in sf"}) == {
@@ -1462,7 +1460,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert (
app.invoke("what is weather in sf", debug=True)
@@ -1562,7 +1560,7 @@ def test_conditional_graph_state(
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
@@ -1914,7 +1912,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
@@ -1980,7 +1978,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke(
{"messages": [HumanMessage(content="what is weather in sf")]}
@@ -2243,7 +2241,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke(
{"messages": [HumanMessage(content="what is weather in sf")]}
@@ -2458,7 +2456,7 @@ def test_message_graph(
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == snapshot
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke(HumanMessage(content="what is weather in sf")) == [
HumanMessage(
@@ -3023,7 +3021,6 @@ def test_start_branch_then(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
)
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_ascii() == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
@@ -3127,7 +3124,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
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().draw_mermaid(with_styles=False) == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
@@ -3300,7 +3297,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(
app = workflow.compile()
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "analyzed: query: what is weather in sf",
@@ -3396,7 +3393,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
app = workflow.compile()
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
"query": "analyzed: query: what is weather in sf",
@@ -3497,7 +3494,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
app = workflow.compile()
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
with pytest.raises(ValidationError):
app.invoke({"query": {}})
@@ -3834,7 +3831,7 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
app = graph.compile()
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
@@ -3863,7 +3860,6 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
app = graph.compile()
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
@@ -3901,7 +3897,7 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
app = graph.compile()
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.get_graph(xray=True).draw_mermaid() == snapshot
assert app.invoke(
{"my_key": "my value", "never_called": never_called}, debug=True