WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-09 11:17:53 +02:00
Code Issues Packages Projects Releases Wiki Activity

checkpoint: switch thread_ts -> checkpoint_id, add checkpoint_ns, change serializer protocol (#1185)

Browse Source 
---------

Co-authored-by: Nuno Campos <nuno@langchain.dev>
This commit is contained in:
Vadym Barda
2024-08-02 01:08:19 +00:00
committed by GitHub
co-authored by Nuno Campos
parent 862afa27de
commit 4b2187c9a3
22 changed files with 1491 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/README.md
+5 -5
View File
@@ -10,21 +10,21 @@ Checkpoint is a snapshot of the graph state at a given point in time. Checkpoint
### Thread
Threads enable the checkpointing of multiple different runs, making them essential for multi-tenant chat applications and other scenarios where maintaining separate states is necessary. A thread is a unique ID assigned to a series of checkpoints saved by a checkpointer. When using a checkpointer, you must specify a `thread_id` or `thread_ts` when running the graph.
Threads enable the checkpointing of multiple different runs, making them essential for multi-tenant chat applications and other scenarios where maintaining separate states is necessary. A thread is a unique ID assigned to a series of checkpoints saved by a checkpointer. When using a checkpointer, you must specify a `thread_id` and optionally `checkpoint_id` when running the graph.
- `thread_id` is simply the ID of a thread. This is always required
- `thread_ts` can optionally be passed. This identifier refers to a specific checkpoint within a thread. This can be used to kick of a run of a graph from some point halfway through a thread.
- `checkpoint_id` can optionally be passed. This identifier refers to a specific checkpoint within a thread. This can be used to kick of a run of a graph from some point halfway through a thread.
You must pass these when invoking the graph as part of the configurable part of the config, e.g.
```python
{"configurable": {"thread_id": "1"}} # valid config
{"configurable": {"thread_id": "1", "thread_ts": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}} # also valid config
{"configurable": {"thread_id": "1", "checkpoint_id": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}} # also valid config
```
### Serde
`langgraph_checkpoint` also defines protocol for serialization/deserialization (serde) and provides an default implementation (`langgraph_checkpoint.serde.jsonplus.JsonPlusSerializer`) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more.
`langgraph_checkpoint` also defines protocol for serialization/deserialization (serde) and provides an default implementation (`langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer`) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more.
### Pending writes
@@ -32,7 +32,7 @@ When a graph node fails mid-execution at a given superstep, LangGraph stores pen
## Interface
Each checkpointer should conform to `langgraph_checkpoint.BaseCheckpointSaver` interface and must implement the following methods:
Each checkpointer should conform to `langgraph.checkpoint.base.BaseCheckpointSaver` interface and must implement the following methods:
- `.put` - Store a checkpoint with its configuration and metadata.
- `.put_writes` - Store intermediate writes linked to a checkpoint (i.e. pending writes).
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+22 -6
View File
@@ -19,7 +19,7 @@ from typing import (
from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import (
ChannelProtocol,
@@ -172,10 +172,19 @@ CheckpointThreadId = ConfigurableFieldSpec(
is_shared=True,
)
CheckpointThreadTs = ConfigurableFieldSpec(
id="thread_ts",
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=None,
is_shared=True,
)
CheckpointId = ConfigurableFieldSpec(
id="checkpoint_id",
annotation=Optional[str],
name="Thread Timestamp",
name="Checkpoint ID",
description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.",
default=None,
is_shared=True,
@@ -203,7 +212,7 @@ class BaseCheckpointSaver(ABC):
*,
serde: Optional[SerializerProtocol] = None,
) -> None:
self.serde = serde or self.serde
self.serde = maybe_add_typed_methods(serde or self.serde)
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
@@ -212,7 +221,7 @@ class BaseCheckpointSaver(ABC):
Returns:
list[ConfigurableFieldSpec]: List of configuration field specs.
"""
return [CheckpointThreadId, CheckpointThreadTs]
return [CheckpointThreadId, CheckpointNS, CheckpointId]
def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
"""Fetch a checkpoint using the given configuration.
@@ -414,3 +423,10 @@ class EmptyChannelError(Exception):
for the first time yet."""
pass
def get_checkpoint_id(config: RunnableConfig) -> Optional[str]:
"""Get checkpoint ID in a backwards-compatible manner (fallback on thread_ts)."""
return config["configurable"].get(
"checkpoint_id", config["configurable"].get("thread_ts")
)
libs/checkpoint/langgraph/checkpoint/memory.py
+72 -40
View File
@@ -11,6 +11,7 @@ from langgraph.checkpoint.base import (
CheckpointMetadata,
CheckpointTuple,
SerializerProtocol,
get_checkpoint_id,
)
@@ -44,7 +45,8 @@ class MemorySaver(BaseCheckpointSaver):
asyncio.run(coro) # Output: 2
"""
storage: defaultdict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]
# thread ID -> checkpoint NS -> checkpoint ID -> checkpoint mapping
storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]]
def __init__(
self,
@@ -52,14 +54,14 @@ class MemorySaver(BaseCheckpointSaver):
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.storage = defaultdict(dict)
self.storage = defaultdict(lambda: defaultdict(dict))
self.writes = defaultdict(list)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the in-memory storage.
This method retrieves a checkpoint tuple from the in-memory storage based on the
provided config. If the config contains a "thread_ts" key, the checkpoint with
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
@@ -70,45 +72,54 @@ class MemorySaver(BaseCheckpointSaver):
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
"""
thread_id = config["configurable"]["thread_id"]
if ts := config["configurable"].get("thread_ts"):
if saved := self.storage[thread_id].get(ts):
checkpoint, metadata, parent_ts = saved
writes = self.writes[(thread_id, ts)]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id := get_checkpoint_id(config):
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)]
return CheckpointTuple(
config=config,
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
checkpoint=self.serde.loads_typed(checkpoint),
metadata=self.serde.loads_typed(metadata),
pending_writes=[
(id, c, self.serde.loads(v)) for id, c, v in writes
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_ts
if parent_checkpoint_id
else None,
)
else:
if checkpoints := self.storage[thread_id]:
ts = max(checkpoints.keys())
checkpoint, metadata, parent_ts = checkpoints[ts]
writes = self.writes[(thread_id, ts)]
if checkpoints := self.storage[thread_id][checkpoint_ns]:
checkpoint_id = max(checkpoints.keys())
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
return CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
checkpoint=self.serde.loads_typed(checkpoint),
metadata=self.serde.loads_typed(metadata),
pending_writes=[
(id, c, self.serde.loads(v)) for id, c, v in writes
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_ts
if parent_checkpoint_id
else None,
)
@@ -135,16 +146,25 @@ class MemorySaver(BaseCheckpointSaver):
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
checkpoint_ns = (
config["configurable"].get("checkpoint_ns", "") if config else ""
)
for thread_id in thread_ids:
for ts, (checkpoint, metadata_b, parent_ts) in sorted(
self.storage[thread_id].items(), key=lambda x: x[0], reverse=True
for checkpoint_id, (checkpoint, metadata_b, parent_checkpoint_id) in sorted(
self.storage[thread_id][checkpoint_ns].items(),
key=lambda x: x[0],
reverse=True,
):
# filter by thread_ts
if before and ts >= before["configurable"]["thread_ts"]:
# filter by checkpoint ID
if (
before
and (before_checkpoint_id := get_checkpoint_id(before))
and checkpoint_id >= before_checkpoint_id
):
continue
# filter by metadata
metadata = self.serde.loads(metadata_b)
metadata = self.serde.loads_typed(metadata_b)
if filter and not all(
query_value == metadata[query_key]
for query_key, query_value in filter.items()
@@ -158,16 +178,23 @@ class MemorySaver(BaseCheckpointSaver):
limit -= 1
yield CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
checkpoint=self.serde.loads_typed(checkpoint),
metadata=metadata,
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_ts
if parent_checkpoint_id
else None,
)
@@ -190,19 +217,22 @@ class MemorySaver(BaseCheckpointSaver):
Returns:
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
"""
self.storage[config["configurable"]["thread_id"]].update(
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
self.storage[thread_id][checkpoint_ns].update(
{
checkpoint["id"]: (
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
config["configurable"].get("thread_ts"), # parent
self.serde.dumps_typed(checkpoint),
self.serde.dumps_typed(metadata),
config["configurable"].get("checkpoint_id"), # parent
)
}
)
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["id"],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
@@ -226,9 +256,11 @@ class MemorySaver(BaseCheckpointSaver):
RunnableConfig: The updated config containing the saved writes' timestamp.
"""
thread_id = config["configurable"]["thread_id"]
ts = config["configurable"]["thread_ts"]
self.writes[(thread_id, ts)].extend(
[(task_id, c, self.serde.dumps(v)) for c, v in writes]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
checkpoint_id = config["configurable"]["checkpoint_id"]
key = (thread_id, checkpoint_ns, checkpoint_id)
self.writes[key].extend(
[(task_id, c, self.serde.dumps_typed(v)) for c, v in writes]
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
libs/checkpoint/langgraph/checkpoint/serde/base.py
+28
View File
@@ -5,7 +5,9 @@ class SerializerProtocol(Protocol):
"""Protocol for serialization and deserialization of objects.
- `dumps`: Serialize an object to bytes.
- `dumps_typed`: Serialize an object to a tuple (type, bytes).
- `loads`: Deserialize an object from bytes.
- `loads_typed`: Deserialize an object from a tuple (type, bytes).
Valid implementations include the `pickle`, `json` and `orjson` modules.
"""
@@ -13,5 +15,31 @@ class SerializerProtocol(Protocol):
def dumps(self, obj: Any) -> bytes:
...
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
...
def loads(self, data: bytes) -> Any:
...
def loads_typed(self, data: tuple[str, bytes]) -> Any:
...
class SerializerCompat(SerializerProtocol):
def __init__(self, serde: SerializerProtocol) -> None:
self.serde = serde
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
return type(obj).__name__, self.serde.dumps(obj)
def loads_typed(self, data: tuple[str, bytes]) -> Any:
return self.serde.loads(data[1])
def maybe_add_typed_methods(serde: SerializerProtocol) -> SerializerProtocol:
"""Wrap serde old serde implementations in a class with loads_typed and dumps_typed for backwards compatibility."""
if not hasattr(serde, "loads_typed") or not hasattr(serde, "dumps_typed"):
return SerializerCompat(serde)
return serde
libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+9
View File
@@ -99,5 +99,14 @@ class JsonPlusSerializer(SerializerProtocol):
"utf-8", "ignore"
)
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
return "json", self.dumps(obj)
def loads(self, data: bytes) -> Any:
return json.loads(data, object_hook=self._reviver)
def loads_typed(self, data: tuple[str, bytes]) -> Any:
type_, data_ = data
if type_ != "json":
raise ValueError("JsonPlusSerializer can only deserialize `json` data")
return self.loads(data_)
libs/checkpoint/langgraph/checkpoint/sqlite/__init__.py
+89 -48
View File
@@ -15,9 +15,11 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
EmptyChannelError,
SerializerProtocol,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import ChannelProtocol
from langgraph.checkpoint.sqlite.utils import JsonPlusSerializerCompat, search_where
from langgraph.checkpoint.sqlite.utils import search_where
_AIO_ERROR_MSG = (
"The SqliteSaver does not support async methods. "
@@ -61,11 +63,9 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
>>> graph.get_state(config)
>>> result = graph.invoke(3, config)
>>> graph.get_state(config)
StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '1', 'thread_ts': '2024-05-04T06:32:42.235444+00:00'}}, parent_config=None)
StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_id': '0c62ca34-ac19-445d-bbb0-5b4984975b2a'}}, parent_config=None)
""" # noqa
serde = JsonPlusSerializerCompat()
conn: sqlite3.Connection
is_setup: bool
@@ -76,6 +76,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.jsonplus_serde = JsonPlusSerializer()
self.conn = conn
self.is_setup = False
self.lock = threading.Lock()
@@ -134,20 +135,24 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
PRAGMA journal_mode=WAL;
CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
thread_ts TEXT NOT NULL,
parent_ts TEXT,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
parent_checkpoint_id TEXT,
type TEXT,
checkpoint BLOB,
metadata BLOB,
PRIMARY KEY (thread_id, thread_ts)
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);
CREATE TABLE IF NOT EXISTS writes (
thread_id TEXT NOT NULL,
thread_ts TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
type TEXT,
value BLOB,
PRIMARY KEY (thread_id, thread_ts, task_id, idx)
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);
"""
)
@@ -180,7 +185,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
"""Get a checkpoint tuple from the database.
This method retrieves a checkpoint tuple from the SQLite database based on the
provided config. If the config contains a "thread_ts" key, the checkpoint with
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
@@ -203,63 +208,76 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
>>> config = {
... "configurable": {
... "thread_id": "1",
... "thread_ts": "2024-05-04T06:32:42.235444+00:00",
... "checkpoint_id": "2024-05-04T06:32:42.235444+00:00",
... }
... }
>>> checkpoint_tuple = memory.get_tuple(config)
>>> print(checkpoint_tuple)
CheckpointTuple(...)
""" # noqa
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
with self.cursor(transaction=False) as cur:
# find the latest checkpoint for the thread_id
if config["configurable"].get("thread_ts"):
if checkpoint_id := get_checkpoint_id(config):
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
checkpoint_ns,
checkpoint_id,
),
)
else:
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(str(config["configurable"]["thread_id"]),),
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1",
(str(config["configurable"]["thread_id"]), checkpoint_ns),
)
# if a checkpoint is found, return it
if value := cur.fetchone():
if not config["configurable"].get("thread_ts"):
(
thread_id,
checkpoint_id,
parent_checkpoint_id,
type,
checkpoint,
metadata,
) = value
if not get_checkpoint_id(config):
config = {
"configurable": {
"thread_id": value[0],
"thread_ts": value[1],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
}
# find any pending writes
cur.execute(
"SELECT task_id, channel, value FROM writes WHERE thread_id = ? AND thread_ts = ?",
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
checkpoint_ns,
str(config["configurable"]["checkpoint_id"]),
),
)
# deserialize the checkpoint and metadata
return CheckpointTuple(
config,
self.serde.loads(value[3]),
self.serde.loads(value[4]) if value[4] is not None else {},
self.serde.loads_typed((type, checkpoint)),
self.jsonplus_serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": value[0],
"thread_ts": value[2],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if value[2]
if parent_checkpoint_id
else None
),
[
(task_id, channel, self.serde.loads(value))
for task_id, channel, value in cur
(task_id, channel, self.serde.loads_typed((type, value)))
for task_id, channel, type, value in cur
],
)
@@ -295,33 +313,48 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
[CheckpointTuple(...), CheckpointTuple(...)]
>>> config = {"configurable": {"thread_id": "1"}}
>>> before = {"configurable": {"thread_ts": "2024-05-04T06:32:42.235444+00:00"}}
>>> before = {"configurable": {"checkpoint_id": "2024-05-04T06:32:42.235444+00:00"}}
>>> checkpoints = list(memory.list(config, before=before))
>>> print(checkpoints)
[CheckpointTuple(...), ...]
"""
where, param_values = search_where(config, filter, before)
query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata
query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
FROM checkpoints
{where}
ORDER BY thread_ts DESC"""
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
with self.cursor(transaction=False) as cur:
cur.execute(query, param_values)
for thread_id, thread_ts, parent_ts, value, metadata in cur:
for (
thread_id,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
type,
checkpoint,
metadata,
) in cur:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
self.serde.loads(metadata) if metadata is not None else {},
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
self.serde.loads_typed((type, checkpoint)),
self.jsonplus_serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_ts
if parent_checkpoint_id
else None
),
)
@@ -354,23 +387,30 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
>>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "data": {"key": "value"}}
>>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}})
>>> print(saved_config)
{"configurable": {"thread_id": "1", "thread_ts": 2024-05-04T06:32:42.235444+00:00"}}
{"configurable": {"thread_id": "1", "checkpoint_id": 2024-05-04T06:32:42.235444+00:00"}}
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
serialized_metadata = self.jsonplus_serde.dumps(metadata)
with self.lock, self.cursor() as cur:
cur.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
checkpoint_ns,
checkpoint["id"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
config["configurable"].get("checkpoint_id"),
type_,
serialized_checkpoint,
serialized_metadata,
),
)
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["id"],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
@@ -391,15 +431,16 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
"""
with self.lock, self.cursor() as cur:
cur.executemany(
"INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)",
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
idx,
channel,
self.serde.dumps(value),
*self.serde.dumps_typed(value),
)
for idx, (channel, value) in enumerate(writes)
],
@@ -463,7 +504,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
current_v = int(current.split(".")[0])
next_v = current_v + 1
try:
next_h = md5(self.serde.dumps(channel.checkpoint())).hexdigest()
next_h = md5(self.serde.dumps_typed(channel.checkpoint())[1]).hexdigest()
except EmptyChannelError:
next_h = ""
return f"{next_v:032}.{next_h}"
libs/checkpoint/langgraph/checkpoint/sqlite/aio.py
+85 -44
View File
@@ -23,8 +23,10 @@ from langgraph.checkpoint.base import (
CheckpointMetadata,
CheckpointTuple,
SerializerProtocol,
get_checkpoint_id,
)
from langgraph.checkpoint.sqlite.utils import JsonPlusSerializerCompat, search_where
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.sqlite.utils import search_where
T = TypeVar("T", bound=callable)
@@ -113,12 +115,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
... saved_config = await saver.aput(config, checkpoint)
... print(saved_config)
>>> asyncio.run(main())
{"configurable": {"thread_id": "1", "thread_ts": "2023-05-03T10:00:00Z"}}
{"configurable": {"thread_id": "1", "checkpoint_id": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}}
```
"""
serde = JsonPlusSerializerCompat()
lock: asyncio.Lock
is_setup: bool
@@ -129,6 +129,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
serde: Optional[SerializerProtocol] = None,
):
super().__init__(serde=serde)
self.jsonplus_serde = JsonPlusSerializer()
self.conn = conn
self.lock = asyncio.Lock()
self.is_setup = False
@@ -209,20 +210,24 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
PRAGMA journal_mode=WAL;
CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
thread_ts TEXT NOT NULL,
parent_ts TEXT,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
parent_checkpoint_id TEXT,
type TEXT,
checkpoint BLOB,
metadata BLOB,
PRIMARY KEY (thread_id, thread_ts)
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);
CREATE TABLE IF NOT EXISTS writes (
thread_id TEXT NOT NULL,
thread_ts TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
type TEXT,
value BLOB,
PRIMARY KEY (thread_id, thread_ts, task_id, idx)
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);
"""
):
@@ -234,7 +239,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"""Get a checkpoint tuple from the database asynchronously.
This method retrieves a checkpoint tuple from the SQLite database based on the
provided config. If the config contains a "thread_ts" key, the checkpoint with
provided config. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
@@ -245,56 +250,69 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
"""
await self.setup()
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
async with self.conn.cursor() as cur:
# find the latest checkpoint for the thread_id
if config["configurable"].get("thread_ts"):
if checkpoint_id := get_checkpoint_id(config):
await cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
checkpoint_ns,
checkpoint_id,
),
)
else:
await cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
(str(config["configurable"]["thread_id"]),),
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1",
(str(config["configurable"]["thread_id"]), checkpoint_ns),
)
# if a checkpoint is found, return it
if value := await cur.fetchone():
if not config["configurable"].get("thread_ts"):
(
thread_id,
checkpoint_id,
parent_checkpoint_id,
type,
checkpoint,
metadata,
) = value
if not get_checkpoint_id(config):
config = {
"configurable": {
"thread_id": value[0],
"thread_ts": value[1],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
}
# find any pending writes
await cur.execute(
"SELECT task_id, channel, value FROM writes WHERE thread_id = ? AND thread_ts = ?",
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
checkpoint_ns,
str(config["configurable"]["checkpoint_id"]),
),
)
# deserialize the checkpoint and metadata
return CheckpointTuple(
config,
self.serde.loads(value[3]),
self.serde.loads(value[4]) if value[4] is not None else {},
self.serde.loads_typed((type, checkpoint)),
self.jsonplus_serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": value[0],
"thread_ts": value[2],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if value[2]
if parent_checkpoint_id
else None
),
[
(task_id, channel, self.serde.loads(value))
async for task_id, channel, value in cur
(task_id, channel, self.serde.loads_typed((type, value)))
async for task_id, channel, type, value in cur
],
)
@@ -322,26 +340,41 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"""
await self.setup()
where, param_values = search_where(config, filter, before)
query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata
query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
FROM checkpoints
{where}
ORDER BY thread_ts DESC"""
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
async with self.conn.execute(query, param_values) as cursor:
async for thread_id, thread_ts, parent_ts, value, metadata in cursor:
async for (
thread_id,
checkpoint_ns,
checkpoint_id,
parent_checkpoint_id,
type,
checkpoint,
metadata,
) in cursor:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
self.serde.loads(metadata) if metadata is not None else {},
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
self.serde.loads_typed((type, checkpoint)),
self.jsonplus_serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_ts
if parent_checkpoint_id
else None
),
)
@@ -366,21 +399,28 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
"""
await self.setup()
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
serialized_metadata = self.jsonplus_serde.dumps(metadata)
async with self.conn.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
(
str(config["configurable"]["thread_id"]),
checkpoint_ns,
checkpoint["id"],
config["configurable"].get("thread_ts"),
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
config["configurable"].get("checkpoint_id"),
type_,
serialized_checkpoint,
serialized_metadata,
),
):
await self.conn.commit()
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["id"],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
@@ -401,15 +441,16 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
"""
await self.setup()
async with self.conn.executemany(
"INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)",
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
idx,
channel,
self.serde.dumps(value),
*self.serde.dumps_typed(value),
)
for idx, (channel, value) in enumerate(writes)
],
libs/checkpoint/langgraph/checkpoint/sqlite/utils.py
+6 -32
View File
@@ -1,38 +1,9 @@
import json
import pickle
from typing import Any, Dict, Optional, Sequence, Tuple
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
class JsonPlusSerializerCompat(JsonPlusSerializer):
"""A serializer that supports loading pickled checkpoints for backwards compatibility.
This serializer extends the JsonPlusSerializer and adds support for loading pickled
checkpoints. If the input data starts with b"\x80" and ends with b".", it is treated
as a pickled checkpoint and loaded using pickle.loads(). Otherwise, the default
JsonPlusSerializer behavior is used.
Examples:
>>> import pickle
>>> from langgraph.checkpoint.sqlite import JsonPlusSerializerCompat
>>>
>>> serializer = JsonPlusSerializerCompat()
>>> pickled_data = pickle.dumps({"key": "value"})
>>> loaded_data = serializer.loads(pickled_data)
>>> print(loaded_data) # Output: {"key": "value"}
>>>
>>> json_data = '{"key": "value"}'.encode("utf-8")
>>> loaded_data = serializer.loads(json_data)
>>> print(loaded_data) # Output: {"key": "value"}
"""
def loads(self, data: bytes) -> Any:
if data.startswith(b"\x80") and data.endswith(b"."):
return pickle.loads(data)
return super().loads(data)
from langgraph.checkpoint.base import get_checkpoint_id
def _metadata_predicate(
@@ -99,6 +70,9 @@ def search_where(
if config is not None:
wheres.append("thread_id = ?")
param_values.append(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
wheres.append("checkpoint_ns = ?")
param_values.append(checkpoint_ns)
# construct predicate for metadata filter
if filter:
@@ -108,7 +82,7 @@ def search_where(
# construct predicate for `before`
if before is not None:
wheres.append("thread_ts < ?")
param_values.append(before["configurable"]["thread_ts"])
wheres.append("checkpoint_id < ?")
param_values.append(get_checkpoint_id(before))
return ("WHERE " + " AND ".join(wheres) if wheres else "", param_values)
libs/checkpoint/tests/test_aiosqlite.py
+49 -3
View File
@@ -17,14 +17,31 @@ class TestAsyncSqliteSaver:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {"thread_id": "thread-1", "thread_ts": "1"}
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {"thread_id": "thread-2", "thread_ts": "2"}
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
@@ -38,12 +55,14 @@ class TestAsyncSqliteSaver:
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async def test_asearch(self):
# set up test
# save checkpoints
await self.sqlite_saver.aput(self.config_1, self.chkpnt_1, self.metadata_1)
await self.sqlite_saver.aput(self.config_2, self.chkpnt_2, self.metadata_2)
await self.sqlite_saver.aput(self.config_3, self.chkpnt_3, self.metadata_3)
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
@@ -70,11 +89,38 @@ class TestAsyncSqliteSaver:
search_results_3 = [
c async for c in sqlite_saver.alist(None, filter=query_3)
]
assert len(search_results_3) == 2
assert len(search_results_3) == 3
search_results_4 = [
c async for c in sqlite_saver.alist(None, filter=query_4)
]
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
search_results_5 = [
c
async for c in self.sqlite_saver.alist(
{"configurable": {"thread_id": "thread-2"}}
)
]
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = [
c
async for c in self.sqlite_saver.alist(
{
"configurable": {
"thread_id": "thread-2",
"checkpoint_ns": "inner",
}
}
)
]
assert len(search_results_6) == 1
assert (
search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
)
# TODO: test before and limit params
libs/checkpoint/tests/test_jsonplus.py
+5 -5
View File
@@ -99,14 +99,14 @@ def test_serde_jsonplus() -> None:
serde = JsonPlusSerializer()
dumped = serde.dumps(to_serialize)
dumped = serde.dumps_typed(to_serialize)
assert (
dumped
== b"""{"uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}"""
assert dumped == (
"json",
b"""{"uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""",
)
assert serde.loads(dumped) == {
assert serde.loads_typed(dumped) == {
**to_serialize,
"text": [v.encode("utf-8", "ignore").decode() for v in to_serialize["text"]],
}
libs/checkpoint/tests/test_memory.py
+37 -2
View File
@@ -17,14 +17,31 @@ class TestMemorySaver:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {"thread_id": "thread-1", "thread_ts": "1"}
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
# for backwards compatibility testing
"thread_ts": "1",
}
}
self.config_2: RunnableConfig = {
"configurable": {"thread_id": "thread-2", "thread_ts": "2"}
"configurable": {
"thread_id": "thread-2",
"checkpoint_ns": "",
"checkpoint_id": "2",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
@@ -38,12 +55,14 @@ class TestMemorySaver:
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async def test_search(self):
# set up test
# save checkpoints
self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1)
self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2)
self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3)
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
@@ -68,6 +87,22 @@ class TestMemorySaver:
search_results_4 = list(self.memory_saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
search_results_5 = list(
self.memory_saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = list(
self.memory_saver.list(
{"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}}
)
)
assert len(search_results_6) == 1
assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
# TODO: test before and limit params
async def test_asearch(self):
libs/checkpoint/tests/test_sqlite.py
+38 -4
View File
@@ -18,14 +18,31 @@ class TestSqliteSaver:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {"thread_id": "thread-1", "thread_ts": "1"}
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {"thread_id": "thread-2", "thread_ts": "2"}
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
@@ -46,6 +63,7 @@ class TestSqliteSaver:
# save checkpoints
self.sqlite_saver.put(self.config_1, self.chkpnt_1, self.metadata_1)
self.sqlite_saver.put(self.config_2, self.chkpnt_2, self.metadata_2)
self.sqlite_saver.put(self.config_3, self.chkpnt_3, self.metadata_3)
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
@@ -65,16 +83,32 @@ class TestSqliteSaver:
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = list(self.sqlite_saver.list(None, filter=query_3))
assert len(search_results_3) == 2
assert len(search_results_3) == 3
search_results_4 = list(self.sqlite_saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
search_results_5 = list(
self.sqlite_saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = list(
self.sqlite_saver.list(
{"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}}
)
)
assert len(search_results_6) == 1
assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
# TODO: test before and limit params
def test_search_where(self):
# call method / assertions
expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND thread_ts < ?"
expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND checkpoint_id < ?"
expected_param_values_1 = ["input", 2, "{}", 1, "1"]
assert search_where(None, self.metadata_1, self.config_1) == (
expected_predicate_1,
libs/langgraph/langgraph/constants.py
+2
View File
@@ -21,6 +21,8 @@ TAG_HIDDEN = "langsmith:hidden"
START = "__start__"
END = "__end__"
CHECKPOINT_NAMESPACE_SEPARATOR = "|"
class Send:
"""A message or packet to send to a specific node in the graph.
libs/langgraph/langgraph/graph/graph.py
+12 -1
View File
@@ -25,7 +25,13 @@ from langchain_core.runnables.graph import Node as DrawableNode
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, START, TAG_HIDDEN, Send
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
END,
START,
TAG_HIDDEN,
Send,
)
from langgraph.errors import InvalidUpdateError
from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import PregelNode
@@ -154,6 +160,11 @@ class Graph:
*,
metadata: Optional[dict[str, Any]] = None,
) -> None:
if isinstance(node, str) and CHECKPOINT_NAMESPACE_SEPARATOR in node:
raise ValueError(
f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names."
)
if self.compiled:
logger.warning(
"Adding a node to a graph that has already been compiled. This will "
libs/langgraph/langgraph/graph/state.py
+7 -1
View File
@@ -30,7 +30,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.named_barrier_value import NamedBarrierValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR, TAG_HIDDEN
from langgraph.errors import InvalidUpdateError
from langgraph.graph.graph import (
END,
@@ -311,6 +311,12 @@ class StateGraph(Graph):
raise ValueError(f"Node `{node}` already present.")
if node == END or node == START:
raise ValueError(f"Node `{node}` is reserved.")
if CHECKPOINT_NAMESPACE_SEPARATOR in node:
raise ValueError(
f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names."
)
try:
if isfunction(action) and (
hints := get_type_hints(action.__call__) or get_type_hints(action)
libs/langgraph/langgraph/pregel/__init__.py
+16 -2
View File
@@ -519,7 +519,14 @@ class Pregel(
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
step = saved.metadata.get("step", -1) if saved else -1
# merge configurable fields with previous checkpoint config
checkpoint_config = config
checkpoint_config = {
**config,
"configurable": {
**config["configurable"],
# TODO: add proper support for updating nested subgraph state
"checkpoint_ns": "",
},
}
if saved:
checkpoint_config = {
"configurable": {
@@ -681,7 +688,14 @@ class Pregel(
step = saved.metadata.get("step", -2) + 1 if saved else -1
# merge configurable fields with previous checkpoint config
checkpoint_config = config
checkpoint_config = {
**config,
"configurable": {
**config["configurable"],
# TODO: add proper support for updating nested subgraph state
"checkpoint_ns": "",
},
}
if saved:
checkpoint_config = {
"configurable": {
libs/langgraph/langgraph/pregel/algo.py
+24 -10
View File
@@ -33,6 +33,7 @@ from langgraph.checkpoint.base import (
create_checkpoint,
)
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
@@ -256,6 +257,7 @@ def prepare_next_tasks(
checkpointer: Optional[BaseCheckpointSaver] = None,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]:
parent_ns = config.get("configurable", {}).get("checkpoint_ns", "")
tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = []
# Consume pending packets
for packet in checkpoint["pending_sends"]:
@@ -275,7 +277,14 @@ def prepare_next_tasks(
"langgraph_triggers": triggers,
"langgraph_task_idx": len(tasks),
}
task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
checkpoint_ns = (
f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}"
if parent_ns
else packet.node
)
task_id = str(
uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata)))
)
writes = deque()
tasks.append(
PregelExecutableTask(
@@ -349,13 +358,18 @@ def prepare_next_tasks(
"langgraph_triggers": triggers,
"langgraph_task_idx": len(tasks),
}
task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
if parent_thread_id := config.get("configurable", {}).get(
"thread_id"
):
thread_id: Optional[str] = f"{parent_thread_id}-{name}"
else:
thread_id = None
checkpoint_ns = (
f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{name}"
if parent_ns
else name
)
task_id = str(
uuid5(
UUID(checkpoint["id"]),
json.dumps((checkpoint_ns, metadata)),
)
)
writes = deque()
tasks.append(
PregelExecutableTask(
@@ -389,8 +403,8 @@ def prepare_next_tasks(
),
CONFIG_KEY_CHECKPOINTER: checkpointer,
CONFIG_KEY_RESUMING: is_resuming,
"thread_id": thread_id,
"thread_ts": checkpoint["id"],
"checkpoint_id": checkpoint["id"],
"checkpoint_ns": checkpoint_ns,
},
),
triggers,
libs/langgraph/langgraph/pregel/debug.py
+4 -2
View File
@@ -71,7 +71,7 @@ def map_debug_tasks(
continue
metadata = config["metadata"].copy()
metadata.pop("thread_ts", None)
metadata.pop("checkpoint_id", None)
yield {
"type": "task",
@@ -97,7 +97,9 @@ def map_debug_task_results(
continue
metadata = config["metadata"].copy()
metadata.pop("thread_ts", None)
metadata.pop("checkpoint_id", None)
# TODO: make task IDs deterministic in tests and reuse task IDs for payload ID
metadata.pop("checkpoint_ns", None)
yield {
"type": "task_result",
libs/langgraph/langgraph/pregel/loop.py
+17 -3
View File
@@ -143,7 +143,10 @@ class PregelLoop:
**self.checkpoint_config,
"configurable": {
**self.checkpoint_config["configurable"],
"thread_ts": self.checkpoint["id"],
"checkpoint_ns": self.config["configurable"].get(
"checkpoint_ns", ""
),
"checkpoint_id": self.checkpoint["id"],
},
},
writes,
@@ -315,8 +318,19 @@ class PregelLoop:
# this is achieved by writing child checkpoints as progress is made
# (so that error recovery / resuming from interrupt don't lose work)
# but doing so always with an id equal to that of the parent checkpoint
id=self.config["configurable"]["thread_ts"] if self.is_nested else None,
id=self.config["configurable"]["checkpoint_id"]
if self.is_nested
else None,
)
self.checkpoint_config = {
**self.checkpoint_config,
"configurable": {
**self.checkpoint_config["configurable"],
"checkpoint_ns": self.config["configurable"].get(
"checkpoint_ns", ""
),
},
}
# save it, without blocking
# if there's a previous checkpoint save in progress, wait for it
# ensuring checkpointers receive checkpoints in order
@@ -331,7 +345,7 @@ class PregelLoop:
**self.checkpoint_config,
"configurable": {
**self.checkpoint_config["configurable"],
"thread_ts": self.checkpoint["id"],
"checkpoint_id": self.checkpoint["id"],
},
}
# produce debug output
libs/langgraph/tests/memory_assert.py
+22 -18
View File
@@ -15,17 +15,17 @@ from langgraph.checkpoint.memory import MemorySaver
class NoopSerializer(SerializerProtocol):
def loads(self, data: bytes) -> Any:
return data
def loads_typed(self, data: tuple[str, bytes]) -> Any:
return data[1]
def dumps(self, obj: Any) -> bytes:
return obj
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
return "type", obj
class MemorySaverAssertImmutable(MemorySaver):
serde = NoopSerializer()
storage_for_copies: defaultdict[str, dict[str, Checkpoint]]
storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]]
def __init__(
self,
@@ -34,7 +34,7 @@ class MemorySaverAssertImmutable(MemorySaver):
put_sleep: Optional[float] = None,
) -> None:
super().__init__(serde=serde)
self.storage_for_copies = defaultdict(dict)
self.storage_for_copies = defaultdict(lambda: defaultdict(dict))
self.put_sleep = put_sleep
def put(
@@ -49,14 +49,17 @@ class MemorySaverAssertImmutable(MemorySaver):
time.sleep(self.put_sleep)
# assert checkpoint hasn't been modified since last written
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
if saved := super().get(config):
assert (
self.serde.loads(self.storage_for_copies[thread_id][saved["id"]])
self.serde.loads_typed(
self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]]
)
== saved
)
self.storage_for_copies[thread_id][checkpoint["id"]] = self.serde.dumps(
copy_checkpoint(checkpoint)
)
self.storage_for_copies[thread_id][checkpoint_ns][
checkpoint["id"]
] = self.serde.dumps_typed(copy_checkpoint(checkpoint))
# call super to write checkpoint
return super().put(config, checkpoint, metadata)
@@ -91,23 +94,24 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
"""
configurable = config["configurable"].copy()
# remove thread_ts to make testing simpler
thread_ts = configurable.pop("thread_ts", None)
self.storage[config["configurable"]["thread_id"]].update(
# remove checkpoint_id to make testing simpler
checkpoint_id = configurable.pop("checkpoint_id", None)
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
self.storage[thread_id][checkpoint_ns].update(
{
checkpoint["id"]: (
self.serde.dumps(checkpoint),
self.serde.dumps_typed(checkpoint),
# merge configurable fields and metadata
self.serde.dumps({**configurable, **metadata}),
thread_ts,
self.serde.dumps_typed({**configurable, **metadata}),
checkpoint_id,
)
}
)
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["id"],
"checkpoint_id": checkpoint["id"],
}
}
libs/langgraph/tests/test_pregel.py
+468 -115
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
+474 -117
View File
File diff suppressed because it is too large Load Diff