WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-28 10:49:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #387 from langchain-ai/nc/3may/checkpoint-metadata

Browse Source 
Add metadata to checkpoints, checkpoint inputs before first step
This commit is contained in:
Nuno Campos
2024-05-06 14:41:29 -07:00
committed by GitHub
parent d1b0be8620 48865daf02
commit 5bee63c7d5
14 changed files with 863 additions and 1485 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/channels/base.py
+3 -1
View File
@@ -125,6 +125,8 @@ def create_checkpoint(
checkpoint: Checkpoint, channels: Mapping[str, BaseChannel]
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
assert ts > checkpoint["ts"], "Timestamps must be monotonically increasing"
values: dict[str, Any] = {}
for k, v in channels.items():
try:
@@ -133,7 +135,7 @@ def create_checkpoint(
pass
return Checkpoint(
v=1,
ts=datetime.now(timezone.utc).isoformat(),
ts=ts,
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
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
+34 -38
View File
@@ -10,7 +10,7 @@ from typing_extensions import Self
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
CheckpointMetadata,
CheckpointTuple,
SerializerProtocol,
)
@@ -80,9 +80,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
@@ -130,6 +129,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
thread_ts TEXT NOT NULL,
parent_ts TEXT,
checkpoint BLOB,
metadata BLOB,
PRIMARY KEY (thread_id, thread_ts)
);
"""
@@ -155,7 +155,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
await self.setup()
if config["configurable"].get("thread_ts"):
async with self.conn.execute(
"SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
"SELECT checkpoint, parent_ts, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
@@ -165,20 +165,19 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
return CheckpointTuple(
config,
self.serde.loads(value[0]),
(
{
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": value[1],
}
self.serde.loads(value[2]) if value[2] is not None else {},
{
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": value[1],
}
if value[1]
else None
),
}
if value[1]
else None,
)
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",
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(str(config["configurable"]["thread_id"]),),
) as cursor:
if value := await cursor.fetchone():
@@ -190,16 +189,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
}
},
self.serde.loads(value[3]),
(
{
"configurable": {
"thread_id": value[0],
"thread_ts": value[2],
}
self.serde.loads(value[4]) if value[4] is not None else {},
{
"configurable": {
"thread_id": value[0],
"thread_ts": value[2],
}
if value[2]
else None
),
}
if value[2]
else None,
)
async def alist(
@@ -224,9 +222,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"""
await self.setup()
query = (
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC"
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC"
if before is None
else "SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts < ? ORDER BY thread_ts DESC"
else "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts < ? ORDER BY thread_ts DESC"
)
if limit:
query += f" LIMIT {limit}"
@@ -241,24 +239,21 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
)
),
) as cursor:
async for thread_id, thread_ts, parent_ts, value in cursor:
async for thread_id, thread_ts, parent_ts, value, metadata in cursor:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None
),
self.serde.loads(metadata) if metadata is not None else {},
{"configurable": {"thread_id": thread_id, "thread_ts": parent_ts}}
if parent_ts
else None,
)
async def aput(
self, config: RunnableConfig, checkpoint: Checkpoint
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
"""Save a checkpoint to the database asynchronously.
@@ -274,12 +269,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"""
await self.setup()
async with self.conn.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
),
):
await self.conn.commit()
langgraph/checkpoint/base.py
+30 -17
View File
@@ -5,6 +5,7 @@ from typing import (
Any,
AsyncIterator,
Iterator,
Literal,
NamedTuple,
Optional,
TypedDict,
@@ -14,7 +15,22 @@ from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
from langgraph.serde.base import SerializerProtocol
from langgraph.serde.jsonplus import JsonPlusSerializer
from langgraph.utils import StrEnum
# Marked as total=False to allow for future expansion.
class CheckpointMetadata(TypedDict, total=False):
source: Literal["input", "loop", "update"]
"""The source of the checkpoint.
- "input": The checkpoint was created from an input to invoke/stream/batch.
- "loop": The checkpoint was created from inside the pregel loop.
- "update": The checkpoint was created from a manual state update.
"""
step: int
"""The step number of the checkpoint.
-1 for the first "input" checkpoint.
0 for the first "loop" checkpoint.
... for the nth checkpoint afterwards.
"""
class Checkpoint(TypedDict):
@@ -71,18 +87,10 @@ 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
metadata: CheckpointMetadata
parent_config: Optional[RunnableConfig] = None
@@ -106,18 +114,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]:
@@ -139,7 +143,12 @@ class BaseCheckpointSaver(ABC):
) -> Iterator[CheckpointTuple]:
raise NotImplementedError
def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
raise NotImplementedError
async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
@@ -149,7 +158,7 @@ class BaseCheckpointSaver(ABC):
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
raise NotImplementedError
async def alist(
def alist(
self,
config: RunnableConfig,
*,
@@ -157,8 +166,12 @@ class BaseCheckpointSaver(ABC):
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
raise NotImplementedError
yield
async def aput(
self, config: RunnableConfig, checkpoint: Checkpoint
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
raise NotImplementedError
langgraph/checkpoint/memory.py
+30 -12
View File
@@ -7,7 +7,7 @@ from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
CheckpointMetadata,
CheckpointTuple,
SerializerProtocol,
)
@@ -39,15 +39,14 @@ class MemorySaver(BaseCheckpointSaver):
asyncio.run(coro) # Output: 2
"""
storage: defaultdict[str, dict[str, Checkpoint]]
storage: defaultdict[str, dict[str, tuple[bytes, bytes]]]
def __init__(
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]:
@@ -66,16 +65,21 @@ class MemorySaver(BaseCheckpointSaver):
"""
thread_id = config["configurable"]["thread_id"]
if ts := config["configurable"].get("thread_ts"):
if checkpoint := self.storage[thread_id].get(ts):
if saved := self.storage[thread_id].get(ts):
checkpoint, metadata = saved
return CheckpointTuple(
config=config, checkpoint=self.serde.loads(checkpoint)
config=config,
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
)
else:
if checkpoints := self.storage[thread_id]:
ts = max(checkpoints.keys())
checkpoint, metadata = checkpoints[ts]
return CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoints[ts]),
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
)
def list(
@@ -99,7 +103,7 @@ class MemorySaver(BaseCheckpointSaver):
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
"""
thread_id = config["configurable"]["thread_id"]
for ts, checkpoint in self.storage[thread_id].items():
for ts, (checkpoint, metadata) in self.storage[thread_id].items():
if before and ts >= before["configurable"]["thread_ts"]:
continue
if limit is not None and limit <= 0:
@@ -108,9 +112,15 @@ class MemorySaver(BaseCheckpointSaver):
yield CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
)
def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
"""Save a checkpoint to the in-memory storage.
This method saves a checkpoint to the in-memory storage. The checkpoint is associated
@@ -124,7 +134,12 @@ class MemorySaver(BaseCheckpointSaver):
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
"""
self.storage[config["configurable"]["thread_id"]].update(
{checkpoint["ts"]: self.serde.dumps(checkpoint)}
{
checkpoint["ts"]: (
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
)
}
)
return {
"configurable": {
@@ -170,8 +185,11 @@ class MemorySaver(BaseCheckpointSaver):
return
async def aput(
self, config: RunnableConfig, checkpoint: Checkpoint
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
return await asyncio.get_running_loop().run_in_executor(
None, self.put, config, checkpoint
None, self.put, config, checkpoint, metadata
)
langgraph/checkpoint/sqlite.py
+51 -39
View File
@@ -1,5 +1,6 @@
import pickle
import sqlite3
import threading
from contextlib import AbstractContextManager, contextmanager
from types import TracebackType
from typing import Any, Iterator, Optional
@@ -10,7 +11,7 @@ from typing_extensions import Self
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
CheckpointMetadata,
CheckpointTuple,
SerializerProtocol,
)
@@ -90,11 +91,11 @@ 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
self.lock = threading.Lock()
@classmethod
def from_conn_string(cls, conn_string: str) -> "SqliteSaver":
@@ -116,7 +117,13 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
memory = SqliteSaver.from_conn_string("checkpoints.sqlite")
"""
return SqliteSaver(conn=sqlite3.connect(conn_string))
return SqliteSaver(
conn=sqlite3.connect(
conn_string,
# https://ricardoanderegg.com/posts/python-sqlite-thread-safety/
check_same_thread=False,
)
)
def __enter__(self) -> Self:
return self
@@ -146,6 +153,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
thread_ts TEXT NOT NULL,
parent_ts TEXT,
checkpoint BLOB,
metadata BLOB,
PRIMARY KEY (thread_id, thread_ts)
);
"""
@@ -211,7 +219,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
with self.cursor(transaction=False) as cur:
if config["configurable"].get("thread_ts"):
cur.execute(
"SELECT checkpoint, parent_ts FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
"SELECT checkpoint, parent_ts, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
@@ -221,20 +229,19 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
return CheckpointTuple(
config,
self.serde.loads(value[0]),
(
{
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": value[1],
}
self.serde.loads(value[2]) if value[2] is not None else {},
{
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": value[1],
}
if value[1]
else None
),
}
if value[1]
else None,
)
else:
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(str(config["configurable"]["thread_id"]),),
)
if value := cur.fetchone():
@@ -246,16 +253,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
}
},
self.serde.loads(value[3]),
(
{
"configurable": {
"thread_id": value[0],
"thread_ts": value[2],
}
self.serde.loads(value[4]) if value[4] is not None else {},
{
"configurable": {
"thread_id": value[0],
"thread_ts": value[2],
}
if value[2]
else None
),
}
if value[2]
else None,
)
def list(
@@ -289,9 +295,9 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
print(checkpoints) # Output: [CheckpointTuple(...), ...]
"""
query = (
"SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC"
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC"
if before is None
else "SELECT thread_id, thread_ts, parent_ts, checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts < ? ORDER BY thread_ts DESC"
else "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts < ? ORDER BY thread_ts DESC"
)
if limit:
query += f" LIMIT {limit}"
@@ -307,23 +313,27 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
)
),
)
for thread_id, thread_ts, parent_ts, value in cur:
for thread_id, thread_ts, parent_ts, value, metadata in cur:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
self.serde.loads(metadata) if metadata is not None else {},
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
if parent_ts
else None
),
}
if parent_ts
else None,
)
def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
"""Save a checkpoint to the database.
This method saves a checkpoint to the SQLite database. The checkpoint is associated
@@ -332,6 +342,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
Args:
config (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (Optional[dict[str, Any]]): Additional metadata to save with the checkpoint. Defaults to None.
Returns:
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
@@ -345,14 +356,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
saved_config
) # Output: {"configurable": {"thread_id": "1", "thread_ts": 2024-05-04T06:32:42.235444+00:00"}}
"""
with self.cursor() as cur:
with self.lock, self.cursor() as cur:
cur.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint) VALUES (?, ?, ?, ?)",
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
checkpoint["ts"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
),
)
return {
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
+145 -110
View File
@@ -58,7 +58,6 @@ from langgraph.channels.base import (
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
CheckpointAt,
copy_checkpoint,
empty_checkpoint,
)
@@ -334,7 +333,6 @@ class Pregel(
saved = self.checkpointer.get_tuple(config)
checkpoint = saved.checkpoint if saved else empty_checkpoint()
config = saved.config if saved else config
with ChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
@@ -342,7 +340,9 @@ class Pregel(
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
saved.config if saved else config,
saved.metadata if saved else None,
saved.parent_config if saved else None,
)
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
@@ -352,7 +352,6 @@ class Pregel(
saved = await self.checkpointer.aget_tuple(config)
checkpoint = saved.checkpoint if saved else empty_checkpoint()
config = saved.config if saved else config
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
@@ -360,7 +359,9 @@ class Pregel(
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
saved.config if saved else config,
saved.metadata if saved else None,
saved.parent_config if saved else None,
)
def get_state_history(
@@ -374,7 +375,7 @@ class Pregel(
if not self.checkpointer:
raise ValueError("No checkpointer set")
for config, checkpoint, parent_config in self.checkpointer.list(
for config, checkpoint, metadata, parent_config in self.checkpointer.list(
config, before=before, limit=limit
):
with ChannelsManager(self.channels, checkpoint) as channels:
@@ -385,6 +386,7 @@ class Pregel(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
metadata,
parent_config,
)
@@ -399,9 +401,12 @@ class Pregel(
if not self.checkpointer:
raise ValueError("No checkpointer set")
async for config, checkpoint, parent_config in self.checkpointer.alist(
config, before=before, limit=limit
):
async for (
config,
checkpoint,
metadata,
parent_config,
) in self.checkpointer.alist(config, before=before, limit=limit):
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, for_execution=False
@@ -410,6 +415,7 @@ class Pregel(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
metadata,
parent_config,
)
@@ -427,8 +433,8 @@ class Pregel(
raise ValueError("No checkpointer set")
# get last checkpoint
checkpoint = self.checkpointer.get(config)
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
saved = self.checkpointer.get_tuple(config)
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
# find last node that updated the state, if not provided
if as_node is None:
last_seen_by_node = sorted(
@@ -476,7 +482,14 @@ class Pregel(
# apply to checkpoint and save
_apply_writes(checkpoint, channels, task.writes)
return self.checkpointer.put(
config, create_checkpoint(checkpoint, channels)
saved.config if saved else config,
create_checkpoint(checkpoint, channels),
{
"source": "update",
"step": saved.metadata.get("step", 0) + 1
if saved.metadata
else None,
},
)
async def aupdate_state(
@@ -489,8 +502,8 @@ class Pregel(
raise ValueError("No checkpointer set")
# get last checkpoint
checkpoint = await self.checkpointer.aget(config)
checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint()
saved = await self.checkpointer.aget_tuple(config)
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
# find last node that updated the state, if not provided
if as_node is None:
last_seen_by_node = sorted(
@@ -538,7 +551,12 @@ class Pregel(
# apply to checkpoint and save
_apply_writes(checkpoint, channels, task.writes)
return await self.checkpointer.aput(
config, create_checkpoint(checkpoint, channels)
saved.config if saved else config,
create_checkpoint(checkpoint, channels),
{
"source": "update",
"step": saved.metadata.get("step", 0) + 1 if saved else None,
},
)
def _defaults(
@@ -605,6 +623,7 @@ class Pregel(
run_id=config.get("run_id"),
)
try:
bg: list[concurrent.futures.Future] = []
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
if self.checkpointer and not config.get("configurable"):
@@ -631,11 +650,10 @@ class Pregel(
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
checkpoint_config = config
checkpoint = (
self.checkpointer.get(checkpoint_config) if self.checkpointer else None
)
checkpoint = checkpoint or empty_checkpoint()
saved = self.checkpointer.get_tuple(config) if self.checkpointer else None
checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint_config = saved.config if saved else config
start = saved.metadata.get("step", -2) + 1 if saved else -1
# create channels from checkpoint
with ChannelsManager(
self.channels, checkpoint
@@ -648,6 +666,27 @@ class Pregel(
)
# apply input writes
_apply_writes(checkpoint, channels, input_writes)
# save input checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
bg.append(
executor.submit(
self.checkpointer.put,
checkpoint_config,
copy_checkpoint(checkpoint),
{"source": "input", "step": start},
)
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# increment start to 0
start += 1
else:
# if received no input, take that as signal to proceed
# past previous interrupt, if any
@@ -661,7 +700,8 @@ class Pregel(
# channel updates from step N are only visible in step N+1
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
for step in range(config["recursion_limit"] + 1):
stop = start + config["recursion_limit"] + 1
for step in range(start, stop):
next_checkpoint, next_tasks = _prepare_next_tasks(
checkpoint, processes, channels, for_execution=True
)
@@ -762,23 +802,31 @@ 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
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
bg.append(
executor.submit(
self.checkpointer.put,
checkpoint_config,
channels,
self.stream_channels_asis,
copy_checkpoint(checkpoint),
{"source": "loop", "step": step},
)
elif stream_mode == "debug":
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# yield debug checkpoint
if stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
step,
checkpoint_config if self.checkpointer else None,
channels,
self.stream_channels_asis,
)
# after execution, check if we should interrupt
@@ -792,33 +840,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
@@ -829,6 +850,12 @@ class Pregel(
task.cancel()
except NameError:
pass
# wait for all background tasks to finish
done, _ = concurrent.futures.wait(
bg, return_when=concurrent.futures.ALL_COMPLETED
)
for task in done:
task.result()
async def astream(
self,
@@ -860,7 +887,7 @@ class Pregel(
None,
)
try:
tasks: list[asyncio.Task] = []
bg: list[asyncio.Task] = []
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
if self.checkpointer and not config.get("configurable"):
@@ -887,13 +914,14 @@ class Pregel(
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
checkpoint_config = config
checkpoint = (
await self.checkpointer.aget(checkpoint_config)
saved = (
await self.checkpointer.aget_tuple(config)
if self.checkpointer
else None
)
checkpoint = checkpoint or empty_checkpoint()
checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint_config = saved.config if saved else config
start = saved.metadata.get("step", -2) + 1 if saved else -1
# create channels from checkpoint
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
# map inputs to channel updates
@@ -904,6 +932,28 @@ class Pregel(
)
# apply input writes
_apply_writes(checkpoint, channels, input_writes)
# save input checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
bg.append(
asyncio.create_task(
self.checkpointer.aput(
checkpoint_config,
copy_checkpoint(checkpoint),
{"source": "input", "step": start},
)
)
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# increment start to 0
start += 1
else:
# if received no input, take that as signal to proceed
# past previous interrupt, if any
@@ -917,7 +967,9 @@ class Pregel(
# channel updates from step N are only visible in step N+1,
# channels are guaranteed to be immutable for the duration of the step,
# channel updates being applied only at the transition between steps
for step in range(config["recursion_limit"] + 1):
start = saved.metadata.get("step", -1) + 1 if saved else 0
stop = start + config["recursion_limit"] + 1
for step in range(start, stop):
next_checkpoint, next_tasks = _prepare_next_tasks(
checkpoint, processes, channels, for_execution=True
)
@@ -1028,23 +1080,32 @@ 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
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
checkpoint_config,
channels,
self.stream_channels_asis,
bg.append(
asyncio.create_task(
self.checkpointer.aput(
checkpoint_config,
checkpoint,
{"source": "loop", "step": step},
)
)
elif stream_mode == "debug":
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# yield debug checkpoint
if stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
step,
checkpoint_config if self.checkpointer else None,
channels,
self.stream_channels_asis,
)
# after execution, check if we should interrupt
@@ -1058,32 +1119,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
@@ -1092,11 +1127,11 @@ class Pregel(
try:
for task in futures:
task.cancel()
tasks.append(task)
bg.append(task)
except NameError:
pass
# wait for all tasks to finish
await asyncio.gather(*tasks, return_exceptions=True)
# wait for all background tasks to finish
await asyncio.gather(*bg)
def invoke(
self,
langgraph/pregel/types.py
+4
View File
@@ -3,6 +3,8 @@ from typing import Any, Literal, NamedTuple, Optional, Union
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.checkpoint.base import CheckpointMetadata
class PregelTaskDescription(NamedTuple):
name: str
@@ -25,6 +27,8 @@ class StateSnapshot(NamedTuple):
"""Nodes to execute in the next step, if any"""
config: RunnableConfig
"""Config used to fetch this snapshot"""
metadata: CheckpointMetadata
"""Metadata associated with this snapshot"""
parent_config: Optional[RunnableConfig] = None
"""Config used to fetch the parent snapshot, if any"""
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
+9 -7
View File
@@ -3,7 +3,7 @@ from typing import Any, Optional
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointAt,
CheckpointMetadata,
SerializerProtocol,
copy_checkpoint,
)
@@ -21,20 +21,22 @@ 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:
def put(
self,
config: dict,
checkpoint: Checkpoint,
metadata: Optional[CheckpointMetadata] = None,
) -> None:
# assert checkpoint hasn't been modified since last written
thread_id = config["configurable"]["thread_id"]
if saved := super().get(config):
@@ -43,4 +45,4 @@ class MemorySaverAssertImmutable(MemorySaver):
checkpoint
)
# call super to write checkpoint
return super().put(config, checkpoint)
return super().put(config, checkpoint, metadata)
tests/test_pregel.py
+264 -344
View File
File diff suppressed because it is too large Load Diff
tests/test_pregel_async.py
+271 -337
View File
@@ -1,6 +1,7 @@
import asyncio
import json
import operator
from collections import Counter
from contextlib import asynccontextmanager, contextmanager
from typing import (
Annotated,
@@ -25,7 +26,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 +270,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 +452,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 +602,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 +617,7 @@ async def test_invoke_checkpoint(
| raise_if_above_10
)
memory = MemorySaverAssertImmutable(at=checkpoint_at)
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one},
@@ -674,12 +658,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 +674,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={
@@ -735,13 +713,21 @@ async def test_invoke_checkpoint_aiosqlite(
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert state.next == ("one",)
"""we checkpoint inputs and it failed on "one", so the next node is one"""
# we can recover from error by sending new inputs
assert await app.ainvoke(2, thread_1) == 9
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
assert state.next == ()
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert await app.ainvoke(5, thread_2) == 5
state = await app.aget_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 7
assert state.values.get("total") == 16
assert state.next == ()
state = await app.aget_state(thread_2)
assert state is not None
@@ -751,8 +737,12 @@ async def test_invoke_checkpoint_aiosqlite(
assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1
# list all checkpoints for thread 1
thread_1_history = [c async for c in app.aget_state_history(thread_1)]
# there are 2: one for each successful ainvoke()
assert len(thread_1_history) == 2
# there are 7 checkpoints
assert len(thread_1_history) == 7
assert Counter(c.metadata["source"] for c in thread_1_history) == {
"input": 4,
"loop": 3,
}
# sorted descending
assert (
thread_1_history[0].config["configurable"]["thread_ts"]
@@ -767,10 +757,10 @@ async def test_invoke_checkpoint_aiosqlite(
]
assert len(cursored) == 1
assert cursored[0].config == thread_1_history[1].config
# the second checkpoint
assert thread_1_history[0].values["total"] == 7
# the first checkpoint
assert thread_1_history[1].values["total"] == 2
# the last checkpoint
assert thread_1_history[0].values["total"] == 16
# the first "loop" checkpoint
assert thread_1_history[-2].values["total"] == 2
# can get each checkpoint using aget with config
assert (await memory.aget(thread_1_history[0].config))[
"ts"
@@ -786,7 +776,14 @@ async def test_invoke_checkpoint_aiosqlite(
> thread_1_history[0].config["configurable"]["thread_ts"]
)
# 1 more checkpoint in history
assert len([h async for h in app.aget_state_history(thread_1)]) == 3
assert len([c async for c in app.aget_state_history(thread_1)]) == 8
assert Counter(
[c.metadata["source"] async for c in app.aget_state_history(thread_1)]
) == {
"update": 1,
"input": 4,
"loop": 3,
}
# the latest checkpoint is the updated one
assert await app.aget_state(thread_1) == await app.aget_state(
thread_1_next_config
@@ -1003,10 +1000,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 +1268,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"}}
@@ -1306,6 +1300,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "loop", "step": 0},
)
await app_w_interrupt.aupdate_state(
@@ -1333,6 +1328,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 1},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1416,12 +1412,13 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 4},
)
# 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"}}
@@ -1454,6 +1451,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "loop", "step": 0},
)
await app_w_interrupt.aupdate_state(
@@ -1481,6 +1479,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 1},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1564,12 +1563,13 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 4},
)
# 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"}}
@@ -1602,6 +1602,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "loop", "step": 0},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1695,10 +1696,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
@@ -1892,7 +1890,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"}}
@@ -1924,6 +1922,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "loop", "step": 1},
)
await app_w_interrupt.aupdate_state(
@@ -1949,6 +1948,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 2},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -2007,12 +2007,13 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 5},
)
# 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"}}
@@ -2043,6 +2044,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "loop", "step": 1},
)
await app_w_interrupt.aupdate_state(
@@ -2068,6 +2070,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
},
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 2},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -2126,6 +2129,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
},
next=(),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "update", "step": 5},
)
@@ -2524,10 +2528,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
@@ -2696,7 +2697,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"}}
@@ -2734,6 +2735,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
],
next=("action",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
metadata={"source": "loop", "step": 1},
)
# modify ai message
@@ -2761,6 +2763,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 2},
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -2813,6 +2816,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 4},
)
await app_w_interrupt.aupdate_state(
@@ -2846,6 +2850,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
],
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 5},
)
@@ -2921,12 +2926,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() -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -2949,7 +2949,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"]
)
@@ -2968,6 +2967,10 @@ async def test_start_branch_then(
values={"my_key": "value", "market": "DE"},
next=("tool_two_slow",),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
metadata={"source": "loop", "step": 0},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread1, debug=1) == {
@@ -2978,6 +2981,10 @@ async def test_start_branch_then(
values={"my_key": "value slow", "market": "DE"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
metadata={"source": "loop", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
thread2 = {"configurable": {"thread_id": "2"}}
@@ -2990,6 +2997,10 @@ async def test_start_branch_then(
values={"my_key": "value", "market": "US"},
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
metadata={"source": "loop", "step": 0},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread2, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread2, debug=1) == {
@@ -3000,15 +3011,55 @@ async def test_start_branch_then(
values={"my_key": "value fast", "market": "US"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
metadata={"source": "loop", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread2, limit=2)
][-1].config,
)
thread3 = {"configurable": {"thread_id": "3"}}
# stop when about to enter node
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == {
"my_key": "value",
"market": "US",
}
assert await tool_two.aget_state(thread3) == StateSnapshot(
values={"my_key": "value", "market": "US"},
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
metadata={"source": "loop", "step": 0},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread3, limit=2)
][-1].config,
)
# update state
await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key
assert await tool_two.aget_state(thread3) == StateSnapshot(
values={"my_key": "valuekey", "market": "US"},
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
metadata={"source": "update", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread3, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread3, debug=1) == {
"my_key": "valuekey fast",
"market": "US",
}
assert await tool_two.aget_state(thread3) == StateSnapshot(
values={"my_key": "valuekey fast", "market": "US"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
metadata={"source": "loop", "step": 2},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread3, limit=2)
][-1].config,
)
@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):
@@ -3039,256 +3090,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"]
@@ -3308,6 +3229,10 @@ async def test_branch_then(
values={"my_key": "value prepared", "market": "DE"},
next=("tool_two_slow",),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
metadata={"source": "loop", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread1, debug=1) == {
@@ -3318,6 +3243,10 @@ async def test_branch_then(
values={"my_key": "value prepared slow finished", "market": "DE"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
metadata={"source": "loop", "step": 3},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
thread2 = {"configurable": {"thread_id": "2"}}
@@ -3330,6 +3259,10 @@ async def test_branch_then(
values={"my_key": "value prepared", "market": "US"},
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
metadata={"source": "loop", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread2, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread2, debug=1) == {
@@ -3340,10 +3273,13 @@ async def test_branch_then(
values={"my_key": "value prepared fast finished", "market": "US"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
metadata={"source": "loop", "step": 3},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread2, limit=2)
][-1].config,
)
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
saver.at = checkpoint_at
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_after=["prepare"]
)
@@ -3362,6 +3298,10 @@ async def test_branch_then(
values={"my_key": "value prepared", "market": "DE"},
next=("tool_two_slow",),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
metadata={"source": "loop", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread1, debug=1) == {
@@ -3372,6 +3312,10 @@ async def test_branch_then(
values={"my_key": "value prepared slow finished", "market": "DE"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
metadata={"source": "loop", "step": 3},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
thread2 = {"configurable": {"thread_id": "2"}}
@@ -3384,6 +3328,10 @@ async def test_branch_then(
values={"my_key": "value prepared", "market": "US"},
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
metadata={"source": "loop", "step": 1},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread2, limit=2)
][-1].config,
)
# resume, for same result as above
assert await tool_two.ainvoke(None, thread2, debug=1) == {
@@ -3394,15 +3342,14 @@ async def test_branch_then(
values={"my_key": "value prepared fast finished", "market": "US"},
next=(),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
metadata={"source": "loop", "step": 3},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread2, limit=2)
][-1].config,
)
@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]:
@@ -3466,7 +3413,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"}}
@@ -3490,12 +3437,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]]]
@@ -3564,7 +3507,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"}}
@@ -3588,12 +3531,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
@@ -3671,7 +3610,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"}}
@@ -3695,12 +3634,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]:
@@ -3769,7 +3703,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"}}