WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-01 22:15:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #365 from langchain-ai/nc/29apr/support-non-str-thread-ids

Browse Source 
Nc/29apr/support non str thread ids
This commit is contained in:
Nuno Campos
2024-04-30 08:54:14 -07:00
committed by GitHub
parent a3c867d688 7554309abf
commit 9505d8cc84
6 changed files with 160 additions and 635 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/time-travel.ipynb
+3 -3
View File
File diff suppressed because one or more lines are too long
langgraph/checkpoint/aiosqlite.py
+5 -5
View File
@@ -75,8 +75,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
async with self.conn.execute(
"SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
(
config["configurable"]["thread_id"],
config["configurable"]["thread_ts"],
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
),
) as cursor:
if value := await cursor.fetchone():
@@ -95,7 +95,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
else:
async with self.conn.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
) as cursor:
if value := await cursor.fetchone():
return CheckpointTuple(
@@ -120,7 +120,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
await self.setup()
async with self.conn.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
) as cursor:
async for thread_id, thread_ts, parent_ts, value in cursor:
yield CheckpointTuple(
@@ -138,7 +138,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
async with self.conn.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
(
config["configurable"]["thread_id"],
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
langgraph/checkpoint/sqlite.py
+5 -5
View File
@@ -93,8 +93,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
cur.execute(
"SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
(
config["configurable"]["thread_id"],
config["configurable"]["thread_ts"],
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
),
)
if value := cur.fetchone():
@@ -113,7 +113,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
else:
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
)
if value := cur.fetchone():
return CheckpointTuple(
@@ -138,7 +138,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
with self.cursor(transaction=False) as cur:
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
(config["configurable"]["thread_id"],),
(str(config["configurable"]["thread_id"]),),
)
for thread_id, thread_ts, parent_ts, value in cur:
yield CheckpointTuple(
@@ -159,7 +159,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
cur.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
(
config["configurable"]["thread_id"],
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
langgraph/graph/graph.py
+1 -1
View File
@@ -1,5 +1,5 @@
import logging
from collections import Counter, defaultdict
from collections import defaultdict
from typing import (
Any,
Awaitable,
tests/__snapshots__/test_pregel.ambr
+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;
tests/test_pregel.py
+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