diff --git a/docs/docs/how-tos/persistence.ipynb b/docs/docs/how-tos/persistence.ipynb index 8a9a6fee2..c68025b9c 100644 --- a/docs/docs/how-tos/persistence.ipynb +++ b/docs/docs/how-tos/persistence.ipynb @@ -739,7 +739,6 @@ " 'id': '1f029ca3-1f5b-6704-8004-820c16b69a5a',\n", " 'channel_versions': {'__start__': '00000000000000000000000000000005.0.5290678567601859', 'messages': '00000000000000000000000000000006.0.3205149138784782', 'branch:to:call_model': '00000000000000000000000000000006.0.14611156755133758'}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000004.0.5736472536395331'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000005.0.1410174088651449'}},\n", " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today?), HumanMessage(content=\"what's my name?\"), AIMessage(content='Your name is Bob.')]},\n", - " 'pending_sends': []\n", " },\n", " metadata={\n", " 'source': 'loop',\n", @@ -856,7 +855,7 @@ " 'id': '1f029ca3-1f5b-6704-8004-820c16b69a5a', \n", " 'channel_versions': {'__start__': '00000000000000000000000000000005.0.5290678567601859', 'messages': '00000000000000000000000000000006.0.3205149138784782', 'branch:to:call_model': '00000000000000000000000000000006.0.14611156755133758'}, \n", " 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000004.0.5736472536395331'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000005.0.1410174088651449'}},\n", - " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\"), AIMessage(content='Your name is Bob.')]}, 'pending_sends': []\n", + " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\"), AIMessage(content='Your name is Bob.')]},\n", " },\n", " metadata={'source': 'loop', 'writes': {'call_model': {'messages': AIMessage(content='Your name is Bob.')}}, 'step': 4, 'parents': {}, 'thread_id': '1'}, \n", " parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1f029ca3-1790-6b0a-8003-baf965b6a38f'}}, \n", @@ -870,8 +869,7 @@ " 'id': '1f029ca3-1790-6b0a-8003-baf965b6a38f', \n", " 'channel_versions': {'__start__': '00000000000000000000000000000005.0.5290678567601859', 'messages': '00000000000000000000000000000005.0.7935064215293443', 'branch:to:call_model': '00000000000000000000000000000005.0.1410174088651449'}, \n", " 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000004.0.5736472536395331'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}}, \n", - " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\")], 'branch:to:call_model': None}, \n", - " 'pending_sends': []\n", + " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\")], 'branch:to:call_model': None}\n", " }, \n", " metadata={'source': 'loop', 'writes': None, 'step': 3, 'parents': {}, 'thread_id': '1'}, \n", " parent_config={...}, \n", @@ -885,8 +883,7 @@ " 'id': '1f029ca3-1790-616e-8002-9e021694a0cd', \n", " 'channel_versions': {'__start__': '00000000000000000000000000000004.0.5736472536395331', 'messages': '00000000000000000000000000000003.0.7056767754077798', 'branch:to:call_model': '00000000000000000000000000000003.0.22059023329132854'}, \n", " 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}}, \n", - " 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"what's my name?\"}]}, 'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}, \n", - " 'pending_sends': []\n", + " 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"what's my name?\"}]}, 'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}\n", " }, \n", " metadata={'source': 'input', 'writes': {'__start__': {'messages': [{'role': 'user', 'content': \"what's my name?\"}]}}, 'step': 2, 'parents': {}, 'thread_id': '1'}, \n", " parent_config={...}, \n", @@ -900,8 +897,7 @@ " 'id': '1f029ca3-178d-6f54-8001-d7b180db0c89', \n", " 'channel_versions': {'__start__': '00000000000000000000000000000002.0.18673090920108737', 'messages': '00000000000000000000000000000003.0.7056767754077798', 'branch:to:call_model': '00000000000000000000000000000003.0.22059023329132854'}, \n", " 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}}, \n", - " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}, \n", - " 'pending_sends': []\n", + " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}\n", " }, \n", " metadata={'source': 'loop', 'writes': {'call_model': {'messages': AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')}}, 'step': 1, 'parents': {}, 'thread_id': '1'}, \n", " parent_config={...}, \n", @@ -915,8 +911,7 @@ " 'id': '1f029ca3-0874-6612-8000-339f2abc83b1', \n", " 'channel_versions': {'__start__': '00000000000000000000000000000002.0.18673090920108737', 'messages': '00000000000000000000000000000002.0.30296526818059655', 'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}, \n", " 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}}, \n", - " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\")], 'branch:to:call_model': None}, \n", - " 'pending_sends': []\n", + " 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\")], 'branch:to:call_model': None}\n", " }, \n", " metadata={'source': 'loop', 'writes': None, 'step': 0, 'parents': {}, 'thread_id': '1'}, \n", " parent_config={...}, \n", @@ -930,8 +925,7 @@ " 'id': '1f029ca3-0870-6ce2-bfff-1f3f14c3e565', \n", " 'channel_versions': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}, \n", " 'versions_seen': {'__input__': {}}, \n", - " 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"hi! I'm bob\"}]}}, \n", - " 'pending_sends': []\n", + " 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"hi! I'm bob\"}]}}\n", " }, \n", " metadata={'source': 'input', 'writes': {'__start__': {'messages': [{'role': 'user', 'content': \"hi! I'm bob\"}]}}, 'step': -1, 'parents': {}, 'thread_id': '1'}, \n", " parent_config=None, \n", diff --git a/libs/checkpoint-postgres/README.md b/libs/checkpoint-postgres/README.md index 2b306960b..449510a6a 100644 --- a/libs/checkpoint-postgres/README.md +++ b/libs/checkpoint-postgres/README.md @@ -25,7 +25,7 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer: # call .setup() the first time you're using the checkpointer checkpointer.setup() checkpoint = { - "v": 2, + "v": 4, "ts": "2024-07-31T20:14:19.804150+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": { @@ -47,7 +47,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer: "start:node": 2 } }, - "pending_sends": [], } # store checkpoint @@ -67,7 +66,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer: checkpoint = { - "v": 2, + "v": 4, "ts": "2024-07-31T20:14:19.804150+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": { @@ -89,7 +88,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer: "start:node": 2 } }, - "pending_sends": [], } # store checkpoint diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 3a2ecbce9..e4f42078a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -1,4 +1,5 @@ import threading +from collections import defaultdict from collections.abc import Iterator, Sequence from contextlib import contextmanager from typing import Any, Optional @@ -142,8 +143,36 @@ class PostgresSaver(BasePostgresSaver): query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor with self._cursor() as cur: - cur.execute(query, args, binary=True) - for value in cur: + cur.execute(query, args) + values = cur.fetchall() + if not values: + return + # migrate pending sends if necessary + if to_migrate := [ + v + for v in values + if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"] + ]: + cur.execute( + self.SELECT_PENDING_SENDS_SQL, + ( + values[0]["thread_id"], + [v["parent_checkpoint_id"] for v in to_migrate], + ), + ) + grouped_by_parent = defaultdict(list) + for value in to_migrate: + grouped_by_parent[value["parent_checkpoint_id"]].append(value) + for sends in cur: + for value in grouped_by_parent[sends["checkpoint_id"]]: + if value["channel_values"] is None: + value["channel_values"] = [] + self._migrate_pending_sends( + sends["sends"], + value["checkpoint"], + value["channel_values"], + ) + for value in values: yield CheckpointTuple( { "configurable": { @@ -155,7 +184,6 @@ class PostgresSaver(BasePostgresSaver): self._load_checkpoint( value["checkpoint"], value["channel_values"], - value["pending_sends"], ), self._load_metadata(value["metadata"]), ( @@ -221,37 +249,52 @@ class PostgresSaver(BasePostgresSaver): cur.execute( self.SELECT_SQL + where, args, - binary=True, ) + value = cur.fetchone() + if value is None: + return None - for value in cur: - return CheckpointTuple( + # migrate pending sends if necessary + if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]: + cur.execute( + self.SELECT_PENDING_SENDS_SQL, + (thread_id, [value["parent_checkpoint_id"]]), + ) + if sends := cur.fetchone(): + if value["channel_values"] is None: + value["channel_values"] = [] + self._migrate_pending_sends( + sends["sends"], + value["checkpoint"], + value["channel_values"], + ) + + return CheckpointTuple( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": value["checkpoint_id"], + } + }, + self._load_checkpoint( + value["checkpoint"], + value["channel_values"], + ), + self._load_metadata(value["metadata"]), + ( { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["checkpoint_id"], + "checkpoint_id": value["parent_checkpoint_id"], } - }, - self._load_checkpoint( - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ), - self._load_metadata(value["metadata"]), - ( - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None - ), - self._load_writes(value["pending_writes"]), - ) + } + if value["parent_checkpoint_id"] + else None + ), + self._load_writes(value["pending_writes"]), + ) def put( self, diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index ddbc2fdbc..4d96e8336 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,4 +1,5 @@ import asyncio +from collections import defaultdict from collections.abc import AsyncIterator, Iterator, Sequence from contextlib import asynccontextmanager from typing import Any, Optional @@ -130,7 +131,35 @@ class AsyncPostgresSaver(BasePostgresSaver): # if we change this to use .stream() we need to make sure to close the cursor async with self._cursor() as cur: await cur.execute(query, args, binary=True) - async for value in cur: + values = await cur.fetchall() + if not values: + return + # migrate pending sends if necessary + if to_migrate := [ + v + for v in values + if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"] + ]: + await cur.execute( + self.SELECT_PENDING_SENDS_SQL, + ( + values[0]["thread_id"], + [v["parent_checkpoint_id"] for v in to_migrate], + ), + ) + grouped_by_parent = defaultdict(list) + for value in to_migrate: + grouped_by_parent[value["parent_checkpoint_id"]].append(value) + async for sends in cur: + for value in grouped_by_parent[sends["checkpoint_id"]]: + if value["channel_values"] is None: + value["channel_values"] = [] + self._migrate_pending_sends( + sends["sends"], + value["checkpoint"], + value["channel_values"], + ) + for value in values: yield CheckpointTuple( { "configurable": { @@ -143,7 +172,6 @@ class AsyncPostgresSaver(BasePostgresSaver): self._load_checkpoint, value["checkpoint"], value["channel_values"], - value["pending_sends"], ), self._load_metadata(value["metadata"]), ( @@ -190,36 +218,52 @@ class AsyncPostgresSaver(BasePostgresSaver): args, binary=True, ) + value = await cur.fetchone() + if value is None: + return None - async for value in cur: - return CheckpointTuple( + # migrate pending sends if necessary + if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]: + await cur.execute( + self.SELECT_PENDING_SENDS_SQL, + (thread_id, [value["parent_checkpoint_id"]]), + ) + if sends := await cur.fetchone(): + if value["channel_values"] is None: + value["channel_values"] = [] + self._migrate_pending_sends( + sends["sends"], + value["checkpoint"], + value["channel_values"], + ) + + return CheckpointTuple( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": value["checkpoint_id"], + } + }, + await asyncio.to_thread( + self._load_checkpoint, + value["checkpoint"], + value["channel_values"], + ), + self._load_metadata(value["metadata"]), + ( { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["checkpoint_id"], + "checkpoint_id": value["parent_checkpoint_id"], } - }, - await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ), - self._load_metadata(value["metadata"]), - ( - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None - ), - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) + } + if value["parent_checkpoint_id"] + else None + ), + await asyncio.to_thread(self._load_writes, value["pending_writes"]), + ) async def aput( self, diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 0d901a78c..785f9989b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -14,7 +14,7 @@ from langgraph.checkpoint.base import ( get_checkpoint_id, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol +from langgraph.checkpoint.serde.types import TASKS MetadataInput = Optional[dict[str, Any]] @@ -72,7 +72,7 @@ MIGRATIONS = [ """ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';""", ] -SELECT_SQL = f""" +SELECT_SQL = """ select thread_id, checkpoint, @@ -96,17 +96,20 @@ select where cw.thread_id = checkpoints.thread_id and cw.checkpoint_ns = checkpoints.checkpoint_ns and cw.checkpoint_id = checkpoints.checkpoint_id - ) as pending_writes, - ( - select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx) - from checkpoint_writes cw - where cw.thread_id = checkpoints.thread_id - and cw.checkpoint_ns = checkpoints.checkpoint_ns - and cw.checkpoint_id = checkpoints.parent_checkpoint_id - and cw.channel = '{TASKS}' - ) as pending_sends + ) as pending_writes from checkpoints """ +SELECT_PENDING_SENDS_SQL = f""" +select + checkpoint_id, + array_agg(array[type::bytea, blob] order by task_path, task_id, idx) as sends +from checkpoint_writes +where thread_id = %s + and checkpoint_id = any(%s) + and channel = '{TASKS}' +group by checkpoint_id +""" + UPSERT_CHECKPOINT_BLOBS_SQL = """ INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob) VALUES (%s, %s, %s, %s, %s, %s) @@ -140,6 +143,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """ class BasePostgresSaver(BaseCheckpointSaver[str]): SELECT_SQL = SELECT_SQL + SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL MIGRATIONS = MIGRATIONS UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL @@ -149,22 +153,38 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): jsonplus_serde = JsonPlusSerializer() supports_pipeline: bool + def _migrate_pending_sends( + self, + pending_sends: list[tuple[bytes, bytes]], + checkpoint: dict[str, Any], + channel_values: list[tuple[bytes, bytes, bytes]], + ) -> None: + if not pending_sends: + return + # add to values + enc, blob = self.serde.dumps_typed( + [self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends], + ) + channel_values.append((TASKS.encode(), enc.encode(), blob)) + # add to versions + checkpoint["channel_versions"][TASKS] = ( + max(checkpoint["channel_versions"].values()) + if checkpoint["channel_versions"] + else self.get_next_version(None) + ) + def _load_checkpoint( self, checkpoint: dict[str, Any], channel_values: list[tuple[bytes, bytes, bytes]], - pending_sends: list[tuple[bytes, bytes]], ) -> Checkpoint: return { **checkpoint, - "pending_sends": [ - self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or [] - ], "channel_values": self._load_blobs(channel_values), } def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]: - return {**checkpoint, "pending_sends": []} + return checkpoint def _load_blobs( self, blob_values: list[tuple[bytes, bytes, bytes]] @@ -249,7 +269,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): # NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing return serialized_metadata.decode().replace("\\u0000", "") - def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: + def get_next_version(self, current: Optional[str]) -> str: if current is None: current_v = 0 elif isinstance(current, int): diff --git a/libs/checkpoint-postgres/tests/checkpoint_utils.py b/libs/checkpoint-postgres/tests/checkpoint_utils.py new file mode 100644 index 000000000..881bef7f4 --- /dev/null +++ b/libs/checkpoint-postgres/tests/checkpoint_utils.py @@ -0,0 +1,51 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Any, Optional, Protocol + +from langgraph.checkpoint.base import Checkpoint, EmptyChannelError +from langgraph.checkpoint.base.id import uuid6 + + +class ChannelProtocol(Protocol): + def checkpoint(self) -> Optional[Any]: ... + + +def empty_checkpoint() -> Checkpoint: + return Checkpoint( + v=1, + id=str(uuid6(clock_seq=-2)), + ts=datetime.now(timezone.utc).isoformat(), + channel_values={}, + channel_versions={}, + versions_seen={}, + ) + + +def create_checkpoint( + checkpoint: Checkpoint, + channels: Optional[Mapping[str, ChannelProtocol]], + step: int, + *, + id: Optional[str] = None, +) -> Checkpoint: + """Create a checkpoint for the given channels.""" + ts = datetime.now(timezone.utc).isoformat() + if channels is None: + values = checkpoint["channel_values"] + else: + values = {} + for k, v in channels.items(): + if k not in checkpoint["channel_versions"]: + continue + try: + values[k] = v.checkpoint() + except EmptyChannelError: + pass + return Checkpoint( + v=1, + ts=ts, + id=id or str(uuid6(clock_seq=step)), + channel_values=values, + channel_versions=checkpoint["channel_versions"], + versions_seen=checkpoint["versions_seen"], + ) diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 45ce0f082..59669d17c 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -14,10 +14,10 @@ from langgraph.checkpoint.base import ( EXCLUDED_METADATA_KEYS, Checkpoint, CheckpointMetadata, - create_checkpoint, - empty_checkpoint, ) from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.checkpoint.serde.types import TASKS +from tests.checkpoint_utils import create_checkpoint, empty_checkpoint from tests.conftest import DEFAULT_POSTGRES_URI @@ -263,3 +263,52 @@ async def test_null_chars(saver_name: str, test_data) -> None: assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][ 0 ].metadata["my_key"] == "abc" + + +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) +async def test_pending_sends_migration(saver_name: str) -> None: + async with _saver(saver_name) as saver: + config = { + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + } + } + + # create the first checkpoint + # and put some pending sends + checkpoint_0 = empty_checkpoint() + config = await saver.aput(config, checkpoint_0, {}, {}) + await saver.aput_writes( + config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1" + ) + await saver.aput_writes(config, [(TASKS, "send-3")], task_id="task-2") + + # check that fetching checkpoint_0 doesn't attach pending sends + # (they should be attached to the next checkpoint) + tuple_0 = await saver.aget_tuple(config) + assert tuple_0.checkpoint["channel_values"] == {} + assert tuple_0.checkpoint["channel_versions"] == {} + + # create the second checkpoint + checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1) + config = await saver.aput(config, checkpoint_1, {}, {}) + + # check that pending sends are attached to checkpoint_1 + tuple_1 = await saver.aget_tuple(config) + assert tuple_1.checkpoint["channel_values"] == { + TASKS: ["send-1", "send-2", "send-3"] + } + assert TASKS in tuple_1.checkpoint["channel_versions"] + + # check that list also applies the migration + search_results = [ + c async for c in saver.alist({"configurable": {"thread_id": "thread-1"}}) + ] + assert len(search_results) == 2 + assert search_results[-1].checkpoint["channel_values"] == {} + assert search_results[-1].checkpoint["channel_versions"] == {} + assert search_results[0].checkpoint["channel_values"] == { + TASKS: ["send-1", "send-2", "send-3"] + } + assert TASKS in search_results[0].checkpoint["channel_versions"] diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index 4409fa2d6..b6eea12bf 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -15,10 +15,10 @@ from langgraph.checkpoint.base import ( EXCLUDED_METADATA_KEYS, Checkpoint, CheckpointMetadata, - create_checkpoint, - empty_checkpoint, ) from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.checkpoint.serde.types import TASKS +from tests.checkpoint_utils import create_checkpoint, empty_checkpoint from tests.conftest import DEFAULT_POSTGRES_URI @@ -258,3 +258,52 @@ def test_nonnull_migrations() -> None: for migration in PostgresSaver.MIGRATIONS: statement = _leading_comment_remover.sub("", migration).split()[0] assert statement.strip() + + +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) +def test_pending_sends_migration(saver_name: str) -> None: + with _saver(saver_name) as saver: + config = { + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + } + } + + # create the first checkpoint + # and put some pending sends + checkpoint_0 = empty_checkpoint() + config = saver.put(config, checkpoint_0, {}, {}) + saver.put_writes( + config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1" + ) + saver.put_writes(config, [(TASKS, "send-3")], task_id="task-2") + + # check that fetching checkpoint_0 doesn't attach pending sends + # (they should be attached to the next checkpoint) + tuple_0 = saver.get_tuple(config) + assert tuple_0.checkpoint["channel_values"] == {} + assert tuple_0.checkpoint["channel_versions"] == {} + + # create the second checkpoint + checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1) + config = saver.put(config, checkpoint_1, {}, {}) + + # check that pending sends are attached to checkpoint_1 + checkpoint_1 = saver.get_tuple(config) + assert checkpoint_1.checkpoint["channel_values"] == { + TASKS: ["send-1", "send-2", "send-3"] + } + assert TASKS in checkpoint_1.checkpoint["channel_versions"] + + # check that list also applies the migration + search_results = [ + c for c in saver.list({"configurable": {"thread_id": "thread-1"}}) + ] + assert len(search_results) == 2 + assert search_results[-1].checkpoint["channel_values"] == {} + assert search_results[-1].checkpoint["channel_versions"] == {} + assert search_results[0].checkpoint["channel_values"] == { + TASKS: ["send-1", "send-2", "send-3"] + } + assert TASKS in search_results[0].checkpoint["channel_versions"] diff --git a/libs/checkpoint-sqlite/README.md b/libs/checkpoint-sqlite/README.md index 8bd948b92..48f4e6690 100644 --- a/libs/checkpoint-sqlite/README.md +++ b/libs/checkpoint-sqlite/README.md @@ -12,7 +12,7 @@ read_config = {"configurable": {"thread_id": "1"}} with SqliteSaver.from_conn_string(":memory:") as checkpointer: checkpoint = { - "v": 2, + "v": 4, "ts": "2024-07-31T20:14:19.804150+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": { @@ -34,7 +34,6 @@ with SqliteSaver.from_conn_string(":memory:") as checkpointer: "start:node": 2 } }, - "pending_sends": [], } # store checkpoint @@ -54,7 +53,7 @@ from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: checkpoint = { - "v": 2, + "v": 4, "ts": "2024-07-31T20:14:19.804150+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": { @@ -76,7 +75,6 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: "start:node": 2 } }, - "pending_sends": [], } # store checkpoint diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index cdfffb681..e13cfbe29 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -19,7 +19,6 @@ from langgraph.checkpoint.base import ( get_checkpoint_metadata, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.serde.types import ChannelProtocol from langgraph.checkpoint.sqlite.utils import search_where _AIO_ERROR_MSG = ( @@ -535,14 +534,13 @@ class SqliteSaver(BaseCheckpointSaver[str]): """ raise NotImplementedError(_AIO_ERROR_MSG) - def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: + def get_next_version(self, current: Optional[str]) -> str: """Generate the next version ID for a channel. This method creates a new version identifier for a channel based on its current version. Args: current (Optional[str]): The current version identifier of the channel. - channel (BaseChannel): The channel being versioned. Returns: str: The next version identifier, which is guaranteed to be monotonically increasing. diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 48907908c..77b499326 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -19,7 +19,6 @@ from langgraph.checkpoint.base import ( get_checkpoint_metadata, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.serde.types import ChannelProtocol from langgraph.checkpoint.sqlite.utils import search_where T = TypeVar("T", bound=Callable) @@ -590,14 +589,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): ) await self.conn.commit() - def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: + def get_next_version(self, current: Optional[str]) -> str: """Generate the next version ID for a channel. This method creates a new version identifier for a channel based on its current version. Args: current (Optional[str]): The current version identifier of the channel. - channel (BaseChannel): The channel being versioned. Returns: str: The next version identifier, which is guaranteed to be monotonically increasing. diff --git a/libs/checkpoint-sqlite/tests/checkpoint_utils.py b/libs/checkpoint-sqlite/tests/checkpoint_utils.py new file mode 100644 index 000000000..881bef7f4 --- /dev/null +++ b/libs/checkpoint-sqlite/tests/checkpoint_utils.py @@ -0,0 +1,51 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Any, Optional, Protocol + +from langgraph.checkpoint.base import Checkpoint, EmptyChannelError +from langgraph.checkpoint.base.id import uuid6 + + +class ChannelProtocol(Protocol): + def checkpoint(self) -> Optional[Any]: ... + + +def empty_checkpoint() -> Checkpoint: + return Checkpoint( + v=1, + id=str(uuid6(clock_seq=-2)), + ts=datetime.now(timezone.utc).isoformat(), + channel_values={}, + channel_versions={}, + versions_seen={}, + ) + + +def create_checkpoint( + checkpoint: Checkpoint, + channels: Optional[Mapping[str, ChannelProtocol]], + step: int, + *, + id: Optional[str] = None, +) -> Checkpoint: + """Create a checkpoint for the given channels.""" + ts = datetime.now(timezone.utc).isoformat() + if channels is None: + values = checkpoint["channel_values"] + else: + values = {} + for k, v in channels.items(): + if k not in checkpoint["channel_versions"]: + continue + try: + values[k] = v.checkpoint() + except EmptyChannelError: + pass + return Checkpoint( + v=1, + ts=ts, + id=id or str(uuid6(clock_seq=step)), + channel_values=values, + channel_versions=checkpoint["channel_versions"], + versions_seen=checkpoint["versions_seen"], + ) diff --git a/libs/checkpoint-sqlite/tests/test_aiosqlite.py b/libs/checkpoint-sqlite/tests/test_aiosqlite.py index 503b7ade2..1e18fbb5e 100644 --- a/libs/checkpoint-sqlite/tests/test_aiosqlite.py +++ b/libs/checkpoint-sqlite/tests/test_aiosqlite.py @@ -6,10 +6,9 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, - create_checkpoint, - empty_checkpoint, ) from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from tests.checkpoint_utils import create_checkpoint, empty_checkpoint class TestAsyncSqliteSaver: diff --git a/libs/checkpoint-sqlite/tests/test_sqlite.py b/libs/checkpoint-sqlite/tests/test_sqlite.py index 2a027fa3b..05bea2907 100644 --- a/libs/checkpoint-sqlite/tests/test_sqlite.py +++ b/libs/checkpoint-sqlite/tests/test_sqlite.py @@ -6,11 +6,10 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, - create_checkpoint, - empty_checkpoint, ) from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.checkpoint.sqlite.utils import _metadata_predicate, search_where +from tests.checkpoint_utils import create_checkpoint, empty_checkpoint class TestSqliteSaver: diff --git a/libs/checkpoint/README.md b/libs/checkpoint/README.md index 2fbfa8fe6..4877a17b9 100644 --- a/libs/checkpoint/README.md +++ b/libs/checkpoint/README.md @@ -51,7 +51,7 @@ read_config = {"configurable": {"thread_id": "1"}} checkpointer = MemorySaver() checkpoint = { - "v": 2, + "v": 4, "ts": "2024-07-31T20:14:19.804150+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": { @@ -73,7 +73,6 @@ checkpoint = { "start:node": 2 } }, - "pending_sends": [], } # store checkpoint diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 9aa051f9a..35aa32fa1 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -1,22 +1,18 @@ -from collections.abc import AsyncIterator, Iterator, Mapping, Sequence -from datetime import datetime, timezone +from collections.abc import AsyncIterator, Iterator, Sequence from typing import ( # noqa: UP035 Any, - Dict, Generic, List, Literal, NamedTuple, Optional, - Tuple, TypedDict, TypeVar, Union, ) -from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig +from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( @@ -24,14 +20,10 @@ from langgraph.checkpoint.serde.types import ( INTERRUPT, RESUME, SCHEDULED, - ChannelProtocol, - SendProtocol, ) V = TypeVar("V", int, float, str) -PendingWrite = Tuple[str, str, Any] -# Kept for backwards compat, newer versions of LangGraph no longer use this. -LATEST_VERSION = 2 +PendingWrite = tuple[str, str, Any] # Marked as total=False to allow for future expansion. @@ -65,10 +57,6 @@ class CheckpointMetadata(TypedDict, total=False): """ -class TaskInfo(TypedDict): - status: Literal["scheduled", "success", "error"] - - ChannelVersions = dict[str, Union[str, int, float]] @@ -96,22 +84,6 @@ class Checkpoint(TypedDict): This keeps track of the versions of the channels that each node has seen. Used to determine which nodes to execute next. """ - pending_sends: List[SendProtocol] - """List of inputs pushed to nodes but not yet processed. - Cleared by the next checkpoint.""" - - -# Kept for backwards compat, newer versions of LangGraph no longer use this. -def empty_checkpoint() -> Checkpoint: - return Checkpoint( - v=LATEST_VERSION, - id=str(uuid6(clock_seq=-2)), - ts=datetime.now(timezone.utc).isoformat(), - channel_values={}, - channel_versions={}, - versions_seen={}, - pending_sends=[], - ) def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: @@ -122,39 +94,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: channel_values=checkpoint["channel_values"].copy(), channel_versions=checkpoint["channel_versions"].copy(), versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()}, - pending_sends=checkpoint.get("pending_sends", []).copy(), - ) - - -# Kept for backwards compat, newer versions of LangGraph no longer use this. -def create_checkpoint( - checkpoint: Checkpoint, - channels: Optional[Mapping[str, ChannelProtocol]], - step: int, - *, - id: Optional[str] = None, -) -> Checkpoint: - """Create a checkpoint for the given channels.""" - ts = datetime.now(timezone.utc).isoformat() - if channels is None: - values = checkpoint["channel_values"] - else: - values = {} - for k, v in channels.items(): - if k not in checkpoint["channel_versions"]: - continue - try: - values[k] = v.checkpoint() - except EmptyChannelError: - pass - return Checkpoint( - v=LATEST_VERSION, - ts=ts, - id=id or str(uuid6(clock_seq=step)), - channel_values=values, - channel_versions=checkpoint["channel_versions"], - versions_seen=checkpoint["versions_seen"], - pending_sends=checkpoint.get("pending_sends", []), ) @@ -168,34 +107,6 @@ class CheckpointTuple(NamedTuple): pending_writes: Optional[List[PendingWrite]] = None -CheckpointThreadId = ConfigurableFieldSpec( - id="thread_id", - annotation=str, - name="Thread ID", - description=None, - default="", - is_shared=True, -) - -CheckpointNS = ConfigurableFieldSpec( - id="checkpoint_ns", - annotation=str, - name="Checkpoint NS", - description='Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', - default="", - is_shared=True, -) - -CheckpointId = ConfigurableFieldSpec( - id="checkpoint_id", - annotation=Optional[str], - name="Checkpoint ID", - description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", - default=None, - is_shared=True, -) - - class BaseCheckpointSaver(Generic[V]): """Base class for creating a graph checkpointer. @@ -219,15 +130,6 @@ class BaseCheckpointSaver(Generic[V]): ) -> None: self.serde = maybe_add_typed_methods(serde or self.serde) - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - """Define the configuration options for the checkpoint saver. - - Returns: - list[ConfigurableFieldSpec]: List of configuration field specs. - """ - return [CheckpointThreadId, CheckpointNS, CheckpointId] - def get(self, config: RunnableConfig) -> Optional[Checkpoint]: """Fetch a checkpoint using the given configuration. @@ -258,7 +160,7 @@ class BaseCheckpointSaver(Generic[V]): self, config: Optional[RunnableConfig], *, - filter: Optional[Dict[str, Any]] = None, + filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: @@ -304,7 +206,7 @@ class BaseCheckpointSaver(Generic[V]): def put_writes( self, config: RunnableConfig, - writes: Sequence[Tuple[str, Any]], + writes: Sequence[tuple[str, Any]], task_id: str, task_path: str = "", ) -> None: @@ -362,7 +264,7 @@ class BaseCheckpointSaver(Generic[V]): self, config: Optional[RunnableConfig], *, - filter: Optional[Dict[str, Any]] = None, + filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> AsyncIterator[CheckpointTuple]: @@ -409,7 +311,7 @@ class BaseCheckpointSaver(Generic[V]): async def aput_writes( self, config: RunnableConfig, - writes: Sequence[Tuple[str, Any]], + writes: Sequence[tuple[str, Any]], task_id: str, task_path: str = "", ) -> None: @@ -437,7 +339,7 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def get_next_version(self, current: Optional[V], channel: ChannelProtocol) -> V: + def get_next_version(self, current: Optional[V]) -> V: """Generate the next version ID for a channel. Default is to use integer versions, incrementing by 1. If you override, you can use str/int/float versions, @@ -445,7 +347,6 @@ class BaseCheckpointSaver(Generic[V]): Args: current: The current version identifier (int, float, or str). - channel: The channel being versioned. Returns: V: The next version identifier, which must be increasing. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 3e916196f..5caf43ff6 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -22,7 +22,6 @@ from langgraph.checkpoint.base import ( get_checkpoint_id, get_checkpoint_metadata, ) -from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol logger = logging.getLogger(__name__) @@ -150,19 +149,6 @@ class InMemorySaver( if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id): checkpoint, metadata, parent_checkpoint_id = saved writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() - if parent_checkpoint_id: - sends = sorted( - ( - (*w, k[1]) - for k, w in self.writes[ - (thread_id, checkpoint_ns, parent_checkpoint_id) - ].items() - if w[1] == TASKS - ), - key=lambda w: (w[3], w[0], w[4]), - ) - else: - sends = [] checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) return CheckpointTuple( config=config, @@ -171,7 +157,6 @@ class InMemorySaver( "channel_values": self._load_blobs( thread_id, checkpoint_ns, checkpoint_["channel_versions"] ), - "pending_sends": [self.serde.loads_typed(s[2]) for s in sends], }, metadata=self.serde.loads_typed(metadata), pending_writes=[ @@ -194,22 +179,7 @@ class InMemorySaver( checkpoint_id = max(checkpoints.keys()) checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id] writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() - if parent_checkpoint_id: - sends = sorted( - ( - (*w, k[1]) - for k, w in self.writes[ - (thread_id, checkpoint_ns, parent_checkpoint_id) - ].items() - if w[1] == TASKS - ), - key=lambda w: (w[3], w[0], w[4]), - ) - else: - sends = [] - checkpoint_ = self.serde.loads_typed(checkpoint) - return CheckpointTuple( config={ "configurable": { @@ -223,7 +193,6 @@ class InMemorySaver( "channel_values": self._load_blobs( thread_id, checkpoint_ns, checkpoint_["channel_versions"] ), - "pending_sends": [self.serde.loads_typed(s[2]) for s in sends], }, metadata=self.serde.loads_typed(metadata), pending_writes=[ @@ -316,20 +285,6 @@ class InMemorySaver( (thread_id, checkpoint_ns, checkpoint_id) ].values() - if parent_checkpoint_id: - sends = sorted( - ( - (*w, k[1]) - for k, w in self.writes[ - (thread_id, checkpoint_ns, parent_checkpoint_id) - ].items() - if w[1] == TASKS - ), - key=lambda w: (w[3], w[0], w[4]), - ) - else: - sends = [] - checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) yield CheckpointTuple( @@ -347,9 +302,6 @@ class InMemorySaver( checkpoint_ns, checkpoint_["channel_versions"], ), - "pending_sends": [ - self.serde.loads_typed(s[2]) for s in sends - ], }, metadata=metadata, parent_config=( @@ -390,7 +342,6 @@ class InMemorySaver( RunnableConfig: The updated config containing the saved checkpoint's timestamp. """ c = checkpoint.copy() - c.pop("pending_sends") # type: ignore[misc] thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"]["checkpoint_ns"] values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc] @@ -561,7 +512,7 @@ class InMemorySaver( """ return self.delete_thread(thread_id) - def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: + def get_next_version(self, current: Optional[str]) -> str: if current is None: current_v = 0 elif isinstance(current, int): diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 9c7f158ed..b612fb6a4 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -1,13 +1,4 @@ -from collections.abc import Sequence -from typing import ( - Any, - Optional, - Protocol, - TypeVar, - runtime_checkable, -) - -from typing_extensions import Self +from typing import Any, Protocol, TypeVar, runtime_checkable ERROR = "__error__" SCHEDULED = "__scheduled__" @@ -20,25 +11,6 @@ Update = TypeVar("Update", contravariant=True) C = TypeVar("C") -class ChannelProtocol(Protocol[Value, Update, C]): - # Mirrors langgraph.channels.base.BaseChannel - @property - def ValueType(self) -> Any: ... - - @property - def UpdateType(self) -> Any: ... - - def checkpoint(self) -> Optional[C]: ... - - def from_checkpoint(self, checkpoint: Optional[C]) -> Self: ... - - def update(self, values: Sequence[Update]) -> bool: ... - - def get(self) -> Value: ... - - def consume(self) -> bool: ... - - @runtime_checkable class SendProtocol(Protocol): # Mirrors langgraph.constants.Send diff --git a/libs/checkpoint/tests/checkpoint_utils.py b/libs/checkpoint/tests/checkpoint_utils.py new file mode 100644 index 000000000..881bef7f4 --- /dev/null +++ b/libs/checkpoint/tests/checkpoint_utils.py @@ -0,0 +1,51 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Any, Optional, Protocol + +from langgraph.checkpoint.base import Checkpoint, EmptyChannelError +from langgraph.checkpoint.base.id import uuid6 + + +class ChannelProtocol(Protocol): + def checkpoint(self) -> Optional[Any]: ... + + +def empty_checkpoint() -> Checkpoint: + return Checkpoint( + v=1, + id=str(uuid6(clock_seq=-2)), + ts=datetime.now(timezone.utc).isoformat(), + channel_values={}, + channel_versions={}, + versions_seen={}, + ) + + +def create_checkpoint( + checkpoint: Checkpoint, + channels: Optional[Mapping[str, ChannelProtocol]], + step: int, + *, + id: Optional[str] = None, +) -> Checkpoint: + """Create a checkpoint for the given channels.""" + ts = datetime.now(timezone.utc).isoformat() + if channels is None: + values = checkpoint["channel_values"] + else: + values = {} + for k, v in channels.items(): + if k not in checkpoint["channel_versions"]: + continue + try: + values[k] = v.checkpoint() + except EmptyChannelError: + pass + return Checkpoint( + v=1, + ts=ts, + id=id or str(uuid6(clock_seq=step)), + channel_values=values, + channel_versions=checkpoint["channel_versions"], + versions_seen=checkpoint["versions_seen"], + ) diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index ad2dbdb1e..975c1e02c 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -6,10 +6,12 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, +) +from langgraph.checkpoint.memory import InMemorySaver +from tests.checkpoint_utils import ( # type: ignore[import-untyped] create_checkpoint, empty_checkpoint, ) -from langgraph.checkpoint.memory import InMemorySaver class TestMemorySaver: diff --git a/libs/langgraph/langgraph/channels/topic.py b/libs/langgraph/langgraph/channels/topic.py index 4cb320ee0..e6cf1b2b6 100644 --- a/libs/langgraph/langgraph/channels/topic.py +++ b/libs/langgraph/langgraph/channels/topic.py @@ -71,12 +71,14 @@ class Topic( return empty def update(self, values: Sequence[Union[Value, list[Value]]]) -> bool: - current = list(self.values) + updated = False if not self.accumulate: + updated = bool(self.values) self.values = list[Value]() - if flat_values := flatten(values): + if flat_values := tuple(flatten(values)): + updated = True self.values.extend(flat_values) - return self.values != current + return updated def get(self) -> Sequence[Value]: if self.values: diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index e361c8f1a..3c89dab4c 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -122,7 +122,6 @@ RESERVED = { ERROR, NO_WRITES, SCHEDULED, - TASKS, # reserved config.configurable keys CONFIG_KEY_SEND, CONFIG_KEY_READ, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 746836e01..1fbbcb96c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -22,17 +22,12 @@ from langchain_core.runnables.config import ( get_callback_manager_for_config, ) from langchain_core.runnables.graph import Graph -from langchain_core.runnables.utils import ( - ConfigurableFieldSpec, - get_unique_config_specs, -) from pydantic import BaseModel from typing_extensions import Self from langgraph.cache.base import BaseCache -from langgraph.channels.base import ( - BaseChannel, -) +from langgraph.channels.base import BaseChannel +from langgraph.channels.topic import Topic from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -49,7 +44,6 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, - CONFIG_KEY_RESUMING, CONFIG_KEY_RUNNER_SUBMIT, CONFIG_KEY_SEND, CONFIG_KEY_STORE, @@ -66,6 +60,7 @@ from langgraph.constants import ( NULL_TASK_ID, PUSH, SCHEDULED, + TASKS, ) from langgraph.errors import ( ErrorCode, @@ -104,6 +99,7 @@ from langgraph.types import ( CachePolicy, Checkpointer, Interrupt, + Send, StateSnapshot, StateUpdate, StreamChunk, @@ -117,8 +113,7 @@ from langgraph.utils.config import ( patch_configurable, recast_checkpoint_ns, ) -from langgraph.utils.fields import get_enhanced_type_hints -from langgraph.utils.pydantic import create_model, is_supported_by_pydantic +from langgraph.utils.pydantic import create_model from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined] from langgraph.utils.runnable import ( Runnable, @@ -706,6 +701,12 @@ class Pregel(PregelProtocol): k: v.build() if isinstance(v, NodeBuilder) else v for k, v in nodes.items() } self.channels = channels or {} + if TASKS in self.channels and not isinstance(self.channels[TASKS], Topic): + raise ValueError( + f"Channel '{TASKS}' is reserved and cannot be used in the graph." + ) + else: + self.channels[TASKS] = Topic(Send, accumulate=False) self.stream_mode = stream_mode self.stream_eager = stream_eager self.output_channels = output_channels @@ -830,54 +831,10 @@ class Pregel(PregelProtocol): self.trigger_to_nodes = _trigger_to_nodes(self.nodes) return self - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - spec - for spec in get_unique_config_specs( - [spec for node in self.nodes.values() for spec in node.config_specs] - + ( - self.checkpointer.config_specs - if isinstance(self.checkpointer, BaseCheckpointSaver) - else [] - ) - + ( - [ - ConfigurableFieldSpec( - id=name, - annotation=typ, - default=default, - description=description, - ) - for name, typ, default, description in get_enhanced_type_hints( - self.config_type - ) - ] - if self.config_type is not None - else [] - ) - ) - # these are provided by the Pregel class - if spec.id - not in [ - CONFIG_KEY_READ, - CONFIG_KEY_SEND, - CONFIG_KEY_CHECKPOINTER, - CONFIG_KEY_RESUMING, - ] - ] - def config_schema(self, *, include: Sequence[str] | None = None) -> type[BaseModel]: - # If the config type is not set explicitly, we will try to infer it. - # If the config type is provided, but isn't directly supported by pydantic - # (e.g., vanilla python class), we will also delegate to the parent class, - # which handles cases where Pydantic doesn't support the type. - if self.config_type is None or not is_supported_by_pydantic(self.config_type): - return super().config_schema(include=include) - include = include or [] fields = { - "configurable": (self.config_type, None), + **({"configurable": (self.config_type, None)} if self.config_type else {}), **{ field_name: (field_type, None) for field_name, field_type in get_type_hints(RunnableConfig).items() @@ -2359,7 +2316,8 @@ class Pregel(PregelProtocol): checkpointer = self.checkpointer if checkpointer and not config.get(CONF): raise ValueError( - f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in checkpointer.config_specs]}" + "Checkpointer requires one or more of the following 'configurable' " + "keys: thread_id, checkpoint_ns, checkpoint_id" ) if CONFIG_KEY_STORE in config.get(CONF, {}): store: BaseStore | None = config[CONF][CONFIG_KEY_STORE] diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 471fcfb99..505928867 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -25,6 +25,7 @@ from langchain_core.runnables.config import RunnableConfig from xxhash import xxh3_128_hexdigest from langgraph.channels.base import BaseChannel +from langgraph.channels.topic import Topic from langgraph.checkpoint.base import ( BaseCheckpointSaver, ChannelVersions, @@ -81,7 +82,7 @@ from langgraph.types import ( ) from langgraph.utils.config import merge_configs, patch_config -GetNextVersion = Callable[[Optional[V], BaseChannel], V] +GetNextVersion = Callable[[Optional[V]], V] SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) @@ -211,7 +212,7 @@ def local_read( return values -def increment(current: Optional[int], channel: BaseChannel) -> int: +def increment(current: Optional[int]) -> int: """Default channel versioning function, increments the current int version.""" return current + 1 if current is not None else 1 @@ -232,6 +233,7 @@ def apply_writes( channels: The channels to update. tasks: The tasks to apply writes from. get_next_version: Optional function to determine the next version of a channel. + trigger_to_nodes: Mapping of channel names to the set of nodes that can be triggered by updates to that channel. Returns: Set of channels that were updated in this step. @@ -255,10 +257,14 @@ def apply_writes( ) # Find the highest version of all channels - if checkpoint["channel_versions"]: - max_version = max(checkpoint["channel_versions"].values()) + if get_next_version is None: + next_version = None else: - max_version = None + next_version = get_next_version( + max(checkpoint["channel_versions"].values()) + if checkpoint["channel_versions"] + else None + ) # Consume all channels that were read for chan in { @@ -267,15 +273,8 @@ def apply_writes( for chan in task.triggers if chan not in RESERVED and chan in channels }: - if channels[chan].consume() and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, - channels[chan], - ) - - # clear pending sends - if checkpoint["pending_sends"] and bump_step: - checkpoint["pending_sends"].clear() + if channels[chan].consume() and next_version is not None: + checkpoint["channel_versions"][chan] = next_version # Group writes by channel pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) @@ -283,8 +282,6 @@ def apply_writes( for chan, val in task.writes: if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): pass - elif chan == TASKS: - checkpoint["pending_sends"].append(val) elif chan in channels: pending_writes_by_channel[chan].append(val) else: @@ -292,21 +289,12 @@ def apply_writes( f"Task {task.name} with path {task.path} wrote to unknown channel {chan}, ignoring it." ) - # Find the highest version of all channels - if checkpoint["channel_versions"]: - max_version = max(checkpoint["channel_versions"].values()) - else: - max_version = None - # Apply writes to channels updated_channels: set[str] = set() for chan, vals in pending_writes_by_channel.items(): if chan in channels: - if channels[chan].update(vals) and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, - channels[chan], - ) + if channels[chan].update(vals) and next_version is not None: + checkpoint["channel_versions"][chan] = next_version # unavailable channels can't trigger tasks, so don't add them if channels[chan].is_available(): updated_channels.add(chan) @@ -315,27 +303,17 @@ def apply_writes( if bump_step: for chan in channels: if channels[chan].is_available() and chan not in updated_channels: - if channels[chan].update(EMPTY_SEQ) and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, - channels[chan], - ) + if channels[chan].update(EMPTY_SEQ) and next_version is not None: + checkpoint["channel_versions"][chan] = next_version # unavailable channels can't trigger tasks, so don't add them if channels[chan].is_available(): updated_channels.add(chan) # If this is (tentatively) the last superstep, notify all channels of finish - if ( - bump_step - and not checkpoint["pending_sends"] - and updated_channels.isdisjoint(trigger_to_nodes) - ): + if bump_step and updated_channels.isdisjoint(trigger_to_nodes): for chan in channels: - if channels[chan].finish() and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, - channels[chan], - ) + if channels[chan].finish() and next_version is not None: + checkpoint["channel_versions"][chan] = next_version # unavailable channels can't trigger tasks, so don't add them if channels[chan].is_available(): updated_channels.add(chan) @@ -344,17 +322,6 @@ def apply_writes( return updated_channels -def has_next_tasks( - trigger_to_nodes: Mapping[str, Sequence[str]], - updated_channels: set[str], - checkpoint: Checkpoint, -) -> bool: - """Check if there are any tasks that should be run in the next step.""" - return bool(checkpoint["pending_sends"]) or not updated_channels.isdisjoint( - trigger_to_nodes - ) - - @overload def prepare_next_tasks( checkpoint: Checkpoint, @@ -448,30 +415,32 @@ def prepare_next_tasks( checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", "")) null_version = checkpoint_null_version(checkpoint) tasks: list[Union[PregelTask, PregelExecutableTask]] = [] - # Consume pending_sends from previous step - for idx, _ in enumerate(checkpoint["pending_sends"]): - if task := prepare_single_task( - (PUSH, idx), - None, - checkpoint=checkpoint, - checkpoint_id_bytes=checkpoint_id_bytes, - checkpoint_null_version=null_version, - pending_writes=pending_writes, - processes=processes, - channels=channels, - managed=managed, - config=config, - step=step, - stop=stop, - for_execution=for_execution, - store=store, - checkpointer=checkpointer, - manager=manager, - input_cache=input_cache, - cache_policy=cache_policy, - retry_policy=retry_policy, - ): - tasks.append(task) + # Consume pending tasks + tasks_channel = cast(Optional[Topic[Send]], channels.get(TASKS)) + if tasks_channel and tasks_channel.is_available(): + for idx, _ in enumerate(tasks_channel.get()): + if task := prepare_single_task( + (PUSH, idx), + None, + checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, + pending_writes=pending_writes, + processes=processes, + channels=channels, + managed=managed, + config=config, + step=step, + stop=stop, + for_execution=for_execution, + store=store, + checkpointer=checkpointer, + manager=manager, + input_cache=input_cache, + cache_policy=cache_policy, + retry_policy=retry_policy, + ): + tasks.append(task) # This section is an optimization that allows which nodes will be active # during the next step. @@ -656,9 +625,12 @@ def prepare_single_task( # SEND tasks, executed in superstep n+1 # (PUSH, idx of pending send) idx = cast(int, task_path[1]) - if idx >= len(checkpoint["pending_sends"]): + if not channels[TASKS].is_available(): return - packet = checkpoint["pending_sends"][idx] + sends: Sequence[Send] = channels[TASKS].get() + if idx < 0 or idx >= len(sends): + return + packet = sends[idx] if not isinstance(packet, Send): logger.warning( f"Ignoring invalid packet type {type(packet)} in pending sends" diff --git a/libs/langgraph/langgraph/pregel/checkpoint.py b/libs/langgraph/langgraph/pregel/checkpoint.py index 47ef04404..0a0f95078 100644 --- a/libs/langgraph/langgraph/pregel/checkpoint.py +++ b/libs/langgraph/langgraph/pregel/checkpoint.py @@ -8,7 +8,7 @@ from langgraph.checkpoint.base.id import uuid6 from langgraph.constants import MISSING from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec -LATEST_VERSION = 3 +LATEST_VERSION = 4 def empty_checkpoint() -> Checkpoint: @@ -19,7 +19,6 @@ def empty_checkpoint() -> Checkpoint: channel_values={}, channel_versions={}, versions_seen={}, - pending_sends=[], ) @@ -49,7 +48,6 @@ def create_checkpoint( channel_values=values, channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], - pending_sends=checkpoint.get("pending_sends", []), ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a07778d8e..872303cab 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -65,7 +65,6 @@ from langgraph.constants import ( RESUME, SCHEDULED, TAG_HIDDEN, - TASKS, ) from langgraph.errors import ( CheckpointNotLatest, @@ -318,9 +317,6 @@ class PregelLoop: """Put writes for a task, to be read by the next tick.""" if not writes: return - # always checkpoint writes containing Send, as they are fetched from the - # parent checkpoint, not the current one - checkpoint_during = self.checkpoint_during or any(w[0] == TASKS for w in writes) # deduplicate writes to special channels, last write wins if all(w[0] in WRITES_IDX_MAP for w in writes): writes = list({w[0]: w for w in writes}.values()) @@ -330,7 +326,7 @@ class PregelLoop: ] # save writes self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) - if checkpoint_during and self.checkpointer_put_writes is not None: + if self.checkpoint_during and self.checkpointer_put_writes is not None: config = patch_configurable( self.checkpoint_config, { @@ -806,7 +802,6 @@ class PregelLoop: else self.stream_keys ), ) - self.checkpoint_id_prev = self.checkpoint["id"] if self.step > -1 else None # do checkpoint? do_checkpoint = self._checkpointer_put_after_previous is not None and ( exiting or self.checkpoint_during @@ -836,7 +831,6 @@ class PregelLoop: CONF: { **self.checkpoint_config[CONF], # this is guaranteed to be set by code above - CONFIG_KEY_CHECKPOINT_ID: self.checkpoint_id_prev, CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" ), diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index dfe892e1d..4dbdffeee 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -1,7 +1,6 @@ -from langgraph.checkpoint.base import empty_checkpoint from langgraph.constants import PULL, PUSH from langgraph.pregel.algo import prepare_next_tasks, task_path_str -from langgraph.pregel.checkpoint import channels_from_checkpoint +from langgraph.pregel.checkpoint import channels_from_checkpoint, empty_checkpoint def test_prepare_next_tasks() -> None: diff --git a/libs/langgraph/tests/test_checkpoint_migration.py b/libs/langgraph/tests/test_checkpoint_migration.py index 942d6b8a8..4f4c406dd 100644 --- a/libs/langgraph/tests/test_checkpoint_migration.py +++ b/libs/langgraph/tests/test_checkpoint_migration.py @@ -350,7 +350,6 @@ SAVED_CHECKPOINTS = { "docs": ["doc1", "doc2", "doc3", "doc4"], "answer": "doc1,doc2,doc3,doc4", }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -413,7 +412,6 @@ SAVED_CHECKPOINTS = { "docs": ["doc1", "doc2", "doc3", "doc4"], "branch:to:qa": None, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -491,7 +489,6 @@ SAVED_CHECKPOINTS = { "branch:to:retriever_one": None, "docs": ["doc3", "doc4"], }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -548,7 +545,6 @@ SAVED_CHECKPOINTS = { "branch:to:analyzer_one": None, "branch:to:retriever_two": None, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -605,7 +601,6 @@ SAVED_CHECKPOINTS = { "query": "what is weather in sf", "branch:to:rewrite_query": None, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -656,7 +651,6 @@ SAVED_CHECKPOINTS = { }, "versions_seen": {"__input__": {}}, "channel_values": {"__start__": {"query": "what is weather in sf"}}, - "pending_sends": [], }, metadata={ "source": "input", @@ -744,7 +738,6 @@ SAVED_CHECKPOINTS = { "answer": "doc1,doc2,doc3,doc4", "qa": "qa", }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -808,7 +801,6 @@ SAVED_CHECKPOINTS = { "docs": ["doc1", "doc2", "doc3", "doc4"], "retriever_one": "retriever_one", }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -889,7 +881,6 @@ SAVED_CHECKPOINTS = { "docs": ["doc3", "doc4"], "retriever_two": "retriever_two", }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -950,7 +941,6 @@ SAVED_CHECKPOINTS = { "rewrite_query": "rewrite_query", "branch:to:retriever_two": "rewrite_query", }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1012,7 +1002,6 @@ SAVED_CHECKPOINTS = { "query": "what is weather in sf", "start:rewrite_query": "__start__", }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1063,7 +1052,6 @@ SAVED_CHECKPOINTS = { }, "versions_seen": {"__input__": {}}, "channel_values": {"__start__": {"query": "what is weather in sf"}}, - "pending_sends": [], }, metadata={ "source": "input", @@ -1151,7 +1139,6 @@ SAVED_CHECKPOINTS = { "retriever_one": "00000000000000000000000000000005.0.222301724202566" }, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1215,7 +1202,6 @@ SAVED_CHECKPOINTS = { "analyzer_one": "00000000000000000000000000000004.0.2684613370070208" }, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1296,7 +1282,6 @@ SAVED_CHECKPOINTS = { "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000003.0.8992241767805405" }, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1357,7 +1342,6 @@ SAVED_CHECKPOINTS = { "start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445" }, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1419,7 +1403,6 @@ SAVED_CHECKPOINTS = { "__start__": "00000000000000000000000000000001.0.6759219622820284" }, }, - "pending_sends": [], }, metadata={ "source": "loop", @@ -1470,7 +1453,6 @@ SAVED_CHECKPOINTS = { "__start__": "00000000000000000000000000000001.0.6759219622820284" }, "versions_seen": {"__input__": {}}, - "pending_sends": [], }, metadata={ "source": "input", diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 0babd6455..8d6ba1909 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -7416,8 +7416,10 @@ def test_send_dedupe_on_resume( if checkpoint_during: assert history == expected_history else: - assert history[0] == expected_history[0] - assert history[1] == expected_history[2] + assert history[0] == expected_history[0]._replace( + parent_config=history[1].config + ) + assert history[1] == expected_history[2]._replace(parent_config=None) def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a652ec66f..8c73daa1b 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -30,7 +30,6 @@ from syrupy import SnapshotAssertion from typing_extensions import TypedDict from langgraph.cache.base import BaseCache -from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue @@ -237,7 +236,7 @@ def test_checkpoint_errors() -> None: raise ValueError("Faulty put_writes") class FaultyVersionCheckpointer(InMemorySaver): - def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: + def get_next_version(self, current: Optional[int]) -> int: raise ValueError("Faulty get_next_version") def logic(inp: str) -> str: @@ -1130,10 +1129,9 @@ def test_pending_writes_resume( } }, checkpoint={ - "v": 3, + "v": 4, "id": AnyStr(), "ts": AnyStr(), - "pending_sends": [], "versions_seen": { "one": { "branch:to:one": AnyVersion(), @@ -1188,10 +1186,9 @@ def test_pending_writes_resume( } }, checkpoint={ - "v": 3, + "v": 4, "id": AnyStr(), "ts": AnyStr(), - "pending_sends": [], "versions_seen": { "__input__": {}, "__start__": { @@ -1223,11 +1220,11 @@ def test_pending_writes_resume( "checkpoint_ns": "", "checkpoint_id": ( checkpoints[2].config["configurable"]["checkpoint_id"] - if checkpoint_during - else AnyStr() ), } - }, + } + if checkpoint_during + else None, pending_writes=( UnsortedSequence( (AnyStr(), "value", 2), @@ -1254,10 +1251,9 @@ def test_pending_writes_resume( } }, checkpoint={ - "v": 3, + "v": 4, "id": AnyStr(), "ts": AnyStr(), - "pending_sends": [], "versions_seen": {"__input__": {}}, "channel_versions": { "__start__": AnyVersion(), @@ -1407,7 +1403,7 @@ def test_imp_task( ) -> None: mapper_calls = 0 - class Configurable: + class Configurable(TypedDict): model: str @task() @@ -1438,25 +1434,9 @@ def test_imp_task( "$defs": { "Configurable": { "properties": { - "model": {"default": None, "title": "Model", "type": "string"}, - "checkpoint_id": { - "anyOf": [{"type": "string"}, {"type": "null"}], - "default": None, - "description": "Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", - "title": "Checkpoint ID", - }, - "checkpoint_ns": { - "default": "", - "description": 'Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', - "title": "Checkpoint NS", - "type": "string", - }, - "thread_id": { - "default": "", - "title": "Thread ID", - "type": "string", - }, + "model": {"title": "Model", "type": "string"}, }, + "required": ["model"], "title": "Configurable", "type": "object", } diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 8aa1f8a41..54759ebf0 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -29,7 +29,6 @@ from syrupy import SnapshotAssertion from typing_extensions import TypedDict from langgraph.cache.base import BaseCache -from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -104,7 +103,7 @@ async def test_checkpoint_errors() -> None: raise ValueError("Faulty put_writes") class FaultyVersionCheckpointer(InMemorySaver): - def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: + def get_next_version(self, current: Optional[int]) -> int: raise ValueError("Faulty get_next_version") def logic(inp: str) -> str: @@ -1940,10 +1939,9 @@ async def test_pending_writes_resume( } }, checkpoint={ - "v": 3, + "v": 4, "id": AnyStr(), "ts": AnyStr(), - "pending_sends": [], "versions_seen": { "one": { "branch:to:one": AnyVersion(), @@ -1998,10 +1996,9 @@ async def test_pending_writes_resume( } }, checkpoint={ - "v": 3, + "v": 4, "id": AnyStr(), "ts": AnyStr(), - "pending_sends": [], "versions_seen": { "__input__": {}, "__start__": { @@ -2031,11 +2028,11 @@ async def test_pending_writes_resume( "configurable": { "thread_id": "1", "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"] - if checkpoint_during - else AnyStr(), + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], } - }, + } + if checkpoint_during + else None, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), @@ -2060,10 +2057,9 @@ async def test_pending_writes_resume( } }, checkpoint={ - "v": 3, + "v": 4, "id": AnyStr(), "ts": AnyStr(), - "pending_sends": [], "versions_seen": {"__input__": {}}, "channel_versions": { "__start__": AnyVersion(), @@ -2936,8 +2932,10 @@ async def test_send_dedupe_on_resume( if checkpoint_during: assert history == expected_history else: - assert history[0] == expected_history[0] - assert history[1] == expected_history[2] + assert history[0] == expected_history[0]._replace( + parent_config=history[1].config + ) + assert history[1] == expected_history[2]._replace(parent_config=None) async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) -> None: