From 1134017d076ff66508f948752dbf5f596b049a84 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 16 Jun 2025 14:57:11 -0700 Subject: [PATCH] Preparation for 0.5 release: langgraph-checkpoint (#5124) Prepare langgraph-checkpoint for 0.5 - Given we have no upper bound on langgraph-checkpoint dep need to undo all changes in langgraph-checkpoint that might break previous versions of langgraph --- .../langgraph/checkpoint/postgres/base.py | 4 +- .../tests/checkpoint_utils.py | 53 ------------ libs/checkpoint-postgres/tests/test_async.py | 3 +- libs/checkpoint-postgres/tests/test_sync.py | 3 +- .../langgraph/checkpoint/sqlite/__init__.py | 2 +- .../langgraph/checkpoint/sqlite/aio.py | 2 +- .../tests/checkpoint_utils.py | 53 ------------ .../checkpoint-sqlite/tests/test_aiosqlite.py | 3 +- libs/checkpoint-sqlite/tests/test_sqlite.py | 3 +- .../langgraph/checkpoint/base/__init__.py | 81 +++++++++++++++---- .../langgraph/checkpoint/memory/__init__.py | 2 +- libs/checkpoint/tests/checkpoint_utils.py | 53 ------------ libs/checkpoint/tests/test_memory.py | 4 +- libs/langgraph/langgraph/pregel/__init__.py | 2 +- libs/langgraph/langgraph/pregel/algo.py | 7 +- libs/langgraph/langgraph/pregel/checkpoint.py | 11 +++ libs/langgraph/langgraph/pregel/loop.py | 19 +---- .../tests/test_checkpoint_migration.py | 7 +- libs/langgraph/tests/test_pregel.py | 2 +- libs/langgraph/tests/test_pregel_async.py | 2 +- .../langgraph/prebuilt/chat_agent_executor.py | 10 +-- libs/prebuilt/tests/memory_assert.py | 2 +- 22 files changed, 109 insertions(+), 219 deletions(-) delete mode 100644 libs/checkpoint-postgres/tests/checkpoint_utils.py delete mode 100644 libs/checkpoint-sqlite/tests/checkpoint_utils.py delete mode 100644 libs/checkpoint/tests/checkpoint_utils.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 8c502a7fe..44b8ee397 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -168,7 +168,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): checkpoint["channel_versions"][TASKS] = ( max(checkpoint["channel_versions"].values()) if checkpoint["channel_versions"] - else self.get_next_version(None) + else self.get_next_version(None, None) ) def _load_blobs( @@ -246,7 +246,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): for idx, (channel, value) in enumerate(writes) ] - def get_next_version(self, current: str | None) -> str: + def get_next_version(self, current: str | None, channel: None) -> 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 deleted file mode 100644 index f38afd740..000000000 --- a/libs/checkpoint-postgres/tests/checkpoint_utils.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from datetime import datetime, timezone -from typing import Any, Protocol - -from langgraph.checkpoint.base import Checkpoint, EmptyChannelError -from langgraph.checkpoint.base.id import uuid6 - - -class ChannelProtocol(Protocol): - def checkpoint(self) -> Any | None: ... - - -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: Mapping[str, ChannelProtocol] | None, - step: int, - *, - id: str | None = 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 f0196f845..905aa8968 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -14,13 +14,14 @@ from langgraph.checkpoint.base import ( EXCLUDED_METADATA_KEYS, Checkpoint, CheckpointMetadata, + create_checkpoint, + empty_checkpoint, ) from langgraph.checkpoint.postgres.aio import ( AsyncPostgresSaver, AsyncShallowPostgresSaver, ) from langgraph.checkpoint.serde.types import TASKS -from tests.checkpoint_utils import create_checkpoint, empty_checkpoint from tests.conftest import DEFAULT_POSTGRES_URI diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index 3ce48c8da..b010b5bbe 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -15,10 +15,11 @@ from langgraph.checkpoint.base import ( EXCLUDED_METADATA_KEYS, Checkpoint, CheckpointMetadata, + create_checkpoint, + empty_checkpoint, ) from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver from langgraph.checkpoint.serde.types import TASKS -from tests.checkpoint_utils import create_checkpoint, empty_checkpoint from tests.conftest import DEFAULT_POSTGRES_URI diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index caf9cdf3a..e716b1f47 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -536,7 +536,7 @@ class SqliteSaver(BaseCheckpointSaver[str]): """ raise NotImplementedError(_AIO_ERROR_MSG) - def get_next_version(self, current: str | None) -> str: + def get_next_version(self, current: str | None, channel: None) -> str: """Generate the next version ID for a channel. This method creates a new version identifier for a channel based on its current version. diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index dd4a61ab9..6ee30b259 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -591,7 +591,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): ) await self.conn.commit() - def get_next_version(self, current: str | None) -> str: + def get_next_version(self, current: str | None, channel: None) -> str: """Generate the next version ID for a channel. This method creates a new version identifier for a channel based on its current version. diff --git a/libs/checkpoint-sqlite/tests/checkpoint_utils.py b/libs/checkpoint-sqlite/tests/checkpoint_utils.py deleted file mode 100644 index f38afd740..000000000 --- a/libs/checkpoint-sqlite/tests/checkpoint_utils.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from datetime import datetime, timezone -from typing import Any, Protocol - -from langgraph.checkpoint.base import Checkpoint, EmptyChannelError -from langgraph.checkpoint.base.id import uuid6 - - -class ChannelProtocol(Protocol): - def checkpoint(self) -> Any | None: ... - - -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: Mapping[str, ChannelProtocol] | None, - step: int, - *, - id: str | None = 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 1e18fbb5e..503b7ade2 100644 --- a/libs/checkpoint-sqlite/tests/test_aiosqlite.py +++ b/libs/checkpoint-sqlite/tests/test_aiosqlite.py @@ -6,9 +6,10 @@ 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 05bea2907..2a027fa3b 100644 --- a/libs/checkpoint-sqlite/tests/test_sqlite.py +++ b/libs/checkpoint-sqlite/tests/test_sqlite.py @@ -6,10 +6,11 @@ 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/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 80b3466ed..e9350a993 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -1,10 +1,8 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Iterator, Sequence -from inspect import signature +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import ( # noqa: UP035 Any, - ClassVar, Generic, Literal, NamedTuple, @@ -15,6 +13,7 @@ from typing import ( # noqa: UP035 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 ( @@ -22,6 +21,7 @@ from langgraph.checkpoint.serde.types import ( INTERRUPT, RESUME, SCHEDULED, + ChannelProtocol, ) V = TypeVar("V", int, float, str) @@ -91,6 +91,7 @@ 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(), ) @@ -118,19 +119,8 @@ class BaseCheckpointSaver(Generic[V]): versions to avoid blocking the main thread. """ - _get_next_version_legacy: ClassVar[bool] = False - """Flag indicating if get_next_version method is legacy (takes two parameters).""" - serde: SerializerProtocol = JsonPlusSerializer() - def __init_subclass__(cls) -> None: - cls._get_next_version_legacy = ( - len(signature(cls.get_next_version).parameters) > 2 # self + current - if hasattr(cls, "get_next_version") - else False - ) - return super().__init_subclass__() - def __init__( self, *, @@ -138,6 +128,15 @@ class BaseCheckpointSaver(Generic[V]): ) -> None: self.serde = maybe_add_typed_methods(serde or self.serde) + @property + def config_specs(self) -> list: + """Define the configuration options for the checkpoint saver. + + Returns: + list: List of configuration field specs. + """ + return [] + def get(self, config: RunnableConfig) -> Checkpoint | None: """Fetch a checkpoint using the given configuration. @@ -347,7 +346,7 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def get_next_version(self, current: V | None) -> V: + def get_next_version(self, current: V | None, channel: None) -> 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, @@ -355,6 +354,7 @@ class BaseCheckpointSaver(Generic[V]): Args: current: The current version identifier (int, float, or str). + channel: Deprecated argument, kept for backwards compatibility. Returns: V: The next version identifier, which must be increasing. @@ -417,3 +417,54 @@ EXCLUDED_METADATA_KEYS = { "checkpoint_ns", "checkpoint_map", } + +# --- below are deprecated utilities used by past versions of LangGraph --- + +LATEST_VERSION = 2 + + +def empty_checkpoint() -> Checkpoint: + from datetime import datetime, timezone + + 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 create_checkpoint( + checkpoint: Checkpoint, + channels: Mapping[str, ChannelProtocol] | None, + step: int, + *, + id: str | None = None, +) -> Checkpoint: + """Create a checkpoint for the given channels.""" + from datetime import datetime, timezone + + 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", []), + ) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index dc6d089e6..14f2a9547 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -512,7 +512,7 @@ class InMemorySaver( """ return self.delete_thread(thread_id) - def get_next_version(self, current: str | None) -> str: + def get_next_version(self, current: str | None, channel: None) -> str: if current is None: current_v = 0 elif isinstance(current, int): diff --git a/libs/checkpoint/tests/checkpoint_utils.py b/libs/checkpoint/tests/checkpoint_utils.py deleted file mode 100644 index f38afd740..000000000 --- a/libs/checkpoint/tests/checkpoint_utils.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from datetime import datetime, timezone -from typing import Any, Protocol - -from langgraph.checkpoint.base import Checkpoint, EmptyChannelError -from langgraph.checkpoint.base.id import uuid6 - - -class ChannelProtocol(Protocol): - def checkpoint(self) -> Any | None: ... - - -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: Mapping[str, ChannelProtocol] | None, - step: int, - *, - id: str | None = 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 b0eeb319b..ad2dbdb1e 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -6,12 +6,10 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, -) -from langgraph.checkpoint.memory import InMemorySaver -from tests.checkpoint_utils import ( create_checkpoint, empty_checkpoint, ) +from langgraph.checkpoint.memory import InMemorySaver class TestMemorySaver: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6fbd1f6b2..5cf977637 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -32,7 +32,6 @@ from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, CheckpointTuple, - copy_checkpoint, ) from langgraph.config import get_config from langgraph.constants import ( @@ -79,6 +78,7 @@ from langgraph.pregel.algo import ( from langgraph.pregel.call import identifier from langgraph.pregel.checkpoint import ( channels_from_checkpoint, + copy_checkpoint, create_checkpoint, empty_checkpoint, ) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index ce2b22411..9d15aff3c 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -83,7 +83,7 @@ from langgraph.types import ( ) from langgraph.utils.config import merge_configs, patch_config -GetNextVersion = Callable[[Optional[V]], V] +GetNextVersion = Callable[[Optional[V], None], V] SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) @@ -214,7 +214,7 @@ def local_read( return values -def increment(current: int | None) -> int: +def increment(current: int | None, channel: None) -> int: """Default channel versioning function, increments the current int version.""" return current + 1 if current is not None else 1 @@ -265,7 +265,8 @@ def apply_writes( next_version = get_next_version( max(checkpoint["channel_versions"].values()) if checkpoint["channel_versions"] - else None + else None, + None, ) # Consume all channels that were read diff --git a/libs/langgraph/langgraph/pregel/checkpoint.py b/libs/langgraph/langgraph/pregel/checkpoint.py index b8ca90db4..b404ee550 100644 --- a/libs/langgraph/langgraph/pregel/checkpoint.py +++ b/libs/langgraph/langgraph/pregel/checkpoint.py @@ -71,3 +71,14 @@ def channels_from_checkpoint( }, managed_specs, ) + + +def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: + return Checkpoint( + v=checkpoint["v"], + ts=checkpoint["ts"], + id=checkpoint["id"], + 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()}, + ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 07d4a97ad..5ff4771b3 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -29,7 +29,6 @@ from typing_extensions import ParamSpec, Self from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel -from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import ( EXCLUDED_METADATA_KEYS, WRITES_IDX_MAP, @@ -39,7 +38,6 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, PendingWrite, - copy_checkpoint, ) from langgraph.constants import ( CONF, @@ -86,6 +84,7 @@ from langgraph.pregel.algo import ( ) from langgraph.pregel.checkpoint import ( channels_from_checkpoint, + copy_checkpoint, create_checkpoint, empty_checkpoint, ) @@ -963,13 +962,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) self.stack = ExitStack() if checkpointer: - if checkpointer._get_next_version_legacy: - empty_channel: LastValue[Any] = LastValue(Any) - self.checkpointer_get_next_version = ( - lambda c: checkpointer.get_next_version(c, empty_channel) # type: ignore[call-arg] - ) - else: - self.checkpointer_get_next_version = checkpointer.get_next_version + self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.put_writes self.checkpointer_put_writes_accepts_task_path = ( signature(checkpointer.put_writes).parameters.get("task_path") @@ -1142,13 +1135,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): ) self.stack = AsyncExitStack() if checkpointer: - if checkpointer._get_next_version_legacy: - empty_channel: LastValue[Any] = LastValue(Any) - self.checkpointer_get_next_version = ( - lambda c: checkpointer.get_next_version(c, empty_channel) # type: ignore[call-arg] - ) - else: - self.checkpointer_get_next_version = checkpointer.get_next_version + self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.aput_writes self.checkpointer_put_writes_accepts_task_path = ( signature(checkpointer.aput_writes).parameters.get("task_path") diff --git a/libs/langgraph/tests/test_checkpoint_migration.py b/libs/langgraph/tests/test_checkpoint_migration.py index e284af3d5..85229c7a8 100644 --- a/libs/langgraph/tests/test_checkpoint_migration.py +++ b/libs/langgraph/tests/test_checkpoint_migration.py @@ -7,12 +7,9 @@ from typing import Annotated, Literal, Optional, Union import pytest from typing_extensions import TypedDict -from langgraph.checkpoint.base import ( - BaseCheckpointSaver, - CheckpointTuple, - copy_checkpoint, -) +from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple from langgraph.graph.state import StateGraph +from langgraph.pregel.checkpoint import copy_checkpoint from langgraph.types import Command, Interrupt, PregelTask, StateSnapshot, interrupt from langgraph.utils.config import patch_configurable from tests.any_int import AnyInt diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 05561bc78..3d38558ab 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -159,7 +159,7 @@ def test_checkpoint_errors() -> None: raise ValueError("Faulty put_writes") class FaultyVersionCheckpointer(InMemorySaver): - def get_next_version(self, current: Optional[int]) -> int: + def get_next_version(self, current: Optional[int], channel: None) -> int: raise ValueError("Faulty get_next_version") def logic(inp: str) -> str: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f7934a9f5..ff5e8496d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -103,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]) -> int: + def get_next_version(self, current: Optional[int], channel: None) -> int: raise ValueError("Faulty get_next_version") def logic(inp: str) -> str: diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 9bf35a941..8f3149663 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -591,7 +591,7 @@ def create_react_agent( workflow = StateGraph(state_schema, config_schema=config_schema) workflow.add_node( "agent", - RunnableCallable(call_model, acall_model), # type: ignore[call-overload] + RunnableCallable(call_model, acall_model), input_schema=input_schema, ) if pre_model_hook is not None: @@ -610,7 +610,7 @@ def create_react_agent( if response_format is not None: workflow.add_node( "generate_structured_response", - RunnableCallable( # type: ignore[call-overload] + RunnableCallable( generate_structured_response, agenerate_structured_response, ), @@ -660,10 +660,10 @@ def create_react_agent( # Define the two nodes we will cycle between workflow.add_node( "agent", - RunnableCallable(call_model, acall_model), # type: ignore[call-overload] + RunnableCallable(call_model, acall_model), input_schema=input_schema, ) - workflow.add_node("tools", tool_node) # type: ignore[call-overload] + workflow.add_node("tools", tool_node) # Optionally add a pre-model hook node that will be called # every time before the "agent" (LLM-calling node) @@ -693,7 +693,7 @@ def create_react_agent( if response_format is not None: workflow.add_node( "generate_structured_response", - RunnableCallable( # type: ignore[call-overload] + RunnableCallable( generate_structured_response, agenerate_structured_response, ), diff --git a/libs/prebuilt/tests/memory_assert.py b/libs/prebuilt/tests/memory_assert.py index f88f0358f..10b93fdbd 100644 --- a/libs/prebuilt/tests/memory_assert.py +++ b/libs/prebuilt/tests/memory_assert.py @@ -13,9 +13,9 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, SerializerProtocol, - copy_checkpoint, ) from langgraph.checkpoint.memory import InMemorySaver, PersistentDict +from langgraph.pregel.checkpoint import copy_checkpoint class NoopSerializer(SerializerProtocol):