From f3049081025f02c6d661584d5f35f50b4c2ba043 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 3 May 2024 15:40:30 -0700 Subject: [PATCH 1/7] Add metadata to checkpoints - not yet used in this PR --- langgraph/checkpoint/aiosqlite.py | 69 ++++++++++++++--------------- langgraph/checkpoint/base.py | 13 +++++- langgraph/checkpoint/memory.py | 32 ++++++++++---- langgraph/checkpoint/sqlite.py | 73 +++++++++++++++++-------------- langgraph/pregel/__init__.py | 15 +++++-- langgraph/pregel/types.py | 2 + tests/test_pregel.py | 39 +++++++++++++++++ tests/test_pregel_async.py | 29 ++++++++++++ 8 files changed, 188 insertions(+), 84 deletions(-) diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index 62818dc16..e0758d84d 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -1,7 +1,7 @@ import asyncio from contextlib import AbstractAsyncContextManager from types import TracebackType -from typing import AsyncIterator, Optional +from typing import Any, AsyncIterator, Optional import aiosqlite from langchain_core.runnables import RunnableConfig @@ -130,6 +130,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 +156,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 +166,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 None, + { + "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 +190,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 None, + { + "configurable": { + "thread_id": value[0], + "thread_ts": value[2], } - if value[2] - else None - ), + } + if value[2] + else None, ) async def alist( @@ -224,9 +223,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 +240,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 None, + {"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: Optional[dict[str, Any]] = None, ) -> RunnableConfig: """Save a checkpoint to the database asynchronously. @@ -274,12 +270,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) if metadata is not None else None, ), ): await self.conn.commit() diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 86d71c17f..6bfc79885 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -83,6 +83,7 @@ class CheckpointAt(StrEnum): class CheckpointTuple(NamedTuple): config: RunnableConfig checkpoint: Checkpoint + metadata: Optional[dict[str, Any]] parent_config: Optional[RunnableConfig] = None @@ -139,7 +140,12 @@ class BaseCheckpointSaver(ABC): ) -> Iterator[CheckpointTuple]: raise NotImplementedError - def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig: + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: dict[str, Any], + ) -> RunnableConfig: raise NotImplementedError async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: @@ -159,6 +165,9 @@ class BaseCheckpointSaver(ABC): raise NotImplementedError async def aput( - self, config: RunnableConfig, checkpoint: Checkpoint + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: dict[str, Any], ) -> RunnableConfig: raise NotImplementedError diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 055f92d54..dab21794f 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -1,6 +1,6 @@ import asyncio from collections import defaultdict -from typing import AsyncIterator, Iterator, Optional +from typing import Any, AsyncIterator, Iterator, Optional from langchain_core.runnables import RunnableConfig @@ -39,7 +39,7 @@ 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, @@ -66,16 +66,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 +104,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 +113,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: dict[str, Any] = None, + ) -> 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 +135,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 or {}), + ) + } ) return { "configurable": { diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index e4c3be2a3..f45c28be4 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -146,6 +146,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 +212,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 +222,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 None, + { + "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 +246,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 None, + { + "configurable": { + "thread_id": value[0], + "thread_ts": value[2], } - if value[2] - else None - ), + } + if value[2] + else None, ) def list( @@ -289,9 +288,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 +306,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 None, + { + "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: dict[str, Any] = None, + ) -> RunnableConfig: """Save a checkpoint to the database. This method saves a checkpoint to the SQLite database. The checkpoint is associated @@ -332,6 +335,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. @@ -347,12 +351,13 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): """ with 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) if metadata is not None else None, ), ) return { diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d83eb906b..c29f55b86 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -343,6 +343,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, + saved.metadata or {}, ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: @@ -361,6 +362,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, + saved.metadata or {}, ) def get_state_history( @@ -374,7 +376,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 +387,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, + metadata or {}, parent_config, ) @@ -399,9 +402,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 +416,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, + metadata or {}, parent_config, ) diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index 7c6075c98..4359b5994 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -25,6 +25,8 @@ class StateSnapshot(NamedTuple): """Nodes to execute in the next step, if any""" config: RunnableConfig """Config used to fetch this snapshot""" + metadata: dict[str, Any] + """Metadata associated with this snapshot""" parent_config: Optional[RunnableConfig] = None """Config used to fetch the parent snapshot, if any""" diff --git a/tests/test_pregel.py b/tests/test_pregel.py index c8f20eda6..2ba1d1d77 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1228,6 +1228,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) 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={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1344,6 +1346,7 @@ def test_conditional_graph( }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # test state get/update methods with interrupt_before @@ -1379,6 +1382,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) 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={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1489,6 +1494,7 @@ def test_conditional_graph( }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # test re-invoke to continue with interrupt_before @@ -1524,6 +1530,7 @@ def test_conditional_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1860,6 +1867,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) app_w_interrupt.update_state( @@ -1885,6 +1893,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1943,6 +1952,7 @@ def test_conditional_state_graph( }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # test state get/update methods with interrupt_before @@ -1977,6 +1987,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) app_w_interrupt.update_state( @@ -2002,6 +2013,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2060,6 +2072,7 @@ def test_conditional_state_graph( }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # test w interrupt before all @@ -2082,6 +2095,7 @@ def test_conditional_state_graph( }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2104,6 +2118,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2142,6 +2157,7 @@ def test_conditional_state_graph( }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2186,6 +2202,7 @@ def test_conditional_state_graph( }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2224,6 +2241,7 @@ def test_conditional_state_graph( }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3042,6 +3060,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # modify ai message @@ -3067,6 +3086,7 @@ def test_message_graph( ], next=("action",), config=next_config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3132,6 +3152,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) app_w_interrupt.update_state( @@ -3167,6 +3188,7 @@ def test_message_graph( ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) app_w_interrupt = workflow.compile( @@ -3212,6 +3234,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # modify ai message @@ -3240,6 +3263,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3305,6 +3329,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) app_w_interrupt.update_state( @@ -3340,6 +3365,7 @@ def test_message_graph( ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) # add an extra message as if it came from "action" node @@ -3375,6 +3401,7 @@ def test_message_graph( ], next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) @@ -3503,6 +3530,7 @@ def test_start_branch_then( values={"my_key": "value", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={}, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3513,6 +3541,7 @@ def test_start_branch_then( values={"my_key": "value slow", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, + metadata={}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3525,6 +3554,7 @@ def test_start_branch_then( values={"my_key": "value", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={}, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3535,6 +3565,7 @@ def test_start_branch_then( values={"my_key": "value fast", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, + metadata={}, ) @@ -3855,6 +3886,7 @@ 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={}, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3865,6 +3897,7 @@ 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={}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3877,6 +3910,7 @@ 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={}, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3887,6 +3921,7 @@ 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={}, ) with SqliteSaver.from_conn_string(":memory:") as saver: @@ -3909,6 +3944,7 @@ 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={}, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3919,6 +3955,7 @@ 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={}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3931,6 +3968,7 @@ 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={}, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3941,6 +3979,7 @@ 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={}, ) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e453460c9..79e4b3e33 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1306,6 +1306,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) await app_w_interrupt.aupdate_state( @@ -1333,6 +1334,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1416,6 +1418,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) # test state get/update methods with interrupt_before @@ -1454,6 +1457,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) await app_w_interrupt.aupdate_state( @@ -1481,6 +1485,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1564,6 +1569,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) # test re-invoke to continue with interrupt_before @@ -1602,6 +1608,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1924,6 +1931,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) await app_w_interrupt.aupdate_state( @@ -1949,6 +1957,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2007,6 +2016,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) # test state get/update methods with interrupt_before @@ -2043,6 +2053,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) await app_w_interrupt.aupdate_state( @@ -2068,6 +2079,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2126,6 +2138,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) @@ -2734,6 +2747,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None: ], next=("action",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + metadata={}, ) # modify ai message @@ -2761,6 +2775,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None: ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2813,6 +2828,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None: ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) await app_w_interrupt.aupdate_state( @@ -2846,6 +2862,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None: ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={}, ) @@ -2968,6 +2985,7 @@ 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={}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread1, debug=1) == { @@ -2978,6 +2996,7 @@ async def test_start_branch_then( values={"my_key": "value slow", "market": "DE"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + metadata={}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -2990,6 +3009,7 @@ 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={}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread2, debug=1) == { @@ -3000,6 +3020,7 @@ async def test_start_branch_then( values={"my_key": "value fast", "market": "US"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + metadata={}, ) @@ -3308,6 +3329,7 @@ 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={}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread1, debug=1) == { @@ -3318,6 +3340,7 @@ 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={}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3330,6 +3353,7 @@ 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={}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread2, debug=1) == { @@ -3340,6 +3364,7 @@ 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={}, ) async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: @@ -3362,6 +3387,7 @@ 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={}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread1, debug=1) == { @@ -3372,6 +3398,7 @@ 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={}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3384,6 +3411,7 @@ 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={}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread2, debug=1) == { @@ -3394,6 +3422,7 @@ 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={}, ) From 3ff3def62b4c655906b77c048cd50dfb8b73caee Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 3 May 2024 15:55:59 -0700 Subject: [PATCH 2/7] Remove option to only checkpoint at end of run - Now that checkpoint at end of each step adds no latency there is not point to keep this - This will make it easier to add future features --- langgraph/checkpoint/__init__.py | 2 - langgraph/checkpoint/aiosqlite.py | 4 +- langgraph/checkpoint/base.py | 14 - langgraph/checkpoint/memory.py | 4 +- langgraph/checkpoint/sqlite.py | 4 +- langgraph/graph/__init__.py | 4 +- langgraph/pregel/__init__.py | 62 +-- tests/__snapshots__/test_pregel.ambr | 518 +-------------------- tests/__snapshots__/test_pregel_async.ambr | 78 +--- tests/memory_assert.py | 6 +- tests/test_pregel.py | 468 ++++++------------- tests/test_pregel_async.py | 466 ++++++------------ 12 files changed, 303 insertions(+), 1327 deletions(-) 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 e0758d84d..314448b01 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -10,7 +10,6 @@ from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, CheckpointTuple, SerializerProtocol, ) @@ -80,9 +79,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): conn: aiosqlite.Connection, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ): - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.conn = conn self.lock = asyncio.Lock() self.is_setup = False diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 6bfc79885..70734f1e6 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -14,7 +14,6 @@ from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer -from langgraph.utils import StrEnum class Checkpoint(TypedDict): @@ -71,15 +70,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: ) -class CheckpointAt(StrEnum): - """When to take a checkpoint.""" - - END_OF_STEP = "end_of_step" - """Take a checkpoint at the end of each step.""" - END_OF_RUN = "end_of_run" - """Take a checkpoint at the end of the run.""" - - class CheckpointTuple(NamedTuple): config: RunnableConfig checkpoint: Checkpoint @@ -107,18 +97,14 @@ CheckpointThreadTs = ConfigurableFieldSpec( class BaseCheckpointSaver(ABC): - at: CheckpointAt = CheckpointAt.END_OF_STEP - serde: SerializerProtocol = JsonPlusSerializer() def __init__( self, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: self.serde = serde or self.serde - self.at = at or self.at @property def config_specs(self) -> list[ConfigurableFieldSpec]: diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index dab21794f..81eba2438 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -7,7 +7,6 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, CheckpointTuple, SerializerProtocol, ) @@ -45,9 +44,8 @@ class MemorySaver(BaseCheckpointSaver): self, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.storage = defaultdict(dict) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index f45c28be4..3c2b63ca3 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -10,7 +10,6 @@ from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, CheckpointTuple, SerializerProtocol, ) @@ -90,9 +89,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): conn: sqlite3.Connection, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.conn = conn self.is_setup = False 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 c29f55b86..ffa14304c 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, ) @@ -769,9 +768,7 @@ class Pregel( yield from map_output_updates(output_keys, next_tasks) # save end of step checkpoint - if self.checkpointer is not None and ( - self.checkpointer.at == CheckpointAt.END_OF_STEP - ): + if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) checkpoint_config = self.checkpointer.put( checkpoint_config, checkpoint @@ -799,33 +796,6 @@ class Pregel( # set final channel values as run output run_manager.on_chain_end(read_channels(channels, output_keys)) - - # save end of run checkpoint - if ( - self.checkpointer is not None - and self.checkpointer.at == CheckpointAt.END_OF_RUN - ): - checkpoint = create_checkpoint(checkpoint, channels) - executor.submit( - self.checkpointer.put(checkpoint_config, checkpoint) - ) - checkpoint_config = { - "configurable": { - "thread_id": checkpoint_config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], - } - } - if stream_mode == "debug": - yield map_debug_checkpoint( - step, - checkpoint_config, - channels, - self.stream_channels_asis, - ) - elif self.checkpointer is None and stream_mode == "debug": - yield map_debug_checkpoint( - step, None, channels, self.stream_channels_asis - ) except BaseException as e: run_manager.on_chain_error(e) raise @@ -1035,9 +1005,7 @@ class Pregel( yield chunk # save end of step checkpoint - if self.checkpointer is not None and ( - self.checkpointer.at == CheckpointAt.END_OF_STEP - ): + if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) checkpoint_config = await self.checkpointer.aput( checkpoint_config, checkpoint @@ -1065,32 +1033,6 @@ class Pregel( # set final channel values as run output await run_manager.on_chain_end(read_channels(channels, output_keys)) - - # save end of run checkpoint - if ( - self.checkpointer is not None - and self.checkpointer.at == CheckpointAt.END_OF_RUN - ): - checkpoint = create_checkpoint(checkpoint, channels) - tasks.append( - asyncio.create_task( - self.checkpointer.aput(checkpoint_config, checkpoint) - ) - ) - checkpoint_config = { - "configurable": { - "thread_id": checkpoint_config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], - } - } - if stream_mode == "debug": - yield map_debug_checkpoint( - step, checkpoint_config, channels, self.stream_channels_asis - ) - elif self.checkpointer is None and stream_mode == "debug": - yield map_debug_checkpoint( - step, None, channels, self.stream_channels_asis - ) except BaseException as e: await run_manager.on_chain_error(e) raise diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 1b1866624..1cf4bd12e 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; @@ -1456,25 +992,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..625b6f12e 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -3,7 +3,6 @@ from typing import Any, Optional from langgraph.checkpoint.base import ( Checkpoint, - CheckpointAt, SerializerProtocol, copy_checkpoint, ) @@ -21,17 +20,14 @@ class NoopSerializer(SerializerProtocol): class MemorySaverAssertImmutable(MemorySaver): serde = NoopSerializer() - at = CheckpointAt.END_OF_STEP - storage_for_copies: defaultdict[str, dict[str, Checkpoint]] def __init__( self, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.storage_for_copies = defaultdict(dict) def put(self, config: dict, checkpoint: Checkpoint) -> None: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 2ba1d1d77..9180e71c9 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -16,7 +16,6 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic -from langgraph.checkpoint.base import CheckpointAt from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.graph import END, Graph from langgraph.graph.message import MessageGraph @@ -292,17 +291,12 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_invoke_two_processes_in_out_interrupt( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one, "two": two}, channels={ @@ -475,12 +469,6 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "step": 1, "payload": {"config": None, "values": {"output": 4, "inbox": []}}, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": {"config": None, "values": {"output": 4, "inbox": []}}, - }, ] @@ -627,10 +615,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) -> None: +def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -645,7 +630,7 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) - | raise_if_above_10 ) - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one}, @@ -686,12 +671,7 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) - assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_invoke_checkpoint_sqlite( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -707,7 +687,6 @@ def test_invoke_checkpoint_sqlite( ) with SqliteSaver.from_conn_string(":memory:") as memory: - memory.at = checkpoint_at app = Pregel( nodes={"one": one}, channels={ @@ -992,12 +971,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_conditional_graph( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_conditional_graph(snapshot: SnapshotAssertion) -> None: from copy import deepcopy from langchain.llms.fake import FakeStreamingListLLM @@ -1199,7 +1173,7 @@ def test_conditional_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1352,7 +1326,7 @@ def test_conditional_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1500,7 +1474,7 @@ def test_conditional_graph( # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1668,12 +1642,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_conditional_state_graph( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish @@ -1840,7 +1809,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1958,7 +1927,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], debug=True, ) @@ -2077,7 +2046,7 @@ def test_conditional_state_graph( # test w interrupt before all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before="*", debug=True, ) @@ -2174,7 +2143,7 @@ def test_conditional_state_graph( # test w interrupt after all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after="*", ) config = {"configurable": {"thread_id": "4"}} @@ -2796,12 +2765,8 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) def test_message_graph( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, deterministic_uuids: MockerFixture, ) -> None: from copy import deepcopy @@ -3020,7 +2985,7 @@ def test_message_graph( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -3192,7 +3157,7 @@ def test_message_graph( ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["action"], ) config = {"configurable": {"thread_id": "2"}} @@ -3472,12 +3437,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_start_branch_then( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_start_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -3511,7 +3471,6 @@ def test_start_branch_then( } with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] ) @@ -3569,10 +3528,7 @@ def test_start_branch_then( ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -> None: +def test_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -3619,254 +3575,125 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - } with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at - # test stream_mode=debug tool_two = tool_two_graph.compile(checkpointer=saver) thread10 = {"configurable": {"thread_id": "10"}} - if checkpoint_at is CheckpointAt.END_OF_RUN: - assert [ - *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": None, - "values": {"my_key": "value", "market": "DE"}, + assert [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", }, }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": None, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": None, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": None, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 4, - "payload": { - "config": { - "configurable": { - "thread_id": "10", - "thread_ts": AnyStr(), - } - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] - else: - assert [ - *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] + }, + ] tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] @@ -3925,7 +3752,6 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - ) with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_after=["prepare"] ) @@ -3983,12 +3809,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_in_one_fan_out_state_graph_waiting_edge( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -4054,7 +3875,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -4075,12 +3896,8 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -4150,7 +3967,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -4171,12 +3988,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError @@ -4254,7 +4067,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -4275,12 +4088,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - checkpoint_at: CheckpointAt, -) -> None: +def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -4349,7 +4157,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 79e4b3e33..32b19b279 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -25,7 +25,6 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver -from langgraph.checkpoint.base import CheckpointAt from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph from langgraph.prebuilt.chat_agent_executor import ( @@ -270,17 +269,12 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_invoke_two_processes_in_out_interrupt( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one, "two": two}, channels={ @@ -457,12 +451,6 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "step": 1, "payload": {"config": None, "values": {"output": 4, "inbox": []}}, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": {"config": None, "values": {"output": 4, "inbox": []}}, - }, ] @@ -613,12 +601,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) assert await app.ainvoke(2) == [3, 3] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_invoke_checkpoint( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +async def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -633,7 +616,7 @@ async def test_invoke_checkpoint( | raise_if_above_10 ) - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one}, @@ -674,12 +657,7 @@ async def test_invoke_checkpoint( assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_invoke_checkpoint_aiosqlite( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -695,7 +673,6 @@ async def test_invoke_checkpoint_aiosqlite( ) async with AsyncSqliteSaver.from_conn_string(":memory:") as memory: - memory.at = checkpoint_at app = Pregel( nodes={"one": one}, channels={ @@ -1003,10 +980,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup_async.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: +async def test_conditional_graph() -> None: from copy import deepcopy from langchain.llms.fake import FakeStreamingListLLM @@ -1274,7 +1248,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1424,7 +1398,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1575,7 +1549,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1702,10 +1676,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: +async def test_conditional_graph_state() -> None: from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish @@ -1899,7 +1870,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2022,7 +1993,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -2537,10 +2508,7 @@ async def test_prebuilt_chat() -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_message_graph(checkpoint_at: CheckpointAt) -> None: +async def test_message_graph() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.agents import AgentAction @@ -2709,7 +2677,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2938,12 +2906,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_start_branch_then( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -2966,7 +2929,6 @@ async def test_start_branch_then( } async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] ) @@ -3024,12 +2986,7 @@ async def test_start_branch_then( ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_branch_then( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +async def test_branch_then() -> None: pass class State(TypedDict): @@ -3060,256 +3017,126 @@ async def test_branch_then( } async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at - # test stream_mode=debug tool_two = tool_two_graph.compile(checkpointer=saver) thread10 = {"configurable": {"thread_id": "10"}} - if checkpoint_at is CheckpointAt.END_OF_RUN: - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": None, - "values": {"my_key": "value", "market": "DE"}, + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", }, }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": None, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": None, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": None, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 4, - "payload": { - "config": { - "configurable": { - "thread_id": "10", - "thread_ts": AnyStr(), - } - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] - else: - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] + }, + ] tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] @@ -3368,7 +3195,6 @@ async def test_branch_then( ) async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_after=["prepare"] ) @@ -3426,12 +3252,7 @@ async def test_branch_then( ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_in_one_fan_out_state_graph_waiting_edge( - checkpoint_at: CheckpointAt, -) -> None: +async def test_in_one_fan_out_state_graph_waiting_edge() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3495,7 +3316,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3519,12 +3340,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -3593,7 +3410,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3617,12 +3434,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError @@ -3700,7 +3513,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3724,12 +3537,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - checkpoint_at: CheckpointAt, -) -> None: +async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3798,7 +3606,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} From f160f8291282b3e7d534b924fad1c32b3b05efc0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 3 May 2024 16:42:51 -0700 Subject: [PATCH 3/7] Add checkpoint metadata fields - source: input, update or loop - step: int - make step counter continue from previous last step --- langgraph/checkpoint/aiosqlite.py | 13 +++--- langgraph/checkpoint/base.py | 23 +++++++-- langgraph/checkpoint/memory.py | 14 ++++-- langgraph/checkpoint/sqlite.py | 11 +++-- langgraph/pregel/__init__.py | 62 ++++++++++++++++-------- tests/memory_assert.py | 10 +++- tests/test_pregel.py | 78 +++++++++++++++---------------- tests/test_pregel_async.py | 58 +++++++++++------------ 8 files changed, 160 insertions(+), 109 deletions(-) diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index 314448b01..c9a322561 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -1,7 +1,7 @@ import asyncio from contextlib import AbstractAsyncContextManager from types import TracebackType -from typing import Any, AsyncIterator, Optional +from typing import AsyncIterator, Optional import aiosqlite from langchain_core.runnables import RunnableConfig @@ -10,6 +10,7 @@ from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, + CheckpointMetadata, CheckpointTuple, SerializerProtocol, ) @@ -164,7 +165,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): return CheckpointTuple( config, self.serde.loads(value[0]), - self.serde.loads(value[2]) if value[2] is not None else None, + self.serde.loads(value[2]) if value[2] is not None else {}, { "configurable": { "thread_id": config["configurable"]["thread_id"], @@ -188,7 +189,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): } }, self.serde.loads(value[3]), - self.serde.loads(value[4]) if value[4] is not None else None, + self.serde.loads(value[4]) if value[4] is not None else {}, { "configurable": { "thread_id": value[0], @@ -242,7 +243,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): yield CheckpointTuple( {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, self.serde.loads(value), - self.serde.loads(metadata) if metadata is not None 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, @@ -252,7 +253,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): self, config: RunnableConfig, checkpoint: Checkpoint, - metadata: Optional[dict[str, Any]] = None, + metadata: CheckpointMetadata, ) -> RunnableConfig: """Save a checkpoint to the database asynchronously. @@ -274,7 +275,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): checkpoint["ts"], config["configurable"].get("thread_ts"), self.serde.dumps(checkpoint), - self.serde.dumps(metadata) if metadata is not None else None, + self.serde.dumps(metadata), ), ): await self.conn.commit() diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 70734f1e6..089d7cc19 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, @@ -16,6 +17,22 @@ from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer +# 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): """State snapshot at a given point in time.""" @@ -73,7 +90,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: class CheckpointTuple(NamedTuple): config: RunnableConfig checkpoint: Checkpoint - metadata: Optional[dict[str, Any]] + metadata: CheckpointMetadata parent_config: Optional[RunnableConfig] = None @@ -130,7 +147,7 @@ class BaseCheckpointSaver(ABC): self, config: RunnableConfig, checkpoint: Checkpoint, - metadata: dict[str, Any], + metadata: CheckpointMetadata, ) -> RunnableConfig: raise NotImplementedError @@ -154,6 +171,6 @@ class BaseCheckpointSaver(ABC): self, config: RunnableConfig, checkpoint: Checkpoint, - metadata: dict[str, Any], + metadata: CheckpointMetadata, ) -> RunnableConfig: raise NotImplementedError diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 81eba2438..fb49bab58 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -1,12 +1,13 @@ import asyncio from collections import defaultdict -from typing import Any, AsyncIterator, Iterator, Optional +from typing import AsyncIterator, Iterator, Optional from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, + CheckpointMetadata, CheckpointTuple, SerializerProtocol, ) @@ -118,7 +119,7 @@ class MemorySaver(BaseCheckpointSaver): self, config: RunnableConfig, checkpoint: Checkpoint, - metadata: dict[str, Any] = None, + metadata: CheckpointMetadata, ) -> RunnableConfig: """Save a checkpoint to the in-memory storage. @@ -136,7 +137,7 @@ class MemorySaver(BaseCheckpointSaver): { checkpoint["ts"]: ( self.serde.dumps(checkpoint), - self.serde.dumps(metadata or {}), + self.serde.dumps(metadata), ) } ) @@ -184,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 3c2b63ca3..54feff576 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -10,6 +10,7 @@ from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, + CheckpointMetadata, CheckpointTuple, SerializerProtocol, ) @@ -220,7 +221,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): return CheckpointTuple( config, self.serde.loads(value[0]), - self.serde.loads(value[2]) if value[2] is not None else None, + self.serde.loads(value[2]) if value[2] is not None else {}, { "configurable": { "thread_id": config["configurable"]["thread_id"], @@ -244,7 +245,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): } }, self.serde.loads(value[3]), - self.serde.loads(value[4]) if value[4] is not None else None, + self.serde.loads(value[4]) if value[4] is not None else {}, { "configurable": { "thread_id": value[0], @@ -308,7 +309,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): yield CheckpointTuple( {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, self.serde.loads(value), - self.serde.loads(metadata) if metadata is not None else None, + self.serde.loads(metadata) if metadata is not None else {}, { "configurable": { "thread_id": thread_id, @@ -323,7 +324,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): self, config: RunnableConfig, checkpoint: Checkpoint, - metadata: dict[str, Any] = None, + metadata: CheckpointMetadata, ) -> RunnableConfig: """Save a checkpoint to the database. @@ -355,7 +356,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): checkpoint["ts"], config["configurable"].get("thread_ts"), self.serde.dumps(checkpoint), - self.serde.dumps(metadata) if metadata is not None else None, + self.serde.dumps(metadata), ), ) return { diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index ffa14304c..de56f8925 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -342,7 +342,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - saved.metadata or {}, + saved.metadata, ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: @@ -361,7 +361,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - saved.metadata or {}, + saved.metadata, ) def get_state_history( @@ -386,7 +386,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - metadata or {}, + metadata, parent_config, ) @@ -415,7 +415,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - metadata or {}, + metadata, parent_config, ) @@ -433,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( @@ -482,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) + config, + create_checkpoint(checkpoint, channels), + { + "source": "update", + "step": saved.metadata.get("step", 0) + 1 + if saved.metadata + else None, + }, ) async def aupdate_state( @@ -495,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( @@ -544,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) + config, + create_checkpoint(checkpoint, channels), + { + "source": "update", + "step": saved.metadata.get("step", 0) + 1 if saved else None, + }, ) def _defaults( @@ -638,10 +650,12 @@ class Pregel( 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 + saved = ( + self.checkpointer.get_tuple(checkpoint_config) + if self.checkpointer + else None ) - checkpoint = checkpoint or empty_checkpoint() + checkpoint = saved.checkpoint if saved else empty_checkpoint() # create channels from checkpoint with ChannelsManager( self.channels, checkpoint @@ -667,7 +681,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, # with channel updates 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 ) @@ -771,7 +787,9 @@ class Pregel( if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) checkpoint_config = self.checkpointer.put( - checkpoint_config, checkpoint + checkpoint_config, + checkpoint, + {"source": "loop", "step": step}, ) if stream_mode == "debug": yield map_debug_checkpoint( @@ -865,12 +883,12 @@ class Pregel( 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(checkpoint_config) if self.checkpointer else None ) - checkpoint = checkpoint or empty_checkpoint() + checkpoint = saved.checkpoint if saved else empty_checkpoint() # create channels from checkpoint async with AsyncChannelsManager(self.channels, checkpoint) as channels: # map inputs to channel updates @@ -894,7 +912,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 ) @@ -1008,7 +1028,9 @@ class Pregel( if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) checkpoint_config = await self.checkpointer.aput( - checkpoint_config, checkpoint + checkpoint_config, + checkpoint, + {"source": "loop", "step": step}, ) if stream_mode == "debug": yield map_debug_checkpoint( diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 625b6f12e..2429cf63d 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -3,6 +3,7 @@ from typing import Any, Optional from langgraph.checkpoint.base import ( Checkpoint, + CheckpointMetadata, SerializerProtocol, copy_checkpoint, ) @@ -30,7 +31,12 @@ class MemorySaverAssertImmutable(MemorySaver): 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): @@ -39,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 9180e71c9..edccfdd78 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1202,7 +1202,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) assert ( app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ @@ -1236,7 +1236,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1320,7 +1320,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 4}, ) # test state get/update methods with interrupt_before @@ -1356,7 +1356,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) app_w_interrupt.update_state( @@ -1384,7 +1384,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1468,7 +1468,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 4}, ) # test re-invoke to continue with interrupt_before @@ -1504,7 +1504,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1836,7 +1836,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) app_w_interrupt.update_state( @@ -1862,7 +1862,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1921,7 +1921,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) # test state get/update methods with interrupt_before @@ -1956,7 +1956,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) app_w_interrupt.update_state( @@ -1982,7 +1982,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2041,7 +2041,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) # test w interrupt before all @@ -2064,7 +2064,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2087,7 +2087,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2126,7 +2126,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2171,7 +2171,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2210,7 +2210,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3025,7 +3025,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # modify ai message @@ -3051,7 +3051,7 @@ def test_message_graph( ], next=("action",), config=next_config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3117,7 +3117,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 4}, ) app_w_interrupt.update_state( @@ -3153,7 +3153,7 @@ def test_message_graph( ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) app_w_interrupt = workflow.compile( @@ -3199,7 +3199,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # modify ai message @@ -3228,7 +3228,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3294,7 +3294,7 @@ def test_message_graph( ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 4}, ) app_w_interrupt.update_state( @@ -3330,7 +3330,7 @@ def test_message_graph( ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) # add an extra message as if it came from "action" node @@ -3366,7 +3366,7 @@ def test_message_graph( ], next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 6}, ) @@ -3489,7 +3489,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3500,7 +3500,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value slow", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3513,7 +3513,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3524,7 +3524,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value fast", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) @@ -3713,7 +3713,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3724,7 +3724,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared slow finished", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3737,7 +3737,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3748,7 +3748,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared fast finished", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) with SqliteSaver.from_conn_string(":memory:") as saver: @@ -3770,7 +3770,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared", "market": "DE"}, next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert tool_two.invoke(None, thread1, debug=1) == { @@ -3781,7 +3781,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared slow finished", "market": "DE"}, next=(), config=tool_two.checkpointer.get_tuple(thread1).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3794,7 +3794,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared", "market": "US"}, next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert tool_two.invoke(None, thread2, debug=1) == { @@ -3805,7 +3805,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value prepared fast finished", "market": "US"}, next=(), config=tool_two.checkpointer.get_tuple(thread2).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 32b19b279..2a79de197 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1280,7 +1280,7 @@ async def test_conditional_graph() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) await app_w_interrupt.aupdate_state( @@ -1308,7 +1308,7 @@ async def test_conditional_graph() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 1}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1392,7 +1392,7 @@ async def test_conditional_graph() -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 4}, ) # test state get/update methods with interrupt_before @@ -1431,7 +1431,7 @@ async def test_conditional_graph() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) await app_w_interrupt.aupdate_state( @@ -1459,7 +1459,7 @@ async def test_conditional_graph() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 1}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1543,7 +1543,7 @@ async def test_conditional_graph() -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 4}, ) # test re-invoke to continue with interrupt_before @@ -1582,7 +1582,7 @@ async def test_conditional_graph() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1902,7 +1902,7 @@ async def test_conditional_graph_state() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) await app_w_interrupt.aupdate_state( @@ -1928,7 +1928,7 @@ async def test_conditional_graph_state() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1987,7 +1987,7 @@ async def test_conditional_graph_state() -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) # test state get/update methods with interrupt_before @@ -2024,7 +2024,7 @@ async def test_conditional_graph_state() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) await app_w_interrupt.aupdate_state( @@ -2050,7 +2050,7 @@ async def test_conditional_graph_state() -> None: }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2109,7 +2109,7 @@ async def test_conditional_graph_state() -> None: }, next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) @@ -2715,7 +2715,7 @@ async def test_message_graph() -> None: ], next=("action",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # modify ai message @@ -2743,7 +2743,7 @@ async def test_message_graph() -> None: ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 2}, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2796,7 +2796,7 @@ async def test_message_graph() -> None: ], next=("action",), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "loop", "step": 4}, ) await app_w_interrupt.aupdate_state( @@ -2830,7 +2830,7 @@ async def test_message_graph() -> None: ], next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={}, + metadata={"source": "update", "step": 5}, ) @@ -2947,7 +2947,7 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value", "market": "DE"}, next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread1, debug=1) == { @@ -2958,7 +2958,7 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value slow", "market": "DE"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -2971,7 +2971,7 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value", "market": "US"}, next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - metadata={}, + metadata={"source": "loop", "step": 0}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread2, debug=1) == { @@ -2982,7 +2982,7 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: values={"my_key": "value fast", "market": "US"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) @@ -3156,7 +3156,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared", "market": "DE"}, next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread1, debug=1) == { @@ -3167,7 +3167,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared slow finished", "market": "DE"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3180,7 +3180,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared", "market": "US"}, next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread2, debug=1) == { @@ -3191,7 +3191,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared fast finished", "market": "US"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: @@ -3213,7 +3213,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared", "market": "DE"}, next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread1, debug=1) == { @@ -3224,7 +3224,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared slow finished", "market": "DE"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) thread2 = {"configurable": {"thread_id": "2"}} @@ -3237,7 +3237,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared", "market": "US"}, next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - metadata={}, + metadata={"source": "loop", "step": 1}, ) # resume, for same result as above assert await tool_two.ainvoke(None, thread2, debug=1) == { @@ -3248,7 +3248,7 @@ async def test_branch_then() -> None: values={"my_key": "value prepared fast finished", "market": "US"}, next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - metadata={}, + metadata={"source": "loop", "step": 3}, ) From 611ecdb1cdac1c6f50f6f19b67db18b0c48d384c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 6 May 2024 08:36:20 -0700 Subject: [PATCH 4/7] Checkpoint inputs before starting the first step for easier error recovery - this enables easier retrying, for any error just do .invoke(None, config) no matter which step the error happened on --- langgraph/checkpoint/sqlite.py | 12 ++- langgraph/pregel/__init__.py | 132 +++++++++++++++++++++++++-------- langgraph/pregel/types.py | 4 +- tests/test_pregel.py | 47 ++++++++---- tests/test_pregel_async.py | 36 +++++++-- 5 files changed, 175 insertions(+), 56 deletions(-) diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index 54feff576..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 @@ -94,6 +95,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): 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": @@ -115,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 @@ -348,7 +356,7 @@ 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, metadata) VALUES (?, ?, ?, ?, ?)", ( diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index de56f8925..4e5cf13a0 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -342,7 +342,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - saved.metadata, + saved.metadata if saved else None, ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: @@ -361,7 +361,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - saved.metadata, + saved.metadata if saved else None, ) def get_state_history( @@ -623,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"): @@ -656,6 +657,7 @@ class Pregel( else None ) checkpoint = saved.checkpoint if saved else empty_checkpoint() + start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint with ChannelsManager( self.channels, checkpoint @@ -668,6 +670,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 @@ -681,7 +704,6 @@ 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 - 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( @@ -786,21 +808,29 @@ class Pregel( # save end of step checkpoint if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) - checkpoint_config = self.checkpointer.put( - checkpoint_config, - checkpoint, - {"source": "loop", "step": step}, - ) - 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 @@ -824,6 +854,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, @@ -855,7 +891,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"): @@ -889,6 +925,7 @@ class Pregel( else None ) checkpoint = saved.checkpoint if saved else empty_checkpoint() + 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 @@ -899,6 +936,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 @@ -1027,21 +1086,30 @@ class Pregel( # save end of step checkpoint if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) - checkpoint_config = await self.checkpointer.aput( - checkpoint_config, - checkpoint, - {"source": "loop", "step": step}, - ) - 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 @@ -1063,11 +1131,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 4359b5994..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,7 +27,7 @@ class StateSnapshot(NamedTuple): """Nodes to execute in the next step, if any""" config: RunnableConfig """Config used to fetch this snapshot""" - metadata: dict[str, Any] + 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/test_pregel.py b/tests/test_pregel.py index edccfdd78..bf295801f 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 @@ -672,7 +673,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: if input > 10: @@ -681,7 +682,7 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: one = ( Channel.subscribe_to(["input"]).join(["total"]) - | add_one + | adder | Channel.write_to("output", "total") | raise_if_above_10 ) @@ -701,10 +702,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: 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 @@ -719,14 +721,22 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: 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 @@ -735,8 +745,12 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: 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"] @@ -748,10 +762,10 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: ) 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"] @@ -769,7 +783,14 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: > thread_1_history[0].config["configurable"]["thread_ts"] ) # 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) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 2a79de197..4a01b1631 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, @@ -712,13 +713,21 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: 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 @@ -728,8 +737,12 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: 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"] @@ -744,10 +757,10 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: ] 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" @@ -763,7 +776,14 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: > 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 From 48df2d10847bd636d2f7ad6657cb3b406d523fb9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 6 May 2024 09:56:16 -0700 Subject: [PATCH 5/7] Add assert --- langgraph/channels/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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"], From 251bd9744de670ced70dc93190ebd9ed750498b0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 6 May 2024 11:08:13 -0700 Subject: [PATCH 6/7] Fix lineage of maual state updates --- langgraph/pregel/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 4e5cf13a0..2d50e4496 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -482,7 +482,7 @@ class Pregel( # apply to checkpoint and save _apply_writes(checkpoint, channels, task.writes) return self.checkpointer.put( - config, + saved.config if saved else config, create_checkpoint(checkpoint, channels), { "source": "update", @@ -551,7 +551,7 @@ class Pregel( # apply to checkpoint and save _apply_writes(checkpoint, channels, task.writes) return await self.checkpointer.aput( - config, + saved.config if saved else config, create_checkpoint(checkpoint, channels), { "source": "update", From 48865daf0220e7c4ca1023696a2d6ee1724d5a4b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 6 May 2024 11:38:28 -0700 Subject: [PATCH 7/7] Fix checkpoint lineage for updates/resumes --- langgraph/checkpoint/base.py | 3 +- langgraph/pregel/__init__.py | 20 ++++----- tests/test_pregel.py | 54 +++++++++++++++++++++++- tests/test_pregel_async.py | 79 +++++++++++++++++++++++++++++++++++- 4 files changed, 141 insertions(+), 15 deletions(-) diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 089d7cc19..5d197e19e 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -158,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, *, @@ -166,6 +166,7 @@ class BaseCheckpointSaver(ABC): limit: Optional[int] = None, ) -> AsyncIterator[CheckpointTuple]: raise NotImplementedError + yield async def aput( self, diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 2d50e4496..75321bf46 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -333,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 @@ -341,8 +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,8 +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( @@ -650,13 +650,9 @@ class Pregel( # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one - checkpoint_config = config - saved = ( - self.checkpointer.get_tuple(checkpoint_config) - if self.checkpointer - else None - ) + 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( @@ -918,13 +914,13 @@ class Pregel( # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one - checkpoint_config = config saved = ( - await self.checkpointer.aget_tuple(checkpoint_config) + await self.checkpointer.aget_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 async with AsyncChannelsManager(self.channels, checkpoint) as channels: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index bf295801f..a4e84956d 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -717,7 +717,7 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: # 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 @@ -782,6 +782,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: 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))) == 8 assert Counter( @@ -3511,6 +3516,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3522,6 +3528,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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"}} @@ -3535,6 +3542,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3546,6 +3554,42 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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, ) @@ -3735,6 +3779,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3746,6 +3791,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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"}} @@ -3759,6 +3805,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3770,6 +3817,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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: @@ -3792,6 +3840,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3803,6 +3852,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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"}} @@ -3816,6 +3866,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3827,6 +3878,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: 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, ) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4a01b1631..d0cf8a161 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2926,7 +2926,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: +async def test_start_branch_then() -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -2968,6 +2968,9 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -2979,6 +2982,9 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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"}} @@ -2992,6 +2998,9 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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) == { @@ -3003,6 +3012,50 @@ async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: 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, ) @@ -3177,6 +3230,9 @@ async def test_branch_then() -> None: 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) == { @@ -3188,6 +3244,9 @@ async def test_branch_then() -> None: 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"}} @@ -3201,6 +3260,9 @@ async def test_branch_then() -> None: 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) == { @@ -3212,6 +3274,9 @@ async def test_branch_then() -> None: 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: @@ -3234,6 +3299,9 @@ async def test_branch_then() -> None: 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) == { @@ -3245,6 +3313,9 @@ async def test_branch_then() -> None: 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"}} @@ -3258,6 +3329,9 @@ async def test_branch_then() -> None: 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) == { @@ -3269,6 +3343,9 @@ async def test_branch_then() -> None: 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, )