diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index c244e750c..368910523 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -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"], diff --git a/langgraph/checkpoint/__init__.py b/langgraph/checkpoint/__init__.py index 5a0f4bd0c..50f9db11b 100644 --- a/langgraph/checkpoint/__init__.py +++ b/langgraph/checkpoint/__init__.py @@ -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", ] diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index 62818dc16..c9a322561 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -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() diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 86d71c17f..5d197e19e 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -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 diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 055f92d54..fb49bab58 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -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 ) diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index e4c3be2a3..bde4a5d0a 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -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 { diff --git a/langgraph/graph/__init__.py b/langgraph/graph/__init__.py index 8fac44cfa..1c0ef3262 100644 --- a/langgraph/graph/__init__.py +++ b/langgraph/graph/__init__.py @@ -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"] diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d83eb906b..75321bf46 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -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, diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index 7c6075c98..d4dc92d40 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -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""" diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 9e8af9433..68e1d3a48 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1,5 +1,5 @@ # serializer version: 1 -# name: test_branch_then[end_of_run] +# name: test_branch_then ''' graph TD; __start__ --> prepare; @@ -11,41 +11,7 @@ ''' # --- -# name: test_branch_then[end_of_run].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - prepare([prepare]):::otherclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - finish([finish]):::otherclass; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_branch_then[end_of_step] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[end_of_step].1 +# name: test_branch_then.1 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; @@ -229,7 +195,7 @@ ''' # --- -# name: test_conditional_graph[end_of_run] +# name: test_conditional_graph ''' { "nodes": [ @@ -294,7 +260,7 @@ } ''' # --- -# name: test_conditional_graph[end_of_run].1 +# name: test_conditional_graph.1 ''' graph TD; __start__ --> agent; @@ -304,7 +270,7 @@ ''' # --- -# name: test_conditional_graph[end_of_run].2 +# name: test_conditional_graph.2 ''' { "nodes": [ @@ -442,7 +408,7 @@ } ''' # --- -# name: test_conditional_graph[end_of_run].3 +# name: test_conditional_graph.3 ''' graph TD; PromptTemplate --> FakeStreamingListLLM; @@ -458,242 +424,13 @@ ''' # --- -# name: test_conditional_graph[end_of_step] - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": "agent", - "type": "runnable", - "data": { - "id": [ - "langchain", - "schema", - "runnable", - "RunnableAssign" - ], - "name": "RunnableAssign" - } - }, - { - "id": "tools", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "tools" - } - } - ], - "edges": [ - { - "source": "__start__", - "target": "agent" - }, - { - "source": "tools", - "target": "agent" - }, - { - "source": "agent", - "target": "tools", - "data": "continue", - "conditional": true - }, - { - "source": "agent", - "target": "__end__", - "data": "exit", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_graph[end_of_step].1 - ''' - graph TD; - __start__ --> agent; - tools --> agent; - agent -. continue .-> tools; - agent -. exit .-> __end__; - - ''' -# --- -# name: test_conditional_graph[end_of_step].2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 3, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_community", - "llms", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 6, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 7, - "type": "runnable", - "data": { - "id": [ - "langchain", - "schema", - "runnable", - "RunnablePassthrough" - ], - "name": "RunnablePassthrough" - } - }, - { - "id": "tools", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "tools" - } - } - ], - "edges": [ - { - "source": 4, - "target": 5 - }, - { - "source": 5, - "target": 6 - }, - { - "source": 2, - "target": 4 - }, - { - "source": 6, - "target": 3 - }, - { - "source": 2, - "target": 7 - }, - { - "source": 7, - "target": 3 - }, - { - "source": "__start__", - "target": 2 - }, - { - "source": "tools", - "target": 2 - }, - { - "source": 3, - "target": "tools", - "data": "continue", - "conditional": true - }, - { - "source": 3, - "target": "__end__", - "data": "exit", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_graph[end_of_step].3 - ''' - graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> Lambda_agent_parser_; - Parallel_agent_outcome_Input --> PromptTemplate; - Lambda_agent_parser_ --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -. continue .-> tools; - Parallel_agent_outcome_Output -. exit .-> __end__; - - ''' -# --- -# name: test_conditional_state_graph[end_of_run] +# name: test_conditional_state_graph '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' # --- -# name: test_conditional_state_graph[end_of_run].1 +# name: test_conditional_state_graph.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- -# name: test_conditional_state_graph[end_of_run].2 +# name: test_conditional_state_graph.2 ''' { "nodes": [ @@ -758,7 +495,7 @@ } ''' # --- -# name: test_conditional_state_graph[end_of_run].3 +# name: test_conditional_state_graph.3 ''' graph TD; __start__ --> agent; @@ -768,88 +505,7 @@ ''' # --- -# name: test_conditional_state_graph[end_of_step] - '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' -# --- -# name: test_conditional_state_graph[end_of_step].1 - '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' -# --- -# name: test_conditional_state_graph[end_of_step].2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": "agent", - "type": "runnable", - "data": { - "id": [ - "langchain", - "schema", - "runnable", - "RunnableSequence" - ], - "name": "RunnableSequence" - } - }, - { - "id": "tools", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "tools" - } - } - ], - "edges": [ - { - "source": "__start__", - "target": "agent" - }, - { - "source": "tools", - "target": "agent" - }, - { - "source": "agent", - "target": "tools", - "data": "continue", - "conditional": true - }, - { - "source": "agent", - "target": "__end__", - "data": "exit", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_state_graph[end_of_step].3 - ''' - graph TD; - __start__ --> agent; - tools --> agent; - agent -. continue .-> tools; - agent -. exit .-> __end__; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[end_of_run] +# name: test_in_one_fan_out_state_graph_waiting_edge ''' graph TD; __start__ --> rewrite_query; @@ -862,20 +518,7 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge[end_of_step] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_run] +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class ''' graph TD; __start__ --> rewrite_query; @@ -888,7 +531,7 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_step] +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' graph TD; __start__ --> rewrite_query; @@ -901,39 +544,13 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_run] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_step] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_message_graph[end_of_run] +# name: test_message_graph '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior, usually passed in as the first of a sequence\\nof input messages.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a function back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}}, "required": ["content", "tool_call_id"]}}}' # --- -# name: test_message_graph[end_of_run].1 +# name: test_message_graph.1 '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior, usually passed in as the first of a sequence\\nof input messages.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a function back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}}, "required": ["content", "tool_call_id"]}}}' # --- -# name: test_message_graph[end_of_run].2 +# name: test_message_graph.2 ''' { "nodes": [ @@ -998,88 +615,7 @@ } ''' # --- -# name: test_message_graph[end_of_run].3 - ''' - graph TD; - __start__ --> agent; - action --> agent; - agent -. continue .-> action; - agent -. end .-> __end__; - - ''' -# --- -# name: test_message_graph[end_of_step] - '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior, usually passed in as the first of a sequence\\nof input messages.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a function back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}}, "required": ["content", "tool_call_id"]}}}' -# --- -# name: test_message_graph[end_of_step].1 - '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior, usually passed in as the first of a sequence\\nof input messages.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a function back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}}, "required": ["content", "tool_call_id"]}}}' -# --- -# name: test_message_graph[end_of_step].2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": "agent", - "type": "runnable", - "data": { - "id": [ - "tests", - "test_pregel", - "FakeFuntionChatModel" - ], - "name": "FakeFuntionChatModel" - } - }, - { - "id": "action", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "prebuilt", - "tool_node", - "ToolNode" - ], - "name": "tools" - } - } - ], - "edges": [ - { - "source": "__start__", - "target": "agent" - }, - { - "source": "action", - "target": "agent" - }, - { - "source": "agent", - "target": "action", - "data": "continue", - "conditional": true - }, - { - "source": "agent", - "target": "__end__", - "data": "end", - "conditional": true - } - ] - } - ''' -# --- -# name: test_message_graph[end_of_step].3 +# name: test_message_graph.3 ''' graph TD; __start__ --> agent; @@ -1457,25 +993,7 @@ ''' # --- -# name: test_start_branch_then[end_of_run] - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - __start__ -.-> tool_two_slow; - tool_two_slow --> __end__; - __start__ -.-> tool_two_fast; - tool_two_fast --> __end__; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_start_branch_then[end_of_step] +# name: test_start_branch_then ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; diff --git a/tests/__snapshots__/test_pregel_async.ambr b/tests/__snapshots__/test_pregel_async.ambr index 2d68a82a5..809ff8731 100644 --- a/tests/__snapshots__/test_pregel_async.ambr +++ b/tests/__snapshots__/test_pregel_async.ambr @@ -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__ | diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 6bceb865e..2429cf63d 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -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) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index aedc89a80..c4f8e1e74 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2,6 +2,7 @@ import json import operator import time import warnings +from collections import Counter from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from typing import Annotated, Any, Generator, Literal, Optional, TypedDict, Union @@ -16,7 +17,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 +292,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 +470,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 +616,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 +631,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,13 +672,8 @@ 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: - add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) +def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: + adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: if input > 10: @@ -701,13 +682,12 @@ def test_invoke_checkpoint_sqlite( one = ( Channel.subscribe_to(["input"]).join(["total"]) - | add_one + | adder | Channel.write_to("output", "total") | raise_if_above_10 ) with SqliteSaver.from_conn_string(":memory:") as memory: - memory.at = checkpoint_at app = Pregel( nodes={"one": one}, channels={ @@ -722,10 +702,11 @@ def test_invoke_checkpoint_sqlite( thread_1 = {"configurable": {"thread_id": "1"}} # total starts out as 0, so output is 0+2=2 - assert app.invoke(2, thread_1) == 2 + assert app.invoke(2, thread_1, debug=1) == 2 state = app.get_state(thread_1) assert state is not None assert state.values.get("total") == 2 + assert state.next == () assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"] # total is now 2, so output is 2+3=5 assert app.invoke(3, thread_1) == 5 @@ -736,18 +717,26 @@ def test_invoke_checkpoint_sqlite( # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): app.invoke(4, thread_1) - # checkpoint is not updated + # checkpoint is updated with new input state = app.get_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 app.invoke(2, thread_1) == 9 + state = app.get_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 app.invoke(5, thread_2) == 5 + assert app.invoke(5, thread_2, debug=True) == 5 state = app.get_state({"configurable": {"thread_id": "1"}}) assert state is not None - assert state.values.get("total") == 7 - assert state.next == () + assert state.values.get("total") == 16 + assert state.next == (), "checkpoint of other thread not touched" state = app.get_state(thread_2) assert state is not None assert state.values.get("total") == 5 @@ -756,8 +745,12 @@ def test_invoke_checkpoint_sqlite( assert len(list(app.get_state_history(thread_1, limit=1))) == 1 # list all checkpoints for thread 1 thread_1_history = [c for c in app.get_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"] @@ -769,10 +762,10 @@ def test_invoke_checkpoint_sqlite( ) 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 ( memory.get(thread_1_history[0].config)["ts"] @@ -789,8 +782,20 @@ def test_invoke_checkpoint_sqlite( thread_1_next_config["configurable"]["thread_ts"] > thread_1_history[0].config["configurable"]["thread_ts"] ) + # update makes new checkpoint child of the previous one + assert ( + app.get_state(thread_1_next_config).parent_config + == thread_1_history[1].config + ) # 1 more checkpoint in history - assert len(list(app.get_state_history(thread_1))) == 3 + assert len(list(app.get_state_history(thread_1))) == 8 + assert Counter( + c.metadata["source"] for c in app.get_state_history(thread_1) + ) == { + "update": 1, + "input": 4, + "loop": 3, + } # the latest checkpoint is the updated one assert app.get_state(thread_1) == app.get_state(thread_1_next_config) @@ -992,12 +997,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 +1199,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"}} @@ -1228,6 +1228,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 0}, ) assert ( app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ @@ -1261,6 +1262,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1344,12 +1346,13 @@ def test_conditional_graph( }, next=(), config=app_w_interrupt.checkpointer.get_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"}} @@ -1379,6 +1382,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 0}, ) app_w_interrupt.update_state( @@ -1406,6 +1410,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1489,12 +1494,13 @@ def test_conditional_graph( }, next=(), config=app_w_interrupt.checkpointer.get_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"}} @@ -1524,6 +1530,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 0}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1661,12 +1668,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 @@ -1833,7 +1835,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"}} @@ -1860,6 +1862,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 1}, ) app_w_interrupt.update_state( @@ -1885,6 +1888,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1943,12 +1947,13 @@ def test_conditional_state_graph( }, next=(), config=app_w_interrupt.checkpointer.get_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"], debug=True, ) @@ -1977,6 +1982,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 1}, ) app_w_interrupt.update_state( @@ -2002,6 +2008,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2060,11 +2067,12 @@ def test_conditional_state_graph( }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 5}, ) # test w interrupt before all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before="*", debug=True, ) @@ -2082,6 +2090,7 @@ def test_conditional_state_graph( }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 0}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2104,6 +2113,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2142,6 +2152,7 @@ def test_conditional_state_graph( }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2158,7 +2169,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"}} @@ -2186,6 +2197,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2224,6 +2236,7 @@ def test_conditional_state_graph( }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2778,12 +2791,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 @@ -3002,7 +3011,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"}} @@ -3042,6 +3051,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 1}, ) # modify ai message @@ -3067,6 +3077,7 @@ def test_message_graph( ], next=("action",), config=next_config, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3132,6 +3143,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 4}, ) app_w_interrupt.update_state( @@ -3167,10 +3179,11 @@ def test_message_graph( ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 5}, ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["action"], ) config = {"configurable": {"thread_id": "2"}} @@ -3212,6 +3225,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 1}, ) # modify ai message @@ -3240,6 +3254,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3305,6 +3320,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "loop", "step": 4}, ) app_w_interrupt.update_state( @@ -3340,6 +3356,7 @@ def test_message_graph( ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 5}, ) # add an extra message as if it came from "action" node @@ -3375,6 +3392,7 @@ def test_message_graph( ], next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={"source": "update", "step": 6}, ) @@ -3445,12 +3463,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 @@ -3484,7 +3497,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"] ) @@ -3503,6 +3515,8 @@ def test_start_branch_then( values={"my_key": "value", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={"source": "loop", "step": 0}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3513,6 +3527,8 @@ def test_start_branch_then( values={"my_key": "value slow", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={"source": "loop", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3525,6 +3541,8 @@ def test_start_branch_then( values={"my_key": "value", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={"source": "loop", "step": 0}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3535,13 +3553,47 @@ def test_start_branch_then( values={"my_key": "value fast", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={"source": "loop", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + + thread3 = {"configurable": {"thread_id": "3"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + metadata={"source": "loop", "step": 0}, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + # update state + tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "valuekey", "market": "US"}, + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + metadata={"source": "update", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + next=(), + config=tool_two.checkpointer.get_tuple(thread3).config, + metadata={"source": "loop", "step": 2}, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, ) -@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 @@ -3588,254 +3640,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"] @@ -3855,6 +3778,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={"source": "loop", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3865,6 +3790,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared slow finished", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={"source": "loop", "step": 3}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3877,6 +3804,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={"source": "loop", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3887,10 +3816,11 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared fast finished", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={"source": "loop", "step": 3}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_after=["prepare"] ) @@ -3909,6 +3839,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={"source": "loop", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3919,6 +3851,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared slow finished", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={"source": "loop", "step": 3}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3931,6 +3865,8 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={"source": "loop", "step": 1}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3941,15 +3877,12 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - values={"my_key": "value prepared fast finished", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={"source": "loop", "step": 3}, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) -@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]: @@ -4015,7 +3948,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"}} @@ -4036,12 +3969,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]]] @@ -4111,7 +4040,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"}} @@ -4132,12 +4061,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 @@ -4215,7 +4140,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"}} @@ -4236,12 +4161,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]: @@ -4310,7 +4230,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"}} diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e453460c9..d0cf8a161 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -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"}}