WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-21 23:22:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove option to only checkpoint at end of run

Browse Source 
- Now that checkpoint at end of each step adds no latency there is not point to keep this
- This will make it easier to add future features
This commit is contained in:
Nuno Campos
2024-05-06 11:52:49 -07:00
parent f304908102
commit 3ff3def62b
12 changed files with 303 additions and 1327 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/checkpoint/__init__.py
-2
View File
@@ -1,7 +1,6 @@
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
SerializerProtocol,
)
from langgraph.checkpoint.memory import MemorySaver
@@ -9,7 +8,6 @@ from langgraph.checkpoint.memory import MemorySaver
__all__ = [
"BaseCheckpointSaver",
"Checkpoint",
"CheckpointAt",
"MemorySaver",
"SerializerProtocol",
]
langgraph/checkpoint/aiosqlite.py
+1 -3
View File
@@ -10,7 +10,6 @@ from typing_extensions import Self
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
CheckpointTuple,
SerializerProtocol,
)
@@ -80,9 +79,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
conn: aiosqlite.Connection,
*,
serde: Optional[SerializerProtocol] = None,
at: Optional[CheckpointAt] = None,
):
super().__init__(serde=serde, at=at)
super().__init__(serde=serde)
self.conn = conn
self.lock = asyncio.Lock()
self.is_setup = False
langgraph/checkpoint/base.py
-14
View File
@@ -14,7 +14,6 @@ from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
from langgraph.serde.base import SerializerProtocol
from langgraph.serde.jsonplus import JsonPlusSerializer
from langgraph.utils import StrEnum
class Checkpoint(TypedDict):
@@ -71,15 +70,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
)
class CheckpointAt(StrEnum):
"""When to take a checkpoint."""
END_OF_STEP = "end_of_step"
"""Take a checkpoint at the end of each step."""
END_OF_RUN = "end_of_run"
"""Take a checkpoint at the end of the run."""
class CheckpointTuple(NamedTuple):
config: RunnableConfig
checkpoint: Checkpoint
@@ -107,18 +97,14 @@ CheckpointThreadTs = ConfigurableFieldSpec(
class BaseCheckpointSaver(ABC):
at: CheckpointAt = CheckpointAt.END_OF_STEP
serde: SerializerProtocol = JsonPlusSerializer()
def __init__(
self,
*,
serde: Optional[SerializerProtocol] = None,
at: Optional[CheckpointAt] = None,
) -> None:
self.serde = serde or self.serde
self.at = at or self.at
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
langgraph/checkpoint/memory.py
+1 -3
View File
@@ -7,7 +7,6 @@ from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
CheckpointTuple,
SerializerProtocol,
)
@@ -45,9 +44,8 @@ class MemorySaver(BaseCheckpointSaver):
self,
*,
serde: Optional[SerializerProtocol] = None,
at: Optional[CheckpointAt] = None,
) -> None:
super().__init__(serde=serde, at=at)
super().__init__(serde=serde)
self.storage = defaultdict(dict)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
langgraph/checkpoint/sqlite.py
+1 -3
View File
@@ -10,7 +10,6 @@ from typing_extensions import Self
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
CheckpointTuple,
SerializerProtocol,
)
@@ -90,9 +89,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
conn: sqlite3.Connection,
*,
serde: Optional[SerializerProtocol] = None,
at: Optional[CheckpointAt] = None,
) -> None:
super().__init__(serde=serde, at=at)
super().__init__(serde=serde)
self.conn = conn
self.is_setup = False
langgraph/graph/__init__.py
+2 -2
View File
@@ -1,5 +1,5 @@
from langgraph.graph.graph import END, Graph
from langgraph.graph.message import MessageGraph
from langgraph.graph.message import MessageGraph, add_messages
from langgraph.graph.state import StateGraph
__all__ = ["END", "Graph", "StateGraph", "MessageGraph"]
__all__ = ["END", "Graph", "StateGraph", "MessageGraph", "add_messages"]
langgraph/pregel/__init__.py
+2 -60
View File
@@ -58,7 +58,6 @@ from langgraph.channels.base import (
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
copy_checkpoint,
empty_checkpoint,
)
@@ -769,9 +768,7 @@ class Pregel(
yield from map_output_updates(output_keys, next_tasks)
# save end of step checkpoint
if self.checkpointer is not None and (
self.checkpointer.at == CheckpointAt.END_OF_STEP
):
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint_config = self.checkpointer.put(
checkpoint_config, checkpoint
@@ -799,33 +796,6 @@ class Pregel(
# set final channel values as run output
run_manager.on_chain_end(read_channels(channels, output_keys))
# save end of run checkpoint
if (
self.checkpointer is not None
and self.checkpointer.at == CheckpointAt.END_OF_RUN
):
checkpoint = create_checkpoint(checkpoint, channels)
executor.submit(
self.checkpointer.put(checkpoint_config, checkpoint)
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"]["thread_id"],
"thread_ts": checkpoint["ts"],
}
}
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
checkpoint_config,
channels,
self.stream_channels_asis,
)
elif self.checkpointer is None and stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
)
except BaseException as e:
run_manager.on_chain_error(e)
raise
@@ -1035,9 +1005,7 @@ class Pregel(
yield chunk
# save end of step checkpoint
if self.checkpointer is not None and (
self.checkpointer.at == CheckpointAt.END_OF_STEP
):
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint_config = await self.checkpointer.aput(
checkpoint_config, checkpoint
@@ -1065,32 +1033,6 @@ class Pregel(
# set final channel values as run output
await run_manager.on_chain_end(read_channels(channels, output_keys))
# save end of run checkpoint
if (
self.checkpointer is not None
and self.checkpointer.at == CheckpointAt.END_OF_RUN
):
checkpoint = create_checkpoint(checkpoint, channels)
tasks.append(
asyncio.create_task(
self.checkpointer.aput(checkpoint_config, checkpoint)
)
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"]["thread_id"],
"thread_ts": checkpoint["ts"],
}
}
if stream_mode == "debug":
yield map_debug_checkpoint(
step, checkpoint_config, channels, self.stream_channels_asis
)
elif self.checkpointer is None and stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
)
except BaseException as e:
await run_manager.on_chain_error(e)
raise
tests/__snapshots__/test_pregel.ambr
+18 -500
View File
File diff suppressed because one or more lines are too long
tests/__snapshots__/test_pregel_async.ambr
+2 -76
View File
@@ -1,5 +1,5 @@
# serializer version: 1
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_run]
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class
'''
+-----------+
| __start__ |
@@ -36,81 +36,7 @@
+---------+
'''
# ---
# 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__ |
+---------+
'''
# ---
# 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__ |
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_step]
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch
'''
+-----------+
| __start__ |
tests/memory_assert.py
+1 -5
View File
@@ -3,7 +3,6 @@ from typing import Any, Optional
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointAt,
SerializerProtocol,
copy_checkpoint,
)
@@ -21,17 +20,14 @@ class NoopSerializer(SerializerProtocol):
class MemorySaverAssertImmutable(MemorySaver):
serde = NoopSerializer()
at = CheckpointAt.END_OF_STEP
storage_for_copies: defaultdict[str, dict[str, Checkpoint]]
def __init__(
self,
*,
serde: Optional[SerializerProtocol] = None,
at: Optional[CheckpointAt] = None,
) -> None:
super().__init__(serde=serde, at=at)
super().__init__(serde=serde)
self.storage_for_copies = defaultdict(dict)
def put(self, config: dict, checkpoint: Checkpoint) -> None:
tests/test_pregel.py
+138 -330
View File
@@ -16,7 +16,6 @@ from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.checkpoint.base import CheckpointAt
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, Graph
from langgraph.graph.message import MessageGraph
@@ -292,17 +291,12 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert step == 2
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_invoke_two_processes_in_out_interrupt(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> None:
def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
memory = MemorySaverAssertImmutable(at=checkpoint_at)
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one, "two": two},
channels={
@@ -475,12 +469,6 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"step": 1,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
]
@@ -627,10 +615,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
assert app.invoke(2) == [3, 3]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) -> None:
def test_invoke_checkpoint(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
@@ -645,7 +630,7 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) -
| raise_if_above_10
)
memory = MemorySaverAssertImmutable(at=checkpoint_at)
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one},
@@ -686,12 +671,7 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) -
assert checkpoint["channel_values"].get("total") == 5
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_invoke_checkpoint_sqlite(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> None:
def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
@@ -707,7 +687,6 @@ def test_invoke_checkpoint_sqlite(
)
with SqliteSaver.from_conn_string(":memory:") as memory:
memory.at = checkpoint_at
app = Pregel(
nodes={"one": one},
channels={
@@ -992,12 +971,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert cleanup.call_count == 1, "Expected cleanup to be called once"
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_conditional_graph(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
from copy import deepcopy
from langchain.llms.fake import FakeStreamingListLLM
@@ -1199,7 +1173,7 @@ def test_conditional_graph(
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -1352,7 +1326,7 @@ def test_conditional_graph(
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
@@ -1500,7 +1474,7 @@ def test_conditional_graph(
# test re-invoke to continue with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
@@ -1668,12 +1642,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_conditional_state_graph(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
@@ -1840,7 +1809,7 @@ def test_conditional_state_graph(
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -1958,7 +1927,7 @@ def test_conditional_state_graph(
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
debug=True,
)
@@ -2077,7 +2046,7 @@ def test_conditional_state_graph(
# test w interrupt before all
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before="*",
debug=True,
)
@@ -2174,7 +2143,7 @@ def test_conditional_state_graph(
# test w interrupt after all
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after="*",
)
config = {"configurable": {"thread_id": "4"}}
@@ -2796,12 +2765,8 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_message_graph(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
deterministic_uuids: MockerFixture,
) -> None:
from copy import deepcopy
@@ -3020,7 +2985,7 @@ def test_message_graph(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -3192,7 +3157,7 @@ def test_message_graph(
)
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["action"],
)
config = {"configurable": {"thread_id": "2"}}
@@ -3472,12 +3437,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_start_branch_then(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -3511,7 +3471,6 @@ def test_start_branch_then(
}
with SqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
)
@@ -3569,10 +3528,7 @@ def test_start_branch_then(
)
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -> None:
def test_branch_then(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -3619,254 +3575,125 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
}
with SqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
# test stream_mode=debug
tool_two = tool_two_graph.compile(checkpointer=saver)
thread10 = {"configurable": {"thread_id": "10"}}
if checkpoint_at is CheckpointAt.END_OF_RUN:
assert [
*tool_two.stream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": None,
"values": {"my_key": "value", "market": "DE"},
assert [
*tool_two.stream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": None,
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": None,
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": None,
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 4,
"payload": {
"config": {
"configurable": {
"thread_id": "10",
"thread_ts": AnyStr(),
}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
else:
assert [
*tool_two.stream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
},
]
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
@@ -3925,7 +3752,6 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
)
with SqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_after=["prepare"]
)
@@ -3983,12 +3809,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -
)
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_in_one_fan_out_state_graph_waiting_edge(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
@@ -4054,7 +3875,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -4075,12 +3896,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
) -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
@@ -4150,7 +3967,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -4171,12 +3988,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
) -> None:
from langchain_core.pydantic_v1 import BaseModel, ValidationError
@@ -4254,7 +4067,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -4275,12 +4088,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
checkpoint_at: CheckpointAt,
) -> None:
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
@@ -4349,7 +4157,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
tests/test_pregel_async.py
+137 -329
View File
@@ -25,7 +25,6 @@ from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver
from langgraph.checkpoint.base import CheckpointAt
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph
from langgraph.prebuilt.chat_agent_executor import (
@@ -270,17 +269,12 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert step == 2
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_invoke_two_processes_in_out_interrupt(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> None:
async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
memory = MemorySaverAssertImmutable(at=checkpoint_at)
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one, "two": two},
channels={
@@ -457,12 +451,6 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"step": 1,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
]
@@ -613,12 +601,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
assert await app.ainvoke(2) == [3, 3]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_invoke_checkpoint(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> None:
async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
@@ -633,7 +616,7 @@ async def test_invoke_checkpoint(
| raise_if_above_10
)
memory = MemorySaverAssertImmutable(at=checkpoint_at)
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one},
@@ -674,12 +657,7 @@ async def test_invoke_checkpoint(
assert checkpoint["channel_values"].get("total") == 5
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_invoke_checkpoint_aiosqlite(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> None:
async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
@@ -695,7 +673,6 @@ async def test_invoke_checkpoint_aiosqlite(
)
async with AsyncSqliteSaver.from_conn_string(":memory:") as memory:
memory.at = checkpoint_at
app = Pregel(
nodes={"one": one},
channels={
@@ -1003,10 +980,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert cleanup_async.call_count == 1, "Expected cleanup to be called once"
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
async def test_conditional_graph() -> None:
from copy import deepcopy
from langchain.llms.fake import FakeStreamingListLLM
@@ -1274,7 +1248,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -1424,7 +1398,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
@@ -1575,7 +1549,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
# test re-invoke to continue with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
@@ -1702,10 +1676,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
async def test_conditional_graph_state() -> None:
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
@@ -1899,7 +1870,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -2022,7 +1993,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
@@ -2537,10 +2508,7 @@ async def test_prebuilt_chat() -> None:
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
async def test_message_graph() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.agents import AgentAction
@@ -2709,7 +2677,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -2938,12 +2906,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_start_branch_then(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
async def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -2966,7 +2929,6 @@ async def test_start_branch_then(
}
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
)
@@ -3024,12 +2986,7 @@ async def test_start_branch_then(
)
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_branch_then(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
async def test_branch_then() -> None:
pass
class State(TypedDict):
@@ -3060,256 +3017,126 @@ async def test_branch_then(
}
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
# test stream_mode=debug
tool_two = tool_two_graph.compile(checkpointer=saver)
thread10 = {"configurable": {"thread_id": "10"}}
if checkpoint_at is CheckpointAt.END_OF_RUN:
assert [
c
async for c in tool_two.astream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": None,
"values": {"my_key": "value", "market": "DE"},
assert [
c
async for c in tool_two.astream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": None,
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": None,
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": None,
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 4,
"payload": {
"config": {
"configurable": {
"thread_id": "10",
"thread_ts": AnyStr(),
}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
else:
assert [
c
async for c in tool_two.astream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
},
]
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
@@ -3368,7 +3195,6 @@ async def test_branch_then(
)
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_after=["prepare"]
)
@@ -3426,12 +3252,7 @@ async def test_branch_then(
)
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge(
checkpoint_at: CheckpointAt,
) -> None:
async def test_in_one_fan_out_state_graph_waiting_edge() -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
@@ -3495,7 +3316,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -3519,12 +3340,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge(
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
) -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
@@ -3593,7 +3410,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -3617,12 +3434,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
) -> None:
from langchain_core.pydantic_v1 import BaseModel, ValidationError
@@ -3700,7 +3513,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
@@ -3724,12 +3537,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
checkpoint_at: CheckpointAt,
) -> None:
async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
@@ -3798,7 +3606,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}